From 9211905ff1a3174bc0fbae1260162fad846c563d Mon Sep 17 00:00:00 2001 From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com> Date: Sat, 5 Jul 2025 23:19:41 +0000 Subject: [PATCH] feat: Handle inline content modification in tool scheduler (#2883) --- .../core/src/core/coreToolScheduler.test.ts | 114 ++++++++++++++++++ packages/core/src/core/coreToolScheduler.ts | 60 ++++++++- packages/core/src/tools/tools.ts | 11 +- 3 files changed, 183 insertions(+), 2 deletions(-) diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index ec01a3af..d6030d6f 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -16,11 +16,14 @@ import { BaseTool, ToolCallConfirmationDetails, ToolConfirmationOutcome, + ToolConfirmationPayload, ToolResult, Config, } from '../index.js'; import { Part, PartListUnion } from '@google/genai'; +import { ModifiableTool, ModifyContext } from '../tools/modifiable-tool.js'; + class MockTool extends BaseTool, ToolResult> { shouldConfirm = false; executeFn = vi.fn(); @@ -54,6 +57,47 @@ class MockTool extends BaseTool, ToolResult> { } } +class MockModifiableTool + extends MockTool + implements ModifiableTool> +{ + constructor(name = 'mockModifiableTool') { + super(name); + this.shouldConfirm = true; + } + + getModifyContext( + _abortSignal: AbortSignal, + ): ModifyContext> { + return { + getFilePath: () => 'test.txt', + getCurrentContent: async () => 'old content', + getProposedContent: async () => 'new content', + createUpdatedParams: ( + _oldContent: string, + modifiedProposedContent: string, + _originalParams: Record, + ) => ({ newContent: modifiedProposedContent }), + }; + } + + async shouldConfirmExecute( + _params: Record, + _abortSignal: AbortSignal, + ): Promise { + if (this.shouldConfirm) { + return { + type: 'edit', + title: 'Confirm Mock Tool', + fileName: 'test.txt', + fileDiff: 'diff', + onConfirm: async () => {}, + }; + } + return false; + } +} + describe('CoreToolScheduler', () => { it('should cancel a tool call if the signal is aborted before confirmation', async () => { const mockTool = new MockTool(); @@ -122,6 +166,76 @@ describe('CoreToolScheduler', () => { }); }); +describe('CoreToolScheduler with payload', () => { + it('should update args and diff and execute tool when payload is provided', async () => { + const mockTool = new MockModifiableTool(); + const toolRegistry = { + getTool: () => mockTool, + getFunctionDeclarations: () => [], + tools: new Map(), + discovery: {} as any, + registerTool: () => {}, + getToolByName: () => mockTool, + getToolByDisplayName: () => mockTool, + 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: 'mockModifiableTool', + args: {}, + isClientInitiated: false, + }; + + await scheduler.schedule([request], abortController.signal); + + const confirmationDetails = await mockTool.shouldConfirmExecute( + {}, + abortController.signal, + ); + + if (confirmationDetails) { + const payload: ToolConfirmationPayload = { newContent: 'final version' }; + await scheduler.handleConfirmationResponse( + '1', + confirmationDetails.onConfirm, + ToolConfirmationOutcome.ProceedOnce, + abortController.signal, + payload, + ); + } + + expect(onAllToolCallsComplete).toHaveBeenCalled(); + const completedCalls = onAllToolCallsComplete.mock + .calls[0][0] as ToolCall[]; + expect(completedCalls[0].status).toBe('success'); + expect(mockTool.executeFn).toHaveBeenCalledWith({ + newContent: 'final version', + }); + }); +}); + describe('convertToFunctionResponse', () => { const toolName = 'testTool'; const callId = 'call1'; diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 81fce558..53596941 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -17,6 +17,7 @@ import { Config, logToolCall, ToolCallEvent, + ToolConfirmationPayload, } from '../index.js'; import { Part, PartListUnion } from '@google/genai'; import { getResponseTextFromParts } from '../utils/generateContentResponseUtilities.js'; @@ -25,6 +26,7 @@ import { ModifyContext, modifyWithEditor, } from '../tools/modifiable-tool.js'; +import * as Diff from 'diff'; export type ValidatingToolCall = { status: 'validating'; @@ -455,12 +457,16 @@ export class CoreToolScheduler { const originalOnConfirm = confirmationDetails.onConfirm; const wrappedConfirmationDetails: ToolCallConfirmationDetails = { ...confirmationDetails, - onConfirm: (outcome: ToolConfirmationOutcome) => + onConfirm: ( + outcome: ToolConfirmationOutcome, + payload?: ToolConfirmationPayload, + ) => this.handleConfirmationResponse( reqInfo.callId, originalOnConfirm, outcome, signal, + payload, ), }; this.setStatusInternal( @@ -492,6 +498,7 @@ export class CoreToolScheduler { originalOnConfirm: (outcome: ToolConfirmationOutcome) => Promise, outcome: ToolConfirmationOutcome, signal: AbortSignal, + payload?: ToolConfirmationPayload, ): Promise { const toolCall = this.toolCalls.find( (c) => c.request.callId === callId && c.status === 'awaiting_approval', @@ -545,11 +552,62 @@ export class CoreToolScheduler { } as ToolCallConfirmationDetails); } } else { + // If the client provided new content, apply it before scheduling. + if (payload?.newContent && toolCall) { + await this._applyInlineModify( + toolCall as WaitingToolCall, + payload, + signal, + ); + } this.setStatusInternal(callId, 'scheduled'); } 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' || + !isModifiableTool(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', { + ...toolCall.confirmationDetails, + fileDiff: updatedDiff, + }); + } + private attemptExecutionOfScheduledCalls(signal: AbortSignal): void { const allCallsFinalOrScheduled = this.toolCalls.every( (call) => diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts index e9fc042b..cff969f8 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -199,12 +199,21 @@ export interface FileDiff { export interface ToolEditConfirmationDetails { type: 'edit'; title: string; - onConfirm: (outcome: ToolConfirmationOutcome) => Promise; + onConfirm: ( + outcome: ToolConfirmationOutcome, + payload?: ToolConfirmationPayload, + ) => Promise; fileName: string; fileDiff: string; isModifying?: boolean; } +export interface ToolConfirmationPayload { + // used to override `modifiedProposedContent` for modifiable tools in the + // inline modify flow + newContent: string; +} + export interface ToolExecuteConfirmationDetails { type: 'exec'; title: string;