Skip to content

Commit 7fec958

Browse files
xtmqxtmq
authored
Merge pull request #3 from JetBrains/a.slapoguzov/LLM-22701-duplicating-messages
LLM-22701 Fix duplicate messages on follow-up prompts
2 parents 741648a + 9a2b7d8 commit 7fec958

6 files changed

Lines changed: 251 additions & 23 deletions

File tree

‎src/CodexAcpClient.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ export class CodexAcpClient {
102102
}
103103

104104
async sendPrompt(request: acp.PromptRequest, eventHandler: (result: ServerNotification) => void): Promise<void> {
105-
this.codexClient.onServerNotification(eventHandler);
105+
this.codexClient.onServerNotification(request.sessionId, eventHandler);
106106

107107
const input = request.prompt.filter(b => b.type === "text")
108108
.map(b => b.text)

‎src/CodexAppServerClient.ts‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -95,19 +95,22 @@ export class CodexAppServerClient {
9595
return await this.sendRequest({ method: "model/list", params });
9696
}
9797

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);
101104
}
102105

103106
private codexEventHandlers: Array<(event: CodexConnectionEvent) => void> = [];
104107
onClientTransportEvent(callback: (event: CodexConnectionEvent) => void){
105108
this.codexEventHandlers.push(callback);
106109
}
107110

108-
private notificationHandlers: Array<(event: ServerNotification) => void> = [];
111+
private notificationHandlers = new Map<string, (event: ServerNotification) => void>();
109112
private notify(notification: ServerNotification) {
110-
for (const notificationHandler of this.notificationHandlers) {
113+
for (const notificationHandler of this.notificationHandlers.values()) {
111114
notificationHandler(notification);
112115
}
113116
}

‎src/__tests__/CodexACPAgent/CodexAcpClient.test.ts‎

Lines changed: 98 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import {describe, expect, it, vi, beforeEach} from 'vitest';
22
import type {CodexAuthRequest} from "../../CodexAuthMethod";
3-
import {createTestFixture, type TestFixture} from "../acp-test-utils";
3+
import {createTestFixture, createCodexMockTestFixture, type TestFixture} from "../acp-test-utils";
44
import type {ServerNotification} from "../../app-server";
55
import type {SessionState} from "../../CodexAcpServer";
66

@@ -64,7 +64,7 @@ describe('ACP server test', { timeout: 40_000 }, () => {
6464
{ method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "ll", }},
6565
{ method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "o!", }},
6666
];
67-
function onServerNotification(callback: (event: ServerNotification) => void){
67+
function onServerNotification(_sessionId: string, callback: (event: ServerNotification) => void){
6868
for (const notification of serverNotifications) {
6969
callback(notification);
7070
}
@@ -95,6 +95,102 @@ describe('ACP server test', { timeout: 40_000 }, () => {
9595

9696
});
9797

98+
it('should not duplicate messages on follow-up prompts', async () => {
99+
const mockFixture = createCodexMockTestFixture();
100+
const codexAcpAgent = mockFixture.getCodexAcpAgent();
101+
102+
mockFixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue(undefined);
103+
mockFixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue(undefined);
104+
105+
const sessionState: SessionState = {
106+
pendingPrompt: null,
107+
sessionMetadata: {
108+
sessionId: "id",
109+
currentModelId: "model-id",
110+
models: [],
111+
}
112+
};
113+
vi.spyOn(codexAcpAgent, "getSessionState").mockReturnValue(sessionState);
114+
115+
// First prompt - registers first notification handler
116+
await codexAcpAgent.prompt({ sessionId: "id", prompt: [{type: "text", text: "First message"}] });
117+
118+
// Follow-up prompt - should NOT accumulate handlers
119+
await codexAcpAgent.prompt({ sessionId: "id", prompt: [{type: "text", text: "Follow-up message"}] });
120+
121+
mockFixture.clearAcpConnectionDump();
122+
123+
// Trigger notifications after both prompts - should produce only 3 events, not 6
124+
const serverNotifications: ServerNotification[] = [
125+
{ method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "He", }},
126+
{ method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "ll", }},
127+
{ method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "o!", }},
128+
];
129+
for (const notification of serverNotifications) {
130+
mockFixture.sendServerNotification(notification);
131+
}
132+
133+
// Wait for async handlers to complete
134+
await vi.waitFor(() => {
135+
const dump = mockFixture.getAcpConnectionDump([]);
136+
expect(dump.length).toBeGreaterThan(0);
137+
});
138+
139+
await expect(mockFixture.getAcpConnectionDump([])).toMatchFileSnapshot("data/follow-up-no-duplicates.json");
140+
});
141+
142+
it('should handle multiple sessions independently', async () => {
143+
const mockFixture = createCodexMockTestFixture();
144+
const codexAcpAgent = mockFixture.getCodexAcpAgent();
145+
146+
mockFixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue(undefined);
147+
mockFixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue(undefined);
148+
149+
const sessionState1: SessionState = {
150+
pendingPrompt: null,
151+
sessionMetadata: {
152+
sessionId: "session-1",
153+
currentModelId: "model-id",
154+
models: [],
155+
}
156+
};
157+
const sessionState2: SessionState = {
158+
pendingPrompt: null,
159+
sessionMetadata: {
160+
sessionId: "session-2",
161+
currentModelId: "model-id",
162+
models: [],
163+
}
164+
};
165+
166+
vi.spyOn(codexAcpAgent, "getSessionState").mockImplementation((sessionId: string) => {
167+
return sessionId === "session-1" ? sessionState1 : sessionState2;
168+
});
169+
170+
// Start prompts for two different sessions
171+
await codexAcpAgent.prompt({ sessionId: "session-1", prompt: [{type: "text", text: "Message to session 1"}] });
172+
await codexAcpAgent.prompt({ sessionId: "session-2", prompt: [{type: "text", text: "Message to session 2"}] });
173+
174+
mockFixture.clearAcpConnectionDump();
175+
176+
// Trigger notifications - both session handlers should receive them
177+
const serverNotifications: ServerNotification[] = [
178+
{ method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "Hello", }},
179+
];
180+
for (const notification of serverNotifications) {
181+
mockFixture.sendServerNotification(notification);
182+
}
183+
184+
// Wait for async handlers to complete
185+
await vi.waitFor(() => {
186+
const dump = mockFixture.getAcpConnectionDump([]);
187+
expect(dump.length).toBeGreaterThan(0);
188+
});
189+
190+
// Should have 2 events - one for each session's handler
191+
await expect(mockFixture.getAcpConnectionDump([])).toMatchFileSnapshot("data/multiple-sessions.json");
192+
});
193+
98194
//dev-time test
99195
it.skip('should convert session notification to acp events', async () => {
100196
fixture.onCodexConnectionEvent((event) => {

‎src/__tests__/CodexACPAgent/data/follow-up-no-duplicates.json‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
{
2+
"method": "sessionUpdate",
3+
"args": [
4+
{
5+
"sessionId": "id",
6+
"update": {
7+
"sessionUpdate": "agent_message_chunk",
8+
"content": {
9+
"type": "text",
10+
"text": "He"
11+
}
12+
}
13+
}
14+
]
15+
}
16+
{
17+
"method": "sessionUpdate",
18+
"args": [
19+
{
20+
"sessionId": "id",
21+
"update": {
22+
"sessionUpdate": "agent_message_chunk",
23+
"content": {
24+
"type": "text",
25+
"text": "ll"
26+
}
27+
}
28+
}
29+
]
30+
}
31+
{
32+
"method": "sessionUpdate",
33+
"args": [
34+
{
35+
"sessionId": "id",
36+
"update": {
37+
"sessionUpdate": "agent_message_chunk",
38+
"content": {
39+
"type": "text",
40+
"text": "o!"
41+
}
42+
}
43+
}
44+
]
45+
}

‎src/__tests__/CodexACPAgent/data/multiple-sessions.json‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
{
2+
"method": "sessionUpdate",
3+
"args": [
4+
{
5+
"sessionId": "session-1",
6+
"update": {
7+
"sessionUpdate": "agent_message_chunk",
8+
"content": {
9+
"type": "text",
10+
"text": "Hello"
11+
}
12+
}
13+
}
14+
]
15+
}
16+
{
17+
"method": "sessionUpdate",
18+
"args": [
19+
{
20+
"sessionId": "session-2",
21+
"update": {
22+
"sessionUpdate": "agent_message_chunk",
23+
"content": {
24+
"type": "text",
25+
"text": "Hello"
26+
}
27+
}
28+
}
29+
]
30+
}

‎src/__tests__/acp-test-utils.ts‎

Lines changed: 69 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,14 @@ import {type CodexConnectionEvent, CodexAppServerClient} from '../CodexAppServer
33
import {startCodexConnection} from "../CodexJsonRpcConnection";
44
import {CodexAcpServer} from "../CodexAcpServer";
55
import type {AgentSideConnection} from "@agentclientprotocol/sdk";
6+
import type {ServerNotification} from "../app-server";
7+
import type {MessageConnection} from "vscode-jsonrpc/node";
68
import path from "node:path";
79
import fs from "node:fs";
810

911
export type MethodCallEvent = { method: string; args: any[] };
1012

11-
function createSmartMock<T extends object>(onCall: (event: MethodCallEvent) => void) {
13+
export function createSmartMock<T extends object>(onCall: (event: MethodCallEvent) => void) {
1214
return new Proxy({} as T, {
1315
get(_, prop) {
1416
return (...args: any[]) => {
@@ -33,25 +35,24 @@ export interface TestFixture {
3335
clearAcpConnectionDump(): void,
3436
}
3537

36-
export function createTestFixture(): TestFixture {
37-
const pathToCodex = path.resolve(process.cwd(), "node_modules", ".bin", process.platform === 'win32' ? "codex.cmd" : "codex");
38-
if (!fs.existsSync(pathToCodex)) {
39-
throw new Error(`Codex binary not found at ${pathToCodex}. Did you run 'npm install'?`);
40-
}
41-
42-
const acpConnectionEvents: MethodCallEvent[] = []
38+
export interface ConnectionConfig {
39+
connection: MessageConnection;
40+
getExitCode: () => number | null;
41+
}
42+
43+
export function createBaseTestFixture(config: ConnectionConfig): TestFixture {
44+
const acpConnectionEvents: MethodCallEvent[] = [];
4345
const acpEventHandlers: ((event: MethodCallEvent) => void)[] = [];
4446
const acpConnection = createSmartMock<AgentSideConnection>((event) => {
4547
acpConnectionEvents.push(event);
4648
acpEventHandlers.forEach(handler => handler(event));
4749
});
48-
const codexConnection = startCodexConnection(pathToCodex);
49-
const codexAppServerClient = new CodexAppServerClient(codexConnection.connection);
5050

51+
const codexAppServerClient = new CodexAppServerClient(config.connection);
5152
const codexAcpClient = new CodexAcpClient(codexAppServerClient);
52-
const codexAcpAgent = new CodexAcpServer(acpConnection, codexAcpClient, undefined, () => codexConnection.process.exitCode);
53+
const codexAcpAgent = new CodexAcpServer(acpConnection, codexAcpClient, undefined, config.getExitCode);
5354

54-
const transportEvents: CodexConnectionEvent[] = []
55+
const transportEvents: CodexConnectionEvent[] = [];
5556
const codexEventHandlers: ((event: CodexConnectionEvent) => void)[] = [];
5657
codexAppServerClient.onClientTransportEvent((event) => {
5758
transportEvents.push(event);
@@ -83,20 +84,73 @@ export function createTestFixture(): TestFixture {
8384
getAcpConnectionDump(ignoredFields: string[]): string {
8485
return createArrayDump(acpConnectionEvents, ignoredFields);
8586
},
86-
clearAcpConnectionDump(){
87+
clearAcpConnectionDump() {
8788
acpConnectionEvents.splice(0, acpConnectionEvents.length);
8889
}
8990
};
9091
}
9192

93+
/**
94+
* Creates a test fixture with a real Codex connection.
95+
* Use for integration tests that need to interact with the actual Codex binary.
96+
*/
97+
export function createTestFixture(): TestFixture {
98+
const pathToCodex = path.resolve(process.cwd(), "node_modules", ".bin", process.platform === 'win32' ? "codex.cmd" : "codex");
99+
if (!fs.existsSync(pathToCodex)) {
100+
throw new Error(`Codex binary not found at ${pathToCodex}. Did you run 'npm install'?`);
101+
}
102+
103+
const codexConnection = startCodexConnection(pathToCodex);
104+
105+
return createBaseTestFixture({
106+
connection: codexConnection.connection,
107+
getExitCode: () => codexConnection.process.exitCode
108+
});
109+
}
110+
111+
export interface CodexMockTestFixture extends TestFixture {
112+
sendServerNotification(notification: ServerNotification): void,
113+
}
114+
115+
/**
116+
* Creates a test fixture with a mock Codex connection.
117+
* Use for unit tests that don't need a real Codex binary.
118+
* Provides `sendServerNotification()` to simulate server notifications.
119+
*/
120+
export function createCodexMockTestFixture(): CodexMockTestFixture {
121+
let unhandledNotificationHandler: ((notification: any) => void) | null = null;
122+
123+
const mockCodexConnection = {
124+
sendRequest: () => Promise.resolve(undefined),
125+
onUnhandledNotification: (handler: (notification: any) => void) => {
126+
unhandledNotificationHandler = handler;
127+
},
128+
onNotification: () => {},
129+
end: () => {},
130+
} as unknown as MessageConnection;
131+
132+
const baseFixture = createBaseTestFixture({
133+
connection: mockCodexConnection,
134+
getExitCode: () => null
135+
});
136+
137+
return {
138+
...baseFixture,
139+
sendServerNotification(notification: ServerNotification): void {
140+
if (unhandledNotificationHandler) {
141+
unhandledNotificationHandler(notification);
142+
}
143+
}
144+
};
145+
}
92146

93-
function createObjectDump(obj: any, anonymizedFields: string[] = []) {
147+
export function createObjectDump(obj: any, anonymizedFields: string[] = []) {
94148
function fieldAnonymizer(key: string, value: any): any {
95149
return anonymizedFields.includes(key) ? key : value;
96150
}
97151
return JSON.stringify(obj, fieldAnonymizer, 2);
98152
}
99153

100-
function createArrayDump(objects: any[], anonymizedFields: string[]): string {
154+
export function createArrayDump(objects: any[], anonymizedFields: string[]): string {
101155
return objects.map(event => createObjectDump(event, anonymizedFields)).join("\n");
102156
}

0 commit comments

Comments
 (0)