refactor: centralize IDE diff interaction in CoreToolScheduler (#2728)

- Move openDiff/confirmation handling from edit.ts and write-file.ts into
  CoreToolScheduler.openIdeDiffIfEnabled(), called after permission hooks
- Use structuredClone in buildInvocation to prevent params mutation leaking
  to LLM history (fixes #2709 token waste)
- Use confirmationDetails as single data source for IDE diff content,
  only rely on ModifyContext.createUpdatedParams() for parameter transform
- Skip inline modify when IDE content unchanged, preserving original tool
  params for multi-edit-on-same-file scenarios (mitigates #2702)
- Remove ideConfirmation field from ToolEditConfirmationDetails
- Remove dead resolveIdeDiffForOutcome from ACP Session.ts
- Fix memory tool scope fallback in createUpdatedParams

Closes #2709
Closes #2673
This commit is contained in:
胡玮文 2026-04-09 14:20:48 +08:00 committed by GitHub
parent d22b7e61ee
commit 32e7b632b8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 374 additions and 365 deletions

View file

@ -711,27 +711,6 @@ export class Session implements SessionContext {
await this.sendUpdate(update);
}
private async resolveIdeDiffForOutcome(
confirmationDetails: ToolCallConfirmationDetails,
outcome: ToolConfirmationOutcome,
): Promise<void> {
if (
confirmationDetails.type !== 'edit' ||
!confirmationDetails.ideConfirmation
) {
return;
}
const { IdeClient } = await import('@qwen-code/qwen-code-core');
const ideClient = await IdeClient.getInstance();
const cliOutcome =
outcome === ToolConfirmationOutcome.Cancel ? 'rejected' : 'accepted';
await ideClient.resolveDiffFromCli(
confirmationDetails.filePath,
cliOutcome as 'accepted' | 'rejected',
);
}
private async runTool(
abortSignal: AbortSignal,
promptId: string,
@ -946,10 +925,6 @@ export class Session implements SessionContext {
hookResult.updatedInput as typeof invocation.params;
}
await this.resolveIdeDiffForOutcome(
confirmationDetails,
ToolConfirmationOutcome.ProceedOnce,
);
await confirmationDetails.onConfirm(
ToolConfirmationOutcome.ProceedOnce,
);
@ -1020,8 +995,6 @@ export class Session implements SessionContext {
.nativeEnum(ToolConfirmationOutcome)
.parse(output.outcome.optionId);
await this.resolveIdeDiffForOutcome(confirmationDetails, outcome);
await confirmationDetails.onConfirm(outcome, {
answers: output.answers,
});

View file

@ -42,11 +42,24 @@ import { MessageBusType } from '../confirmation-bus/types.js';
import type { HookExecutionResponse } from '../confirmation-bus/types.js';
import { type NotificationType } from '../hooks/types.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { IdeClient } from '../ide/ide-client.js';
vi.mock('fs/promises', () => ({
writeFile: vi.fn(),
}));
vi.mock('../ide/ide-client.js', () => ({
IdeClient: {
getInstance: vi.fn(),
},
}));
const mockIdeClient = {
openDiff: vi.fn(),
isDiffingEnabled: vi.fn(),
closeDiff: vi.fn(),
};
class TestApprovalTool extends BaseDeclarativeTool<{ id: string }, ToolResult> {
static readonly Name = 'testApprovalTool';
@ -3219,3 +3232,288 @@ describe('Fire hook functions integration', () => {
});
});
});
describe('CoreToolScheduler IDE interaction', () => {
function createIdeMockConfig(
overrides: {
approvalMode?: ApprovalMode;
ideMode?: boolean;
} = {},
) {
const mockModifiableTool = new MockModifiableTool();
mockModifiableTool.executeFn = vi.fn();
const mockToolRegistry = {
getTool: () => mockModifiableTool,
getFunctionDeclarations: () => [],
tools: new Map(),
discovery: {},
registerTool: () => {},
getToolByName: () => mockModifiableTool,
getToolByDisplayName: () => mockModifiableTool,
getTools: () => [],
discoverTools: async () => {},
getAllTools: () => [],
getToolsByServer: () => [],
} as unknown as ToolRegistry;
const mockConfig = {
getSessionId: () => 'test-session-id',
getUsageStatisticsEnabled: () => true,
getDebugMode: () => false,
getApprovalMode: () => overrides.approvalMode ?? ApprovalMode.DEFAULT,
getPermissionsAllow: () => [],
getContentGeneratorConfig: () => ({
model: 'test-model',
authType: 'gemini',
}),
getShellExecutionConfig: () => ({
terminalWidth: 90,
terminalHeight: 30,
}),
storage: {
getProjectTempDir: () => '/tmp',
},
getTruncateToolOutputThreshold: () =>
DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD,
getTruncateToolOutputLines: () => DEFAULT_TRUNCATE_TOOL_OUTPUT_LINES,
getToolRegistry: () => mockToolRegistry,
getUseModelRouter: () => false,
getGeminiClient: () => null,
isInteractive: () => true,
getIdeMode: () => overrides.ideMode ?? true,
getExperimentalZedIntegration: () => false,
getChatRecordingService: () => undefined,
getMessageBus: vi.fn().mockReturnValue(undefined),
getDisableAllHooks: vi.fn().mockReturnValue(true),
setApprovalMode: vi.fn(),
} as unknown as Config;
return { mockConfig, mockModifiableTool, mockToolRegistry };
}
beforeEach(() => {
vi.mocked(IdeClient.getInstance).mockResolvedValue(
mockIdeClient as unknown as IdeClient,
);
mockIdeClient.isDiffingEnabled.mockReturnValue(true);
mockIdeClient.openDiff.mockReset();
});
it('should safely update args via _applyInlineModify when IDE returns modified content (#2709)', async () => {
const { mockConfig, mockModifiableTool } = createIdeMockConfig({
ideMode: true,
});
// IDE returns accepted with modified content
mockIdeClient.openDiff.mockResolvedValue({
status: 'accepted',
content: 'IDE-modified content',
});
const onAllToolCallsComplete = vi.fn();
const onToolCallsUpdate = vi.fn();
const scheduler = new CoreToolScheduler({
config: mockConfig,
onAllToolCallsComplete,
onToolCallsUpdate,
getPreferredEditor: () => 'vscode',
onEditorClose: vi.fn(),
});
const originalArgs = { param: 'original-value' };
const request = {
callId: 'ide-1',
name: 'mockModifiableTool',
args: originalArgs,
isClientInitiated: false,
prompt_id: 'prompt-ide-1',
};
const abortController = new AbortController();
await scheduler.schedule([request], abortController.signal);
// Wait for the tool to complete (IDE auto-confirms)
await vi.waitFor(() => {
expect(onAllToolCallsComplete).toHaveBeenCalled();
});
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('success');
// The tool should have been executed with the IDE-modified content
// via _applyInlineModify -> createUpdatedParams -> setArgsInternal
expect(mockModifiableTool.executeFn).toHaveBeenCalledWith({
newContent: 'IDE-modified content',
});
// CRITICAL: The original args object should NOT have been mutated (#2709)
expect(originalArgs).toEqual({ param: 'original-value' });
// The request.args (which is what goes into history) should also be safe.
// structuredClone in buildInvocation ensures the tool gets its own copy.
expect(request.args).toEqual({ param: 'original-value' });
});
it('should NOT call openDiff when AUTO_EDIT mode is active (#2673)', async () => {
const { mockConfig, mockModifiableTool } = createIdeMockConfig({
approvalMode: ApprovalMode.AUTO_EDIT,
ideMode: true,
});
mockModifiableTool.shouldConfirm = false; // AUTO_EDIT returns 'allow'
const onAllToolCallsComplete = vi.fn();
const onToolCallsUpdate = vi.fn();
const scheduler = new CoreToolScheduler({
config: mockConfig,
onAllToolCallsComplete,
onToolCallsUpdate,
getPreferredEditor: () => 'vscode',
onEditorClose: vi.fn(),
});
const request = {
callId: 'auto-edit-1',
name: 'mockModifiableTool',
args: { param: 'value' },
isClientInitiated: false,
prompt_id: 'prompt-auto-edit-1',
};
const abortController = new AbortController();
await scheduler.schedule([request], abortController.signal);
await vi.waitFor(() => {
expect(onAllToolCallsComplete).toHaveBeenCalled();
});
// openDiff should NOT have been called since AUTO_EDIT auto-approves
expect(mockIdeClient.openDiff).not.toHaveBeenCalled();
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('success');
});
it('should execute normally when IDE accepts without modifying content', async () => {
const { mockConfig, mockModifiableTool } = createIdeMockConfig({
ideMode: true,
});
// IDE returns accepted without content (no modifications)
mockIdeClient.openDiff.mockResolvedValue({
status: 'accepted',
content: undefined,
});
const onAllToolCallsComplete = vi.fn();
const onToolCallsUpdate = vi.fn();
const scheduler = new CoreToolScheduler({
config: mockConfig,
onAllToolCallsComplete,
onToolCallsUpdate,
getPreferredEditor: () => 'vscode',
onEditorClose: vi.fn(),
});
const request = {
callId: 'ide-no-mod-1',
name: 'mockModifiableTool',
args: { param: 'keep-this' },
isClientInitiated: false,
prompt_id: 'prompt-ide-no-mod-1',
};
const abortController = new AbortController();
await scheduler.schedule([request], abortController.signal);
await vi.waitFor(() => {
expect(onAllToolCallsComplete).toHaveBeenCalled();
});
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('success');
// Tool should execute with original params (no _applyInlineModify call)
// executeFn receives the params object from the invocation
expect(mockModifiableTool.executeFn).toHaveBeenCalled();
});
it('should cancel tool when IDE rejects the diff', async () => {
const { mockConfig } = createIdeMockConfig({
ideMode: true,
});
// IDE rejects the diff
mockIdeClient.openDiff.mockResolvedValue({
status: 'rejected',
});
const onAllToolCallsComplete = vi.fn();
const onToolCallsUpdate = vi.fn();
const scheduler = new CoreToolScheduler({
config: mockConfig,
onAllToolCallsComplete,
onToolCallsUpdate,
getPreferredEditor: () => 'vscode',
onEditorClose: vi.fn(),
});
const request = {
callId: 'ide-reject-1',
name: 'mockModifiableTool',
args: { param: 'value' },
isClientInitiated: false,
prompt_id: 'prompt-ide-reject-1',
};
const abortController = new AbortController();
await scheduler.schedule([request], abortController.signal);
await vi.waitFor(() => {
expect(onAllToolCallsComplete).toHaveBeenCalled();
});
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('cancelled');
});
it('should not call openDiff when IDE mode is disabled', async () => {
const { mockConfig } = createIdeMockConfig({
ideMode: false,
});
const onAllToolCallsComplete = vi.fn();
const onToolCallsUpdate = vi.fn();
const scheduler = new CoreToolScheduler({
config: mockConfig,
onAllToolCallsComplete,
onToolCallsUpdate,
getPreferredEditor: () => 'vscode',
onEditorClose: vi.fn(),
});
const request = {
callId: 'no-ide-1',
name: 'mockModifiableTool',
args: { param: 'value' },
isClientInitiated: false,
prompt_id: 'prompt-no-ide-1',
};
const abortController = new AbortController();
await scheduler.schedule([request], abortController.signal);
// Tool should be awaiting approval but openDiff was never called
await waitForStatus(onToolCallsUpdate, 'awaiting_approval');
expect(mockIdeClient.openDiff).not.toHaveBeenCalled();
});
});

View file

@ -65,6 +65,7 @@ import * as Diff from 'diff';
import levenshtein from 'fast-levenshtein';
import { getPlanModeSystemReminder } from './prompts.js';
import { ShellToolInvocation } from '../tools/shell.js';
import { IdeClient } from '../ide/ide-client.js';
const TRUNCATION_PARAM_GUIDANCE =
'Note: Your previous response was truncated due to max_tokens limit, ' +
@ -592,7 +593,7 @@ export class CoreToolScheduler {
args: object,
): AnyToolInvocation | Error {
try {
return tool.build(args);
return tool.build(structuredClone(args));
} catch (e) {
if (e instanceof Error) {
return e;
@ -974,41 +975,6 @@ export class CoreToolScheduler {
continue;
}
// Allow IDE to resolve confirmation
if (
confirmationDetails.type === 'edit' &&
confirmationDetails.ideConfirmation
) {
confirmationDetails.ideConfirmation.then((resolution) => {
// Guard: skip if the tool was already handled (e.g. by CLI
// confirmation). Without this check, resolveDiffFromCli
// triggers this handler AND the CLI's onConfirm, causing a
// race where ProceedOnce overwrites ProceedAlways.
const still = this.toolCalls.find(
(c) =>
c.request.callId === reqInfo.callId &&
c.status === 'awaiting_approval',
);
if (!still) return;
if (resolution.status === 'accepted') {
this.handleConfirmationResponse(
reqInfo.callId,
confirmationDetails!.onConfirm,
ToolConfirmationOutcome.ProceedOnce,
signal,
);
} else {
this.handleConfirmationResponse(
reqInfo.callId,
confirmationDetails!.onConfirm,
ToolConfirmationOutcome.Cancel,
signal,
);
}
});
}
// Fire PermissionRequest hook before showing the permission dialog.
const messageBus = this.config.getMessageBus() as
| MessageBus
@ -1074,6 +1040,13 @@ export class CoreToolScheduler {
}
}
// Allow IDE to resolve confirmation
this.openIdeDiffIfEnabled(
confirmationDetails,
reqInfo.callId,
signal,
);
const originalOnConfirm = confirmationDetails.onConfirm;
const wrappedConfirmationDetails: ToolCallConfirmationDetails = {
...confirmationDetails,
@ -1229,6 +1202,65 @@ export class CoreToolScheduler {
await this.attemptExecutionOfScheduledCalls(signal);
}
/**
* Opens an IDE diff view for edit-type tools when IDE mode is active.
* The IDE resolution is handled asynchronously if the user accepts or
* rejects from the IDE, it triggers handleConfirmationResponse.
*
* Uses confirmationDetails.filePath / newContent (the same data shown in
* CLI diff) rather than ModifyContext so that the IDE diff is always
* consistent with the CLI and with resolveDiffFromCli.
*/
private async openIdeDiffIfEnabled(
confirmationDetails: ToolCallConfirmationDetails,
callId: string,
signal: AbortSignal,
) {
if (confirmationDetails.type !== 'edit' || !this.config.getIdeMode()) {
return;
}
const ideClient = await IdeClient.getInstance();
if (!ideClient.isDiffingEnabled()) return;
const resolution = await ideClient.openDiff(
confirmationDetails.filePath,
confirmationDetails.newContent,
);
// Guard: skip if the tool was already handled (e.g. by CLI
// confirmation). Without this check, resolveDiffFromCli
// triggers this handler AND the CLI's onConfirm, causing a
// race where ProceedOnce overwrites ProceedAlways.
const still = this.toolCalls.find(
(c) => c.request.callId === callId && c.status === 'awaiting_approval',
);
if (!still) return;
if (resolution.status === 'accepted') {
// When content is unchanged, skip the inline modify path so that
// the original tool params (e.g. partial old_string for edit tool)
// are preserved. Mitigate the multi-edit-on-same-file issue (#2702)
// for the common accept-without-edit case.
const userEdited =
resolution.content != null &&
resolution.content !== confirmationDetails.newContent;
this.handleConfirmationResponse(
callId,
confirmationDetails.onConfirm,
ToolConfirmationOutcome.ProceedOnce,
signal,
userEdited ? { newContent: resolution.content } : undefined,
);
} else {
this.handleConfirmationResponse(
callId,
confirmationDetails.onConfirm,
ToolConfirmationOutcome.Cancel,
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
@ -1240,18 +1272,17 @@ export class CoreToolScheduler {
payload: ToolConfirmationPayload,
signal: AbortSignal,
): Promise<void> {
const confirmDetails = toolCall.confirmationDetails;
if (
toolCall.confirmationDetails.type !== 'edit' ||
confirmDetails.type !== 'edit' ||
!isModifiableDeclarativeTool(toolCall.tool) ||
!payload.newContent
) {
return;
}
const currentContent = confirmDetails.originalContent ?? '';
const modifyContext = toolCall.tool.getModifyContext(signal);
const currentContent = await modifyContext.getCurrentContent(
toolCall.request.args,
);
const updatedParams = modifyContext.createUpdatedParams(
currentContent,
@ -1259,7 +1290,7 @@ export class CoreToolScheduler {
toolCall.request.args,
);
const updatedDiff = Diff.createPatch(
modifyContext.getFilePath(toolCall.request.args),
confirmDetails.filePath,
currentContent,
payload.newContent,
'Current',
@ -1268,7 +1299,7 @@ export class CoreToolScheduler {
this.setArgsInternal(toolCall.request.callId, updatedParams);
this.setStatusInternal(toolCall.request.callId, 'awaiting_approval', {
...toolCall.confirmationDetails,
...confirmDetails,
fileDiff: updatedDiff,
});
}

View file

@ -349,7 +349,7 @@ class AgentToolInvocation extends BaseToolInvocation<AgentParams, ToolResult> {
// When a tool result arrives for the tool that had a pending
// confirmation, clear the stale prompt. This handles the case where
// the IDE diff-tab accept resolved the tool via CoreToolScheduler's
// ideConfirmation.then path, which bypasses the UI's onConfirm wrapper.
// IDE confirmation handler, which bypasses the UI's onConfirm wrapper.
const clearPending =
pendingConfirmationCallId === event.callId
? { pendingConfirmation: undefined }

View file

@ -7,18 +7,9 @@
/* eslint-disable @typescript-eslint/no-explicit-any */
const mockGenerateJson = vi.hoisted(() => vi.fn());
const mockOpenDiff = vi.hoisted(() => vi.fn());
import { IdeClient } from '../ide/ide-client.js';
vi.mock('../ide/ide-client.js', () => ({
IdeClient: {
getInstance: vi.fn(),
},
}));
vi.mock('../utils/editor.js', () => ({
openDiff: mockOpenDiff,
openDiff: vi.fn(),
}));
vi.mock('../telemetry/loggers.js', () => ({
@ -30,7 +21,6 @@ import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import type { EditToolParams } from './edit.js';
import { applyReplacement, EditTool } from './edit.js';
import type { FileDiff } from './tools.js';
import { ToolConfirmationOutcome } from './tools.js';
import { ToolErrorType } from './tool-error.js';
import path from 'node:path';
import fs from 'node:fs';
@ -877,93 +867,4 @@ describe('EditTool', () => {
expect(error).toBeNull();
});
});
describe('IDE mode', () => {
const testFile = 'edit_me.txt';
let filePath: string;
let ideClient: any;
beforeEach(() => {
filePath = path.join(rootDir, testFile);
ideClient = {
openDiff: vi.fn(),
isDiffingEnabled: vi.fn().mockReturnValue(true),
};
vi.mocked(IdeClient.getInstance).mockResolvedValue(ideClient);
(mockConfig as any).getIdeMode = () => true;
});
it('should call ideClient.openDiff and update params on confirmation', async () => {
const initialContent = 'some old content here';
const newContent = 'some new content here';
const modifiedContent = 'some modified content here';
fs.writeFileSync(filePath, initialContent);
const params: EditToolParams = {
file_path: filePath,
old_string: 'old',
new_string: 'new',
};
ideClient.openDiff.mockResolvedValueOnce({
status: 'accepted',
content: modifiedContent,
});
const invocation = tool.build(params);
const confirmation = await invocation.getConfirmationDetails(
new AbortController().signal,
);
expect(ideClient.openDiff).toHaveBeenCalledWith(filePath, newContent);
if (confirmation && 'onConfirm' in confirmation) {
await confirmation.onConfirm(ToolConfirmationOutcome.ProceedOnce);
}
expect(params.old_string).toBe(initialContent);
expect(params.new_string).toBe(modifiedContent);
});
it('should not call ideClient.openDiff in AUTO_EDIT mode', async () => {
const initialContent = 'some old content here';
fs.writeFileSync(filePath, initialContent);
const params: EditToolParams = {
file_path: filePath,
old_string: 'old',
new_string: 'new',
};
(mockConfig.getApprovalMode as Mock).mockReturnValueOnce(
ApprovalMode.AUTO_EDIT,
);
const invocation = tool.build(params);
const confirmation = await invocation.getConfirmationDetails(
new AbortController().signal,
);
expect(ideClient.openDiff).not.toHaveBeenCalled();
expect(confirmation).toBeDefined();
expect((confirmation as any).ideConfirmation).toBeUndefined();
});
it('should not call ideClient.openDiff in YOLO mode', async () => {
const initialContent = 'some old content here';
fs.writeFileSync(filePath, initialContent);
const params: EditToolParams = {
file_path: filePath,
old_string: 'old',
new_string: 'new',
};
(mockConfig.getApprovalMode as Mock).mockReturnValueOnce(
ApprovalMode.YOLO,
);
const invocation = tool.build(params);
const confirmation = await invocation.getConfirmationDetails(
new AbortController().signal,
);
expect(ideClient.openDiff).not.toHaveBeenCalled();
expect((confirmation as any).ideConfirmation).toBeUndefined();
});
});
});

View file

@ -42,7 +42,6 @@ import type {
ModifiableDeclarativeTool,
ModifyContext,
} from './modifiable-tool.js';
import { IdeClient } from '../ide/ide-client.js';
import { safeLiteralReplace } from '../utils/textUtils.js';
import {
countOccurrences,
@ -308,16 +307,6 @@ class EditToolInvocation implements ToolInvocation<EditToolParams, ToolResult> {
'Proposed',
DEFAULT_DIFF_OPTIONS,
);
const approvalMode = this.config.getApprovalMode();
const ideClient = await IdeClient.getInstance();
const ideConfirmation =
this.config.getIdeMode() &&
ideClient.isDiffingEnabled() &&
approvalMode !== ApprovalMode.AUTO_EDIT &&
approvalMode !== ApprovalMode.YOLO
? ideClient.openDiff(this.params.file_path, editData.newContent)
: undefined;
const confirmationDetails: ToolEditConfirmationDetails = {
type: 'edit',
title: `Confirm Edit: ${shortenPath(makeRelative(this.params.file_path, this.config.getTargetDir()))}`,
@ -330,16 +319,7 @@ class EditToolInvocation implements ToolInvocation<EditToolParams, ToolResult> {
if (outcome === ToolConfirmationOutcome.ProceedAlways) {
this.config.setApprovalMode(ApprovalMode.AUTO_EDIT);
}
if (ideConfirmation) {
const result = await ideConfirmation;
if (result.status === 'accepted' && result.content) {
this.params.old_string = editData.currentContent ?? '';
this.params.new_string = result.content;
}
}
},
ideConfirmation,
};
return confirmationDetails;
}

View file

@ -519,7 +519,7 @@ export class MemoryTool
);
const scope = scopeMatch
? (scopeMatch[1].toLowerCase() as 'global' | 'project')
: 'global';
: originalParams.scope || 'global';
// Strip out the scope directive and instruction lines, keep only the actual memory content
const contentWithoutScope = modifiedProposedContent.replace(

View file

@ -6,7 +6,6 @@
import type { FunctionDeclaration, Part, PartListUnion } from '@google/genai';
import { ToolErrorType } from './tool-error.js';
import type { DiffUpdateResult } from '../ide/ide-client.js';
import type { ShellExecutionConfig } from '../services/shellExecutionService.js';
import { SchemaValidator } from '../utils/schemaValidator.js';
import { type AgentStatsSummary } from '../agents/runtime/agent-statistics.js';
@ -592,7 +591,6 @@ export interface ToolEditConfirmationDetails {
originalContent: string | null;
newContent: string;
isModifying?: boolean;
ideConfirmation?: Promise<DiffUpdateResult>;
}
export interface ToolConfirmationPayload {

View file

@ -27,28 +27,13 @@ import os from 'node:os';
import { GeminiClient } from '../core/client.js';
import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js';
import { StandardFileSystemService } from '../services/fileSystemService.js';
import type { DiffUpdateResult } from '../ide/ide-client.js';
import { IdeClient } from '../ide/ide-client.js';
const rootDir = path.resolve(os.tmpdir(), 'qwen-code-test-root');
// --- MOCKS ---
vi.mock('../core/client.js');
vi.mock('../ide/ide-client.js', () => ({
IdeClient: {
getInstance: vi.fn(),
},
}));
let mockGeminiClientInstance: Mocked<GeminiClient>;
const mockIdeClient = {
openDiff: vi.fn(),
isDiffingEnabled: vi.fn(),
};
vi.mocked(IdeClient.getInstance).mockResolvedValue(
mockIdeClient as unknown as IdeClient,
);
// Mock Config
const fsService = new StandardFileSystemService();
@ -59,7 +44,6 @@ const mockConfigInternal = {
getGeminiClient: vi.fn(), // Initialize as a plain mock function
getBaseLlmClient: vi.fn(), // Initialize as a plain mock function
getFileSystemService: () => fsService,
getIdeMode: vi.fn(() => false),
getWorkspaceContext: () => createMockWorkspaceContext(rootDir),
getApiKey: () => 'test-key',
getModel: () => 'test-model',
@ -269,148 +253,11 @@ describe('WriteFileTool', () => {
originalContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
);
});
describe('with IDE integration', () => {
beforeEach(() => {
// Enable IDE mode and set connection status for these tests
mockConfigInternal.getIdeMode.mockReturnValue(true);
mockIdeClient.isDiffingEnabled.mockReturnValue(true);
mockIdeClient.openDiff.mockResolvedValue({
status: 'accepted',
content: 'ide-modified-content',
});
});
it('should call openDiff and await it when in IDE mode and connected', async () => {
const filePath = path.join(rootDir, 'ide_confirm_file.txt');
const params = { file_path: filePath, content: 'test' };
const invocation = tool.build(params);
const confirmation = (await invocation.getConfirmationDetails(
abortSignal,
)) as ToolEditConfirmationDetails;
expect(mockIdeClient.openDiff).toHaveBeenCalledWith(
filePath,
'test', // The corrected content
);
// Ensure the promise is awaited by checking the result
expect(confirmation.ideConfirmation).toBeDefined();
await confirmation.ideConfirmation; // Should resolve
});
it('should not call openDiff if not in IDE mode', async () => {
mockConfigInternal.getIdeMode.mockReturnValue(false);
const filePath = path.join(rootDir, 'ide_disabled_file.txt');
const params = { file_path: filePath, content: 'test' };
const invocation = tool.build(params);
await invocation.getConfirmationDetails(abortSignal);
expect(mockIdeClient.openDiff).not.toHaveBeenCalled();
});
it('should not call openDiff if IDE is not connected', async () => {
mockIdeClient.isDiffingEnabled.mockReturnValue(false);
const filePath = path.join(rootDir, 'ide_disconnected_file.txt');
const params = { file_path: filePath, content: 'test' };
const invocation = tool.build(params);
await invocation.getConfirmationDetails(abortSignal);
expect(mockIdeClient.openDiff).not.toHaveBeenCalled();
});
it('should not call openDiff in AUTO_EDIT mode', async () => {
mockConfigInternal.getApprovalMode.mockReturnValue(
ApprovalMode.AUTO_EDIT,
);
const filePath = path.join(rootDir, 'ide_auto_edit_file.txt');
const params = { file_path: filePath, content: 'test' };
const invocation = tool.build(params);
const confirmation = (await invocation.getConfirmationDetails(
abortSignal,
)) as ToolEditConfirmationDetails;
expect(mockIdeClient.openDiff).not.toHaveBeenCalled();
expect(confirmation.ideConfirmation).toBeUndefined();
});
it('should not call openDiff in YOLO mode', async () => {
mockConfigInternal.getApprovalMode.mockReturnValue(ApprovalMode.YOLO);
const filePath = path.join(rootDir, 'ide_yolo_file.txt');
const params = { file_path: filePath, content: 'test' };
const invocation = tool.build(params);
const confirmation = (await invocation.getConfirmationDetails(
abortSignal,
)) as ToolEditConfirmationDetails;
expect(mockIdeClient.openDiff).not.toHaveBeenCalled();
expect(confirmation.ideConfirmation).toBeUndefined();
});
it('should update params.content with IDE content when onConfirm is called', async () => {
const filePath = path.join(rootDir, 'ide_onconfirm_file.txt');
const params = { file_path: filePath, content: 'original-content' };
const invocation = tool.build(params);
// This is the key part: get the confirmation details
const confirmation = (await invocation.getConfirmationDetails(
abortSignal,
)) as ToolEditConfirmationDetails;
// The `onConfirm` function should exist on the details object
expect(confirmation.onConfirm).toBeDefined();
// Call `onConfirm` to trigger the logic that updates the content
await confirmation.onConfirm!(ToolConfirmationOutcome.ProceedOnce);
// Now, check if the original `params` object (captured by the invocation) was modified
expect(invocation.params.content).toBe('ide-modified-content');
});
it('should not await ideConfirmation promise', async () => {
const filePath = path.join(rootDir, 'ide_no_await_file.txt');
const params = { file_path: filePath, content: 'test' };
const invocation = tool.build(params);
let diffPromiseResolved = false;
const diffPromise = new Promise<DiffUpdateResult>((resolve) => {
setTimeout(() => {
diffPromiseResolved = true;
resolve({ status: 'accepted', content: 'ide-modified-content' });
}, 50); // A small delay to ensure the check happens before resolution
});
mockIdeClient.openDiff.mockReturnValue(diffPromise);
const confirmation = (await invocation.getConfirmationDetails(
abortSignal,
)) as ToolEditConfirmationDetails;
// This is the key check: the confirmation details should be returned
// *before* the diffPromise is resolved.
expect(diffPromiseResolved).toBe(false);
expect(confirmation).toBeDefined();
expect(confirmation.ideConfirmation).toBe(diffPromise);
// Now, we can await the promise to let the test finish cleanly.
await diffPromise;
expect(diffPromiseResolved).toBe(true);
});
});
});
describe('execute', () => {
const abortSignal = new AbortController().signal;
beforeEach(() => {
// Ensure IDE mode is disabled for these tests
mockConfigInternal.getIdeMode.mockReturnValue(false);
mockIdeClient.isDiffingEnabled.mockReturnValue(false);
});
it('should return error if _getCorrectedFileContent returns an error during execute', async () => {
const filePath = path.join(rootDir, 'execute_error_file.txt');
const params = { file_path: filePath, content: 'test content' };

View file

@ -39,7 +39,6 @@ import type {
ModifiableDeclarativeTool,
ModifyContext,
} from './modifiable-tool.js';
import { IdeClient } from '../ide/ide-client.js';
import { logFileOperation } from '../telemetry/loggers.js';
import { FileOperationEvent } from '../telemetry/types.js';
import { FileOperation } from '../telemetry/metrics.js';
@ -143,16 +142,6 @@ class WriteFileToolInvocation extends BaseToolInvocation<
DEFAULT_DIFF_OPTIONS,
);
const approvalMode = this.config.getApprovalMode();
const ideClient = await IdeClient.getInstance();
const ideConfirmation =
this.config.getIdeMode() &&
ideClient.isDiffingEnabled() &&
approvalMode !== ApprovalMode.AUTO_EDIT &&
approvalMode !== ApprovalMode.YOLO
? ideClient.openDiff(this.params.file_path, this.params.content)
: undefined;
const confirmationDetails: ToolEditConfirmationDetails = {
type: 'edit',
title: `Confirm Write: ${shortenPath(relativePath)}`,
@ -165,15 +154,7 @@ class WriteFileToolInvocation extends BaseToolInvocation<
if (outcome === ToolConfirmationOutcome.ProceedAlways) {
this.config.setApprovalMode(ApprovalMode.AUTO_EDIT);
}
if (ideConfirmation) {
const result = await ideConfirmation;
if (result.status === 'accepted' && result.content) {
this.params.content = result.content;
}
}
},
ideConfirmation,
};
return confirmationDetails;
}