From 4dac8f47ed46c8382cf97f88acceef65bb176bd6 Mon Sep 17 00:00:00 2001
From: Vincent Koc <25068+vincentkoc@users.noreply.github.com>
Date: Tue, 23 Jun 2026 15:33:50 +0800
Subject: [PATCH] perf(agents): index displaced tool results

---
 src/agents/session-transcript-repair.ts | 35 ++++++++++++++-----------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/src/agents/session-transcript-repair.ts b/src/agents/session-transcript-repair.ts
index 5243bdccbea..fe6669216de 100644
--- a/src/agents/session-transcript-repair.ts
+++ b/src/agents/session-transcript-repair.ts
@@ -554,14 +554,15 @@ function assistantHasToolCalls(message: AgentMessage): boolean {
   return extractToolCallsFromAssistant(message).length > 0;
 }
 
-function findLaterMatchingToolResult(params: {
+function collectLaterMatchingToolResults(params: {
   messages: AgentMessage[];
   startIndex: number;
-  toolCallId: string;
-  toolName?: string;
   toolCalls: Array<{ id: string; name?: string }>;
+  toolNamesById: Map<string, string>;
   seenToolResultIds: Set<string>;
-}): Extract<AgentMessage, { role: "toolResult" }> | undefined {
+}): Map<string, Extract<AgentMessage, { role: "toolResult" }>> {
+  const resultsById = new Map<string, Extract<AgentMessage, { role: "toolResult" }>>();
+  const toolCallIds = new Set(params.toolCalls.map((toolCall) => toolCall.id));
   for (let index = params.startIndex; index < params.messages.length; index += 1) {
     const candidate = params.messages[index];
     if (!candidate || typeof candidate !== "object" || candidate.role !== "toolResult") {
@@ -569,12 +570,15 @@ function findLaterMatchingToolResult(params: {
     }
     const normalizedLegacyResult = normalizeLegacyToolResultId(candidate, params.toolCalls);
     const id = extractToolResultId(normalizedLegacyResult);
-    if (!id || id !== params.toolCallId || params.seenToolResultIds.has(id)) {
+    if (!id || !toolCallIds.has(id) || params.seenToolResultIds.has(id) || resultsById.has(id)) {
       continue;
     }
-    return normalizeToolResultName(normalizedLegacyResult, params.toolName);
+    resultsById.set(
+      id,
+      normalizeToolResultName(normalizedLegacyResult, params.toolNamesById.get(id)),
+    );
   }
-  return undefined;
+  return resultsById;
 }
 
 export function repairToolUseResultPairing(
@@ -769,20 +773,21 @@ export function repairToolUseResultPairing(
       changed = true;
     }
 
+    const laterResultsById = collectLaterMatchingToolResults({
+      messages,
+      startIndex: j,
+      toolCalls,
+      toolNamesById: toolCallNamesById,
+      seenToolResultIds,
+    });
     for (const call of toolCalls) {
       const existing = spanResultsById.get(call.id);
       if (existing) {
         pushToolResult(existing);
       } else {
-        const laterResult = findLaterMatchingToolResult({
-          messages,
-          startIndex: j,
-          toolCallId: call.id,
-          toolName: call.name,
-          toolCalls,
-          seenToolResultIds,
-        });
+        const laterResult = laterResultsById.get(call.id);
         if (laterResult) {
+          laterResultsById.delete(call.id);
           moved = true;
           changed = true;
           pushToolResult(laterResult);