Normalize parallel tool result order

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: stephentoub

test/harness/replayingCapiProxy.test.ts

@@ -730,6 +730,108 @@ describe("ReplayingCapiProxy", () => {
       }
     });
 
+    test("matches parallel tool results regardless of arrival order", async () => {
+      const cachePath = path.join(tempDir, "cache.yaml");
+      const cacheContent = yaml.stringify({
+        models: ["test-model"],
+        conversations: [
+          {
+            messages: [
+              { role: "system", content: "${system}" },
+              { role: "user", content: "Lookup city and country" },
+              {
+                role: "assistant",
+                tool_calls: [
+                  {
+                    id: "toolcall_0",
+                    type: "function",
+                    function: {
+                      name: "lookup_city",
+                      arguments: '{"city":"Paris"}',
+                    },
+                  },
+                  {
+                    id: "toolcall_1",
+                    type: "function",
+                    function: {
+                      name: "lookup_country",
+                      arguments: '{"country":"France"}',
+                    },
+                  },
+                ],
+              },
+              {
+                role: "tool",
+                tool_call_id: "toolcall_0",
+                content: "CITY_PARIS",
+              },
+              {
+                role: "tool",
+                tool_call_id: "toolcall_1",
+                content: "COUNTRY_FRANCE",
+              },
+              { role: "assistant", content: "Paris is in France." },
+            ],
+          },
+        ],
+      } satisfies NormalizedData);
+      await writeFile(cachePath, cacheContent);
+
+      const proxy = new ReplayingCapiProxy(
+        "http://localhost:9999",
+        cachePath,
+        workDir,
+      );
+      const proxyUrl = await proxy.start();
+
+      try {
+        const response = await makeRequest(proxyUrl, "/chat/completions", {
+          body: {
+            model: "test-model",
+            messages: [
+              { role: "system", content: "Be helpful" },
+              { role: "user", content: "Lookup city and country" },
+              {
+                role: "assistant",
+                tool_calls: [
+                  {
+                    id: "city-id",
+                    type: "function",
+                    function: {
+                      name: "lookup_city",
+                      arguments: '{"city":"Paris"}',
+                    },
+                  },
+                  {
+                    id: "country-id",
+                    type: "function",
+                    function: {
+                      name: "lookup_country",
+                      arguments: '{"country":"France"}',
+                    },
+                  },
+                ],
+              },
+              {
+                role: "tool",
+                tool_call_id: "country-id",
+                content: "COUNTRY_FRANCE",
+              },
+              { role: "tool", tool_call_id: "city-id", content: "CITY_PARIS" },
+            ],
+          },
+        });
+
+        expect(response.status).toBe(200);
+        expect(
+          (JSON.parse(response.body) as ChatCompletion).choices[0].message
+            .content,
+        ).toBe("Paris is in France.");
+      } finally {
+        await proxy.stop();
+      }
+    });
+
     test("returns streaming response when stream: true", async () => {
       const cachePath = path.join(tempDir, "cache.yaml");
       const cacheContent = yaml.stringify({

test/harness/replayingCapiProxy.ts

@@ -677,6 +677,7 @@ async function transformHttpExchanges(
   );
 
   normalizeToolCalls(dedupedExchanges, toolResultNormalizers);
+  normalizeToolResultOrder(dedupedExchanges);
   normalizeFilenames(dedupedExchanges, workDir);
   return { models: Array.from(dedupedModels), conversations: dedupedExchanges };
 }
@@ -782,6 +783,51 @@ function normalizeToolCalls(
   }
 }
 
+function normalizeToolResultOrder(conversations: NormalizedConversation[]) {
+  for (const conv of conversations) {
+    for (let start = 0; start < conv.messages.length; ) {
+      if (conv.messages[start].role !== "tool") {
+        start++;
+        continue;
+      }
+
+      let end = start + 1;
+      while (end < conv.messages.length && conv.messages[end].role === "tool") {
+        end++;
+      }
+
+      conv.messages
+        .slice(start, end)
+        .sort(compareToolResultMessages)
+        .forEach((message, index) => {
+          conv.messages[start + index] = message;
+        });
+      start = end;
+    }
+  }
+}
+
+function compareToolResultMessages(
+  left: NormalizedMessage,
+  right: NormalizedMessage,
+) {
+  return compareToolCallIds(left.tool_call_id, right.tool_call_id);
+}
+
+function compareToolCallIds(left?: string, right?: string) {
+  const leftNumber = parseNormalizedToolCallId(left);
+  const rightNumber = parseNormalizedToolCallId(right);
+  if (leftNumber !== undefined && rightNumber !== undefined) {
+    return leftNumber - rightNumber;
+  }
+  return (left ?? "").localeCompare(right ?? "");
+}
+
+function parseNormalizedToolCallId(id?: string) {
+  const match = id?.match(/^toolcall_(\d+)$/);
+  return match ? Number(match[1]) : undefined;
+}
+
 // As we capture LLM calls, we see:
 // - Request A, response AB
 // - Request ABC, response ABCD