fix: add mutex to session-tool-result-guard to prevent async race condition

Ivan Somov · Ivan Somov · commit 135fdf395d79 · 2026-03-21T19:25:24.000-07:00
diff --git a/src/agents/session-tool-result-guard.test.ts b/src/agents/session-tool-result-guard.test.ts
@@ -519,3 +519,117 @@ describe("installSessionToolResultGuard", () => {
     expect(syntheticForError).toHaveLength(0);
   });
 });
+
+describe("concurrent append race condition", () => {
+  it("does not synthesize when toolResult arrives concurrently with next assistant message", async () => {
+    // Simulate the exact race: tool_calls appended, then toolResult and a
+    // completion message are dispatched concurrently (both microtasks queued).
+    const sm = SessionManager.inMemory();
+    installSessionToolResultGuard(sm);
+
+    sm.appendMessage(
+      asAppendMessage({
+        role: "assistant",
+        content: [{ type: "toolCall", id: "race_1", name: "read", arguments: {} }],
+      }),
+    );
+
+    // Fire both concurrently — do NOT await between them
+    const p1 = Promise.resolve().then(() =>
+      sm.appendMessage(
+        asAppendMessage({
+          role: "toolResult",
+          toolCallId: "race_1",
+          content: [{ type: "text", text: "real result" }],
+          isError: false,
+        }),
+      ),
+    );
+    const p2 = Promise.resolve().then(() =>
+      sm.appendMessage(
+        asAppendMessage({
+          role: "assistant",
+          content: [{ type: "text", text: "completion" }],
+          stopReason: "endTurn",
+        }),
+      ),
+    );
+
+    await Promise.all([p1, p2]);
+
+    const messages = expectPersistedRoles(sm, ["assistant", "toolResult", "assistant"]);
+    // The toolResult must be the real one, not synthetic
+    const tr = messages[1] as { isError?: boolean; content?: Array<{ text?: string }> };
+    expect(tr.isError).toBe(false);
+    expect(tr.content?.[0]?.text).toBe("real result");
+  });
+
+  it("serializes multiple concurrent appends without duplicating messages", async () => {
+    const sm = SessionManager.inMemory();
+    installSessionToolResultGuard(sm);
+
+    sm.appendMessage(
+      asAppendMessage({
+        role: "assistant",
+        content: [
+          { type: "toolCall", id: "c1", name: "a", arguments: {} },
+          { type: "toolCall", id: "c2", name: "b", arguments: {} },
+        ],
+      }),
+    );
+
+    // Fire all results concurrently
+    await Promise.all([
+      Promise.resolve().then(() =>
+        sm.appendMessage(
+          asAppendMessage({
+            role: "toolResult",
+            toolCallId: "c1",
+            content: [{ type: "text", text: "r1" }],
+            isError: false,
+          }),
+        ),
+      ),
+      Promise.resolve().then(() =>
+        sm.appendMessage(
+          asAppendMessage({
+            role: "toolResult",
+            toolCallId: "c2",
+            content: [{ type: "text", text: "r2" }],
+            isError: false,
+          }),
+        ),
+      ),
+    ]);
+
+    const messages = expectPersistedRoles(sm, ["assistant", "toolResult", "toolResult"]);
+    // Both results must be the real ones, not synthetic
+    const results = messages.slice(1) as Array<{ isError?: boolean }>;
+    expect(results.every((r) => r.isError === false)).toBe(true);
+  });
+
+  it("getPendingIds() returns empty after concurrent resolution", async () => {
+    const sm = SessionManager.inMemory();
+    const guard = installSessionToolResultGuard(sm);
+
+    sm.appendMessage(
+      asAppendMessage({
+        role: "assistant",
+        content: [{ type: "toolCall", id: "drain_1", name: "x", arguments: {} }],
+      }),
+    );
+
+    await Promise.resolve().then(() =>
+      sm.appendMessage(
+        asAppendMessage({
+          role: "toolResult",
+          toolCallId: "drain_1",
+          content: [{ type: "text", text: "ok" }],
+          isError: false,
+        }),
+      ),
+    );
+
+    expect(guard.getPendingIds()).toEqual([]);
+  });
+});
diff --git a/src/agents/session-tool-result-guard.ts b/src/agents/session-tool-result-guard.ts
@@ -22,6 +22,30 @@ type SessionManagerWithRawAppend = SessionManager & {
   [RAW_APPEND_MESSAGE]?: SessionManager["appendMessage"];
 };
 
+/**
+ * Creates a simple serializing async mutex.
+ * Each call to `run` queues `fn` behind all previously-queued work so that
+ * concurrent callers execute sequentially in arrival order.
+ *
+ * This is intentionally minimal — no timeout, no fairness guarantees beyond
+ * FIFO promise resolution. Sufficient for low-throughput session appends.
+ */
+function createMutex() {
+  let queue: Promise<void> = Promise.resolve();
+  return {
+    run<T>(fn: () => T | Promise<T>): Promise<T> {
+      const result = queue.then(() => fn());
+      // Advance the queue regardless of whether fn() throws, so a failed
+      // append does not permanently block subsequent calls.
+      queue = result.then(
+        () => undefined,
+        () => undefined,
+      );
+      return result;
+    },
+  };
+}
+
 /**
  * Truncate oversized text content blocks in a tool result message.
  * Returns the original message if under the limit, or a new message with
@@ -127,6 +151,8 @@ export function installSessionToolResultGuard(
   const originalAppend = getRawSessionAppendMessage(sessionManager);
   (sessionManager as SessionManagerWithRawAppend)[RAW_APPEND_MESSAGE] = originalAppend;
   const pendingState = createPendingToolCallState();
+  const appendMutex = createMutex();
+
   const persistMessage = (message: AgentMessage) => {
     const transformer = opts?.transformMessageForPersistence;
     return transformer ? transformer(message) : message;
@@ -187,7 +213,7 @@ export function installSessionToolResultGuard(
     pendingState.clear();
   };
 
-  const guardedAppend = (message: AgentMessage) => {
+  const guardedAppendInner = (message: AgentMessage) => {
     let nextMessage = message;
     const role = (message as { role?: unknown }).role;
     if (role === "assistant") {
@@ -253,6 +279,12 @@ export function installSessionToolResultGuard(
       flushPendingToolResults();
     }
 
+    // Register tool calls BEFORE appending so that a racing toolResult handler
+    // (which may arrive during an awaited originalAppend) finds the IDs in pending.
+    if (toolCalls.length > 0) {
+      pendingState.trackToolCalls(toolCalls);
+    }
+
     const finalMessage = applyBeforeWriteHook(persistMessage(nextMessage));
     if (!finalMessage) {
       return undefined;
@@ -271,13 +303,15 @@ export function installSessionToolResultGuard(
       });
     }
 
-    if (toolCalls.length > 0) {
-      pendingState.trackToolCalls(toolCalls);
-    }
-
     return result;
   };
 
+  const guardedAppend = (message: AgentMessage): ReturnType<typeof originalAppend> => {
+    return appendMutex.run(() => guardedAppendInner(message)) as unknown as ReturnType<
+      typeof originalAppend
+    >;
+  };
+
   // Monkey-patch appendMessage with our guarded version.
   sessionManager.appendMessage = guardedAppend as SessionManager["appendMessage"];
 
