From cb43bb9ca41f6e5b5850eb41699e356d5fd0094a Mon Sep 17 00:00:00 2001 From: Tommaso Sciortino Date: Thu, 4 Sep 2025 09:32:09 -0700 Subject: [PATCH] Use IdeClient directly instead of config.ideClient (#7627) --- .../a2a-server/src/utils/testing_utils.ts | 1 - packages/cli/src/gemini.tsx | 8 - .../src/services/BuiltinCommandLoader.test.ts | 6 +- .../cli/src/services/BuiltinCommandLoader.ts | 2 +- packages/cli/src/ui/App.test.tsx | 17 +- packages/cli/src/ui/App.tsx | 23 ++- .../cli/src/ui/commands/aboutCommand.test.ts | 24 ++- packages/cli/src/ui/commands/aboutCommand.ts | 16 +- .../cli/src/ui/commands/bugCommand.test.ts | 23 ++- packages/cli/src/ui/commands/bugCommand.ts | 15 +- .../cli/src/ui/commands/ideCommand.test.ts | 182 ++++++++++-------- packages/cli/src/ui/commands/ideCommand.ts | 34 +++- .../messages/ToolConfirmationMessage.tsx | 4 +- .../ui/hooks/slashCommandProcessor.test.ts | 7 - .../cli/src/ui/hooks/slashCommandProcessor.ts | 12 +- .../cli/src/ui/hooks/useIdeTrustListener.ts | 23 +-- packages/core/src/config/config.ts | 23 +-- packages/core/src/ide/ide-client.ts | 28 +-- packages/core/src/tools/edit.test.ts | 15 +- packages/core/src/tools/edit.ts | 4 +- packages/core/src/tools/smart-edit.test.ts | 15 +- packages/core/src/tools/smart-edit.ts | 4 +- packages/core/src/tools/write-file.test.ts | 15 +- packages/core/src/tools/write-file.ts | 4 +- 24 files changed, 288 insertions(+), 217 deletions(-) diff --git a/packages/a2a-server/src/utils/testing_utils.ts b/packages/a2a-server/src/utils/testing_utils.ts index 07cfe783f1..e175002ff6 100644 --- a/packages/a2a-server/src/utils/testing_utils.ts +++ b/packages/a2a-server/src/utils/testing_utils.ts @@ -24,7 +24,6 @@ export function createMockConfig( getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT), getIdeMode: vi.fn().mockReturnValue(false), getAllowedTools: vi.fn().mockReturnValue([]), - getIdeClient: vi.fn(), getWorkspaceContext: vi.fn().mockReturnValue({ isPathWithinWorkspace: () => true, }), diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 79bf14270f..603c1cbd76 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -36,9 +36,6 @@ import { logUserPrompt, AuthType, getOauthClient, - logIdeConnection, - IdeConnectionEvent, - IdeConnectionType, uiTelemetryService, } from '@google/gemini-cli-core'; import { validateAuthMethod } from './config/auth.js'; @@ -289,11 +286,6 @@ export async function main() { spinnerInstance.unmount(); } - if (config.getIdeMode()) { - await config.getIdeClient().connect(); - logIdeConnection(config, new IdeConnectionEvent(IdeConnectionType.START)); - } - // Load custom themes from settings themeManager.loadCustomThemes(settings.merged.ui?.customThemes); diff --git a/packages/cli/src/services/BuiltinCommandLoader.test.ts b/packages/cli/src/services/BuiltinCommandLoader.test.ts index 0ae31e10ae..9a1993fece 100644 --- a/packages/cli/src/services/BuiltinCommandLoader.test.ts +++ b/packages/cli/src/services/BuiltinCommandLoader.test.ts @@ -64,7 +64,7 @@ describe('BuiltinCommandLoader', () => { vi.clearAllMocks(); mockConfig = { some: 'config' } as unknown as Config; - ideCommandMock.mockReturnValue({ + ideCommandMock.mockResolvedValue({ name: 'ide', description: 'IDE command', kind: CommandKind.BUILT_IN, @@ -81,7 +81,7 @@ describe('BuiltinCommandLoader', () => { await loader.loadCommands(new AbortController().signal); expect(ideCommandMock).toHaveBeenCalledTimes(1); - expect(ideCommandMock).toHaveBeenCalledWith(mockConfig); + expect(ideCommandMock).toHaveBeenCalledWith(); expect(restoreCommandMock).toHaveBeenCalledTimes(1); expect(restoreCommandMock).toHaveBeenCalledWith(mockConfig); }); @@ -105,7 +105,7 @@ describe('BuiltinCommandLoader', () => { const loader = new BuiltinCommandLoader(null); await loader.loadCommands(new AbortController().signal); expect(ideCommandMock).toHaveBeenCalledTimes(1); - expect(ideCommandMock).toHaveBeenCalledWith(null); + expect(ideCommandMock).toHaveBeenCalledWith(); expect(restoreCommandMock).toHaveBeenCalledTimes(1); expect(restoreCommandMock).toHaveBeenCalledWith(null); }); diff --git a/packages/cli/src/services/BuiltinCommandLoader.ts b/packages/cli/src/services/BuiltinCommandLoader.ts index 9dafcd205b..e4cf330f4b 100644 --- a/packages/cli/src/services/BuiltinCommandLoader.ts +++ b/packages/cli/src/services/BuiltinCommandLoader.ts @@ -64,7 +64,7 @@ export class BuiltinCommandLoader implements ICommandLoader { editorCommand, extensionsCommand, helpCommand, - ideCommand(this.config), + await ideCommand(), initCommand, mcpCommand, memoryCommand, diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index 2b7169a0d2..ef3dd4a9fe 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -93,7 +93,6 @@ interface MockServerConfig { getAllGeminiMdFilenames: Mock<() => string[]>; getGeminiClient: Mock<() => GeminiClient | undefined>; getUserTier: Mock<() => Promise>; - getIdeClient: Mock<() => { getCurrentIde: Mock<() => string | undefined> }>; getScreenReader: Mock<() => boolean>; } @@ -183,13 +182,6 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { getWorkspaceContext: vi.fn(() => ({ getDirectories: vi.fn(() => []), })), - getIdeClient: vi.fn(() => ({ - getCurrentIde: vi.fn(() => 'vscode'), - getDetectedIdeDisplayName: vi.fn(() => 'VSCode'), - addStatusChangeListener: vi.fn(), - removeStatusChangeListener: vi.fn(), - getConnectionStatus: vi.fn(() => 'connected'), - })), isTrustedFolder: vi.fn(() => true), getScreenReader: vi.fn(() => false), getFolderTrustFeature: vi.fn(() => false), @@ -208,6 +200,15 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { MCPServerConfig: actualCore.MCPServerConfig, getAllGeminiMdFilenames: vi.fn(() => ['GEMINI.md']), ideContext: ideContextMock, + IdeClient: { + getInstance: vi.fn().mockResolvedValue({ + getCurrentIde: vi.fn(() => 'vscode'), + getDetectedIdeDisplayName: vi.fn(() => 'VSCode'), + addStatusChangeListener: vi.fn(), + removeStatusChangeListener: vi.fn(), + getConnectionStatus: vi.fn(() => 'connected'), + }), + }, isGitRepository: vi.fn(), }; }); diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index c83cbc8a85..6f0ac09aab 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -59,7 +59,12 @@ import { ContextSummaryDisplay } from './components/ContextSummaryDisplay.js'; import { useHistory } from './hooks/useHistoryManager.js'; import { useInputHistoryStore } from './hooks/useInputHistoryStore.js'; import process from 'node:process'; -import type { EditorType, Config, IdeContext } from '@google/gemini-cli-core'; +import type { + EditorType, + Config, + IdeContext, + DetectedIde, +} from '@google/gemini-cli-core'; import { ApprovalMode, getAllGeminiMdFilenames, @@ -73,6 +78,7 @@ import { isGenericQuotaExceededError, UserTierId, DEFAULT_GEMINI_FLASH_MODEL, + IdeClient, } from '@google/gemini-cli-core'; import type { IdeIntegrationNudgeResult } from './IdeIntegrationNudge.js'; import { IdeIntegrationNudge } from './IdeIntegrationNudge.js'; @@ -161,10 +167,19 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { const { history, addItem, clearItems, loadHistory } = useHistory(); const [idePromptAnswered, setIdePromptAnswered] = useState(false); - const currentIDE = config.getIdeClient().getCurrentIde(); + const [currentIDE, setCurrentIDE] = useState(); + useEffect(() => { - registerCleanup(() => config.getIdeClient().disconnect()); + (async () => { + const ideClient = await IdeClient.getInstance(); + setCurrentIDE(ideClient.getCurrentIde()); + })(); + registerCleanup(async () => { + const ideClient = await IdeClient.getInstance(); + ideClient.disconnect(); + }); }, [config]); + const shouldShowIdePrompt = currentIDE && !config.getIdeMode() && @@ -306,7 +321,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { const { isFolderTrustDialogOpen, handleFolderTrustSelect, isRestarting } = useFolderTrust(settings, setIsTrustedFolder); - const { needsRestart: ideNeedsRestart } = useIdeTrustListener(config); + const { needsRestart: ideNeedsRestart } = useIdeTrustListener(); useEffect(() => { if (ideNeedsRestart) { // IDE trust changed, force a restart. diff --git a/packages/cli/src/ui/commands/aboutCommand.test.ts b/packages/cli/src/ui/commands/aboutCommand.test.ts index eeab10d83c..5ae1577d82 100644 --- a/packages/cli/src/ui/commands/aboutCommand.test.ts +++ b/packages/cli/src/ui/commands/aboutCommand.test.ts @@ -10,8 +10,20 @@ import { type CommandContext } from './types.js'; import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; import * as versionUtils from '../../utils/version.js'; import { MessageType } from '../types.js'; +import { IdeClient } from '@google/gemini-cli-core'; -import type { IdeClient } from '../../../../core/src/ide/ide-client.js'; +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + IdeClient: { + getInstance: vi.fn().mockResolvedValue({ + getDetectedIdeDisplayName: vi.fn().mockReturnValue('test-ide'), + }), + }, + }; +}); vi.mock('../../utils/version.js', () => ({ getCliVersion: vi.fn(), @@ -27,7 +39,6 @@ describe('aboutCommand', () => { services: { config: { getModel: vi.fn(), - getIdeClient: vi.fn(), getIdeMode: vi.fn().mockReturnValue(true), }, settings: { @@ -53,9 +64,6 @@ describe('aboutCommand', () => { Object.defineProperty(process, 'platform', { value: 'test-os', }); - vi.spyOn(mockContext.services.config!, 'getIdeClient').mockReturnValue({ - getDetectedIdeDisplayName: vi.fn().mockReturnValue('test-ide'), - } as Partial as IdeClient); }); afterEach(() => { @@ -129,11 +137,11 @@ describe('aboutCommand', () => { }); it('should not show ide client when it is not detected', async () => { - vi.spyOn(mockContext.services.config!, 'getIdeClient').mockReturnValue({ + vi.mocked(IdeClient.getInstance).mockResolvedValue({ getDetectedIdeDisplayName: vi.fn().mockReturnValue(undefined), - } as Partial as IdeClient); + } as unknown as IdeClient); - process.env.SANDBOX = ''; + process.env['SANDBOX'] = ''; if (!aboutCommand.action) { throw new Error('The about command must have an action.'); } diff --git a/packages/cli/src/ui/commands/aboutCommand.ts b/packages/cli/src/ui/commands/aboutCommand.ts index 9de430003e..6b2cf16ec9 100644 --- a/packages/cli/src/ui/commands/aboutCommand.ts +++ b/packages/cli/src/ui/commands/aboutCommand.ts @@ -5,10 +5,11 @@ */ import { getCliVersion } from '../../utils/version.js'; -import type { SlashCommand } from './types.js'; +import type { CommandContext, SlashCommand } from './types.js'; import { CommandKind } from './types.js'; import process from 'node:process'; import { MessageType, type HistoryItemAbout } from '../types.js'; +import { IdeClient } from '@google/gemini-cli-core'; export const aboutCommand: SlashCommand = { name: 'about', @@ -29,10 +30,7 @@ export const aboutCommand: SlashCommand = { const selectedAuthType = context.services.settings.merged.security?.auth?.selectedType || ''; const gcpProject = process.env['GOOGLE_CLOUD_PROJECT'] || ''; - const ideClient = - (context.services.config?.getIdeMode() && - context.services.config?.getIdeClient()?.getDetectedIdeDisplayName()) || - ''; + const ideClient = await getIdeClientName(context); const aboutItem: Omit = { type: MessageType.ABOUT, @@ -48,3 +46,11 @@ export const aboutCommand: SlashCommand = { context.ui.addItem(aboutItem, Date.now()); }, }; + +async function getIdeClientName(context: CommandContext) { + if (!context.services.config?.getIdeMode()) { + return ''; + } + const ideClient = await IdeClient.getInstance(); + return ideClient?.getDetectedIdeDisplayName() ?? ''; +} diff --git a/packages/cli/src/ui/commands/bugCommand.test.ts b/packages/cli/src/ui/commands/bugCommand.test.ts index 2dfb3fe8f8..a050620e9a 100644 --- a/packages/cli/src/ui/commands/bugCommand.test.ts +++ b/packages/cli/src/ui/commands/bugCommand.test.ts @@ -16,7 +16,19 @@ import { formatMemoryUsage } from '../utils/formatters.js'; vi.mock('open'); vi.mock('../../utils/version.js'); vi.mock('../utils/formatters.js'); -vi.mock('@google/gemini-cli-core'); +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + IdeClient: { + getInstance: () => ({ + getDetectedIdeDisplayName: vi.fn().mockReturnValue('VSCode'), + }), + }, + sessionId: 'test-session-id', + }; +}); vi.mock('node:process', () => ({ default: { platform: 'test-platform', @@ -31,9 +43,6 @@ describe('bugCommand', () => { beforeEach(() => { vi.mocked(getCliVersion).mockResolvedValue('0.1.0'); vi.mocked(formatMemoryUsage).mockReturnValue('100 MB'); - vi.mock('@google/gemini-cli-core', () => ({ - sessionId: 'test-session-id', - })); vi.stubEnv('SANDBOX', 'gemini-test'); }); @@ -48,9 +57,6 @@ describe('bugCommand', () => { config: { getModel: () => 'gemini-pro', getBugCommand: () => undefined, - getIdeClient: () => ({ - getDetectedIdeDisplayName: () => 'VSCode', - }), getIdeMode: () => true, }, }, @@ -84,9 +90,6 @@ describe('bugCommand', () => { config: { getModel: () => 'gemini-pro', getBugCommand: () => ({ urlTemplate: customTemplate }), - getIdeClient: () => ({ - getDetectedIdeDisplayName: () => 'VSCode', - }), getIdeMode: () => true, }, }, diff --git a/packages/cli/src/ui/commands/bugCommand.ts b/packages/cli/src/ui/commands/bugCommand.ts index c5f559b678..e792d45886 100644 --- a/packages/cli/src/ui/commands/bugCommand.ts +++ b/packages/cli/src/ui/commands/bugCommand.ts @@ -15,7 +15,7 @@ import { MessageType } from '../types.js'; import { GIT_COMMIT_INFO } from '../../generated/git-commit.js'; import { formatMemoryUsage } from '../utils/formatters.js'; import { getCliVersion } from '../../utils/version.js'; -import { sessionId } from '@google/gemini-cli-core'; +import { IdeClient, sessionId } from '@google/gemini-cli-core'; export const bugCommand: SlashCommand = { name: 'bug', @@ -37,10 +37,7 @@ export const bugCommand: SlashCommand = { const modelVersion = config?.getModel() || 'Unknown'; const cliVersion = await getCliVersion(); const memoryUsage = formatMemoryUsage(process.memoryUsage().rss); - const ideClient = - (context.services.config?.getIdeMode() && - context.services.config?.getIdeClient()?.getDetectedIdeDisplayName()) || - ''; + const ideClient = await getIdeClientName(context); let info = ` * **CLI Version:** ${cliVersion} @@ -90,3 +87,11 @@ export const bugCommand: SlashCommand = { } }, }; + +async function getIdeClientName(context: CommandContext) { + if (!context.services.config?.getIdeMode()) { + return ''; + } + const ideClient = await IdeClient.getInstance(); + return ideClient.getDetectedIdeDisplayName() ?? ''; +} diff --git a/packages/cli/src/ui/commands/ideCommand.test.ts b/packages/cli/src/ui/commands/ideCommand.test.ts index 7de1bd7da5..bb2cd19e6a 100644 --- a/packages/cli/src/ui/commands/ideCommand.test.ts +++ b/packages/cli/src/ui/commands/ideCommand.test.ts @@ -8,26 +8,43 @@ import type { MockInstance } from 'vitest'; import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest'; import { ideCommand } from './ideCommand.js'; import { type CommandContext } from './types.js'; -import { type Config, DetectedIde } from '@google/gemini-cli-core'; +import { DetectedIde } from '@google/gemini-cli-core'; import * as core from '@google/gemini-cli-core'; -vi.mock('child_process'); -vi.mock('glob'); vi.mock('@google/gemini-cli-core', async (importOriginal) => { const original = await importOriginal(); return { ...original, getOauthClient: vi.fn(original.getOauthClient), getIdeInstaller: vi.fn(original.getIdeInstaller), + IdeClient: { + getInstance: vi.fn(), + }, + ideContext: { + getIdeContext: vi.fn(), + }, }; }); describe('ideCommand', () => { let mockContext: CommandContext; - let mockConfig: Config; + let mockIdeClient: core.IdeClient; let platformSpy: MockInstance; beforeEach(() => { + vi.resetAllMocks(); + + mockIdeClient = { + reconnect: vi.fn(), + disconnect: vi.fn(), + connect: vi.fn(), + getCurrentIde: vi.fn(), + getDetectedIdeDisplayName: vi.fn(), + getConnectionStatus: vi.fn(), + } as unknown as core.IdeClient; + + vi.mocked(core.IdeClient.getInstance).mockResolvedValue(mockIdeClient); + mockContext = { ui: { addItem: vi.fn(), @@ -36,22 +53,14 @@ describe('ideCommand', () => { settings: { setValue: vi.fn(), }, + config: { + getIdeMode: vi.fn(), + setIdeMode: vi.fn(), + getUsageStatisticsEnabled: vi.fn().mockReturnValue(false), + }, }, } as unknown as CommandContext; - mockConfig = { - getIdeMode: vi.fn(), - getIdeClient: vi.fn(() => ({ - reconnect: vi.fn(), - disconnect: vi.fn(), - getCurrentIde: vi.fn(), - getDetectedIdeDisplayName: vi.fn(), - getConnectionStatus: vi.fn(), - })), - setIdeModeAndSyncConnection: vi.fn(), - setIdeMode: vi.fn(), - } as unknown as Config; - platformSpy = vi.spyOn(process, 'platform', 'get'); }); @@ -59,64 +68,57 @@ describe('ideCommand', () => { vi.restoreAllMocks(); }); - it('should return null if config is not provided', () => { - const command = ideCommand(null); - expect(command).toBeNull(); + it('should return the ide command', async () => { + vi.mocked(mockIdeClient.getCurrentIde).mockReturnValue(DetectedIde.VSCode); + vi.mocked(mockIdeClient.getDetectedIdeDisplayName).mockReturnValue( + 'VS Code', + ); + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ + status: core.IDEConnectionStatus.Disconnected, + }); + const command = await ideCommand(); + expect(command).not.toBeNull(); + expect(command.name).toBe('ide'); + expect(command.subCommands).toHaveLength(3); + expect(command.subCommands?.[0].name).toBe('enable'); + expect(command.subCommands?.[1].name).toBe('status'); + expect(command.subCommands?.[2].name).toBe('install'); }); - it('should return the ide command', () => { - vi.mocked(mockConfig.getIdeMode).mockReturnValue(true); - vi.mocked(mockConfig.getIdeClient).mockReturnValue({ - getCurrentIde: () => DetectedIde.VSCode, - getDetectedIdeDisplayName: () => 'VS Code', - getConnectionStatus: () => ({ - status: core.IDEConnectionStatus.Disconnected, - }), - } as ReturnType); - const command = ideCommand(mockConfig); + it('should show disable command when connected', async () => { + vi.mocked(mockIdeClient.getCurrentIde).mockReturnValue(DetectedIde.VSCode); + vi.mocked(mockIdeClient.getDetectedIdeDisplayName).mockReturnValue( + 'VS Code', + ); + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ + status: core.IDEConnectionStatus.Connected, + }); + const command = await ideCommand(); expect(command).not.toBeNull(); - expect(command?.name).toBe('ide'); - expect(command?.subCommands).toHaveLength(3); - expect(command?.subCommands?.[0].name).toBe('enable'); - expect(command?.subCommands?.[1].name).toBe('status'); - expect(command?.subCommands?.[2].name).toBe('install'); - }); - - it('should show disable command when connected', () => { - vi.mocked(mockConfig.getIdeMode).mockReturnValue(true); - vi.mocked(mockConfig.getIdeClient).mockReturnValue({ - getCurrentIde: () => DetectedIde.VSCode, - getDetectedIdeDisplayName: () => 'VS Code', - getConnectionStatus: () => ({ - status: core.IDEConnectionStatus.Connected, - }), - } as ReturnType); - const command = ideCommand(mockConfig); - expect(command).not.toBeNull(); - const subCommandNames = command?.subCommands?.map((cmd) => cmd.name); + const subCommandNames = command.subCommands?.map((cmd) => cmd.name); expect(subCommandNames).toContain('disable'); expect(subCommandNames).not.toContain('enable'); }); describe('status subcommand', () => { - const mockGetConnectionStatus = vi.fn(); beforeEach(() => { - vi.mocked(mockConfig.getIdeClient).mockReturnValue({ - getConnectionStatus: mockGetConnectionStatus, - getCurrentIde: () => DetectedIde.VSCode, - getDetectedIdeDisplayName: () => 'VS Code', - } as unknown as ReturnType); + vi.mocked(mockIdeClient.getCurrentIde).mockReturnValue( + DetectedIde.VSCode, + ); + vi.mocked(mockIdeClient.getDetectedIdeDisplayName).mockReturnValue( + 'VS Code', + ); }); it('should show connected status', async () => { - mockGetConnectionStatus.mockReturnValue({ + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ status: core.IDEConnectionStatus.Connected, }); - const command = ideCommand(mockConfig); + const command = await ideCommand(); const result = await command!.subCommands!.find( (c) => c.name === 'status', )!.action!(mockContext, ''); - expect(mockGetConnectionStatus).toHaveBeenCalled(); + expect(vi.mocked(mockIdeClient.getConnectionStatus)).toHaveBeenCalled(); expect(result).toEqual({ type: 'message', messageType: 'info', @@ -125,14 +127,14 @@ describe('ideCommand', () => { }); it('should show connecting status', async () => { - mockGetConnectionStatus.mockReturnValue({ + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ status: core.IDEConnectionStatus.Connecting, }); - const command = ideCommand(mockConfig); + const command = await ideCommand(); const result = await command!.subCommands!.find( (c) => c.name === 'status', )!.action!(mockContext, ''); - expect(mockGetConnectionStatus).toHaveBeenCalled(); + expect(vi.mocked(mockIdeClient.getConnectionStatus)).toHaveBeenCalled(); expect(result).toEqual({ type: 'message', messageType: 'info', @@ -140,14 +142,14 @@ describe('ideCommand', () => { }); }); it('should show disconnected status', async () => { - mockGetConnectionStatus.mockReturnValue({ + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ status: core.IDEConnectionStatus.Disconnected, }); - const command = ideCommand(mockConfig); + const command = await ideCommand(); const result = await command!.subCommands!.find( (c) => c.name === 'status', )!.action!(mockContext, ''); - expect(mockGetConnectionStatus).toHaveBeenCalled(); + expect(vi.mocked(mockIdeClient.getConnectionStatus)).toHaveBeenCalled(); expect(result).toEqual({ type: 'message', messageType: 'error', @@ -157,15 +159,15 @@ describe('ideCommand', () => { it('should show disconnected status with details', async () => { const details = 'Something went wrong'; - mockGetConnectionStatus.mockReturnValue({ + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ status: core.IDEConnectionStatus.Disconnected, details, }); - const command = ideCommand(mockConfig); + const command = await ideCommand(); const result = await command!.subCommands!.find( (c) => c.name === 'status', )!.action!(mockContext, ''); - expect(mockGetConnectionStatus).toHaveBeenCalled(); + expect(vi.mocked(mockIdeClient.getConnectionStatus)).toHaveBeenCalled(); expect(result).toEqual({ type: 'message', messageType: 'error', @@ -177,32 +179,40 @@ describe('ideCommand', () => { describe('install subcommand', () => { const mockInstall = vi.fn(); beforeEach(() => { - vi.mocked(mockConfig.getIdeMode).mockReturnValue(true); - vi.mocked(mockConfig.getIdeClient).mockReturnValue({ - getCurrentIde: () => DetectedIde.VSCode, - getConnectionStatus: () => ({ - status: core.IDEConnectionStatus.Disconnected, - }), - getDetectedIdeDisplayName: () => 'VS Code', - } as unknown as ReturnType); + vi.mocked(mockIdeClient.getCurrentIde).mockReturnValue( + DetectedIde.VSCode, + ); + vi.mocked(mockIdeClient.getDetectedIdeDisplayName).mockReturnValue( + 'VS Code', + ); + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ + status: core.IDEConnectionStatus.Disconnected, + }); vi.mocked(core.getIdeInstaller).mockReturnValue({ install: mockInstall, - isInstalled: vi.fn(), }); platformSpy.mockReturnValue('linux'); }); it('should install the extension', async () => { + vi.useFakeTimers(); mockInstall.mockResolvedValue({ success: true, message: 'Successfully installed.', }); - const command = ideCommand(mockConfig); - await command!.subCommands!.find((c) => c.name === 'install')!.action!( - mockContext, - '', - ); + const command = await ideCommand(); + + // For the polling loop inside the action. + vi.mocked(mockIdeClient.getConnectionStatus).mockReturnValue({ + status: core.IDEConnectionStatus.Connected, + }); + + const actionPromise = command!.subCommands!.find( + (c) => c.name === 'install', + )!.action!(mockContext, ''); + await vi.runAllTimersAsync(); + await actionPromise; expect(core.getIdeInstaller).toHaveBeenCalledWith('vscode'); expect(mockInstall).toHaveBeenCalled(); @@ -220,6 +230,14 @@ describe('ideCommand', () => { }), expect.any(Number), ); + expect(mockContext.ui.addItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'info', + text: '🟢 Connected to VS Code', + }), + expect.any(Number), + ); + vi.useRealTimers(); }, 10000); it('should show an error if installation fails', async () => { @@ -228,7 +246,7 @@ describe('ideCommand', () => { message: 'Installation failed.', }); - const command = ideCommand(mockConfig); + const command = await ideCommand(); await command!.subCommands!.find((c) => c.name === 'install')!.action!( mockContext, '', diff --git a/packages/cli/src/ui/commands/ideCommand.ts b/packages/cli/src/ui/commands/ideCommand.ts index 6af30472f7..bf1f1e4231 100644 --- a/packages/cli/src/ui/commands/ideCommand.ts +++ b/packages/cli/src/ui/commands/ideCommand.ts @@ -4,7 +4,14 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { Config, IdeClient, File } from '@google/gemini-cli-core'; +import { + type Config, + IdeClient, + type File, + logIdeConnection, + IdeConnectionEvent, + IdeConnectionType, +} from '@google/gemini-cli-core'; import { getIdeInstaller, IDEConnectionStatus, @@ -111,11 +118,22 @@ async function getIdeStatusMessageWithFiles(ideClient: IdeClient): Promise<{ } } -export const ideCommand = (config: Config | null): SlashCommand | null => { - if (!config) { - return null; +async function setIdeModeAndSyncConnection( + config: Config, + value: boolean, +): Promise { + config.setIdeMode(value); + const ideClient = await IdeClient.getInstance(); + if (value) { + await ideClient.connect(); + logIdeConnection(config, new IdeConnectionEvent(IdeConnectionType.SESSION)); + } else { + await ideClient.disconnect(); } - const ideClient = config.getIdeClient(); +} + +export const ideCommand = async (): Promise => { + const ideClient = await IdeClient.getInstance(); const currentIDE = ideClient.getCurrentIde(); if (!currentIDE || !ideClient.getDetectedIdeDisplayName()) { return { @@ -194,7 +212,7 @@ export const ideCommand = (config: Config | null): SlashCommand | null => { ); // Poll for up to 5 seconds for the extension to activate. for (let i = 0; i < 10; i++) { - await config.setIdeModeAndSyncConnection(true); + await setIdeModeAndSyncConnection(context.services.config!, true); if ( ideClient.getConnectionStatus().status === IDEConnectionStatus.Connected @@ -236,7 +254,7 @@ export const ideCommand = (config: Config | null): SlashCommand | null => { 'ide.enabled', true, ); - await config.setIdeModeAndSyncConnection(true); + await setIdeModeAndSyncConnection(context.services.config!, true); const { messageType, content } = getIdeStatusMessage(ideClient); context.ui.addItem( { @@ -258,7 +276,7 @@ export const ideCommand = (config: Config | null): SlashCommand | null => { 'ide.enabled', false, ); - await config.setIdeModeAndSyncConnection(false); + await setIdeModeAndSyncConnection(context.services.config!, false); const { messageType, content } = getIdeStatusMessage(ideClient); context.ui.addItem( { diff --git a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx index e3c48d9415..a440e59e85 100644 --- a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx +++ b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx @@ -15,7 +15,7 @@ import type { ToolMcpConfirmationDetails, Config, } from '@google/gemini-cli-core'; -import { ToolConfirmationOutcome } from '@google/gemini-cli-core'; +import { IdeClient, ToolConfirmationOutcome } from '@google/gemini-cli-core'; import type { RadioSelectItem } from '../shared/RadioButtonSelect.js'; import { RadioButtonSelect } from '../shared/RadioButtonSelect.js'; import { MaxSizedBox } from '../shared/MaxSizedBox.js'; @@ -43,10 +43,10 @@ export const ToolConfirmationMessage: React.FC< const handleConfirm = async (outcome: ToolConfirmationOutcome) => { if (confirmationDetails.type === 'edit') { - const ideClient = config.getIdeClient(); if (config.getIdeMode()) { const cliOutcome = outcome === ToolConfirmationOutcome.Cancel ? 'rejected' : 'accepted'; + const ideClient = await IdeClient.getInstance(); await ideClient?.resolveDiffFromCli( confirmationDetails.filePath, cliOutcome, diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index f04caf1fb4..5210ed546c 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -86,8 +86,6 @@ import { SlashCommandStatus, ToolConfirmationOutcome, makeFakeConfig, - ToolConfirmationOutcome, - type IdeClient, } from '@google/gemini-cli-core'; function createTestCommand( @@ -111,11 +109,6 @@ describe('useSlashCommandProcessor', () => { const mockSetQuittingMessages = vi.fn(); const mockConfig = makeFakeConfig({}); - vi.spyOn(mockConfig, 'getIdeClient').mockReturnValue({ - addStatusChangeListener: vi.fn(), - removeStatusChangeListener: vi.fn(), - } as unknown as IdeClient); - const mockSettings = {} as LoadedSettings; beforeEach(() => { diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 350df660d2..a7c3ce583f 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -17,6 +17,7 @@ import { SlashCommandStatus, ToolConfirmationOutcome, Storage, + IdeClient, } from '@google/gemini-cli-core'; import { useSessionStats } from '../contexts/SessionContext.js'; import { runExitCleanup } from '../../utils/cleanup.js'; @@ -215,15 +216,20 @@ export const useSlashCommandProcessor = ( return; } - const ideClient = config.getIdeClient(); const listener = () => { reloadCommands(); }; - ideClient.addStatusChangeListener(listener); + (async () => { + const ideClient = await IdeClient.getInstance(); + ideClient.addStatusChangeListener(listener); + })(); return () => { - ideClient.removeStatusChangeListener(listener); + (async () => { + const ideClient = await IdeClient.getInstance(); + ideClient.removeStatusChangeListener(listener); + })(); }; }, [config, reloadCommands]); diff --git a/packages/cli/src/ui/hooks/useIdeTrustListener.ts b/packages/cli/src/ui/hooks/useIdeTrustListener.ts index bcd40200b6..88b3d5ff97 100644 --- a/packages/cli/src/ui/hooks/useIdeTrustListener.ts +++ b/packages/cli/src/ui/hooks/useIdeTrustListener.ts @@ -5,25 +5,26 @@ */ import { useCallback, useEffect, useState, useSyncExternalStore } from 'react'; -import type { Config } from '@google/gemini-cli-core'; -import { ideContext } from '@google/gemini-cli-core'; +import { IdeClient, ideContext } from '@google/gemini-cli-core'; /** * This hook listens for trust status updates from the IDE companion extension. * It provides the current trust status from the IDE and a flag indicating * if a restart is needed because the trust state has changed. */ -export function useIdeTrustListener(config: Config) { - const subscribe = useCallback( - (onStoreChange: () => void) => { - const ideClient = config.getIdeClient(); +export function useIdeTrustListener() { + const subscribe = useCallback((onStoreChange: () => void) => { + (async () => { + const ideClient = await IdeClient.getInstance(); ideClient.addTrustChangeListener(onStoreChange); - return () => { + })(); + return () => { + (async () => { + const ideClient = await IdeClient.getInstance(); ideClient.removeTrustChangeListener(onStoreChange); - }; - }, - [config], - ); + })(); + }; + }, []); const getSnapshot = () => ideContext.getIdeContext()?.workspaceState?.isTrusted; diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 543199274f..7adf182f64 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -260,7 +260,7 @@ export class Config { private readonly folderTrustFeature: boolean; private readonly folderTrust: boolean; private ideMode: boolean; - private ideClient!: IdeClient; + private inFallbackMode = false; private readonly maxSessionTurns: number; private readonly listExtensions: boolean; @@ -383,7 +383,12 @@ export class Config { throw Error('Config was already initialized'); } this.initialized = true; - this.ideClient = await IdeClient.getInstance(); + + if (this.getIdeMode()) { + await (await IdeClient.getInstance()).connect(); + logIdeConnection(this, new IdeConnectionEvent(IdeConnectionType.START)); + } + // Initialize centralized FileDiscoveryService this.getFileService(); if (this.getCheckpointingEnabled()) { @@ -762,20 +767,6 @@ export class Config { this.ideMode = value; } - async setIdeModeAndSyncConnection(value: boolean): Promise { - this.ideMode = value; - if (value) { - await this.ideClient.connect(); - logIdeConnection(this, new IdeConnectionEvent(IdeConnectionType.SESSION)); - } else { - await this.ideClient.disconnect(); - } - } - - getIdeClient(): IdeClient { - return this.ideClient; - } - /** * Get the current FileSystemService */ diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts index fc87967217..0c86059405 100644 --- a/packages/core/src/ide/ide-client.ts +++ b/packages/core/src/ide/ide-client.ts @@ -65,7 +65,7 @@ function getRealPath(path: string): string { * Manages the connection to and interaction with the IDE server. */ export class IdeClient { - private static instance: IdeClient; + private static instancePromise: Promise | null = null; private client: Client | undefined = undefined; private state: IDEConnectionState = { status: IDEConnectionStatus.Disconnected, @@ -81,19 +81,21 @@ export class IdeClient { private constructor() {} - static async getInstance(): Promise { - if (!IdeClient.instance) { - const client = new IdeClient(); - client.ideProcessInfo = await getIdeProcessInfo(); - client.currentIde = detectIde(client.ideProcessInfo); - if (client.currentIde) { - client.currentIdeDisplayName = getIdeInfo( - client.currentIde, - ).displayName; - } - IdeClient.instance = client; + static getInstance(): Promise { + if (!IdeClient.instancePromise) { + IdeClient.instancePromise = (async () => { + const client = new IdeClient(); + client.ideProcessInfo = await getIdeProcessInfo(); + client.currentIde = detectIde(client.ideProcessInfo); + if (client.currentIde) { + client.currentIdeDisplayName = getIdeInfo( + client.currentIde, + ).displayName; + } + return client; + })(); } - return IdeClient.instance; + return IdeClient.instancePromise; } addStatusChangeListener(listener: (state: IDEConnectionState) => void) { diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index 391ad12d67..03627dc085 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -10,7 +10,17 @@ const mockEnsureCorrectEdit = vi.hoisted(() => vi.fn()); const mockGenerateJson = vi.hoisted(() => vi.fn()); const mockOpenDiff = vi.hoisted(() => vi.fn()); -import { IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; + +vi.mock('../ide/ide-client.js', () => ({ + IdeClient: { + getInstance: vi.fn(), + }, + IDEConnectionStatus: { + Connected: 'connected', + Disconnected: 'disconnected', + }, +})); vi.mock('../utils/editCorrector.js', () => ({ ensureCorrectEdit: mockEnsureCorrectEdit, @@ -70,7 +80,6 @@ describe('EditTool', () => { setApprovalMode: vi.fn(), getWorkspaceContext: () => createMockWorkspaceContext(rootDir), getFileSystemService: () => new StandardFileSystemService(), - getIdeClient: () => undefined, getIdeMode: () => false, // getGeminiConfig: () => ({ apiKey: 'test-api-key' }), // This was not a real Config method // Add other properties/methods of Config if EditTool uses them @@ -878,8 +887,8 @@ describe('EditTool', () => { status: IDEConnectionStatus.Connected, }), }; + vi.mocked(IdeClient.getInstance).mockResolvedValue(ideClient); (mockConfig as any).getIdeMode = () => true; - (mockConfig as any).getIdeClient = () => ideClient; }); it('should call ideClient.openDiff and update params on confirmation', async () => { diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index c9c67a0ab1..f3ed36c798 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -33,7 +33,7 @@ import type { ModifiableDeclarativeTool, ModifyContext, } from './modifiable-tool.js'; -import { IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; export function applyReplacement( currentContent: string | null, @@ -267,7 +267,7 @@ class EditToolInvocation implements ToolInvocation { 'Proposed', DEFAULT_DIFF_OPTIONS, ); - const ideClient = this.config.getIdeClient(); + const ideClient = await IdeClient.getInstance(); const ideConfirmation = this.config.getIdeMode() && ideClient?.getConnectionStatus().status === IDEConnectionStatus.Connected diff --git a/packages/core/src/tools/smart-edit.test.ts b/packages/core/src/tools/smart-edit.test.ts index 51ab29a57f..132d993306 100644 --- a/packages/core/src/tools/smart-edit.test.ts +++ b/packages/core/src/tools/smart-edit.test.ts @@ -10,7 +10,17 @@ const mockFixLLMEditWithInstruction = vi.hoisted(() => vi.fn()); const mockGenerateJson = vi.hoisted(() => vi.fn()); const mockOpenDiff = vi.hoisted(() => vi.fn()); -import { IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; + +vi.mock('../ide/ide-client.js', () => ({ + IdeClient: { + getInstance: vi.fn(), + }, + IDEConnectionStatus: { + Connected: 'connected', + Disconnected: 'disconnected', + }, +})); vi.mock('../utils/llm-edit-fixer.js', () => ({ FixLLMEditWithInstruction: mockFixLLMEditWithInstruction, @@ -75,7 +85,6 @@ describe('SmartEditTool', () => { setApprovalMode: vi.fn(), getWorkspaceContext: () => createMockWorkspaceContext(rootDir), getFileSystemService: () => new StandardFileSystemService(), - getIdeClient: () => undefined, getIdeMode: () => false, getApiKey: () => 'test-api-key', getModel: () => 'test-model', @@ -449,8 +458,8 @@ describe('SmartEditTool', () => { status: IDEConnectionStatus.Connected, }), }; + vi.mocked(IdeClient.getInstance).mockResolvedValue(ideClient); (mockConfig as any).getIdeMode = () => true; - (mockConfig as any).getIdeClient = () => ideClient; }); it('should call ideClient.openDiff and update params on confirmation', async () => { diff --git a/packages/core/src/tools/smart-edit.ts b/packages/core/src/tools/smart-edit.ts index 6d0342c80b..3647b73246 100644 --- a/packages/core/src/tools/smart-edit.ts +++ b/packages/core/src/tools/smart-edit.ts @@ -28,7 +28,7 @@ import { type ModifiableDeclarativeTool, type ModifyContext, } from './modifiable-tool.js'; -import { IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; import { FixLLMEditWithInstruction } from '../utils/llm-edit-fixer.js'; export function applyReplacement( @@ -526,7 +526,7 @@ class EditToolInvocation implements ToolInvocation { 'Proposed', DEFAULT_DIFF_OPTIONS, ); - const ideClient = this.config.getIdeClient(); + const ideClient = await IdeClient.getInstance(); const ideConfirmation = this.config.getIdeMode() && ideClient?.getConnectionStatus().status === IDEConnectionStatus.Connected diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index bb0aa6d8ce..df9a33ef63 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -39,7 +39,11 @@ const rootDir = path.resolve(os.tmpdir(), 'gemini-cli-test-root'); // --- MOCKS --- vi.mock('../core/client.js'); vi.mock('../utils/editCorrector.js'); - +vi.mock('../ide/ide-client.js', () => ({ + IdeClient: { + getInstance: vi.fn(), + }, +})); let mockGeminiClientInstance: Mocked; const mockEnsureCorrectEdit = vi.fn(); const mockEnsureCorrectFileContent = vi.fn(); @@ -58,7 +62,6 @@ const mockConfigInternal = { setApprovalMode: vi.fn(), getGeminiClient: vi.fn(), // Initialize as a plain mock function getFileSystemService: () => fsService, - getIdeClient: vi.fn(), getIdeMode: vi.fn(() => false), getWorkspaceContext: () => createMockWorkspaceContext(rootDir), getApiKey: () => 'test-key', @@ -120,14 +123,6 @@ 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); diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index f72bf72db9..597043f05b 100644 --- a/packages/core/src/tools/write-file.ts +++ b/packages/core/src/tools/write-file.ts @@ -35,7 +35,7 @@ import type { ModifiableDeclarativeTool, ModifyContext, } from './modifiable-tool.js'; -import { IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; import { logFileOperation } from '../telemetry/loggers.js'; import { FileOperationEvent } from '../telemetry/types.js'; import { FileOperation } from '../telemetry/metrics.js'; @@ -193,7 +193,7 @@ class WriteFileToolInvocation extends BaseToolInvocation< DEFAULT_DIFF_OPTIONS, ); - const ideClient = this.config.getIdeClient(); + const ideClient = await IdeClient.getInstance(); const ideConfirmation = this.config.getIdeMode() && ideClient.getConnectionStatus().status === IDEConnectionStatus.Connected