mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-10 22:21:22 -07:00
refactor(core): extract ToolModificationHandler from scheduler (#16118)
This commit is contained in:
@@ -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();
|
||||
|
||||
@@ -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> {
|
||||
|
||||
252
packages/core/src/scheduler/tool-modifier.test.ts
Normal file
252
packages/core/src/scheduler/tool-modifier.test.ts
Normal 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',
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
105
packages/core/src/scheduler/tool-modifier.ts
Normal file
105
packages/core/src/scheduler/tool-modifier.ts
Normal 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,
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -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")',
|
||||
],
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user