feat(core): display declined confirmation code diff (#4440)

2 files changed, 138 insertions, 1 deletions

diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts
index 94d4f7c1..7b6a130c 100644
--- a/packages/core/src/core/coreToolScheduler.test.ts
+++ b/packages/core/src/core/coreToolScheduler.test.ts
@@ -407,3 +407,123 @@ describe('convertToFunctionResponse', () => {
     });
   });
 });
+
+describe('CoreToolScheduler edit cancellation', () => {
+  it('should preserve diff when an edit is cancelled', async () => {
+    class MockEditTool extends BaseTool<Record<string, unknown>, ToolResult> {
+      constructor() {
+        super(
+          'mockEditTool',
+          'mockEditTool',
+          'A mock edit tool',
+          Icon.Pencil,
+          {},
+        );
+      }
+
+      async shouldConfirmExecute(
+        _params: Record<string, unknown>,
+        _abortSignal: AbortSignal,
+      ): Promise<ToolCallConfirmationDetails | false> {
+        return {
+          type: 'edit',
+          title: 'Confirm Edit',
+          fileName: 'test.txt',
+          fileDiff:
+            '--- test.txt\n+++ test.txt\n@@ -1,1 +1,1 @@\n-old content\n+new content',
+          originalContent: 'old content',
+          newContent: 'new content',
+          onConfirm: async () => {},
+        };
+      }
+
+      async execute(
+        _params: Record<string, unknown>,
+        _abortSignal: AbortSignal,
+      ): Promise<ToolResult> {
+        return {
+          llmContent: 'Edited successfully',
+          returnDisplay: 'Edited successfully',
+        };
+      }
+    }
+
+    const mockEditTool = new MockEditTool();
+    const toolRegistry = {
+      getTool: () => mockEditTool,
+      getFunctionDeclarations: () => [],
+      tools: new Map(),
+      discovery: {} as any,
+      registerTool: () => {},
+      getToolByName: () => mockEditTool,
+      getToolByDisplayName: () => mockEditTool,
+      getTools: () => [],
+      discoverTools: async () => {},
+      getAllTools: () => [],
+      getToolsByServer: () => [],
+    };
+
+    const onAllToolCallsComplete = vi.fn();
+    const onToolCallsUpdate = vi.fn();
+
+    const mockConfig = {
+      getSessionId: () => 'test-session-id',
+      getUsageStatisticsEnabled: () => true,
+      getDebugMode: () => false,
+    } as unknown as Config;
+
+    const scheduler = new CoreToolScheduler({
+      config: mockConfig,
+      toolRegistry: Promise.resolve(toolRegistry as any),
+      onAllToolCallsComplete,
+      onToolCallsUpdate,
+      getPreferredEditor: () => 'vscode',
+    });
+
+    const abortController = new AbortController();
+    const request = {
+      callId: '1',
+      name: 'mockEditTool',
+      args: {},
+      isClientInitiated: false,
+      prompt_id: 'prompt-id-1',
+    };
+
+    await scheduler.schedule([request], abortController.signal);
+
+    // Wait for the tool to reach awaiting_approval state
+    const awaitingCall = onToolCallsUpdate.mock.calls.find(
+      (call) => call[0][0].status === 'awaiting_approval',
+    )?.[0][0];
+
+    expect(awaitingCall).toBeDefined();
+
+    // Cancel the edit
+    const confirmationDetails = await mockEditTool.shouldConfirmExecute(
+      {},
+      abortController.signal,
+    );
+    if (confirmationDetails) {
+      await scheduler.handleConfirmationResponse(
+        '1',
+        confirmationDetails.onConfirm,
+        ToolConfirmationOutcome.Cancel,
+        abortController.signal,
+      );
+    }
+
+    expect(onAllToolCallsComplete).toHaveBeenCalled();
+    const completedCalls = onAllToolCallsComplete.mock
+      .calls[0][0] as ToolCall[];
+
+    expect(completedCalls[0].status).toBe('cancelled');
+
+    // Check that the diff is preserved
+    const cancelledCall = completedCalls[0] as any;
+    expect(cancelledCall.response.resultDisplay).toBeDefined();
+    expect(cancelledCall.response.resultDisplay.fileDiff).toBe(
+      '--- test.txt\n+++ test.txt\n@@ -1,1 +1,1 @@\n-old content\n+new content',
+    );
+    expect(cancelledCall.response.resultDisplay.fileName).toBe('test.txt');
+  });
+});
diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts
index 8f9ec1e2..0d7d5923 100644
--- a/packages/core/src/core/coreToolScheduler.ts
+++ b/packages/core/src/core/coreToolScheduler.ts
@@ -11,6 +11,7 @@ import {
   Tool,
   ToolCallConfirmationDetails,
   ToolResult,
+  ToolResultDisplay,
   ToolRegistry,
   ApprovalMode,
   EditorType,
@@ -335,6 +336,22 @@ export class CoreToolScheduler {
           const durationMs = existingStartTime
             ? Date.now() - existingStartTime
             : undefined;
+
+          // Preserve diff for cancelled edit operations
+          let resultDisplay: ToolResultDisplay | undefined = undefined;
+          if (currentCall.status === 'awaiting_approval') {
+            const waitingCall = currentCall as WaitingToolCall;
+            if (waitingCall.confirmationDetails.type === 'edit') {
+              resultDisplay = {
+                fileDiff: waitingCall.confirmationDetails.fileDiff,
+                fileName: waitingCall.confirmationDetails.fileName,
+                originalContent:
+                  waitingCall.confirmationDetails.originalContent,
+                newContent: waitingCall.confirmationDetails.newContent,
+              };
+            }
+          }
+
           return {
             request: currentCall.request,
             tool: toolInstance,
@@ -350,7 +367,7 @@ export class CoreToolScheduler {
                   },
                 },
               },
-              resultDisplay: undefined,
+              resultDisplay,
               error: undefined,
             },
             durationMs,