From 30b13c63576417fac8cd3bb53545ff68a253838b Mon Sep 17 00:00:00 2001 From: HugoMurillo Date: Wed, 10 Sep 2025 12:21:46 -0600 Subject: [PATCH] fix(#7039): broken IDE integration for multi-edit/multi-write flows (#8159) --- packages/core/src/ide/ide-client.ts | 49 ++++++++- packages/core/src/tools/write-file.test.ts | 119 +++++++++++++++++++++ 2 files changed, 166 insertions(+), 2 deletions(-) diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts index 7da863c564..6c6454b781 100644 --- a/packages/core/src/ide/ide-client.ts +++ b/packages/core/src/ide/ide-client.ts @@ -78,6 +78,12 @@ export class IdeClient { private diffResponses = new Map void>(); private statusListeners = new Set<(state: IDEConnectionState) => void>(); private trustChangeListeners = new Set<(isTrusted: boolean) => void>(); + /** + * A mutex to ensure that only one diff view is open in the IDE at a time. + * This prevents race conditions and UI issues in IDEs like VSCode that + * can't handle multiple diff views being opened simultaneously. + */ + private diffMutex = Promise.resolve(); private constructor() {} @@ -202,10 +208,16 @@ export class IdeClient { filePath: string, newContent?: string, ): Promise { - return new Promise((resolve, reject) => { + const release = await this.acquireMutex(); + + const promise = new Promise((resolve, reject) => { + if (!this.client) { + // The promise will be rejected, and the finally block below will release the mutex. + return reject(new Error('IDE client is not connected.')); + } this.diffResponses.set(filePath, resolve); this.client - ?.callTool({ + .callTool({ name: `openDiff`, arguments: { filePath, @@ -214,9 +226,42 @@ export class IdeClient { }) .catch((err) => { logger.debug(`callTool for ${filePath} failed:`, err); + this.diffResponses.delete(filePath); reject(err); }); }); + + // Ensure the mutex is released only after the diff interaction is complete. + promise.finally(release); + + return promise; + } + + /** + * Acquires a lock to ensure sequential execution of critical sections. + * + * This method implements a promise-based mutex. It works by chaining promises. + * Each call to `acquireMutex` gets the current `diffMutex` promise. It then + * creates a *new* promise (`newMutex`) that will be resolved when the caller + * invokes the returned `release` function. The `diffMutex` is immediately + * updated to this `newMutex`. + * + * The method returns a promise that resolves with the `release` function only + * *after* the *previous* `diffMutex` promise has resolved. This creates a + * queue where each subsequent operation must wait for the previous one to release + * the lock. + * + * @returns A promise that resolves to a function that must be called to + * release the lock. + */ + private acquireMutex(): Promise<() => void> { + let release: () => void; + const newMutex = new Promise((resolve) => { + release = resolve; + }); + const oldMutex = this.diffMutex; + this.diffMutex = newMutex; + return oldMutex.then(() => release); } async closeDiff( diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index df9a33ef63..5dcafe6249 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -33,6 +33,8 @@ import { } from '../utils/editCorrector.js'; import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js'; import { StandardFileSystemService } from '../services/fileSystemService.js'; +import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; +import type { DiffUpdateResult } from '../ide/ideContext.js'; const rootDir = path.resolve(os.tmpdir(), 'gemini-cli-test-root'); @@ -43,16 +45,28 @@ vi.mock('../ide/ide-client.js', () => ({ IdeClient: { getInstance: vi.fn(), }, + IDEConnectionStatus: { + Connected: 'connected', + Disconnected: 'disconnected', + Connecting: 'connecting', + }, })); let mockGeminiClientInstance: Mocked; const mockEnsureCorrectEdit = vi.fn(); const mockEnsureCorrectFileContent = vi.fn(); +const mockIdeClient = { + getConnectionStatus: vi.fn(), + openDiff: vi.fn(), +}; // Wire up the mocked functions to be used by the actual module imports vi.mocked(ensureCorrectEdit).mockImplementation(mockEnsureCorrectEdit); vi.mocked(ensureCorrectFileContent).mockImplementation( mockEnsureCorrectFileContent, ); +vi.mocked(IdeClient.getInstance).mockResolvedValue( + mockIdeClient as unknown as IdeClient, +); // Mock Config const fsService = new StandardFileSystemService(); @@ -437,6 +451,111 @@ 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.getConnectionStatus.mockReturnValue({ + status: IDEConnectionStatus.Connected, + }); + 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.shouldConfirmExecute( + 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.shouldConfirmExecute(abortSignal); + + expect(mockIdeClient.openDiff).not.toHaveBeenCalled(); + }); + + it('should not call openDiff if IDE is not connected', async () => { + mockIdeClient.getConnectionStatus.mockReturnValue({ + status: IDEConnectionStatus.Disconnected, + }); + const filePath = path.join(rootDir, 'ide_disconnected_file.txt'); + const params = { file_path: filePath, content: 'test' }; + const invocation = tool.build(params); + + await invocation.shouldConfirmExecute(abortSignal); + + expect(mockIdeClient.openDiff).not.toHaveBeenCalled(); + }); + + 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.shouldConfirmExecute( + 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((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.shouldConfirmExecute( + 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', () => {