fix(gateway): clean up aborted queued turns (#103194)

Co-authored-by: Ted Li <tl2493@columbia.edu>

src/gateway/chat-queued-turns.test.ts

@@ -14,6 +14,21 @@ function emptyMap(): QueuedChatTurnMap {
   return new Map();
 }
 
+function registerTurn(
+  map: QueuedChatTurnMap,
+  runId: string,
+  controller: AbortController,
+  sessionId = runId,
+): boolean {
+  return registerQueuedChatTurn({
+    chatQueuedTurns: map,
+    runId,
+    controller,
+    sessionId,
+    sessionKey: "main",
+  });
+}
+
 describe("chat-queued-turns", () => {
   it("registers and completes a queued turn", () => {
     const map = emptyMap();
@@ -30,10 +45,96 @@ describe("chat-queued-turns", () => {
       }),
     ).toBe(true);
     expect(getQueuedChatTurn(map, "run-a")?.sessionKey).toBe("main");
-    expect(completeQueuedChatTurn(map, "run-a")).toBe(true);
+    expect(completeQueuedChatTurn(map, "run-a", controller)).toBe(true);
     expect(getQueuedChatTurn(map, "run-a")).toBeUndefined();
   });
 
+  it("removes the queued entry when its controller aborts", () => {
+    const map = emptyMap();
+    const controller = new AbortController();
+    expect(registerTurn(map, "run-abort", controller, "sess-abort")).toBe(true);
+
+    controller.abort();
+
+    expect(getQueuedChatTurn(map, "run-abort")).toBeUndefined();
+  });
+
+  it("does not let a stale abort listener remove a reused run id", () => {
+    const map = emptyMap();
+    const first = new AbortController();
+    const second = new AbortController();
+    expect(registerTurn(map, "run-reused", first, "sess-a")).toBe(true);
+    expect(completeQueuedChatTurn(map, "run-reused", first)).toBe(true);
+    expect(registerTurn(map, "run-reused", second, "sess-b")).toBe(true);
+
+    first.abort();
+
+    expect(getQueuedChatTurn(map, "run-reused")?.controller).toBe(second);
+    second.abort();
+    expect(getQueuedChatTurn(map, "run-reused")).toBeUndefined();
+  });
+
+  it("does not let stale lifecycle callbacks mutate a reused run id", () => {
+    const map = emptyMap();
+    const first = new AbortController();
+    const second = new AbortController();
+    expect(registerTurn(map, "run-reused", first, "sess-a")).toBe(true);
+
+    first.abort();
+    expect(registerTurn(map, "run-reused", second, "sess-b")).toBe(true);
+
+    expect(retireQueuedChatTurnCancellation(map, "run-reused", first)).toBe(false);
+    expect(completeQueuedChatTurn(map, "run-reused", first)).toBe(false);
+    const current = getQueuedChatTurn(map, "run-reused");
+    expect(current?.controller).toBe(second);
+    expect(current?.abortable).toBeUndefined();
+  });
+
+  it.each(["single", "bulk"] as const)(
+    "preserves a synchronous replacement during %s abort cleanup",
+    (mode) => {
+      const map = emptyMap();
+      const first = new AbortController();
+      const second = new AbortController();
+      expect(registerTurn(map, "run-replaced", first, "sess-a")).toBe(true);
+      const firstEntry = getQueuedChatTurn(map, "run-replaced");
+      expect(firstEntry).toBeDefined();
+      first.signal.addEventListener(
+        "abort",
+        () => {
+          expect(registerTurn(map, "run-replaced", second, "sess-b")).toBe(true);
+        },
+        { once: true },
+      );
+
+      const aborted =
+        mode === "single"
+          ? abortQueuedChatTurnById(map, {
+              runId: "run-replaced",
+              sessionKey: "main",
+            }).aborted
+          : abortQueuedChatTurns(map, [{ runId: "run-replaced", entry: firstEntry! }]).includes(
+              "run-replaced",
+            );
+
+      expect(aborted).toBe(true);
+      expect(getQueuedChatTurn(map, "run-replaced")?.controller).toBe(second);
+    },
+  );
+
+  it("keeps retired collect identities until completion after abort", () => {
+    const map = emptyMap();
+    const controller = new AbortController();
+    expect(registerTurn(map, "run-retired", controller, "sess-retired")).toBe(true);
+    expect(retireQueuedChatTurnCancellation(map, "run-retired", controller)).toBe(true);
+
+    controller.abort();
+
+    expect(getQueuedChatTurn(map, "run-retired")?.abortable).toBe(false);
+    expect(completeQueuedChatTurn(map, "run-retired", controller)).toBe(true);
+    expect(getQueuedChatTurn(map, "run-retired")).toBeUndefined();
+  });
+
   it("rejects re-register with a different controller", () => {
     const map = emptyMap();
     const first = new AbortController();
@@ -117,7 +218,7 @@ describe("chat-queued-turns", () => {
       sessionKey: "main",
     });
 
-    expect(retireQueuedChatTurnCancellation(map, "run-collected")).toBe(true);
+    expect(retireQueuedChatTurnCancellation(map, "run-collected", controller)).toBe(true);
     expect(
       abortQueuedChatTurnById(map, { runId: "run-collected", sessionKey: "main" }).aborted,
     ).toBe(false);
@@ -126,7 +227,7 @@ describe("chat-queued-turns", () => {
     expect(listQueuedChatTurnsForSession({ chatQueuedTurns: map, sessionKeys: ["main"] })).toEqual(
       [],
     );
-    expect(completeQueuedChatTurn(map, "run-collected")).toBe(true);
+    expect(completeQueuedChatTurn(map, "run-collected", controller)).toBe(true);
   });
 
   it("refuses abort when sessionKey mismatches unless allowed", () => {

src/gateway/chat-queued-turns.ts

@@ -38,6 +38,19 @@ function resolveExactRunId(runId: string): string | undefined {
   return runId.length > 0 ? runId : undefined;
 }
 
+// Queue callbacks can outlive their map entry, and protocol run IDs may be reused.
+// Mutate only the exact entry captured by the callback or abort operation.
+function deleteQueuedChatTurnEntry(
+  chatQueuedTurns: QueuedChatTurnMap,
+  runId: string,
+  entry: QueuedChatTurnEntry,
+): boolean {
+  if (chatQueuedTurns.get(runId) !== entry) {
+    return false;
+  }
+  return chatQueuedTurns.delete(runId);
+}
+
 export function registerQueuedChatTurn(params: RegisterQueuedChatTurnParams): boolean {
   const runId = resolveExactRunId(params.runId);
   const sessionKey = normalizeOptionalString(params.sessionKey);
@@ -63,15 +76,33 @@ export function registerQueuedChatTurn(params: RegisterQueuedChatTurnParams): bo
     ownerDeviceId: normalizeOptionalString(params.ownerDeviceId),
   };
   params.chatQueuedTurns.set(runId, entry);
+  params.controller.signal.addEventListener(
+    "abort",
+    () => {
+      // Queued entries can outlive active-run cleanup. Retired collect entries
+      // stay as idempotency guards until aggregate completion removes them.
+      if (entry.abortable !== false) {
+        deleteQueuedChatTurnEntry(params.chatQueuedTurns, runId, entry);
+      }
+    },
+    { once: true },
+  );
   return true;
 }
 
-export function completeQueuedChatTurn(chatQueuedTurns: QueuedChatTurnMap, runId: string): boolean {
+export function completeQueuedChatTurn(
+  chatQueuedTurns: QueuedChatTurnMap,
+  runId: string,
+  controller: AbortController,
+): boolean {
   const key = resolveExactRunId(runId);
   if (!key) {
     return false;
   }
-  return chatQueuedTurns.delete(key);
+  const entry = chatQueuedTurns.get(key);
+  return entry?.controller === controller
+    ? deleteQueuedChatTurnEntry(chatQueuedTurns, key, entry)
+    : false;
 }
 
 /**
@@ -81,9 +112,10 @@ export function completeQueuedChatTurn(chatQueuedTurns: QueuedChatTurnMap, runId
 export function retireQueuedChatTurnCancellation(
   chatQueuedTurns: QueuedChatTurnMap,
   runId: string,
+  controller: AbortController,
 ): boolean {
   const entry = getQueuedChatTurn(chatQueuedTurns, runId);
-  if (!entry) {
+  if (!entry || entry.controller !== controller) {
     return false;
   }
   entry.abortable = false;
@@ -132,7 +164,7 @@ export function abortQueuedChatTurnById(
       params.stopReason ? new Error(`queued turn aborted: ${params.stopReason}`) : undefined,
     );
   }
-  chatQueuedTurns.delete(runId);
+  deleteQueuedChatTurnEntry(chatQueuedTurns, runId, entry);
   return { aborted: true };
 }
 
@@ -194,7 +226,7 @@ export function abortQueuedChatTurns(
 ): string[] {
   const runIds: string[] = [];
   for (const { runId, entry } of matches) {
-    if (!chatQueuedTurns.has(runId)) {
+    if (chatQueuedTurns.get(runId) !== entry) {
       continue;
     }
     if (!entry.controller.signal.aborted) {
@@ -202,7 +234,7 @@ export function abortQueuedChatTurns(
         stopReason ? new Error(`queued turn aborted: ${stopReason}`) : undefined,
       );
     }
-    chatQueuedTurns.delete(runId);
+    deleteQueuedChatTurnEntry(chatQueuedTurns, runId, entry);
     runIds.push(runId);
   }
   return runIds;

src/gateway/server-methods/chat.ts

@@ -4841,10 +4841,18 @@ export const chatHandlers: GatewayRequestHandlers = {
                       return queuedFollowupEnqueued;
                     },
                     onCancellationRetired: () => {
-                      retireQueuedChatTurnCancellation(ensureChatQueuedTurns(context), clientRunId);
+                      retireQueuedChatTurnCancellation(
+                        ensureChatQueuedTurns(context),
+                        clientRunId,
+                        activeRunAbort.controller,
+                      );
                     },
                     onComplete: () => {
-                      completeQueuedChatTurn(ensureChatQueuedTurns(context), clientRunId);
+                      completeQueuedChatTurn(
+                        ensureChatQueuedTurns(context),
+                        clientRunId,
+                        activeRunAbort.controller,
+                      );
                     },
                   },
                   images: replyOptionImages,

src/gateway/server.chat.gateway-server-chat-b.test.ts

@@ -2567,6 +2567,11 @@ describe("gateway server chat", () => {
       );
       expect(dispatchInboundMessageMock).toHaveBeenCalledTimes(1);
 
+      const queuedEntry = context.chatQueuedTurns.get("idem-queued-followup");
+      expect(queuedEntry).toBeDefined();
+      queuedEntry?.controller.abort();
+      expect(context.chatQueuedTurns.has("idem-queued-followup")).toBe(false);
+
       queuedLifecycle?.onComplete?.();
       expect(context.chatQueuedTurns.has("idem-queued-followup")).toBe(false);
       await vi.waitFor(