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, ToolResult> { + constructor() { + super( + 'mockEditTool', + 'mockEditTool', + 'A mock edit tool', + Icon.Pencil, + {}, + ); + } + + async shouldConfirmExecute( + _params: Record, + _abortSignal: AbortSignal, + ): Promise { + 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, + _abortSignal: AbortSignal, + ): Promise { + 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,