LLM-22701 Fix duplicate messages on follow-up prompts

Aleksandr Slapoguzov · Aleksandr Slapoguzov · commit 9a2b7d85c799 · 2025-12-20T22:59:03.000+01:00
Notification handlers were accumulating in an array, causing each
  handler to process every notification. Changed to a Map keyed by
  session ID so each session has exactly one handler that gets replaced
  on follow-up prompts
diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts
@@ -102,7 +102,7 @@ export class CodexAcpClient {
     }
 
     async sendPrompt(request: acp.PromptRequest, eventHandler: (result: ServerNotification) => void): Promise<void> {
-        this.codexClient.onServerNotification(eventHandler);
+        this.codexClient.onServerNotification(request.sessionId, eventHandler);
 
         const input = request.prompt.filter(b => b.type === "text")
             .map(b => b.text)
diff --git a/src/CodexAppServerClient.ts b/src/CodexAppServerClient.ts
@@ -95,19 +95,22 @@ export class CodexAppServerClient {
         return await this.sendRequest({ method: "model/list", params });
     }
 
-    //TODO support removal (leads to duplicated processing of follow-ups)
-    onServerNotification(callback: (event: ServerNotification) => void){
-        this.notificationHandlers.push(callback);
+    /**
+     * Registers a notification handler for a specific session.
+     * Replaces any existing handler for the same session, preventing handler accumulation.
+     */
+    onServerNotification(sessionId: string, callback: (event: ServerNotification) => void) {
+        this.notificationHandlers.set(sessionId, callback);
     }
 
     private codexEventHandlers: Array<(event: CodexConnectionEvent) => void> = [];
     onClientTransportEvent(callback: (event: CodexConnectionEvent) => void){
         this.codexEventHandlers.push(callback);
     }
 
-    private notificationHandlers: Array<(event: ServerNotification) => void> = [];
+    private notificationHandlers = new Map<string, (event: ServerNotification) => void>();
     private notify(notification: ServerNotification) {
-        for (const notificationHandler of this.notificationHandlers) {
+        for (const notificationHandler of this.notificationHandlers.values()) {
             notificationHandler(notification);
         }
     }
diff --git a/src/__tests__/CodexACPAgent/CodexAcpClient.test.ts b/src/__tests__/CodexACPAgent/CodexAcpClient.test.ts
@@ -1,6 +1,6 @@
 import {describe, expect, it, vi, beforeEach} from 'vitest';
 import type {CodexAuthRequest} from "../../CodexAuthMethod";
-import {createTestFixture, type TestFixture} from "../acp-test-utils";
+import {createTestFixture, createCodexMockTestFixture, type TestFixture} from "../acp-test-utils";
 import type {ServerNotification} from "../../app-server";
 import type {SessionState} from "../../CodexAcpServer";
 
@@ -64,7 +64,7 @@ describe('ACP server test', { timeout: 40_000 }, () => {
             { method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "ll", }},
             { method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "o!", }},
         ];
-        function onServerNotification(callback: (event: ServerNotification) => void){
+        function onServerNotification(_sessionId: string, callback: (event: ServerNotification) => void){
             for (const notification of serverNotifications) {
                 callback(notification);
             }
@@ -95,6 +95,102 @@ describe('ACP server test', { timeout: 40_000 }, () => {
 
     });
 
+    it('should not duplicate messages on follow-up prompts', async () => {
+        const mockFixture = createCodexMockTestFixture();
+        const codexAcpAgent = mockFixture.getCodexAcpAgent();
+
+        mockFixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue(undefined);
+        mockFixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue(undefined);
+
+        const sessionState: SessionState = {
+            pendingPrompt: null,
+            sessionMetadata: {
+                sessionId: "id",
+                currentModelId: "model-id",
+                models: [],
+            }
+        };
+        vi.spyOn(codexAcpAgent, "getSessionState").mockReturnValue(sessionState);
+
+        // First prompt - registers first notification handler
+        await codexAcpAgent.prompt({ sessionId: "id", prompt: [{type: "text", text: "First message"}] });
+
+        // Follow-up prompt - should NOT accumulate handlers
+        await codexAcpAgent.prompt({ sessionId: "id", prompt: [{type: "text", text: "Follow-up message"}] });
+
+        mockFixture.clearAcpConnectionDump();
+
+        // Trigger notifications after both prompts - should produce only 3 events, not 6
+        const serverNotifications: ServerNotification[] = [
+            { method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "He", }},
+            { method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "ll", }},
+            { method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "o!", }},
+        ];
+        for (const notification of serverNotifications) {
+            mockFixture.sendServerNotification(notification);
+        }
+
+        // Wait for async handlers to complete
+        await vi.waitFor(() => {
+            const dump = mockFixture.getAcpConnectionDump([]);
+            expect(dump.length).toBeGreaterThan(0);
+        });
+
+        await expect(mockFixture.getAcpConnectionDump([])).toMatchFileSnapshot("data/follow-up-no-duplicates.json");
+    });
+
+    it('should handle multiple sessions independently', async () => {
+        const mockFixture = createCodexMockTestFixture();
+        const codexAcpAgent = mockFixture.getCodexAcpAgent();
+
+        mockFixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue(undefined);
+        mockFixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue(undefined);
+
+        const sessionState1: SessionState = {
+            pendingPrompt: null,
+            sessionMetadata: {
+                sessionId: "session-1",
+                currentModelId: "model-id",
+                models: [],
+            }
+        };
+        const sessionState2: SessionState = {
+            pendingPrompt: null,
+            sessionMetadata: {
+                sessionId: "session-2",
+                currentModelId: "model-id",
+                models: [],
+            }
+        };
+
+        vi.spyOn(codexAcpAgent, "getSessionState").mockImplementation((sessionId: string) => {
+            return sessionId === "session-1" ? sessionState1 : sessionState2;
+        });
+
+        // Start prompts for two different sessions
+        await codexAcpAgent.prompt({ sessionId: "session-1", prompt: [{type: "text", text: "Message to session 1"}] });
+        await codexAcpAgent.prompt({ sessionId: "session-2", prompt: [{type: "text", text: "Message to session 2"}] });
+
+        mockFixture.clearAcpConnectionDump();
+
+        // Trigger notifications - both session handlers should receive them
+        const serverNotifications: ServerNotification[] = [
+            { method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "Hello", }},
+        ];
+        for (const notification of serverNotifications) {
+            mockFixture.sendServerNotification(notification);
+        }
+
+        // Wait for async handlers to complete
+        await vi.waitFor(() => {
+            const dump = mockFixture.getAcpConnectionDump([]);
+            expect(dump.length).toBeGreaterThan(0);
+        });
+
+        // Should have 2 events - one for each session's handler
+        await expect(mockFixture.getAcpConnectionDump([])).toMatchFileSnapshot("data/multiple-sessions.json");
+    });
+
     //dev-time test
     it.skip('should convert session notification to acp events', async () => {
         fixture.onCodexConnectionEvent((event) => {
diff --git a/src/__tests__/CodexACPAgent/data/follow-up-no-duplicates.json b/src/__tests__/CodexACPAgent/data/follow-up-no-duplicates.json
@@ -0,0 +1,45 @@
+{
+  "method": "sessionUpdate",
+  "args": [
+    {
+      "sessionId": "id",
+      "update": {
+        "sessionUpdate": "agent_message_chunk",
+        "content": {
+          "type": "text",
+          "text": "He"
+        }
+      }
+    }
+  ]
+}
+{
+  "method": "sessionUpdate",
+  "args": [
+    {
+      "sessionId": "id",
+      "update": {
+        "sessionUpdate": "agent_message_chunk",
+        "content": {
+          "type": "text",
+          "text": "ll"
+        }
+      }
+    }
+  ]
+}
+{
+  "method": "sessionUpdate",
+  "args": [
+    {
+      "sessionId": "id",
+      "update": {
+        "sessionUpdate": "agent_message_chunk",
+        "content": {
+          "type": "text",
+          "text": "o!"
+        }
+      }
+    }
+  ]
+}
diff --git a/src/__tests__/CodexACPAgent/data/multiple-sessions.json b/src/__tests__/CodexACPAgent/data/multiple-sessions.json
@@ -0,0 +1,30 @@
+{
+  "method": "sessionUpdate",
+  "args": [
+    {
+      "sessionId": "session-1",
+      "update": {
+        "sessionUpdate": "agent_message_chunk",
+        "content": {
+          "type": "text",
+          "text": "Hello"
+        }
+      }
+    }
+  ]
+}
+{
+  "method": "sessionUpdate",
+  "args": [
+    {
+      "sessionId": "session-2",
+      "update": {
+        "sessionUpdate": "agent_message_chunk",
+        "content": {
+          "type": "text",
+          "text": "Hello"
+        }
+      }
+    }
+  ]
+}
diff --git a/src/__tests__/acp-test-utils.ts b/src/__tests__/acp-test-utils.ts
@@ -3,12 +3,14 @@ import {type CodexConnectionEvent, CodexAppServerClient} from '../CodexAppServer
 import {startCodexConnection} from "../CodexJsonRpcConnection";
 import {CodexAcpServer} from "../CodexAcpServer";
 import type {AgentSideConnection} from "@agentclientprotocol/sdk";
+import type {ServerNotification} from "../app-server";
+import type {MessageConnection} from "vscode-jsonrpc/node";
 import path from "node:path";
 import fs from "node:fs";
 
 export type MethodCallEvent = { method: string; args: any[] };
 
-function createSmartMock<T extends object>(onCall: (event: MethodCallEvent) => void) {
+export function createSmartMock<T extends object>(onCall: (event: MethodCallEvent) => void) {
     return new Proxy({} as T, {
         get(_, prop) {
             return (...args: any[]) => {
@@ -33,25 +35,24 @@ export interface TestFixture {
     clearAcpConnectionDump(): void,
 }
 
-export function createTestFixture(): TestFixture {
-    const pathToCodex = path.resolve(process.cwd(), "node_modules", ".bin", process.platform === 'win32' ? "codex.cmd" : "codex");
-    if (!fs.existsSync(pathToCodex)) {
-        throw new Error(`Codex binary not found at ${pathToCodex}. Did you run 'npm install'?`);
-    }
-    
-    const acpConnectionEvents: MethodCallEvent[] = []
+export interface ConnectionConfig {
+    connection: MessageConnection;
+    getExitCode: () => number | null;
+}
+
+export function createBaseTestFixture(config: ConnectionConfig): TestFixture {
+    const acpConnectionEvents: MethodCallEvent[] = [];
     const acpEventHandlers: ((event: MethodCallEvent) => void)[] = [];
     const acpConnection = createSmartMock<AgentSideConnection>((event) => {
         acpConnectionEvents.push(event);
         acpEventHandlers.forEach(handler => handler(event));
     });
-    const codexConnection = startCodexConnection(pathToCodex);
-    const codexAppServerClient = new CodexAppServerClient(codexConnection.connection);
 
+    const codexAppServerClient = new CodexAppServerClient(config.connection);
     const codexAcpClient = new CodexAcpClient(codexAppServerClient);
-    const codexAcpAgent = new CodexAcpServer(acpConnection, codexAcpClient, undefined, () => codexConnection.process.exitCode);
+    const codexAcpAgent = new CodexAcpServer(acpConnection, codexAcpClient, undefined, config.getExitCode);
 
-    const transportEvents: CodexConnectionEvent[] = []
+    const transportEvents: CodexConnectionEvent[] = [];
     const codexEventHandlers: ((event: CodexConnectionEvent) => void)[] = [];
     codexAppServerClient.onClientTransportEvent((event) => {
         transportEvents.push(event);
@@ -83,20 +84,73 @@ export function createTestFixture(): TestFixture {
         getAcpConnectionDump(ignoredFields: string[]): string {
             return createArrayDump(acpConnectionEvents, ignoredFields);
         },
-        clearAcpConnectionDump(){
+        clearAcpConnectionDump() {
             acpConnectionEvents.splice(0, acpConnectionEvents.length);
         }
     };
 }
 
+/**
+ * Creates a test fixture with a real Codex connection.
+ * Use for integration tests that need to interact with the actual Codex binary.
+ */
+export function createTestFixture(): TestFixture {
+    const pathToCodex = path.resolve(process.cwd(), "node_modules", ".bin", process.platform === 'win32' ? "codex.cmd" : "codex");
+    if (!fs.existsSync(pathToCodex)) {
+        throw new Error(`Codex binary not found at ${pathToCodex}. Did you run 'npm install'?`);
+    }
+
+    const codexConnection = startCodexConnection(pathToCodex);
+
+    return createBaseTestFixture({
+        connection: codexConnection.connection,
+        getExitCode: () => codexConnection.process.exitCode
+    });
+}
+
+export interface CodexMockTestFixture extends TestFixture {
+    sendServerNotification(notification: ServerNotification): void,
+}
+
+/**
+ * Creates a test fixture with a mock Codex connection.
+ * Use for unit tests that don't need a real Codex binary.
+ * Provides `sendServerNotification()` to simulate server notifications.
+ */
+export function createCodexMockTestFixture(): CodexMockTestFixture {
+    let unhandledNotificationHandler: ((notification: any) => void) | null = null;
+
+    const mockCodexConnection = {
+        sendRequest: () => Promise.resolve(undefined),
+        onUnhandledNotification: (handler: (notification: any) => void) => {
+            unhandledNotificationHandler = handler;
+        },
+        onNotification: () => {},
+        end: () => {},
+    } as unknown as MessageConnection;
+
+    const baseFixture = createBaseTestFixture({
+        connection: mockCodexConnection,
+        getExitCode: () => null
+    });
+
+    return {
+        ...baseFixture,
+        sendServerNotification(notification: ServerNotification): void {
+            if (unhandledNotificationHandler) {
+                unhandledNotificationHandler(notification);
+            }
+        }
+    };
+}
 
-function createObjectDump(obj: any, anonymizedFields: string[] = []) {
+export function createObjectDump(obj: any, anonymizedFields: string[] = []) {
     function fieldAnonymizer(key: string, value: any): any {
         return anonymizedFields.includes(key) ? key : value;
     }
     return JSON.stringify(obj, fieldAnonymizer, 2);
 }
 
-function createArrayDump(objects: any[], anonymizedFields: string[]): string {
+export function createArrayDump(objects: any[], anonymizedFields: string[]): string {
     return objects.map(event => createObjectDump(event, anonymizedFields)).join("\n");
 }

Original file line number	Diff line number	Diff line change
`@@ -102,7 +102,7 @@ export class CodexAcpClient {`
`102`	`102`	`}`
`103`	`103`
`104`	`104`	`async sendPrompt(request: acp.PromptRequest, eventHandler: (result: ServerNotification) => void): Promise<void> {`
`105`		`- this.codexClient.onServerNotification(eventHandler);`
	`105`	`+ this.codexClient.onServerNotification(request.sessionId, eventHandler);`
`106`	`106`
`107`	`107`	`const input = request.prompt.filter(b => b.type === "text")`
`108`	`108`	`.map(b => b.text)`
Original file line number	Diff line number	Diff line change
`@@ -95,19 +95,22 @@ export class CodexAppServerClient {`
`95`	`95`	`return await this.sendRequest({ method: "model/list", params });`
`96`	`96`	`}`
`97`	`97`
`98`		`- //TODO support removal (leads to duplicated processing of follow-ups)`
`99`		`- onServerNotification(callback: (event: ServerNotification) => void){`
`100`		`- this.notificationHandlers.push(callback);`
	`98`	`+ /**`
	`99`	`+ * Registers a notification handler for a specific session.`
	`100`	`+ * Replaces any existing handler for the same session, preventing handler accumulation.`
	`101`	`+ */`
	`102`	`+ onServerNotification(sessionId: string, callback: (event: ServerNotification) => void) {`
	`103`	`+ this.notificationHandlers.set(sessionId, callback);`
`101`	`104`	`}`
`102`	`105`
`103`	`106`	`private codexEventHandlers: Array<(event: CodexConnectionEvent) => void> = [];`
`104`	`107`	`onClientTransportEvent(callback: (event: CodexConnectionEvent) => void){`
`105`	`108`	`this.codexEventHandlers.push(callback);`
`106`	`109`	`}`
`107`	`110`
`108`		`- private notificationHandlers: Array<(event: ServerNotification) => void> = [];`
	`111`	`+ private notificationHandlers = new Map<string, (event: ServerNotification) => void>();`
`109`	`112`	`private notify(notification: ServerNotification) {`
`110`		`- for (const notificationHandler of this.notificationHandlers) {`
	`113`	`+ for (const notificationHandler of this.notificationHandlers.values()) {`
`111`	`114`	`notificationHandler(notification);`
`112`	`115`	`}`
`113`	`116`	`}`