From 9a2b7d85c79948ad0f2c5e52eb20d689c35dc76a Mon Sep 17 00:00:00 2001 From: Aleksandr Slapoguzov Date: Sat, 20 Dec 2025 22:59:03 +0100 Subject: [PATCH] LLM-22701 Fix duplicate messages on follow-up prompts 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 --- src/CodexAcpClient.ts | 2 +- src/CodexAppServerClient.ts | 13 ++- .../CodexACPAgent/CodexAcpClient.test.ts | 100 +++++++++++++++++- .../data/follow-up-no-duplicates.json | 45 ++++++++ .../CodexACPAgent/data/multiple-sessions.json | 30 ++++++ src/__tests__/acp-test-utils.ts | 84 ++++++++++++--- 6 files changed, 251 insertions(+), 23 deletions(-) create mode 100644 src/__tests__/CodexACPAgent/data/follow-up-no-duplicates.json create mode 100644 src/__tests__/CodexACPAgent/data/multiple-sessions.json diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index 5d342205..459046ff 100644 --- 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 { - 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 index 3caccea6..f4f763f2 100644 --- a/src/CodexAppServerClient.ts +++ b/src/CodexAppServerClient.ts @@ -95,9 +95,12 @@ 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> = []; @@ -105,9 +108,9 @@ export class CodexAppServerClient { this.codexEventHandlers.push(callback); } - private notificationHandlers: Array<(event: ServerNotification) => void> = []; + private notificationHandlers = new Map 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 index 0345290e..57cf8e2c 100644 --- 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 new file mode 100644 index 00000000..efd040d7 --- /dev/null +++ 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!" + } + } + } + ] +} \ No newline at end of file diff --git a/src/__tests__/CodexACPAgent/data/multiple-sessions.json b/src/__tests__/CodexACPAgent/data/multiple-sessions.json new file mode 100644 index 00000000..45136b60 --- /dev/null +++ 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" + } + } + } + ] +} \ No newline at end of file diff --git a/src/__tests__/acp-test-utils.ts b/src/__tests__/acp-test-utils.ts index 8d88bcf7..4b4c9c0a 100644 --- 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(onCall: (event: MethodCallEvent) => void) { +export function createSmartMock(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((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"); } \ No newline at end of file