@illuma-ai/agents 1.1.25 → 1.3.0

package/src/common/spawnPath.ts

@@ -0,0 +1,101 @@
+/**
+ * spawnPath — hierarchical invocation identity for nested multi-agent orchestration.
+ *
+ * A spawnPath is a slash-separated chain of spawnKeys from the root of the current
+ * agent invocation to the current spawn. The root agent has an empty spawnPath;
+ * each handoff/transfer/sequence that spawns a new subgraph appends a new spawnKey.
+ *
+ * Examples:
+ *   ""                               → primary agent (no spawn)
+ *   "call_abc"                       → first-level handoff child
+ *   "call_abc/call_def"              → grandchild (depth 2)
+ *   "call_abc/call_def/call_ghi"     → depth 3
+ *
+ * These utilities are the single source of truth for path manipulation across:
+ *   - @illuma-ai/agents (MultiAgentGraph, HandoffRegistry, callbacks)
+ *   - host api (initialize.js, callbacks.js, ExecutionTrace writes)
+ *   - host client (subagent store, sidebar rendering)
+ *
+ * See docs/multi-agent-nesting-architecture.md for the full design.
+ */
+
+/** Separator between spawnKeys in a spawnPath. Chosen so that the path looks
+ * like a filesystem/URL path, which makes it easy to read in logs and traces. */
+export const SPAWN_PATH_SEP = '/';
+
+/** Hard cap on nested multi-agent invocations. Prevents runaway recursion.
+ * Can be overridden at the host api layer via MULTI_AGENT_MAX_NESTING_DEPTH. */
+export const MAX_NESTING_DEPTH = 5;
+
+/**
+ * Append a spawnKey to a parent spawnPath.
+ *
+ * @param parent - Parent spawnPath (may be undefined/null/empty for root)
+ * @param key    - spawnKey to append
+ * @returns New spawnPath string
+ */
+export function buildSpawnPath(
+  parent: string | undefined | null,
+  key: string
+): string {
+  if (!key) {
+    throw new Error('[spawnPath] buildSpawnPath called with empty key');
+  }
+  if (parent == null || parent === '') return key;
+  return `${parent}${SPAWN_PATH_SEP}${key}`;
+}
+
+/**
+ * Compute the depth of a spawnPath.
+ * Root (empty) → 0; single-segment → 1; etc.
+ */
+export function spawnPathDepth(path: string | undefined | null): number {
+  if (path == null || path === '') return 0;
+  return path.split(SPAWN_PATH_SEP).filter(Boolean).length;
+}
+
+/**
+ * Return the parent spawnPath, or null if the input is already root.
+ *
+ * - parentSpawnPath("a/b/c") === "a/b"
+ * - parentSpawnPath("a") === ""
+ * - parentSpawnPath("") === null
+ */
+export function parentSpawnPath(
+  path: string | undefined | null
+): string | null {
+  if (path == null || path === '') return null;
+  const parts = path.split(SPAWN_PATH_SEP).filter(Boolean);
+  if (parts.length <= 1) return '';
+  return parts.slice(0, -1).join(SPAWN_PATH_SEP);
+}
+
+/** Split a spawnPath into its constituent spawnKey segments. */
+export function spawnPathParts(path: string | undefined | null): string[] {
+  if (path == null || path === '') return [];
+  return path.split(SPAWN_PATH_SEP).filter(Boolean);
+}
+
+/**
+ * Return the last spawnKey in a spawnPath (the "current" spawn).
+ * Returns null for root.
+ */
+export function leafSpawnKey(path: string | undefined | null): string | null {
+  const parts = spawnPathParts(path);
+  return parts.length === 0 ? null : parts[parts.length - 1];
+}
+
+/**
+ * True if `ancestor` is a strict ancestor of `descendant`. Root ("") is
+ * ancestor of everything except itself.
+ */
+export function isAncestorSpawnPath(
+  ancestor: string | undefined | null,
+  descendant: string | undefined | null
+): boolean {
+  const a = ancestor ?? '';
+  const d = descendant ?? '';
+  if (a === d) return false;
+  if (a === '') return d !== '';
+  return d.startsWith(a + SPAWN_PATH_SEP);
+}

package/src/graphs/Graph.ts

@@ -79,6 +79,9 @@ import { getChatModelClass, manualToolStreamProviders } from '@/llm/providers';
 import { ToolNode as CustomToolNode, toolsCondition } from '@/tools/ToolNode';
 import { ChatOpenAI, AzureChatOpenAI } from '@/llm/openai';
 import { safeDispatchCustomEvent } from '@/utils/events';
+import { mlog, mwarn } from '@/utils/logging';
+import { normalizeMessageToolCalls } from '@/utils/toolCallNormalization';
+import { isTruncationReason } from '@/utils/finishReasons';
 import {
   detectDocuments,
   shouldInjectMultiDocHint,
@@ -1144,10 +1147,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
       const resolved = agentContext.resolveStructuredOutputMode();
       method = resolved.method;
       if (resolved.warnings.length > 0) {
-        console.warn(
-          '[Graph] Structured output mode warnings:',
-          resolved.warnings
-        );
+        mwarn('[Graph] Structured output mode warnings:', resolved.warnings);
       }
     } else {
       // Legacy fallback: use the old mode-based resolution
@@ -1172,7 +1172,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
       );
       preparedSchema = prepared;
       if (warnings.length > 0) {
-        console.warn('[Graph] Schema preparation warnings:', warnings);
+        mwarn('[Graph] Schema preparation warnings:', warnings);
       }
     }
 
@@ -1264,7 +1264,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
             ? handleErrors
             : `The response did not match the expected schema. Error: ${lastError.message}. Please try again with a valid response.`;
 
-        console.warn(
+        mwarn(
           `[Graph] Structured output attempt ${attempts} failed: ${lastError.message}. Retrying...`
         );
 
@@ -1467,7 +1467,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         this._toolDiscoveryCache.getNewDiscoveries(messages);
       if (cachedDiscoveries.length > 0) {
         agentContext.markToolsAsDiscovered(cachedDiscoveries);
-        console.debug(
+        mlog(
           `[Graph:ToolDiscovery] Cached ${cachedDiscoveries.length} new tools (total: ${this._toolDiscoveryCache.size})`
         );
       }
@@ -1497,6 +1497,14 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
           clientOptions: effectiveClientOptions,
         });
 
+      // DEBUG: Log which model and tools each agent uses during handoff
+      mlog(
+        `[createCallModel] Agent "${agentId}" invoking LLM | provider=${agentContext.provider} | ` +
+          `model=${(effectiveClientOptions as Record<string, unknown>).model ?? 'default'} | ` +
+          `toolsForBinding=${toolsForBinding?.length ?? 0} | ` +
+          `toolNames=[${(toolsForBinding ?? []).map((t) => (t as { name?: string }).name ?? 'unknown').join(', ')}]`
+      );
+
       if (agentContext.systemRunnable) {
         model = agentContext.systemRunnable.pipe(model as Runnable);
       }
@@ -1507,7 +1515,15 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
       if (!config.signal) {
         config.signal = this.signal;
       }
-      this.config = config;
+      // First-writer-wins: `this.config` is used ONLY as a "has a run started"
+      // existence flag by the dispatch* methods (they never read its value —
+      // they read the current RunnableConfig from LangChain AsyncLocalStorage).
+      // Unconditionally reassigning here races across concurrent child
+      // subgraph.invoke() calls under parallel multi-agent handoffs; the last
+      // writer wins, and any dispatch firing between writes would historically
+      // have been tagged with the wrong child's metadata. Keeping the first
+      // write pinned makes this a true flag, eliminating the race.
+      this.config ??= config;
 
       let messagesToUse = messages;
 
@@ -1611,7 +1627,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
 
           if (oldMessages.length > 0) {
             this._summaryInFlight = true;
-            console.debug(
+            mlog(
               `[Graph:ProactiveSummary] Context at ${utilization.toFixed(1)}% (threshold ${threshold}%) — summarizing ${oldMessages.length} older msgs in background`
             );
 
@@ -1620,7 +1636,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
               .then((updated) => {
                 if (updated != null && updated !== '') {
                   this._cachedRunSummary = updated;
-                  console.debug(
+                  mlog(
                     `[Graph:ProactiveSummary] Background summary ready (len=${updated.length})`
                   );
                 }
@@ -1830,7 +1846,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         }
         agentContext.indexTokenCountMap = viewTokenMap;
 
-        console.debug(
+        mlog(
           `[Graph:Compaction] ${messages.length}→${viewParts.length} msgs | ` +
             `compacted=${compactedMessages.length} window=${recentMessages.length} | ` +
             `summary=${summarySource} | budget=${usedTokens}/${recentBudget}` +
@@ -1854,7 +1870,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
           if (shouldSummarize) {
             if (this._summaryInFlight) {
               this._pendingMessagesToRefine.push(...compactedMessages);
-              console.debug(
+              mlog(
                 `[Graph:Compaction] Summary in-flight, queued ${compactedMessages.length} msgs (pending=${this._pendingMessagesToRefine.length})`
               );
             } else {
@@ -1907,7 +1923,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         deduplicateSystemMessages(messagesToUse);
       if (removedCount > 0) {
         messagesToUse = dedupedMessages;
-        console.debug(
+        mlog(
           `[Graph:Dedup] Removed ${removedCount} duplicate system message(s)`
         );
       }
@@ -2031,7 +2047,6 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         );
       }
 
-
       // Get model info for analytics
       const bedrockOpts = agentContext.clientOptions as
         | t.BedrockAnthropicClientOptions
@@ -2145,11 +2160,11 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
 
         // Log when we detect the error
         if (isInputTooLongError) {
-          console.warn(
+          mwarn(
             '[Graph] Detected input too long error:',
             errorMessage.substring(0, 200)
           );
-          console.warn('[Graph] Checking emergency pruning conditions:', {
+          mwarn('[Graph] Checking emergency pruning conditions:', {
             hasPruneMessages: !!agentContext.pruneMessages,
             hasTokenCounter: !!agentContext.tokenCounter,
             maxContextTokens: agentContext.maxContextTokens,
@@ -2174,7 +2189,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
             const reducedMaxTokens = Math.floor(
               agentContext.maxContextTokens! * reductionFactor
             );
-            console.warn(
+            mwarn(
               `[Graph] Input too long. Retrying with ${reductionFactor * 100}% context (${reducedMaxTokens} tokens)...`
             );
 
@@ -2182,7 +2197,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
             // This is needed when messages were dynamically added without updating the token map
             let tokenMapForPruning = agentContext.indexTokenCountMap;
             if (Object.keys(tokenMapForPruning).length < messages.length) {
-              console.warn(
+              mwarn(
                 '[Graph] Building fresh token count map for emergency pruning...'
               );
               tokenMapForPruning = {};
@@ -2207,7 +2222,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
 
             // Skip if we can't fit any messages
             if (reducedMessages.length === 0) {
-              console.warn(
+              mwarn(
                 `[Graph] Cannot fit any messages at ${reductionFactor * 100}% reduction, trying next level...`
               );
               continue;
@@ -2290,7 +2305,7 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
                 retryErrorMsg.includes('validationexception');
 
               if (stillTooLong && reductionFactor > 0.1) {
-                console.warn(
+                mwarn(
                   `[Graph] Still too long at ${reductionFactor * 100}%, trying more aggressive pruning...`
                 );
               } else {
@@ -2362,6 +2377,27 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
        * handled everything — both paths become no-ops.
        */
       const responseMessage = result.messages?.[0];
+
+      // Tool-call name normalization — catches LLM output that names tools
+      // with wrong delimiters (outlook/operations), prefixes
+      // (functions.outlook_operations), case drift, counter suffixes, or
+      // empty names recoverable from the tool_call id. Mutates in place so
+      // the downstream ToolNode dispatch sees the corrected names.
+      if (responseMessage && agentContext.toolMap) {
+        const allowedNames = new Set(Object.keys(agentContext.toolMap));
+        if (allowedNames.size > 0) {
+          const rewrote = normalizeMessageToolCalls(
+            responseMessage,
+            allowedNames
+          );
+          if (rewrote) {
+            mlog(
+              `[Graph] normalized tool_call names on agent "${agentId}" response`
+            );
+          }
+        }
+      }
+
       const toolCalls = (responseMessage as AIMessageChunk | undefined)
         ?.tool_calls;
       const hasToolCalls = Array.isArray(toolCalls) && toolCalls.length > 0;
@@ -2477,12 +2513,22 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
         const messageStop = meta.messageStop as
           | Record<string, unknown>
           | undefined;
-        this.lastFinishReason =
+        const nextReason =
           (meta.finish_reason as string | undefined) ?? // OpenAI/Azure
           (meta.stop_reason as string | undefined) ?? // Anthropic direct API
           (meta.stopReason as string | undefined) ?? // Bedrock invoke (non-streaming)
           (messageStop?.stopReason as string | undefined) ?? // Bedrock streaming
           (meta.finishReason as string | undefined); // VertexAI/Google
+
+        // Sticky on truncation: a single Graph instance is reused across
+        // every scoped-subgraph inner node invocation (see MultiAgentGraph
+        // buildScopedSubgraph). If an earlier inner node hit max_tokens
+        // but a later inner node finished cleanly, the host's continuation layer
+        // would miss the truncation signal unless we preserve it. Keep the
+        // truncation reason pinned so the outer caller can retry.
+        if (!isTruncationReason(this.lastFinishReason)) {
+          this.lastFinishReason = nextReason;
+        }
       }
 
       this.cleanupSignalListener();
@@ -2545,7 +2591,7 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
               '[Graph] Deferred structured output failed after successful tool use:',
               structuredError
             );
-            console.warn(
+            mwarn(
               '[Graph] Falling back to unstructured response from tool-use phase'
             );
             return result;
@@ -2570,7 +2616,10 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
       state: t.BaseGraphState,
       config?: RunnableConfig
     ): string => {
-      this.config = config;
+      // First-writer-wins — see note in createCallModel. `this.config` is an
+      // existence flag only; assigning unconditionally would race under
+      // parallel child subgraph.invoke().
+      this.config ??= config;
       return toolsCondition(state, toolNode, this.invokedToolIds);
     };
 
@@ -2615,10 +2664,16 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
         default: () => [],
       }),
     });
+    // Pass compileOptions (including the HITL checkpointer) to the OUTER
+    // workflow — not just the inner agent subgraph. hasInterrupts() calls
+    // getState() on the outer compiled graph; without a checkpointer here,
+    // getState reports zero tasks and the HITL interrupt/resume loop breaks
+    // out immediately even though interrupt() fired correctly inside the
+    // agent subgraph.
     const workflow = new StateGraph(StateAnnotation)
       .addNode(this.defaultAgentId, agentNode, { ends: [END] })
       .addEdge(START, this.defaultAgentId)
-      .compile();
+      .compile(this.compileOptions as unknown as never);
 
     return workflow;
   }
@@ -2701,7 +2756,7 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
         }
       } catch (_e) {
         /** If we can't get agent context, that's okay - agentId remains undefined */
-        console.debug(
+        mlog(
           `[dispatchRunStep] Could not resolve agentId from metadata.langgraph_node="${(metadata as Record<string, unknown>).langgraph_node}": ${(_e as Error).message}`
         );
       }
@@ -2709,11 +2764,11 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
 
     this.contentData.push(runStep);
     this.contentIndexMap.set(stepId, runStep.index);
-    await safeDispatchCustomEvent(
-      GraphEvents.ON_RUN_STEP,
-      runStep,
-      this.config
-    );
+    // Pass undefined so safeDispatchCustomEvent resolves the runnable config
+    // from LangChain's AsyncLocalStorage. Using the shared `this.config` would
+    // race across concurrent child subgraph.invoke calls under parallel
+    // multi-agent handoffs and tag events with the wrong child's spawnKey.
+    await safeDispatchCustomEvent(GraphEvents.ON_RUN_STEP, runStep);
     return stepId;
   }
 
@@ -2854,7 +2909,7 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
     }
 
     if (!data.id) {
-      console.warn('No Tool ID provided for Tool Error');
+      mwarn('No Tool ID provided for Tool Error');
       return;
     }
 
@@ -2919,11 +2974,10 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
       id,
       delta,
     };
-    await safeDispatchCustomEvent(
-      GraphEvents.ON_RUN_STEP_DELTA,
-      runStepDelta,
-      this.config
-    );
+    // See dispatchRunStep note: do not pass `this.config`. The implicit
+    // AsyncLocalStorage config is the correct per-async-branch source under
+    // parallel handoffs.
+    await safeDispatchCustomEvent(GraphEvents.ON_RUN_STEP_DELTA, runStepDelta);
   }
 
   async dispatchMessageDelta(id: string, delta: t.MessageDelta): Promise<void> {
@@ -2934,11 +2988,8 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
       id,
       delta,
     };
-    await safeDispatchCustomEvent(
-      GraphEvents.ON_MESSAGE_DELTA,
-      messageDelta,
-      this.config
-    );
+    // See dispatchRunStep note.
+    await safeDispatchCustomEvent(GraphEvents.ON_MESSAGE_DELTA, messageDelta);
   }
 
   dispatchReasoningDelta = async (
@@ -2952,10 +3003,10 @@ If I seem to be missing something we discussed earlier, just give me a quick rem
       id: stepId,
       delta,
     };
+    // See dispatchRunStep note.
     await safeDispatchCustomEvent(
       GraphEvents.ON_REASONING_DELTA,
-      reasoningDelta,
-      this.config
+      reasoningDelta
     );
   };
 }

package/src/graphs/HandoffRegistry.ts

@@ -0,0 +1,199 @@
+import type { BaseMessage } from '@langchain/core/messages';
+import type * as t from '@/types';
+
+/**
+ * Tracks the lifecycle of a spawned handoff child agent.
+ */
+export type HandoffRecord = {
+  /** Agent identity (destination agentId) — stable across multiple spawns of the same agent */
+  id: string;
+  /** Unique internal key for this specific spawn (agentId + monotonic counter) */
+  spawnKey: string;
+  /** Display name of the child agent */
+  name: string;
+  /** Task description / instructions passed to child */
+  task: string;
+  /** When the handoff was spawned */
+  spawnedAt: number;
+  /** Current status */
+  status: 'pending' | 'running' | 'completed' | 'failed';
+  /** The background promise executing the child subgraph */
+  promise: Promise<t.BaseGraphState>;
+  /** Resolved result text (populated on completion) */
+  resultText?: string;
+  /** Error message (populated on failure) */
+  error?: string;
+  /** Duration in ms (populated on completion/failure) */
+  durationMs?: number;
+  /** Number of messages in child's output */
+  resultMessageCount?: number;
+};
+
+/**
+ * Registry for async handoff execution.
+ *
+ * Enables the autonomous orchestration pattern:
+ * 1. Orchestrator spawns children (non-blocking)
+ * 2. Orchestrator stays alive to reason, spawn more, or check status
+ * 3. Orchestrator collects results when ready
+ *
+ * Scoped per MultiAgentGraph instance — each orchestrator graph gets its own registry.
+ */
+export class HandoffRegistry {
+  private records: Map<string, HandoffRecord> = new Map();
+  /** Monotonically increasing counter for unique spawn IDs */
+  private spawnCounter = 0;
+
+  /**
+   * Register a spawned handoff child.
+   * The promise runs in the background — not awaited here.
+   * Uses a unique key per spawn so the same agent can be spawned multiple times
+   * across rounds without overwriting prior records.
+   */
+  spawn(params: {
+    id: string;
+    name: string;
+    task: string;
+    promise: Promise<t.BaseGraphState>;
+    extractResult: (messages: BaseMessage[], agentId: string) => string;
+    truncateResult: (text: string, maxChars: number) => string;
+    maxResultChars: number;
+    /** Callback when child completes (for SSE events) */
+    onComplete?: (record: HandoffRecord) => void;
+  }): void {
+    // Unique internal key: agentId + counter to support multiple spawns of same agent.
+    // record.id stays as the agent identity so callers/observability can attribute
+    // results to the agent regardless of which spawn round they came from.
+    const spawnKey = `${params.id}__${this.spawnCounter++}`;
+    const record: HandoffRecord = {
+      id: params.id,
+      spawnKey,
+      name: params.name,
+      task: params.task,
+      spawnedAt: Date.now(),
+      status: 'running',
+      promise: params.promise,
+    };
+
+    // Wire up the promise to update the record on completion
+    params.promise
+      .then((result) => {
+        const resultText = params.extractResult(result.messages, params.id);
+        const truncated = params.truncateResult(
+          resultText,
+          params.maxResultChars
+        );
+        record.status = 'completed';
+        record.resultText = truncated;
+        record.durationMs = Date.now() - record.spawnedAt;
+        record.resultMessageCount = result.messages.length;
+        params.onComplete?.(record);
+      })
+      .catch((err) => {
+        record.status = 'failed';
+        record.error = err instanceof Error ? err.message : String(err);
+        record.durationMs = Date.now() - record.spawnedAt;
+        params.onComplete?.(record);
+      });
+
+    this.records.set(spawnKey, record);
+  }
+
+  /** List all pending (running) handoffs */
+  listPending(): HandoffRecord[] {
+    return Array.from(this.records.values()).filter(
+      (r) => r.status === 'running'
+    );
+  }
+
+  /** List all completed handoffs (not yet collected) */
+  listCompleted(): HandoffRecord[] {
+    return Array.from(this.records.values()).filter(
+      (r) => r.status === 'completed' || r.status === 'failed'
+    );
+  }
+
+  /** List all handoffs regardless of status */
+  listAll(): HandoffRecord[] {
+    return Array.from(this.records.values());
+  }
+
+  /**
+   * Get a handoff record by either its unique spawnKey or by agentId.
+   * - Exact spawnKey match wins (O(1)).
+   * - Falls back to the most recently spawned record for that agentId — matching
+   *   lookups by callers that only know the agent identity, not the spawn round.
+   */
+  get(idOrSpawnKey: string): HandoffRecord | undefined {
+    const direct = this.records.get(idOrSpawnKey);
+    if (direct) return direct;
+    let latest: HandoffRecord | undefined;
+    for (const record of this.records.values()) {
+      if (record.id !== idOrSpawnKey) continue;
+      if (!latest || record.spawnedAt >= latest.spawnedAt) latest = record;
+    }
+    return latest;
+  }
+
+  /** Check if any handoffs are still running */
+  hasPending(): boolean {
+    return this.listPending().length > 0;
+  }
+
+  /**
+   * Wait for ALL pending handoffs to complete and return all records.
+   * Records are NOT auto-cleared — caller removes collected records via remove().
+   */
+  async waitForAll(): Promise<HandoffRecord[]> {
+    const pending = this.listPending();
+    if (pending.length > 0) {
+      await Promise.allSettled(pending.map((r) => r.promise));
+    }
+    const results = this.listAll();
+    return results;
+  }
+
+  /**
+   * Wait for ANY pending handoff to complete.
+   * Returns the newly completed record(s).
+   */
+  async waitForAny(): Promise<HandoffRecord[]> {
+    const pending = this.listPending();
+    if (pending.length === 0) {
+      return this.listCompleted();
+    }
+
+    // Race all pending promises — at least one will resolve
+    await Promise.race(
+      pending.map((r) => r.promise.then(() => r).catch(() => r))
+    );
+
+    // Small yield to let promise handlers update records
+    await new Promise((resolve) => setTimeout(resolve, 0));
+
+    return this.listCompleted();
+  }
+
+  /**
+   * Remove record(s) by spawnKey or agentId.
+   * - Exact spawnKey match removes only that record.
+   * - agentId match removes ALL records for that agent (covers callers that
+   *   want to forget everything tied to a given agent).
+   */
+  remove(idOrSpawnKey: string): void {
+    if (this.records.delete(idOrSpawnKey)) return;
+    for (const [key, record] of this.records) {
+      if (record.id === idOrSpawnKey) this.records.delete(key);
+    }
+  }
+
+  /** Clear all records (for cleanup between graph invocations) */
+  clear(): void {
+    this.records.clear();
+  }
+
+  /** Number of total tracked handoffs */
+  get size(): number {
+    return this.records.size;
+  }
+}