Skip to content

Commit 4c9030c

Browse files
i1broi1bro
committed
Serialize session updates before tool-call completion
1 parent 6c83fc4 commit 4c9030c

9 files changed

Lines changed: 242 additions & 82 deletions

‎src/CodexAcpClient.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -359,7 +359,7 @@ export class CodexAcpClient {
359359

360360
async subscribeToSessionEvents(
361361
sessionId: string,
362-
eventHandler: (result: ServerNotification) => void,
362+
eventHandler: (result: ServerNotification) => void | Promise<void>,
363363
approvalHandler: ApprovalHandler
364364
) {
365365
this.codexClient.onServerNotification(sessionId, eventHandler);
@@ -393,7 +393,9 @@ export class CodexAcpClient {
393393

394394
// Wait for turn completion
395395
// If turnInterrupt() was called, Codex will send turn/completed event with status "interrupted"
396-
return await this.codexClient.awaitTurnCompleted();
396+
const turnCompleted = await this.codexClient.awaitTurnCompleted();
397+
await this.codexClient.flushServerNotifications(request.sessionId);
398+
return turnCompleted;
397399
}
398400

399401
async listSkills(params?: SkillsListParams): Promise<SkillsListResponse> {

‎src/CodexAppServerClient.ts‎

Lines changed: 41 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import type {
3535
ListMcpServerStatusParams,
3636
ListMcpServerStatusResponse, ConfigReadParams, ConfigReadResponse,
3737
} from "./app-server/v2";
38+
import { logger } from "./Logger";
3839

3940
export interface ApprovalHandler {
4041
handleCommandExecution(params: CommandExecutionRequestApprovalParams): Promise<CommandExecutionRequestApprovalResponse>;
@@ -60,6 +61,8 @@ const FileChangeApprovalRequest = new RequestType<
6061
export class CodexAppServerClient {
6162
readonly connection: MessageConnection;
6263
private approvalHandlers = new Map<string, ApprovalHandler>();
64+
private readonly notificationHandlers = new Map<string, (event: ServerNotification) => void | Promise<void>>();
65+
private readonly notificationQueues = new Map<string, Promise<void> | null>();
6366
private mcpStartupCompleteVersion = 0;
6467
private lastMcpStartupComplete: McpStartupCompleteEvent | null = null;
6568
private readonly mcpStartupCompleteResolvers: Array<SignalResolver<McpStartupCompleteEvent>> = [];
@@ -190,22 +193,56 @@ export class CodexAppServerClient {
190193
* Registers a notification handler for a specific session.
191194
* Replaces any existing handler for the same session, preventing handler accumulation.
192195
*/
193-
onServerNotification(sessionId: string, callback: (event: ServerNotification) => void) {
196+
onServerNotification(sessionId: string, callback: (event: ServerNotification) => void | Promise<void>) {
194197
this.notificationHandlers.set(sessionId, callback);
198+
this.notificationQueues.set(sessionId, null);
195199
}
196200

197201
private codexEventHandlers: Array<(event: CodexConnectionEvent) => void> = [];
198202
onClientTransportEvent(callback: (event: CodexConnectionEvent) => void){
199203
this.codexEventHandlers.push(callback);
200204
}
201205

202-
private notificationHandlers = new Map<string, (event: ServerNotification) => void>();
203206
private notify(notification: ServerNotification) {
204-
for (const notificationHandler of this.notificationHandlers.values()) {
205-
notificationHandler(notification);
207+
for (const [sessionId, notificationHandler] of this.notificationHandlers.entries()) {
208+
const queue = this.notificationQueues.get(sessionId);
209+
if (queue) {
210+
const next = queue
211+
.then(() => notificationHandler(notification))
212+
.catch((error) => {
213+
logger.error("Error handling server notification", error);
214+
});
215+
this.notificationQueues.set(sessionId, this.trackNotificationQueue(sessionId, next));
216+
continue;
217+
}
218+
219+
try {
220+
const result = notificationHandler(notification);
221+
if (result instanceof Promise) {
222+
const next = result.catch((error) => {
223+
logger.error("Error handling server notification", error);
224+
});
225+
this.notificationQueues.set(sessionId, this.trackNotificationQueue(sessionId, next));
226+
}
227+
} catch (error) {
228+
logger.error("Error handling server notification", error);
229+
}
206230
}
207231
}
208232

233+
async flushServerNotifications(sessionId: string): Promise<void> {
234+
await (this.notificationQueues.get(sessionId) ?? Promise.resolve());
235+
}
236+
237+
private trackNotificationQueue(sessionId: string, queue: Promise<void>): Promise<void> {
238+
const trackedQueue = queue.finally(() => {
239+
if (this.notificationQueues.get(sessionId) === trackedQueue) {
240+
this.notificationQueues.set(sessionId, null);
241+
}
242+
});
243+
return trackedQueue;
244+
}
245+
209246
private resolveSignal<T>(
210247
event: T,
211248
version: number,

‎src/CodexEventHandler.ts‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import {toTokenCount} from "./TokenCount";
2525
import {
2626
createCommandExecutionUpdate,
2727
createDynamicToolCallUpdate,
28+
createFileChangeCompletionUpdate,
2829
createFileChangeUpdate,
2930
createMcpRawInput,
3031
createMcpRawOutput,
@@ -88,8 +89,6 @@ export class CodexEventHandler {
8889
return null;
8990
case "turn/completed":
9091
this.sessionState.currentTurnId = null;
91-
this.approvalContext.fileChangesByItemId.clear();
92-
this.approvalContext.turnDiffsByTurnId.clear();
9392
return null;
9493
case "thread/tokenUsage/updated":
9594
return this.createUsageUpdate(notification.params);
@@ -224,11 +223,7 @@ export class CodexEventHandler {
224223
switch (event.item.type) {
225224
case "fileChange":
226225
this.approvalContext.fileChangesByItemId.set(event.item.id, event.item);
227-
return {
228-
sessionUpdate: "tool_call_update",
229-
toolCallId: event.item.id,
230-
status: event.item.status === "completed" ? "completed" : "failed",
231-
}
226+
return createFileChangeCompletionUpdate(event.item);
232227
case "dynamicToolCall":
233228
return {
234229
sessionUpdate: "tool_call_update",

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -983,6 +983,10 @@ describe('ACP server test', { timeout: 40_000 }, () => {
983983
}
984984
});
985985

986+
await vi.waitFor(() => {
987+
expect(sessionState.rateLimits?.size).toBe(2);
988+
});
989+
986990
expect(sessionState.rateLimits).not.toBeNull();
987991
expect(sessionState.rateLimits!.size).toBe(2);
988992
expect(sessionState.rateLimits!.get("standard-limit")).toEqual({

‎src/__tests__/CodexACPAgent/data/file-change-completion-update.json‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
{
2+
"method": "sessionUpdate",
3+
"args": [
4+
{
5+
"sessionId": "test-session-id",
6+
"update": {
7+
"sessionUpdate": "tool_call",
8+
"toolCallId": "file-change-complete",
9+
"title": "/test/project/OldFile.kt",
10+
"kind": "edit",
11+
"status": "in_progress",
12+
"content": [
13+
{
14+
"type": "diff",
15+
"oldText": "package test.project\n\nclass OldFile {}",
16+
"newText": "package test.project\n\nclass OldFile { fun hello() = \"Hello\" }",
17+
"path": "/test/project/OldFile.kt",
18+
"_meta": {
19+
"kind": "update"
20+
}
21+
}
22+
],
23+
"locations": [
24+
{
25+
"path": "/test/project/OldFile.kt"
26+
}
27+
],
28+
"rawInput": {
29+
"changes": [
30+
{
31+
"path": "/test/project/OldFile.kt",
32+
"kind": {
33+
"type": "update"
34+
},
35+
"diff": "--- /test/project/OldFile.kt\n+++ /test/project/OldFile.kt\n@@ -1,3 +1,3 @@\n package test.project\n \n-class OldFile {}\n+class OldFile { fun hello() = \"Hello\" }"
36+
}
37+
]
38+
}
39+
}
40+
}
41+
]
42+
}
43+
{
44+
"method": "sessionUpdate",
45+
"args": [
46+
{
47+
"sessionId": "test-session-id",
48+
"update": {
49+
"sessionUpdate": "tool_call_update",
50+
"toolCallId": "file-change-complete",
51+
"status": "completed"
52+
}
53+
}
54+
]
55+
}

‎src/__tests__/CodexACPAgent/data/mcp-tool-completed-with-logs.json‎

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,19 @@
44
{
55
"sessionId": "test-session-id",
66
"update": {
7-
"sessionUpdate": "tool_call_update",
7+
"sessionUpdate": "tool_call",
88
"toolCallId": "call-id",
9-
"_meta": {
10-
"mcp_output_delta": {
11-
"data": "File /Users/aleksandr.slapoguzov/Projects/ultimate/.ai/local.md doesn't exist or can't be opened"
9+
"kind": "execute",
10+
"title": "mcp.ijproxy.read_file",
11+
"status": "in_progress",
12+
"rawInput": {
13+
"server": "ijproxy",
14+
"tool": "read_file",
15+
"arguments": {
16+
"file_path": ".ai/local.md",
17+
"mode": "slice",
18+
"start_line": 1,
19+
"max_lines": 200
1220
}
1321
}
1422
}
@@ -23,21 +31,9 @@
2331
"update": {
2432
"sessionUpdate": "tool_call_update",
2533
"toolCallId": "call-id",
26-
"status": "failed",
27-
"rawInput": {
28-
"server": "ijproxy",
29-
"tool": "read_file",
30-
"arguments": {
31-
"file_path": ".ai/local.md",
32-
"mode": "slice",
33-
"start_line": 1,
34-
"max_lines": 200
35-
}
36-
},
37-
"rawOutput": {
38-
"result": null,
39-
"error": {
40-
"message": "File /Users/aleksandr.slapoguzov/Projects/ultimate/.ai/local.md doesn't exist or can't be opened"
34+
"_meta": {
35+
"mcp_output_delta": {
36+
"data": "File /Users/aleksandr.slapoguzov/Projects/ultimate/.ai/local.md doesn't exist or can't be opened"
4137
}
4238
}
4339
}
@@ -50,11 +46,9 @@
5046
{
5147
"sessionId": "test-session-id",
5248
"update": {
53-
"sessionUpdate": "tool_call",
49+
"sessionUpdate": "tool_call_update",
5450
"toolCallId": "call-id",
55-
"kind": "execute",
56-
"title": "mcp.ijproxy.read_file",
57-
"status": "in_progress",
51+
"status": "failed",
5852
"rawInput": {
5953
"server": "ijproxy",
6054
"tool": "read_file",
@@ -64,6 +58,12 @@
6458
"start_line": 1,
6559
"max_lines": 200
6660
}
61+
},
62+
"rawOutput": {
63+
"result": null,
64+
"error": {
65+
"message": "File /Users/aleksandr.slapoguzov/Projects/ultimate/.ai/local.md doesn't exist or can't be opened"
66+
}
6767
}
6868
}
6969
}

‎src/__tests__/CodexACPAgent/data/mcp-tool-repeated-progress.json‎

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,16 @@
44
{
55
"sessionId": "test-session-id",
66
"update": {
7-
"sessionUpdate": "tool_call_update",
7+
"sessionUpdate": "tool_call",
88
"toolCallId": "call-id",
9-
"_meta": {
10-
"mcp_output_delta": {
11-
"data": "Polling for status"
9+
"kind": "execute",
10+
"title": "mcp.server-name.tool-name",
11+
"status": "in_progress",
12+
"rawInput": {
13+
"server": "server-name",
14+
"tool": "tool-name",
15+
"arguments": {
16+
"argument": "example"
1217
}
1318
}
1419
}
@@ -40,18 +45,9 @@
4045
"update": {
4146
"sessionUpdate": "tool_call_update",
4247
"toolCallId": "call-id",
43-
"status": "failed",
44-
"rawInput": {
45-
"server": "server-name",
46-
"tool": "tool-name",
47-
"arguments": {
48-
"argument": "example"
49-
}
50-
},
51-
"rawOutput": {
52-
"result": null,
53-
"error": {
54-
"message": "Polling for status"
48+
"_meta": {
49+
"mcp_output_delta": {
50+
"data": "Polling for status"
5551
}
5652
}
5753
}
@@ -64,17 +60,21 @@
6460
{
6561
"sessionId": "test-session-id",
6662
"update": {
67-
"sessionUpdate": "tool_call",
63+
"sessionUpdate": "tool_call_update",
6864
"toolCallId": "call-id",
69-
"kind": "execute",
70-
"title": "mcp.server-name.tool-name",
71-
"status": "in_progress",
65+
"status": "failed",
7266
"rawInput": {
7367
"server": "server-name",
7468
"tool": "tool-name",
7569
"arguments": {
7670
"argument": "example"
7771
}
72+
},
73+
"rawOutput": {
74+
"result": null,
75+
"error": {
76+
"message": "Polling for status"
77+
}
7878
}
7979
}
8080
}

0 commit comments

Comments
 (0)