Fix/edit tool calls

src/CodexAcpClient.ts

@@ -359,7 +359,7 @@ export class CodexAcpClient {
 
     async subscribeToSessionEvents(
         sessionId: string,
-        eventHandler: (result: ServerNotification) => void,
+        eventHandler: (result: ServerNotification) => void | Promise<void>,
         approvalHandler: ApprovalHandler
     ) {
         this.codexClient.onServerNotification(sessionId, eventHandler);
@@ -393,7 +393,9 @@ export class CodexAcpClient {
 
         // 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.codexClient.flushServerNotifications(request.sessionId);
+        return turnCompleted;
     }
 
     async listSkills(params?: SkillsListParams): Promise<SkillsListResponse> {

src/CodexAcpServer.ts

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

src/CodexAppServerClient.ts

@@ -35,6 +35,7 @@ import type {
     ListMcpServerStatusParams,
     ListMcpServerStatusResponse, ConfigReadParams, ConfigReadResponse,
 } from "./app-server/v2";
+import { logger } from "./Logger";
 
 export interface ApprovalHandler {
     handleCommandExecution(params: CommandExecutionRequestApprovalParams): Promise<CommandExecutionRequestApprovalResponse>;
@@ -60,6 +61,8 @@ const FileChangeApprovalRequest = new RequestType<
 export class CodexAppServerClient {
     readonly connection: MessageConnection;
     private approvalHandlers = new Map<string, ApprovalHandler>();
+    private readonly notificationHandlers = new Map<string, (event: ServerNotification) => void | Promise<void>>();
+    private readonly notificationQueues = new Map<string, Promise<void> | null>();
     private mcpStartupCompleteVersion = 0;
     private lastMcpStartupComplete: McpStartupCompleteEvent | null = null;
     private readonly mcpStartupCompleteResolvers: Array<SignalResolver<McpStartupCompleteEvent>> = [];
@@ -190,22 +193,56 @@ export class CodexAppServerClient {
      * Registers a notification handler for a specific session.
      * Replaces any existing handler for the same session, preventing handler accumulation.
      */
-    onServerNotification(sessionId: string, callback: (event: ServerNotification) => void) {
+    onServerNotification(sessionId: string, callback: (event: ServerNotification) => void | Promise<void>) {
         this.notificationHandlers.set(sessionId, callback);
+        this.notificationQueues.set(sessionId, null);
     }
 
     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);
+        for (const [sessionId, notificationHandler] of this.notificationHandlers.entries()) {
+            const queue = this.notificationQueues.get(sessionId);
+            if (queue) {
+                const next = queue
+                    .then(() => notificationHandler(notification))
+                    .catch((error) => {
+                        logger.error("Error handling server notification", error);
+                    });
+                this.notificationQueues.set(sessionId, this.trackNotificationQueue(sessionId, next));
+                continue;
+            }
+
+            try {
+                const result = notificationHandler(notification);
+                if (result instanceof Promise) {
+                    const next = result.catch((error) => {
+                        logger.error("Error handling server notification", error);
+                    });
+                    this.notificationQueues.set(sessionId, this.trackNotificationQueue(sessionId, next));
+                }
+            } catch (error) {
+                logger.error("Error handling server notification", error);
+            }
         }
     }
 
+    async flushServerNotifications(sessionId: string): Promise<void> {
+        await (this.notificationQueues.get(sessionId) ?? Promise.resolve());
+    }
+
+    private trackNotificationQueue(sessionId: string, queue: Promise<void>): Promise<void> {
+        const trackedQueue = queue.finally(() => {
+            if (this.notificationQueues.get(sessionId) === trackedQueue) {
+                this.notificationQueues.set(sessionId, null);
+            }
+        });
+        return trackedQueue;
+    }
+
     private resolveSignal<T>(
         event: T,
         version: number,