diff options
Diffstat (limited to 'packages/core/src/tools')
| -rw-r--r-- | packages/core/src/tools/edit.ts | 1 | ||||
| -rw-r--r-- | packages/core/src/tools/memoryTool.ts | 1 | ||||
| -rw-r--r-- | packages/core/src/tools/tools.ts | 3 | ||||
| -rw-r--r-- | packages/core/src/tools/write-file.test.ts | 29 | ||||
| -rw-r--r-- | packages/core/src/tools/write-file.ts | 18 |
5 files changed, 49 insertions, 3 deletions
diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index 25da2292..0d129e42 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -332,6 +332,7 @@ Expectation for required parameters: type: 'edit', title: `Confirm Edit: ${shortenPath(makeRelative(params.file_path, this.config.getTargetDir()))}`, fileName, + filePath: params.file_path, fileDiff, originalContent: editData.currentContent, newContent: editData.newContent, diff --git a/packages/core/src/tools/memoryTool.ts b/packages/core/src/tools/memoryTool.ts index 96509f79..847ea5cf 100644 --- a/packages/core/src/tools/memoryTool.ts +++ b/packages/core/src/tools/memoryTool.ts @@ -220,6 +220,7 @@ export class MemoryTool type: 'edit', title: `Confirm Memory Save: ${tildeifyPath(memoryFilePath)}`, fileName: memoryFilePath, + filePath: memoryFilePath, fileDiff, originalContent: currentContent, newContent, diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts index 5d9d9253..3404093f 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -6,6 +6,7 @@ import { FunctionDeclaration, PartListUnion, Schema } from '@google/genai'; import { ToolErrorType } from './tool-error.js'; +import { DiffUpdateResult } from '../ide/ideContext.js'; /** * Interface representing the base Tool functionality @@ -330,10 +331,12 @@ export interface ToolEditConfirmationDetails { payload?: ToolConfirmationPayload, ) => Promise<void>; fileName: string; + filePath: string; fileDiff: string; originalContent: string | null; newContent: string; isModifying?: boolean; + ideConfirmation?: Promise<DiffUpdateResult>; } export interface ToolConfirmationPayload { diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index fe662a02..563579bb 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -55,6 +55,9 @@ const mockConfigInternal = { getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT), setApprovalMode: vi.fn(), getGeminiClient: vi.fn(), // Initialize as a plain mock function + getIdeClient: vi.fn(), + getIdeMode: vi.fn(() => false), + getIdeModeFeature: vi.fn(() => false), getWorkspaceContext: () => createMockWorkspaceContext(rootDir), getApiKey: () => 'test-key', getModel: () => 'test-model', @@ -110,6 +113,14 @@ describe('WriteFileTool', () => { mockConfigInternal.getGeminiClient.mockReturnValue( mockGeminiClientInstance, ); + mockConfigInternal.getIdeClient.mockReturnValue({ + openDiff: vi.fn(), + closeDiff: vi.fn(), + getIdeContext: vi.fn(), + subscribeToIdeContext: vi.fn(), + isCodeTrackerEnabled: vi.fn(), + getTrackedCode: vi.fn(), + }); tool = new WriteFileTool(mockConfig); @@ -500,7 +511,11 @@ describe('WriteFileTool', () => { params, abortSignal, ); - if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) { + if ( + typeof confirmDetails === 'object' && + 'onConfirm' in confirmDetails && + confirmDetails.onConfirm + ) { await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce); } @@ -554,7 +569,11 @@ describe('WriteFileTool', () => { params, abortSignal, ); - if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) { + if ( + typeof confirmDetails === 'object' && + 'onConfirm' in confirmDetails && + confirmDetails.onConfirm + ) { await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce); } @@ -595,7 +614,11 @@ describe('WriteFileTool', () => { params, abortSignal, ); - if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) { + if ( + typeof confirmDetails === 'object' && + 'onConfirm' in confirmDetails && + confirmDetails.onConfirm + ) { await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce); } diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index 1cb1a917..32ecc068 100644 --- a/packages/core/src/tools/write-file.ts +++ b/packages/core/src/tools/write-file.ts @@ -32,6 +32,7 @@ import { recordFileOperationMetric, FileOperation, } from '../telemetry/metrics.js'; +import { IDEConnectionStatus } from '../ide/ide-client.js'; /** * Parameters for the WriteFile tool @@ -184,10 +185,19 @@ export class WriteFileTool DEFAULT_DIFF_OPTIONS, ); + const ideClient = this.config.getIdeClient(); + const ideConfirmation = + this.config.getIdeModeFeature() && + this.config.getIdeMode() && + ideClient.getConnectionStatus().status === IDEConnectionStatus.Connected + ? ideClient.openDiff(params.file_path, correctedContent) + : undefined; + const confirmationDetails: ToolEditConfirmationDetails = { type: 'edit', title: `Confirm Write: ${shortenPath(relativePath)}`, fileName, + filePath: params.file_path, fileDiff, originalContent, newContent: correctedContent, @@ -195,7 +205,15 @@ export class WriteFileTool if (outcome === ToolConfirmationOutcome.ProceedAlways) { this.config.setApprovalMode(ApprovalMode.AUTO_EDIT); } + + if (ideConfirmation) { + const result = await ideConfirmation; + if (result.status === 'accepted' && result.content) { + params.content = result.content; + } + } }, + ideConfirmation, }; return confirmationDetails; } |
