diff --git a/src/gateway/chat-queued-turns.test.ts b/src/gateway/chat-queued-turns.test.ts index 2a55cb9ae6c5..1650ad705dff 100644 --- a/src/gateway/chat-queued-turns.test.ts +++ b/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", () => { diff --git a/src/gateway/chat-queued-turns.ts b/src/gateway/chat-queued-turns.ts index 11ba265b1f9c..f1b78672888a 100644 --- a/src/gateway/chat-queued-turns.ts +++ b/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; diff --git a/src/gateway/server-methods/chat.ts b/src/gateway/server-methods/chat.ts index 0322ad2cc1a4..516dd2631958 100644 --- a/src/gateway/server-methods/chat.ts +++ b/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, diff --git a/src/gateway/server.chat.gateway-server-chat-b.test.ts b/src/gateway/server.chat.gateway-server-chat-b.test.ts index d30deb643aba..a46fb0268a43 100644 --- a/src/gateway/server.chat.gateway-server-chat-b.test.ts +++ b/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(