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

@@ -35,6 +35,7 @@ import {
     createFileChangeUpdate,
     createMcpToolCallUpdate,
 } from "./CodexToolCallMapper";
+import { createApprovalContextStore } from "./CodexApprovalContext";
 
 export interface SessionState {
     sessionId: string,
@@ -705,8 +706,9 @@ export class CodexAcpServer implements acp.Agent {
         sessionState.lastTokenUsage = null;
 
         try {
-            const eventHandler = new CodexEventHandler(this.connection, sessionState);
-            const approvalHandler = new CodexApprovalHandler(this.connection, sessionState);
+            const approvalContext = createApprovalContextStore();
+            const eventHandler = new CodexEventHandler(this.connection, sessionState, approvalContext);
+            const approvalHandler = new CodexApprovalHandler(this.connection, sessionState, approvalContext);
             await this.codexAcpClient.subscribeToSessionEvents(params.sessionId,
                 (event) => eventHandler.handleNotification(event),
                 approvalHandler);

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,

src/CodexApprovalContext.ts

@@ -0,0 +1,13 @@
+import type { ThreadItem } from "./app-server/v2";
+
+export interface ApprovalContextStore {
+    readonly fileChangesByItemId: Map<string, ThreadItem & { type: "fileChange" }>;
+    readonly turnDiffsByTurnId: Map<string, string>;
+}
+
+export function createApprovalContextStore(): ApprovalContextStore {
+    return {
+        fileChangesByItemId: new Map(),
+        turnDiffsByTurnId: new Map(),
+    };
+}

src/CodexApprovalHandler.ts

@@ -5,11 +5,19 @@ import type {
     CommandExecutionRequestApprovalParams,
     CommandExecutionRequestApprovalResponse,
     FileChangeRequestApprovalParams,
-    FileChangeRequestApprovalResponse
+    FileChangeRequestApprovalResponse,
+    FileUpdateChange,
 } from "./app-server/v2";
 import type {ToolCallContent} from "@agentclientprotocol/sdk/dist/schema/types.gen";
 import {logger} from "./Logger";
 import {stripShellPrefix} from "./CodexEventHandler";
+import type {ApprovalContextStore} from "./CodexApprovalContext";
+import {
+    createFileChangeContents,
+    createFileChangeLocations,
+    createRawFileChangeInput,
+    parseUnifiedDiffChanges,
+} from "./CodexToolCallMapper";
 
 const APPROVAL_OPTIONS: acp.PermissionOption[] = [
     { optionId: "allow_once", name: "Allow Once", kind: "allow_once" },
@@ -20,13 +28,16 @@ const APPROVAL_OPTIONS: acp.PermissionOption[] = [
 export class CodexApprovalHandler implements ApprovalHandler {
     private readonly connection: acp.AgentSideConnection;
     private readonly sessionState: SessionState;
+    private readonly approvalContext: ApprovalContextStore;
 
     constructor(
         connection: acp.AgentSideConnection,
-        sessionState: SessionState
+        sessionState: SessionState,
+        approvalContext: ApprovalContextStore,
     ) {
         this.connection = connection;
         this.sessionState = sessionState;
+        this.approvalContext = approvalContext;
     }
 
     async handleCommandExecution(
@@ -48,7 +59,7 @@ export class CodexApprovalHandler implements ApprovalHandler {
     ): Promise<FileChangeRequestApprovalResponse> {
         try {
             const sessionId = this.sessionState.sessionId;
-            const acpRequest = this.buildFileChangePermissionRequest(sessionId, params);
+            const acpRequest = await this.buildFileChangePermissionRequest(sessionId, params);
             const response = await this.connection.requestPermission(acpRequest);
             return this.convertFileChangeResponse(response);
         } catch (error) {
@@ -61,7 +72,7 @@ export class CodexApprovalHandler implements ApprovalHandler {
         sessionId: string,
         params: CommandExecutionRequestApprovalParams
     ): acp.RequestPermissionRequest {
-        const reasonContent = this.createContentFromReason(params.reason ?? null);
+        const reasonContent = this.createTextContent(params.reason ?? null);
         return {
             sessionId,
             toolCall: {
@@ -75,32 +86,55 @@ export class CodexApprovalHandler implements ApprovalHandler {
         };
     }
 
-    private createContentFromReason(reason: string | null): ToolCallContent | null {
-        if (reason === null || reason === "") {
+    private createTextContent(text: string | null): ToolCallContent | null {
+        if (text === null || text === "") {
             return null;
         }
         return {
             type: "content",
             content: {
                 type: "text",
-                text: reason
+                text
             }
         }
     }
 
-    private buildFileChangePermissionRequest(
+    private async buildFileChangePermissionRequest(
         sessionId: string,
         params: FileChangeRequestApprovalParams
-    ): acp.RequestPermissionRequest {
-        const reasonContent = this.createContentFromReason(params.reason ?? null);
+    ): Promise<acp.RequestPermissionRequest> {
+        const reasonContent = this.createTextContent(params.reason ?? null);
+        const fileChange = this.approvalContext.fileChangesByItemId.get(params.itemId);
+        const content: ToolCallContent[] = reasonContent ? [reasonContent] : [];
+        const toolCall: acp.ToolCallUpdate = {
+            toolCallId: params.itemId,
+            kind: "edit",
+            status: "pending",
+        };
+        if (fileChange) {
+            content.push(...await createFileChangeContents(fileChange.changes));
+            toolCall.locations = createFileChangeLocations(fileChange.changes);
+            toolCall.rawInput = createRawFileChangeInput(fileChange.changes);
+        } else {
+            const turnDiff = this.approvalContext.turnDiffsByTurnId.get(params.turnId);
+            if (turnDiff) {
+                const parsedChanges = parseUnifiedDiffChanges(turnDiff);
+                content.push(...await createFileChangeContents(parsedChanges));
+                const locations = createFileChangeLocations(parsedChanges);
+                if (locations.length > 0) {
+                    toolCall.locations = locations;
+                }
+                toolCall.rawInput = parsedChanges.length > 0
+                    ? { unifiedDiff: turnDiff, ...createRawFileChangeInput(parsedChanges) }
+                    : { unifiedDiff: turnDiff };
+            }
+        }
+        if (content.length > 0) {
+            toolCall.content = content;
+        }
         return {
             sessionId,
-            toolCall: {
-                toolCallId: params.itemId,
-                kind: "edit",
-                status: "pending",
-                content: reasonContent ? [reasonContent] : null,
-            },
+            toolCall,
             options: APPROVAL_OPTIONS,
         };
     }

src/CodexEventHandler.ts

@@ -25,6 +25,7 @@ import {toTokenCount} from "./TokenCount";
 import {
     createCommandExecutionUpdate,
     createDynamicToolCallUpdate,
+    createFileChangeCompletionUpdate,
     createFileChangeUpdate,
     createMcpRawInput,
     createMcpRawOutput,
@@ -34,19 +35,22 @@ import {
     fuzzyFileSearchToolCallId,
 } from "./CodexToolCallMapper";
 import { stripShellPrefix } from "./CommandUtils";
+import type { ApprovalContextStore } from "./CodexApprovalContext";
 
 export { stripShellPrefix };
 
 export class CodexEventHandler {
 
     private readonly connection: acp.AgentSideConnection;
     private readonly sessionState: SessionState;
+    private readonly approvalContext: ApprovalContextStore;
     private failure: RequestError | null = null;
     private readonly activeFuzzyFileSearchSessions = new Set<string>();
 
-    constructor(connection: acp.AgentSideConnection, sessionState: SessionState) {
+    constructor(connection: acp.AgentSideConnection, sessionState: SessionState, approvalContext: ApprovalContextStore) {
         this.connection = connection;
         this.sessionState = sessionState;
+        this.approvalContext = approvalContext;
     }
 
     getFailure(): RequestError | null {
@@ -100,14 +104,16 @@ export class CodexEventHandler {
             case "item/reasoning/summaryPartAdded":
             //skipped events
             case "item/reasoning/textDelta": //for raw output
-            case "turn/diff/updated":
             case "item/commandExecution/terminalInteraction":
             case "item/fileChange/outputDelta":
             case "serverRequest/resolved":
             case "account/updated":
             case "fs/changed":
             case "mcpServer/startupStatus/updated":
                 return null;
+            case "turn/diff/updated":
+                this.approvalContext.turnDiffsByTurnId.set(notification.params.turnId, notification.params.diff);
+                return null;
             case "item/mcpToolCall/progress":
                 return this.createMcpToolProgressEvent(notification.params);
             case "account/rateLimits/updated":
@@ -189,6 +195,7 @@ export class CodexEventHandler {
     private async createItemEvent(event: ItemStartedNotification): Promise<UpdateSessionEvent | null> {
         switch (event.item.type) {
             case "fileChange":
+                this.approvalContext.fileChangesByItemId.set(event.item.id, event.item);
                 return await createFileChangeUpdate(event.item);
             case "commandExecution":
                 return await createCommandExecutionUpdate(event.item);
@@ -215,6 +222,8 @@ export class CodexEventHandler {
     private async completeItemEvent(event: ItemCompletedNotification): Promise<UpdateSessionEvent | null> {
         switch (event.item.type) {
             case "fileChange":
+                this.approvalContext.fileChangesByItemId.set(event.item.id, event.item);
+                return createFileChangeCompletionUpdate(event.item);
             case "dynamicToolCall":
                 return {
                     sessionUpdate: "tool_call_update",