From 16da6918cb56b04141583ad7253270afc7a77e34 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Thu, 8 Jan 2026 14:52:56 -0500 Subject: [PATCH] refactor(core): extract ToolModificationHandler from scheduler (#16118) --- .../core/src/core/coreToolScheduler.test.ts | 11 + packages/core/src/core/coreToolScheduler.ts | 121 +++------ .../core/src/scheduler/tool-modifier.test.ts | 252 ++++++++++++++++++ packages/core/src/scheduler/tool-modifier.ts | 105 ++++++++ packages/core/src/utils/editor.test.ts | 13 +- packages/core/src/utils/editor.ts | 13 +- 6 files changed, 427 insertions(+), 88 deletions(-) create mode 100644 packages/core/src/scheduler/tool-modifier.test.ts create mode 100644 packages/core/src/scheduler/tool-modifier.ts diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 22ef939a62..9cfacc2358 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -750,6 +750,17 @@ describe('CoreToolScheduler with payload', () => { ); } + // After internal update, the tool should be awaiting approval again with the NEW content. + const updatedAwaitingCall = (await waitForStatus( + onToolCallsUpdate, + 'awaiting_approval', + )) as WaitingToolCall; + + // Now confirm for real to execute. + await updatedAwaitingCall.confirmationDetails.onConfirm( + ToolConfirmationOutcome.ProceedOnce, + ); + // Wait for the tool execution to complete await vi.waitFor(() => { expect(onAllToolCallsComplete).toHaveBeenCalled(); diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 1120074248..bec4eacd53 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -19,12 +19,7 @@ import { logToolCall } from '../telemetry/loggers.js'; import { ToolErrorType } from '../tools/tool-error.js'; import { ToolCallEvent } from '../telemetry/types.js'; import { runInDevTraceSpan } from '../telemetry/trace.js'; -import type { ModifyContext } from '../tools/modifiable-tool.js'; -import { - isModifiableDeclarativeTool, - modifyWithEditor, -} from '../tools/modifiable-tool.js'; -import * as Diff from 'diff'; +import { ToolModificationHandler } from '../scheduler/tool-modifier.js'; import { getToolSuggestion } from '../utils/tool-utils.js'; import type { ToolConfirmationRequest } from '../confirmation-bus/types.js'; import { MessageBusType } from '../confirmation-bus/types.js'; @@ -124,6 +119,7 @@ export class CoreToolScheduler { private toolCallQueue: ToolCall[] = []; private completedToolCallsForBatch: CompletedToolCall[] = []; private toolExecutor: ToolExecutor; + private toolModifier: ToolModificationHandler; constructor(options: CoreToolSchedulerOptions) { this.config = options.config; @@ -132,6 +128,7 @@ export class CoreToolScheduler { this.onToolCallsUpdate = options.onToolCallsUpdate; this.getPreferredEditor = options.getPreferredEditor; this.toolExecutor = new ToolExecutor(this.config); + this.toolModifier = new ToolModificationHandler(); // Subscribe to message bus for ASK_USER policy decisions // Use a static WeakMap to ensure we only subscribe ONCE per MessageBus instance @@ -749,105 +746,61 @@ export class CoreToolScheduler { return; // `cancelAll` calls `checkAndNotifyCompletion`, so we can exit here. } else if (outcome === ToolConfirmationOutcome.ModifyWithEditor) { const waitingToolCall = toolCall as WaitingToolCall; - if (isModifiableDeclarativeTool(waitingToolCall.tool)) { - const modifyContext = waitingToolCall.tool.getModifyContext(signal); - const editorType = this.getPreferredEditor(); - if (!editorType) { - return; - } + const editorType = this.getPreferredEditor(); + if (!editorType) { + return; + } + + this.setStatusInternal(callId, 'awaiting_approval', signal, { + ...waitingToolCall.confirmationDetails, + isModifying: true, + } as ToolCallConfirmationDetails); + + const result = await this.toolModifier.handleModifyWithEditor( + waitingToolCall, + editorType, + signal, + ); + + // Restore status (isModifying: false) and update diff if result exists + if (result) { + this.setArgsInternal(callId, result.updatedParams); this.setStatusInternal(callId, 'awaiting_approval', signal, { ...waitingToolCall.confirmationDetails, - isModifying: true, + fileDiff: result.updatedDiff, + isModifying: false, } as ToolCallConfirmationDetails); - - const contentOverrides = - waitingToolCall.confirmationDetails.type === 'edit' - ? { - currentContent: - waitingToolCall.confirmationDetails.originalContent, - proposedContent: waitingToolCall.confirmationDetails.newContent, - } - : undefined; - - const { updatedParams, updatedDiff } = await modifyWithEditor< - typeof waitingToolCall.request.args - >( - waitingToolCall.request.args, - modifyContext as ModifyContext, - editorType, - signal, - contentOverrides, - ); - this.setArgsInternal(callId, updatedParams); + } else { this.setStatusInternal(callId, 'awaiting_approval', signal, { ...waitingToolCall.confirmationDetails, - fileDiff: updatedDiff, isModifying: false, } as ToolCallConfirmationDetails); } } else { - // If the client provided new content, apply it before scheduling. + // If the client provided new content, apply it and wait for + // re-confirmation. if (payload?.newContent && toolCall) { - await this._applyInlineModify( + const result = await this.toolModifier.applyInlineModify( toolCall as WaitingToolCall, payload, signal, ); + if (result) { + this.setArgsInternal(callId, result.updatedParams); + this.setStatusInternal(callId, 'awaiting_approval', signal, { + ...(toolCall as WaitingToolCall).confirmationDetails, + fileDiff: result.updatedDiff, + } as ToolCallConfirmationDetails); + // After an inline modification, wait for another user confirmation. + return; + } } this.setStatusInternal(callId, 'scheduled', signal); } await this.attemptExecutionOfScheduledCalls(signal); } - /** - * Applies user-provided content changes to a tool call that is awaiting confirmation. - * This method updates the tool's arguments and refreshes the confirmation prompt with a new diff - * before the tool is scheduled for execution. - * @private - */ - private async _applyInlineModify( - toolCall: WaitingToolCall, - payload: ToolConfirmationPayload, - signal: AbortSignal, - ): Promise { - if ( - toolCall.confirmationDetails.type !== 'edit' || - !isModifiableDeclarativeTool(toolCall.tool) - ) { - return; - } - - const modifyContext = toolCall.tool.getModifyContext(signal); - const currentContent = await modifyContext.getCurrentContent( - toolCall.request.args, - ); - - const updatedParams = modifyContext.createUpdatedParams( - currentContent, - payload.newContent, - toolCall.request.args, - ); - const updatedDiff = Diff.createPatch( - modifyContext.getFilePath(toolCall.request.args), - currentContent, - payload.newContent, - 'Current', - 'Proposed', - ); - - this.setArgsInternal(toolCall.request.callId, updatedParams); - this.setStatusInternal( - toolCall.request.callId, - 'awaiting_approval', - signal, - { - ...toolCall.confirmationDetails, - fileDiff: updatedDiff, - }, - ); - } - private async attemptExecutionOfScheduledCalls( signal: AbortSignal, ): Promise { diff --git a/packages/core/src/scheduler/tool-modifier.test.ts b/packages/core/src/scheduler/tool-modifier.test.ts new file mode 100644 index 0000000000..8107e4c901 --- /dev/null +++ b/packages/core/src/scheduler/tool-modifier.test.ts @@ -0,0 +1,252 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { ToolModificationHandler } from './tool-modifier.js'; +import type { WaitingToolCall, ToolCallRequestInfo } from './types.js'; +import * as modifiableToolModule from '../tools/modifiable-tool.js'; +import * as Diff from 'diff'; +import { MockModifiableTool, MockTool } from '../test-utils/mock-tool.js'; +import type { + ToolResult, + ToolInvocation, + ToolConfirmationPayload, +} from '../tools/tools.js'; +import type { ModifyContext } from '../tools/modifiable-tool.js'; +import type { Mock } from 'vitest'; + +// Mock the modules that export functions we need to control +vi.mock('diff', () => ({ + createPatch: vi.fn(), + diffLines: vi.fn(), +})); + +vi.mock('../tools/modifiable-tool.js', () => ({ + isModifiableDeclarativeTool: vi.fn(), + modifyWithEditor: vi.fn(), +})); + +type MockModifyContext = { + [K in keyof ModifyContext>]: Mock; +}; + +function createMockWaitingToolCall( + overrides: Partial = {}, +): WaitingToolCall { + return { + status: 'awaiting_approval', + request: { + callId: 'test-call-id', + name: 'test-tool', + args: {}, + isClientInitiated: false, + prompt_id: 'test-prompt-id', + } as ToolCallRequestInfo, + tool: new MockTool({ name: 'test-tool' }), + invocation: {} as ToolInvocation, ToolResult>, // We generally don't check invocation details in these tests + confirmationDetails: { + type: 'edit', + title: 'Test Confirmation', + fileName: 'test.txt', + filePath: '/path/to/test.txt', + fileDiff: 'diff', + originalContent: 'original', + newContent: 'new', + onConfirm: async () => {}, + }, + ...overrides, + }; +} + +describe('ToolModificationHandler', () => { + let handler: ToolModificationHandler; + let mockModifiableTool: MockModifiableTool; + let mockPlainTool: MockTool; + let mockModifyContext: MockModifyContext; + + beforeEach(() => { + vi.clearAllMocks(); + handler = new ToolModificationHandler(); + mockModifiableTool = new MockModifiableTool(); + mockPlainTool = new MockTool({ name: 'plainTool' }); + + mockModifyContext = { + getCurrentContent: vi.fn(), + getFilePath: vi.fn(), + createUpdatedParams: vi.fn(), + getProposedContent: vi.fn(), + }; + + vi.spyOn(mockModifiableTool, 'getModifyContext').mockReturnValue( + mockModifyContext as unknown as ModifyContext>, + ); + }); + + describe('handleModifyWithEditor', () => { + it('should return undefined if tool is not modifiable', async () => { + vi.mocked( + modifiableToolModule.isModifiableDeclarativeTool, + ).mockReturnValue(false); + + const mockWaitingToolCall = createMockWaitingToolCall({ + tool: mockPlainTool, + request: { + callId: 'call-1', + name: 'plainTool', + args: { path: 'foo.txt' }, + isClientInitiated: false, + prompt_id: 'p1', + }, + }); + + const result = await handler.handleModifyWithEditor( + mockWaitingToolCall, + 'vscode', + new AbortController().signal, + ); + + expect(result).toBeUndefined(); + }); + + it('should call modifyWithEditor and return updated params', async () => { + vi.mocked( + modifiableToolModule.isModifiableDeclarativeTool, + ).mockReturnValue(true); + + vi.mocked(modifiableToolModule.modifyWithEditor).mockResolvedValue({ + updatedParams: { path: 'foo.txt', content: 'new' }, + updatedDiff: 'diff', + }); + + const mockWaitingToolCall = createMockWaitingToolCall({ + tool: mockModifiableTool, + request: { + callId: 'call-1', + name: 'mockModifiableTool', + args: { path: 'foo.txt' }, + isClientInitiated: false, + prompt_id: 'p1', + }, + confirmationDetails: { + type: 'edit', + title: 'Confirm', + fileName: 'foo.txt', + filePath: 'foo.txt', + fileDiff: 'diff', + originalContent: 'old', + newContent: 'new', + onConfirm: async () => {}, + }, + }); + + const result = await handler.handleModifyWithEditor( + mockWaitingToolCall, + 'vscode', + new AbortController().signal, + ); + + expect(modifiableToolModule.modifyWithEditor).toHaveBeenCalledWith( + mockWaitingToolCall.request.args, + mockModifyContext, + 'vscode', + expect.any(AbortSignal), + { currentContent: 'old', proposedContent: 'new' }, + ); + + expect(result).toEqual({ + updatedParams: { path: 'foo.txt', content: 'new' }, + updatedDiff: 'diff', + }); + }); + }); + + describe('applyInlineModify', () => { + it('should return undefined if tool is not modifiable', async () => { + vi.mocked( + modifiableToolModule.isModifiableDeclarativeTool, + ).mockReturnValue(false); + + const mockWaitingToolCall = createMockWaitingToolCall({ + tool: mockPlainTool, + }); + + const result = await handler.applyInlineModify( + mockWaitingToolCall, + { newContent: 'foo' }, + new AbortController().signal, + ); + + expect(result).toBeUndefined(); + }); + + it('should return undefined if payload has no new content', async () => { + vi.mocked( + modifiableToolModule.isModifiableDeclarativeTool, + ).mockReturnValue(true); + + const mockWaitingToolCall = createMockWaitingToolCall({ + tool: mockModifiableTool, + }); + + const result = await handler.applyInlineModify( + mockWaitingToolCall, + { newContent: undefined } as unknown as ToolConfirmationPayload, + new AbortController().signal, + ); + + expect(result).toBeUndefined(); + }); + + it('should calculate diff and return updated params', async () => { + vi.mocked( + modifiableToolModule.isModifiableDeclarativeTool, + ).mockReturnValue(true); + (Diff.createPatch as unknown as Mock).mockReturnValue('mock-diff'); + + mockModifyContext.getCurrentContent.mockResolvedValue('old content'); + mockModifyContext.getFilePath.mockReturnValue('test.txt'); + mockModifyContext.createUpdatedParams.mockReturnValue({ + content: 'new content', + }); + + const mockWaitingToolCall = createMockWaitingToolCall({ + tool: mockModifiableTool, + request: { + callId: 'call-1', + name: 'mockModifiableTool', + args: { content: 'original' }, + isClientInitiated: false, + prompt_id: 'p1', + }, + }); + + const result = await handler.applyInlineModify( + mockWaitingToolCall, + { newContent: 'new content' }, + new AbortController().signal, + ); + + expect(mockModifyContext.getCurrentContent).toHaveBeenCalled(); + expect(mockModifyContext.createUpdatedParams).toHaveBeenCalledWith( + 'old content', + 'new content', + { content: 'original' }, + ); + expect(Diff.createPatch).toHaveBeenCalledWith( + 'test.txt', + 'old content', + 'new content', + 'Current', + 'Proposed', + ); + + expect(result).toEqual({ + updatedParams: { content: 'new content' }, + updatedDiff: 'mock-diff', + }); + }); + }); +}); diff --git a/packages/core/src/scheduler/tool-modifier.ts b/packages/core/src/scheduler/tool-modifier.ts new file mode 100644 index 0000000000..c7d9c93c67 --- /dev/null +++ b/packages/core/src/scheduler/tool-modifier.ts @@ -0,0 +1,105 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as Diff from 'diff'; +import type { EditorType } from '../utils/editor.js'; +import { + isModifiableDeclarativeTool, + modifyWithEditor, + type ModifyContext, +} from '../tools/modifiable-tool.js'; +import type { ToolConfirmationPayload } from '../tools/tools.js'; +import type { WaitingToolCall } from './types.js'; + +export interface ModificationResult { + updatedParams: Record; + updatedDiff?: string; +} + +export class ToolModificationHandler { + /** + * Handles the "Modify with Editor" flow where an external editor is launched + * to modify the tool's parameters. + */ + async handleModifyWithEditor( + toolCall: WaitingToolCall, + editorType: EditorType, + signal: AbortSignal, + ): Promise { + if (!isModifiableDeclarativeTool(toolCall.tool)) { + return undefined; + } + + const confirmationDetails = toolCall.confirmationDetails; + const modifyContext = toolCall.tool.getModifyContext(signal); + + const contentOverrides = + confirmationDetails.type === 'edit' + ? { + currentContent: confirmationDetails.originalContent, + proposedContent: confirmationDetails.newContent, + } + : undefined; + + const { updatedParams, updatedDiff } = await modifyWithEditor< + typeof toolCall.request.args + >( + toolCall.request.args, + modifyContext as ModifyContext, + editorType, + signal, + contentOverrides, + ); + + return { + updatedParams, + updatedDiff, + }; + } + + /** + * Applies user-provided inline content updates (e.g. from the chat UI). + */ + async applyInlineModify( + toolCall: WaitingToolCall, + payload: ToolConfirmationPayload, + signal: AbortSignal, + ): Promise { + if ( + toolCall.confirmationDetails.type !== 'edit' || + !payload.newContent || + !isModifiableDeclarativeTool(toolCall.tool) + ) { + return undefined; + } + + const modifyContext = toolCall.tool.getModifyContext( + signal, + ) as ModifyContext; + const currentContent = await modifyContext.getCurrentContent( + toolCall.request.args, + ); + + const updatedParams = modifyContext.createUpdatedParams( + currentContent, + payload.newContent, + toolCall.request.args, + ); + + const updatedDiff = Diff.createPatch( + modifyContext.getFilePath(toolCall.request.args), + currentContent, + payload.newContent, + 'Current', + 'Proposed', + ); + + return { + updatedParams, + updatedDiff, + }; + } +} diff --git a/packages/core/src/utils/editor.test.ts b/packages/core/src/utils/editor.test.ts index 82b886f366..78035c4cc9 100644 --- a/packages/core/src/utils/editor.test.ts +++ b/packages/core/src/utils/editor.test.ts @@ -311,11 +311,18 @@ describe('editor utils', () => { }); } - it('should return the correct command for emacs', () => { - const command = getDiffCommand('old.txt', 'new.txt', 'emacs'); + it('should return the correct command for emacs with escaped paths', () => { + const command = getDiffCommand( + 'old file "quote".txt', + 'new file \\back\\slash.txt', + 'emacs', + ); expect(command).toEqual({ command: 'emacs', - args: ['--eval', '(ediff "old.txt" "new.txt")'], + args: [ + '--eval', + '(ediff "old file \\"quote\\".txt" "new file \\\\back\\\\slash.txt")', + ], }); }); diff --git a/packages/core/src/utils/editor.ts b/packages/core/src/utils/editor.ts index b71a0b23eb..742d1157fb 100644 --- a/packages/core/src/utils/editor.ts +++ b/packages/core/src/utils/editor.ts @@ -60,6 +60,14 @@ function isValidEditorType(editor: string): editor is EditorType { return EDITORS_SET.has(editor); } +/** + * Escapes a string for use in an Emacs Lisp string literal. + * Wraps in double quotes and escapes backslashes and double quotes. + */ +function escapeELispString(str: string): string { + return `"${str.replace(/\\/g, '\\\\').replace(/"/g, '\\"')}"`; +} + interface DiffCommand { command: string; args: string[]; @@ -182,7 +190,10 @@ export function getDiffCommand( case 'emacs': return { command: 'emacs', - args: ['--eval', `(ediff "${oldPath}" "${newPath}")`], + args: [ + '--eval', + `(ediff ${escapeELispString(oldPath)} ${escapeELispString(newPath)})`, + ], }; case 'hx': return {