refactor(core): extract ToolModificationHandler from scheduler (#16118)

This commit is contained in:
Abhi
2026-01-08 14:52:56 -05:00
committed by GitHub
parent fbfad06307
commit 16da6918cb
6 changed files with 427 additions and 88 deletions

View File

@@ -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();

View File

@@ -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<typeof waitingToolCall.request.args>,
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<void> {
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<void> {

View File

@@ -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<Record<string, unknown>>]: Mock;
};
function createMockWaitingToolCall(
overrides: Partial<WaitingToolCall> = {},
): 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<Record<string, unknown>, 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<Record<string, unknown>>,
);
});
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',
});
});
});
});

View File

@@ -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<string, unknown>;
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<ModificationResult | undefined> {
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<typeof toolCall.request.args>,
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<ModificationResult | undefined> {
if (
toolCall.confirmationDetails.type !== 'edit' ||
!payload.newContent ||
!isModifiableDeclarativeTool(toolCall.tool)
) {
return undefined;
}
const modifyContext = toolCall.tool.getModifyContext(
signal,
) as ModifyContext<typeof toolCall.request.args>;
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,
};
}
}

View File

@@ -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")',
],
});
});

View File

@@ -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 {