Fix/edit tool calls

src/CodexAcpClient.ts

@@ -48,6 +48,7 @@
     private gatewayConfig: GatewayConfig | null;
     private pendingLoginCompleted: Promise<AccountLoginCompletedNotification> | null = null;
     private pendingAccountUpdated: Promise<AccountUpdatedNotification> | null = null;
+    private readonly sessionNotificationStates = new Map<string, SessionNotificationState>();
 
 
     constructor(codexClient: CodexAppServerClient, codexConfig?: JsonObject, modelProvider?: string) {
@@ -359,14 +360,16 @@
         return ModelId.create(selectedModel.id, reasoningEffort ?? selectedModel.defaultReasoningEffort);
     }
 
-    async subscribeToSessionEvents(
+    subscribeToSessionEvents(
         sessionId: string,
-        eventHandler: (result: ServerNotification) => void,
-        approvalHandler: ApprovalHandler,
+        eventHandler: SessionEventHandler,
         elicitationHandler: ElicitationHandler
-    ) {
-        this.codexClient.onServerNotification(sessionId, eventHandler);
-        this.codexClient.onApprovalRequest(sessionId, approvalHandler);
+    ): void {
+        this.replaceSessionNotificationHandler(sessionId, async (event) => {
+            elicitationHandler.handleNotification(event);
+            await eventHandler.handleNotification(event);
+        });
+        this.codexClient.onApprovalRequest(sessionId, eventHandler);
         this.codexClient.onElicitationRequest(sessionId, elicitationHandler);
     }
 
@@ -397,7 +400,9 @@
 
         // Wait for turn completion
         // If turnInterrupt() was called, Codex will send turn/completed event with status "interrupted"
-        return await this.codexClient.awaitTurnCompleted();
+        const turnCompleted = await this.codexClient.awaitTurnCompleted();
+        await this.awaitPendingSessionNotifications(request.sessionId);
+        return turnCompleted;
     }
 
     async listSkills(params?: SkillsListParams): Promise<SkillsListResponse> {
@@ -447,6 +452,30 @@
         });
     }
 
+    private replaceSessionNotificationHandler(
+        sessionId: string,
+        eventHandler: (event: ServerNotification) => Promise<void>,
+    ): void {
+        this.sessionNotificationStates.get(sessionId)?.subscription.dispose();
+        const state: SessionNotificationState = {
+            pending: Promise.resolve(),
+            subscription: { dispose() {} },
+        };
+        state.subscription = this.codexClient.onServerNotification(async (event) => {
+            state.pending = state.pending
+                .then(() => eventHandler(event))
+                .catch((error) => {
+                    logger.error("Error handling server notification", error);
+                });
+            await state.pending;
+        });
+        this.sessionNotificationStates.set(sessionId, state);
+    }
+
+    private async awaitPendingSessionNotifications(sessionId: string): Promise<void> {
+        await (this.sessionNotificationStates.get(sessionId)?.pending ?? Promise.resolve());
+    }
+
     async listMcpServers(params: ListMcpServerStatusParams = { cursor: null, limit: null }): Promise<ListMcpServerStatusResponse> {
         return this.codexClient.listMcpServerStatus(params);
     }
@@ -586,6 +615,15 @@
 
 export type JsonObject = { [key in string]?: JsonValue }
 
+type SessionEventHandler = ApprovalHandler & {
+    handleNotification(notification: ServerNotification): Promise<void>;
+}
+
+type SessionNotificationState = {
+    pending: Promise<void>;
+    subscription: Disposable;
+}
+
 export type SessionMetadata = {
     sessionId: string,
     currentModelId: string,

src/CodexAcpServer.ts

@@ -6,7 +6,6 @@ import {
     type SessionModeState
 } from "@agentclientprotocol/sdk";
 import {CodexEventHandler} from "./CodexEventHandler";
-import {CodexApprovalHandler} from "./CodexApprovalHandler";
 import {CodexElicitationHandler} from "./CodexElicitationHandler";
 import {CodexAuthMethods, type CodexAuthRequest} from "./CodexAuthMethod";
 import {CodexAcpClient, type SessionMetadata, type SessionMetadataWithThread} from "./CodexAcpClient";
@@ -718,15 +717,8 @@ export class CodexAcpServer implements acp.Agent {
 
         try {
             const eventHandler = new CodexEventHandler(this.connection, sessionState);
-            const approvalHandler = new CodexApprovalHandler(this.connection, sessionState);
             const elicitationHandler = new CodexElicitationHandler(this.connection, sessionState);
-            await this.codexAcpClient.subscribeToSessionEvents(params.sessionId,
-                (event) => {
-                    elicitationHandler.handleNotification(event);
-                    return eventHandler.handleNotification(event);
-                },
-                approvalHandler,
-                elicitationHandler);
+            this.codexAcpClient.subscribeToSessionEvents(params.sessionId, eventHandler, elicitationHandler);
 
             if (await this.availableCommands.tryHandle(params.prompt, sessionState)) {
                 logger.log("Prompt handled by a command");

src/CodexAppServerClient.ts

@@ -1,3 +1,4 @@
+import type {Disposable} from "vscode-jsonrpc";
 import {type MessageConnection, RequestType} from "vscode-jsonrpc/node";
 import type {
     ClientRequest,
@@ -44,6 +45,7 @@ import type {
     FileChangeRequestApprovalParams,
     FileChangeRequestApprovalResponse,
 } from "./app-server/v2";
+import { logger } from "./Logger";
 
 export interface ApprovalHandler {
     handleCommandExecution(params: CommandExecutionRequestApprovalParams): Promise<CommandExecutionRequestApprovalResponse>;
@@ -90,6 +92,7 @@ const McpServerElicitationRequest = new RequestType<
 export class CodexAppServerClient {
     readonly connection: MessageConnection;
     private approvalHandlers = new Map<string, ApprovalHandler>();
+    private readonly notificationHandlers = new Set<(event: ServerNotification) => Promise<void>>();
     private elicitationHandlers = new Map<string, ElicitationHandler>();
     private mcpServerStartupVersion = 0;
     private readonly mcpServerStartupStates = new Map<string, McpServerStartupSnapshot>();
@@ -108,7 +111,7 @@ export class CodexAppServerClient {
                 });
                 this.resolveMcpServerStartupResolvers();
             }
-            this.notify(serverNotification);
+            this.notifyServerNotificationHandlers(serverNotification);
             for (const callback of this.codexEventHandlers) {
                 callback({ eventType: "notification", ...serverNotification });
             }
@@ -241,22 +244,27 @@ export class CodexAppServerClient {
     }
 
     /**
-     * Registers a notification handler for a specific session.
-     * Replaces any existing handler for the same session, preventing handler accumulation.
+     * Registers a notification handler for server notifications.
      */
-    onServerNotification(sessionId: string, callback: (event: ServerNotification) => void) {
-        this.notificationHandlers.set(sessionId, callback);
+    onServerNotification(callback: (event: ServerNotification) => Promise<void>): Disposable {
+        this.notificationHandlers.add(callback);
+        return {
+            dispose: () => {
+                this.notificationHandlers.delete(callback);
+            }
+        };
     }
 
     private codexEventHandlers: Array<(event: CodexConnectionEvent) => void> = [];
     onClientTransportEvent(callback: (event: CodexConnectionEvent) => void){
         this.codexEventHandlers.push(callback);
     }
 
-    private notificationHandlers = new Map<string, (event: ServerNotification) => void>();
-    private notify(notification: ServerNotification) {
-        for (const notificationHandler of this.notificationHandlers.values()) {
-            notificationHandler(notification);
+    private notifyServerNotificationHandlers(notification: ServerNotification): void {
+        for (const notificationHandler of this.notificationHandlers) {
+            void notificationHandler(notification).catch((error) => {
+                logger.error("Error handling server notification", error);
+            });
         }
     }