Suppress stale turn notifications after routing them

Author: benbrandt

src/CodexAppServerClient.ts

@@ -144,17 +144,13 @@ export class CodexAppServerClient {
                 this.recordCompactionCompleted(serverNotification);
             }
             const routing = extractTurnRouting(serverNotification);
-            const staleTurnNotification = this.isStaleTurn(routing.threadId, routing.turnId);
-            if (staleTurnNotification) {
-                if (isTurnCompletedNotification(serverNotification) && routing.threadId !== null && routing.turnId !== null) {
-                    this.clearStaleTurn(routing.threadId, routing.turnId);
-                }
-                for (const callback of this.codexEventHandlers) {
-                    callback({ eventType: "notification", ...serverNotification });
-                }
+            if (this.handleStaleTurnNotification(serverNotification, routing)) {
                 return;
             }
             this.recordTurnRouting(routing);
+            if (this.handleStaleTurnNotification(serverNotification, routing)) {
+                return;
+            }
             this.notify(serverNotification);
             for (const callback of this.codexEventHandlers) {
                 callback({ eventType: "notification", ...serverNotification });
@@ -554,6 +550,22 @@ export class CodexAppServerClient {
         }
     }
 
+    private handleStaleTurnNotification(
+        notification: ServerNotification,
+        routing: { threadId: string | null, turnId: string | null },
+    ): boolean {
+        if (!this.isStaleTurn(routing.threadId, routing.turnId)) {
+            return false;
+        }
+        if (isTurnCompletedNotification(notification) && routing.threadId !== null && routing.turnId !== null) {
+            this.clearStaleTurn(routing.threadId, routing.turnId);
+        }
+        for (const callback of this.codexEventHandlers) {
+            callback({ eventType: "notification", ...notification });
+        }
+        return true;
+    }
+
     private isStaleTurn(threadId: string | null, turnId: string | null): boolean {
         if (threadId === null || turnId === null) {
             return false;

src/__tests__/CodexACPAgent/CodexAcpClient.test.ts

@@ -1536,6 +1536,52 @@ describe('ACP server test', { timeout: 40_000 }, () => {
         });
     });
 
+    it('suppresses the first routed goal notification after cancellation marks the turn stale', async () => {
+        const { mockFixture } = setupPromptFixture();
+        const codexAppServerClient = mockFixture.getCodexAppServerClient();
+        const threadGoalSetSpy = vi.spyOn(codexAppServerClient, "threadGoalSet")
+            .mockResolvedValue({ goal: createThreadGoal() });
+        const turnInterruptSpy = vi.spyOn(mockFixture.getCodexAcpClient(), "turnInterrupt")
+            .mockResolvedValue(undefined);
+        const controller = new AbortController();
+
+        const promptPromise = mockFixture.getCodexAcpAgent().prompt({
+            sessionId: "session-id",
+            prompt: [{ type: "text", text: "/goal Ship the migration and keep tests green" }],
+        }, controller.signal);
+
+        await vi.waitFor(() => {
+            expect(threadGoalSetSpy).toHaveBeenCalledWith({
+                threadId: "session-id",
+                objective: "Ship the migration and keep tests green",
+                status: "active",
+            });
+        });
+
+        controller.abort();
+        await expect(promptPromise).resolves.toMatchObject({stopReason: "cancelled"});
+        mockFixture.clearAcpConnectionDump();
+
+        mockFixture.sendServerNotification({
+            method: "item/agentMessage/delta",
+            params: {
+                threadId: "session-id",
+                turnId: "goal-turn-id",
+                itemId: "goal-message-id",
+                delta: "leaked goal output",
+            },
+        });
+
+        await vi.waitFor(() => {
+            expect(turnInterruptSpy).toHaveBeenCalledWith({
+                threadId: "session-id",
+                turnId: "goal-turn-id",
+            });
+        });
+        await flushAsyncWork();
+        expect(mockFixture.getAcpConnectionDump([])).not.toContain("leaked goal output");
+    });
+
     it('does not hang when goal set starts no continuation turn', async () => {
         const mockFixture = createCodexMockTestFixture();
         const codexAppServerClient = mockFixture.getCodexAppServerClient();