Skip to content

Commit 398d936

Browse files
i1broi1bro
committed
Serialize session notifications before finishing prompts
1 parent e0f1e37 commit 398d936

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
@@ -362,7 +362,7 @@ export class CodexAcpClient {
362362

363363
async subscribeToSessionEvents(
364364
sessionId: string,
365-
eventHandler: (result: ServerNotification) => void,
365+
eventHandler: (result: ServerNotification) => void | Promise<void>,
366366
approvalHandler: ApprovalHandler
367367
) {
368368
this.codexClient.onServerNotification(sessionId, eventHandler);
@@ -396,7 +396,9 @@ export class CodexAcpClient {
396396

397397
// Wait for turn completion
398398
// If turnInterrupt() was called, Codex will send turn/completed event with status "interrupted"
399-
return await this.codexClient.awaitTurnCompleted();
399+
const turnCompleted = await this.codexClient.awaitTurnCompleted();
400+
await this.codexClient.flushServerNotifications(request.sessionId);
401+
return turnCompleted;
400402
}
401403

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

‎src/CodexAppServerClient.ts‎

Lines changed: 41 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ import type {
3737
McpServerElicitationRequestParams,
3838
McpServerElicitationRequestResponse,
3939
} from "./app-server/v2";
40+
import { logger } from "./Logger";
4041

4142
export interface ApprovalHandler {
4243
handleCommandExecution(params: CommandExecutionRequestApprovalParams): Promise<CommandExecutionRequestApprovalResponse>;
@@ -69,6 +70,8 @@ const McpServerElicitationRequest = new RequestType<
6970
export class CodexAppServerClient {
7071
readonly connection: MessageConnection;
7172
private approvalHandlers = new Map<string, ApprovalHandler>();
73+
private readonly notificationHandlers = new Map<string, (event: ServerNotification) => void | Promise<void>>();
74+
private readonly notificationQueues = new Map<string, Promise<void> | null>();
7275
private mcpStartupCompleteVersion = 0;
7376
private lastMcpStartupComplete: McpStartupCompleteEvent | null = null;
7477
private readonly mcpStartupCompleteResolvers: Array<SignalResolver<McpStartupCompleteEvent>> = [];
@@ -207,22 +210,56 @@ export class CodexAppServerClient {
207210
* Registers a notification handler for a specific session.
208211
* Replaces any existing handler for the same session, preventing handler accumulation.
209212
*/
210-
onServerNotification(sessionId: string, callback: (event: ServerNotification) => void) {
213+
onServerNotification(sessionId: string, callback: (event: ServerNotification) => void | Promise<void>) {
211214
this.notificationHandlers.set(sessionId, callback);
215+
this.notificationQueues.set(sessionId, null);
212216
}
213217

214218
private codexEventHandlers: Array<(event: CodexConnectionEvent) => void> = [];
215219
onClientTransportEvent(callback: (event: CodexConnectionEvent) => void){
216220
this.codexEventHandlers.push(callback);
217221
}
218222

219-
private notificationHandlers = new Map<string, (event: ServerNotification) => void>();
220223
private notify(notification: ServerNotification) {
221-
for (const notificationHandler of this.notificationHandlers.values()) {
222-
notificationHandler(notification);
224+
for (const [sessionId, notificationHandler] of this.notificationHandlers.entries()) {
225+
const queue = this.notificationQueues.get(sessionId);
226+
if (queue) {
227+
const next = queue
228+
.then(() => notificationHandler(notification))
229+
.catch((error) => {
230+
logger.error("Error handling server notification", error);
231+
});
232+
this.notificationQueues.set(sessionId, this.trackNotificationQueue(sessionId, next));
233+
continue;
234+
}
235+
236+
try {
237+
const result = notificationHandler(notification);
238+
if (result instanceof Promise) {
239+
const next = result.catch((error) => {
240+
logger.error("Error handling server notification", error);
241+
});
242+
this.notificationQueues.set(sessionId, this.trackNotificationQueue(sessionId, next));
243+
}
244+
} catch (error) {
245+
logger.error("Error handling server notification", error);
246+
}
223247
}
224248
}
225249

250+
async flushServerNotifications(sessionId: string): Promise<void> {
251+
await (this.notificationQueues.get(sessionId) ?? Promise.resolve());
252+
}
253+
254+
private trackNotificationQueue(sessionId: string, queue: Promise<void>): Promise<void> {
255+
const trackedQueue = queue.finally(() => {
256+
if (this.notificationQueues.get(sessionId) === trackedQueue) {
257+
this.notificationQueues.set(sessionId, null);
258+
}
259+
});
260+
return trackedQueue;
261+
}
262+
226263
private resolveSignal<T>(
227264
event: T,
228265
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)