[Refactor] Decouple Tool-Specific Cancellation Logic from SchedulerStateManager

packages/core/src/confirmation-bus/types.ts

@@ -9,6 +9,7 @@ import type {
   ToolConfirmationOutcome,
   ToolConfirmationPayload,
   DiffStat,
+  ToolResultDisplay,
 } from '../tools/tools.js';
 import type { ToolCall } from '../scheduler/types.js';
 import type { SandboxPermissions } from '../services/sandboxManager.js';
@@ -86,13 +87,15 @@ export type SerializableConfirmationDetails =
       rootCommand: string;
       additionalPermissions: SandboxPermissions;
       systemMessage?: string;
+      cancelResult?: ToolResultDisplay;
     }
   | {
       type: 'info';
       title: string;
       systemMessage?: string;
       prompt: string;
       urls?: string[];
+      cancelResult?: ToolResultDisplay;
     }
   | {
       type: 'edit';
@@ -105,6 +108,7 @@ export type SerializableConfirmationDetails =
       newContent: string;
       isModifying?: boolean;
       diffStat?: DiffStat;
+      cancelResult?: ToolResultDisplay;
     }
   | {
       type: 'exec';
@@ -114,6 +118,7 @@ export type SerializableConfirmationDetails =
       rootCommand: string;
       rootCommands: string[];
       commands?: string[];
+      cancelResult?: ToolResultDisplay;
     }
   | {
       type: 'mcp';
@@ -125,18 +130,21 @@ export type SerializableConfirmationDetails =
       toolArgs?: Record<string, unknown>;
       toolDescription?: string;
       toolParameterSchema?: unknown;
+      cancelResult?: ToolResultDisplay;
     }
   | {
       type: 'ask_user';
       title: string;
       systemMessage?: string;
       questions: Question[];
+      cancelResult?: ToolResultDisplay;
     }
   | {
       type: 'exit_plan_mode';
       title: string;
       systemMessage?: string;
       planPath: string;
+      cancelResult?: ToolResultDisplay;
     };
 
 export interface UpdatePolicy {

packages/core/src/scheduler/state-manager.test.ts

@@ -360,20 +360,40 @@ describe('SchedulerStateManager', () => {
       expect(active.confirmationDetails).toEqual(details);
     });
 
-    it('should preserve diff and derive stats when cancelling an edit tool call', () => {
+    it('should preserve diff when cancelling an edit tool call with cancelResult', () => {
       const call = createValidatingCall();
       stateManager.enqueue([call]);
       stateManager.dequeue();
 
+      const diffContent = '@@ -1,1 +1,1 @@\n-old line\n+new line';
+      const cancelResult: FileDiff = {
+        fileDiff: diffContent,
+        fileName: 'test.txt',
+        filePath: '/path/to/test.txt',
+        originalContent: 'old line',
+        newContent: 'new line',
+        diffStat: {
+          model_added_lines: 1,
+          model_removed_lines: 1,
+          model_added_chars: 8,
+          model_removed_chars: 8,
+          user_added_lines: 0,
+          user_removed_lines: 0,
+          user_added_chars: 0,
+          user_removed_chars: 0,
+        },
+      };
+
       const details = {
         type: 'edit' as const,
         title: 'Edit',
         fileName: 'test.txt',
         filePath: '/path/to/test.txt',
-        fileDiff: '@@ -1,1 +1,1 @@\n-old line\n+new line',
+        fileDiff: diffContent,
         originalContent: 'old line',
         newContent: 'new line',
         onConfirm: vi.fn(),
+        cancelResult,
       };
 
       stateManager.updateStatus(
@@ -391,7 +411,7 @@ describe('SchedulerStateManager', () => {
       const completed = stateManager.completedBatch[0] as CancelledToolCall;
       expect(completed.status).toBe(CoreToolCallStatus.Cancelled);
       const result = completed.response.resultDisplay as FileDiff;
-      expect(result.fileDiff).toBe(details.fileDiff);
+      expect(result.fileDiff).toBe(diffContent);
       expect(result.diffStat).toEqual(
         expect.objectContaining({
           model_added_lines: 1,
@@ -400,6 +420,38 @@ describe('SchedulerStateManager', () => {
       );
     });
 
+    it('should use cancelResult from confirmation details for any tool type', () => {
+      const call = createValidatingCall();
+      stateManager.enqueue([call]);
+      stateManager.dequeue();
+
+      const cancelResult = 'Operation was cancelled before completion';
+
+      const details = {
+        type: 'info' as const,
+        title: 'Info',
+        prompt: 'Some info prompt',
+        onConfirm: vi.fn(),
+        cancelResult,
+      };
+
+      stateManager.updateStatus(
+        call.request.callId,
+        CoreToolCallStatus.AwaitingApproval,
+        details,
+      );
+      stateManager.updateStatus(
+        call.request.callId,
+        CoreToolCallStatus.Cancelled,
+        'User cancelled',
+      );
+      stateManager.finalizeCall(call.request.callId);
+
+      const completed = stateManager.completedBatch[0] as CancelledToolCall;
+      expect(completed.status).toBe(CoreToolCallStatus.Cancelled);
+      expect(completed.response.resultDisplay).toBe(cancelResult);
+    });
+
     it('should ignore status updates for non-existent callIds', () => {
       stateManager.updateStatus('unknown', CoreToolCallStatus.Scheduled);
       expect(onUpdate).not.toHaveBeenCalled();

packages/core/src/scheduler/state-manager.ts

@@ -32,7 +32,6 @@ import {
   type SerializableConfirmationDetails,
 } from '../confirmation-bus/types.js';
 import { isToolCallResponseInfo } from '../utils/tool-utils.js';
-import { getDiffStatFromPatch } from '../tools/diffOptions.js';
 
 /**
  * Handler for terminal tool calls.
@@ -455,28 +454,12 @@ export class SchedulerStateManager {
     this.validateHasToolAndInvocation(call, CoreToolCallStatus.Cancelled);
     const startTime = 'startTime' in call ? call.startTime : undefined;
 
-    // TODO: Refactor this tool-specific logic into the confirmation details payload.
-    // See: https://github.com/google-gemini/gemini-cli/issues/16716
+    // Use the tool-provided cancelResult if available
     let resultDisplay: ToolResultDisplay | undefined = undefined;
     if (this.isWaitingToolCall(call)) {
       const details = call.confirmationDetails;
-      if (
-        details.type === 'edit' &&
-        'fileDiff' in details &&
-        'fileName' in details &&
-        'filePath' in details &&
-        'originalContent' in details &&
-        'newContent' in details
-      ) {
-        resultDisplay = {
-          fileDiff: details.fileDiff,
-          fileName: details.fileName,
-          filePath: details.filePath,
-          originalContent: details.originalContent,
-          newContent: details.newContent,
-          // Derive stats from the patch if they aren't already present
-          diffStat: details.diffStat ?? getDiffStatFromPatch(details.fileDiff),
-        };
+      if ('cancelResult' in details && details.cancelResult) {
+        resultDisplay = details.cancelResult;
       }
     }
 

packages/core/src/tools/tools.ts

@@ -984,6 +984,7 @@ export interface ToolEditConfirmationDetails {
   isModifying?: boolean;
   diffStat?: DiffStat;
   ideConfirmation?: Promise<DiffUpdateResult>;
+  cancelResult?: ToolResultDisplay;
 }
 
 export interface ToolEditConfirmationPayload {

packages/core/src/tools/write-file.test.ts

@@ -643,6 +643,37 @@ describe('WriteFileTool', () => {
         await diffPromise;
         expect(diffPromiseResolved).toBe(true);
       });
+
+      it('should populate cancelResult field in confirmation details', async () => {
+        const filePath = path.join(rootDir, 'cancel_result_file.txt');
+        const originalContent = 'Original content';
+        const proposedContent = 'Proposed new content';
+        const correctedContent = 'Corrected new content';
+        fs.writeFileSync(filePath, originalContent, 'utf8');
+
+        mockEnsureCorrectFileContent.mockResolvedValue(correctedContent);
+
+        const params = { file_path: filePath, content: proposedContent };
+        const invocation = tool.build(params);
+        const confirmation = (await invocation.shouldConfirmExecute(
+          abortSignal,
+        )) as ToolEditConfirmationDetails;
+
+        // Verify that cancelResult exists and contains the expected fields
+        expect(confirmation.cancelResult).toBeDefined();
+        expect(confirmation.cancelResult).toEqual(
+          expect.objectContaining({
+            fileDiff: expect.any(String),
+            fileName: 'cancel_result_file.txt',
+            filePath,
+            originalContent,
+            newContent: correctedContent,
+          }),
+        );
+        // Verify the diff is properly generated
+        expect(confirmation.cancelResult.fileDiff).toContain(originalContent);
+        expect(confirmation.cancelResult.fileDiff).toContain(correctedContent);
+      });
     });
   });
 

packages/core/src/tools/write-file.ts

@@ -257,6 +257,13 @@ class WriteFileToolInvocation extends BaseToolInvocation<
         }
       },
       ideConfirmation,
+      cancelResult: {
+        fileDiff,
+        fileName,
+        filePath: this.resolvedPath,
+        originalContent,
+        newContent: correctedContent,
+      },
     };
     return confirmationDetails;
   }