Skip to content

Commit 1ee6ac7

Browse files
i1broi1bro
committed
Move session notification sequencing out of app-server client
1 parent 291feb1 commit 1ee6ac7

4 files changed

Lines changed: 59 additions & 58 deletions

File tree

‎src/CodexAcpClient.ts‎

Lines changed: 40 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ export class CodexAcpClient {
4444
private gatewayConfig: GatewayConfig | null;
4545
private pendingLoginCompleted: Promise<AccountLoginCompletedNotification> | null = null;
4646
private pendingAccountUpdated: Promise<AccountUpdatedNotification> | null = null;
47+
private readonly sessionNotificationStates = new Map<string, SessionNotificationState>();
4748

4849

4950
constructor(codexClient: CodexAppServerClient, codexConfig?: JsonObject, modelProvider?: string) {
@@ -357,13 +358,12 @@ export class CodexAcpClient {
357358
return ModelId.create(selectedModel.id, reasoningEffort ?? selectedModel.defaultReasoningEffort);
358359
}
359360

360-
async subscribeToSessionEvents(
361+
subscribeToSessionEvents(
361362
sessionId: string,
362-
eventHandler: (result: ServerNotification) => void | Promise<void>,
363-
approvalHandler: ApprovalHandler
364-
) {
365-
this.codexClient.onServerNotification(sessionId, eventHandler);
366-
this.codexClient.onApprovalRequest(sessionId, approvalHandler);
363+
eventHandler: SessionEventHandler
364+
): void {
365+
this.replaceSessionNotificationHandler(sessionId, eventHandler.handleNotification.bind(eventHandler));
366+
this.codexClient.onApprovalRequest(sessionId, eventHandler);
367367
}
368368

369369
async sendPrompt(
@@ -394,7 +394,7 @@ export class CodexAcpClient {
394394
// Wait for turn completion
395395
// If turnInterrupt() was called, Codex will send turn/completed event with status "interrupted"
396396
const turnCompleted = await this.codexClient.awaitTurnCompleted();
397-
await this.codexClient.flushServerNotifications(request.sessionId);
397+
await this.awaitPendingSessionNotifications(request.sessionId);
398398
return turnCompleted;
399399
}
400400

@@ -445,6 +445,30 @@ export class CodexAcpClient {
445445
});
446446
}
447447

448+
private replaceSessionNotificationHandler(
449+
sessionId: string,
450+
eventHandler: (event: ServerNotification) => Promise<void>,
451+
): void {
452+
this.sessionNotificationStates.get(sessionId)?.subscription.dispose();
453+
const state: SessionNotificationState = {
454+
pending: Promise.resolve(),
455+
subscription: { dispose() {} },
456+
};
457+
state.subscription = this.codexClient.onServerNotification(async (event) => {
458+
state.pending = state.pending
459+
.then(() => eventHandler(event))
460+
.catch((error) => {
461+
logger.error("Error handling server notification", error);
462+
});
463+
await state.pending;
464+
});
465+
this.sessionNotificationStates.set(sessionId, state);
466+
}
467+
468+
private async awaitPendingSessionNotifications(sessionId: string): Promise<void> {
469+
await (this.sessionNotificationStates.get(sessionId)?.pending ?? Promise.resolve());
470+
}
471+
448472
async listMcpServers(params: ListMcpServerStatusParams = { cursor: null, limit: null }): Promise<ListMcpServerStatusResponse> {
449473
return this.codexClient.listMcpServerStatus(params);
450474
}
@@ -584,6 +608,15 @@ export class CodexAcpClient {
584608

585609
export type JsonObject = { [key in string]?: JsonValue }
586610

611+
type SessionEventHandler = ApprovalHandler & {
612+
handleNotification(notification: ServerNotification): Promise<void>;
613+
}
614+
615+
type SessionNotificationState = {
616+
pending: Promise<void>;
617+
subscription: Disposable;
618+
}
619+
587620
export type SessionMetadata = {
588621
sessionId: string,
589622
currentModelId: string,

‎src/CodexAcpServer.ts‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -705,11 +705,7 @@ export class CodexAcpServer implements acp.Agent {
705705

706706
try {
707707
const eventHandler = new CodexEventHandler(this.connection, sessionState);
708-
this.codexAcpClient.subscribeToSessionEvents(
709-
params.sessionId,
710-
(event) => eventHandler.handleNotification(event),
711-
eventHandler
712-
);
708+
this.codexAcpClient.subscribeToSessionEvents(params.sessionId, eventHandler);
713709

714710
if (await this.availableCommands.tryHandle(params.prompt, sessionState)) {
715711
logger.log("Prompt handled by a command");

‎src/CodexAppServerClient.ts‎

Lines changed: 15 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import type {Disposable} from "vscode-jsonrpc";
12
import {type MessageConnection, RequestType} from "vscode-jsonrpc/node";
23
import type {
34
ClientRequest,
@@ -61,8 +62,7 @@ const FileChangeApprovalRequest = new RequestType<
6162
export class CodexAppServerClient {
6263
readonly connection: MessageConnection;
6364
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>();
65+
private readonly notificationHandlers = new Set<(event: ServerNotification) => Promise<void>>();
6666
private mcpStartupCompleteVersion = 0;
6767
private lastMcpStartupComplete: McpStartupCompleteEvent | null = null;
6868
private readonly mcpStartupCompleteResolvers: Array<SignalResolver<McpStartupCompleteEvent>> = [];
@@ -80,7 +80,7 @@ export class CodexAppServerClient {
8080
return;
8181
}
8282
const serverNotification = data as ServerNotification;
83-
this.notify(serverNotification);
83+
this.notifyServerNotificationHandlers(serverNotification);
8484
for (const callback of this.codexEventHandlers) {
8585
callback({ eventType: "notification", ...serverNotification });
8686
}
@@ -190,59 +190,30 @@ export class CodexAppServerClient {
190190
}
191191

192192
/**
193-
* Registers a notification handler for a specific session.
194-
* Replaces any existing handler for the same session, preventing handler accumulation.
193+
* Registers a notification handler for server notifications.
195194
*/
196-
onServerNotification(sessionId: string, callback: (event: ServerNotification) => void | Promise<void>) {
197-
this.notificationHandlers.set(sessionId, callback);
198-
this.notificationQueues.set(sessionId, null);
195+
onServerNotification(callback: (event: ServerNotification) => Promise<void>): Disposable {
196+
this.notificationHandlers.add(callback);
197+
return {
198+
dispose: () => {
199+
this.notificationHandlers.delete(callback);
200+
}
201+
};
199202
}
200203

201204
private codexEventHandlers: Array<(event: CodexConnectionEvent) => void> = [];
202205
onClientTransportEvent(callback: (event: CodexConnectionEvent) => void){
203206
this.codexEventHandlers.push(callback);
204207
}
205208

206-
private notify(notification: ServerNotification) {
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) {
209+
private notifyServerNotificationHandlers(notification: ServerNotification): void {
210+
for (const notificationHandler of this.notificationHandlers) {
211+
void notificationHandler(notification).catch((error) => {
228212
logger.error("Error handling server notification", error);
229-
}
213+
});
230214
}
231215
}
232216

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-
246217
private resolveSignal<T>(
247218
event: T,
248219
version: number,

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -387,10 +387,11 @@ describe('ACP server test', { timeout: 40_000 }, () => {
387387
{ method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "ll", }},
388388
{ method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "o!", }},
389389
];
390-
function onServerNotification(_sessionId: string, callback: (event: ServerNotification) => void){
390+
function onServerNotification(callback: (event: ServerNotification) => Promise<void>){
391391
for (const notification of serverNotifications) {
392-
callback(notification);
392+
void callback(notification);
393393
}
394+
return { dispose() {} };
394395
}
395396
return onServerNotification;
396397
}

0 commit comments

Comments
 (0)