From d892cde0b04739c3c802fc93c2333fe03bef4ea6 Mon Sep 17 00:00:00 2001 From: Tommaso Sciortino Date: Thu, 11 Sep 2025 11:22:20 -0700 Subject: [PATCH] Refactor IdeContextStore (#8278) --- packages/cli/src/ui/AppContainer.tsx | 6 +- packages/cli/src/ui/commands/ideCommand.ts | 4 +- .../cli/src/ui/hooks/useIdeTrustListener.ts | 5 +- packages/core/src/config/config.ts | 4 +- packages/core/src/core/client.test.ts | 34 ++++---- packages/core/src/core/client.ts | 4 +- packages/core/src/ide/ide-client.ts | 6 +- packages/core/src/ide/ideContext.test.ts | 83 +++++++++---------- packages/core/src/ide/ideContext.ts | 56 +++++-------- packages/core/src/utils/ide-trust.ts | 4 +- 10 files changed, 92 insertions(+), 114 deletions(-) diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 4bef9d7531..b2781648a0 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -29,7 +29,7 @@ import { type UserTierId, DEFAULT_GEMINI_FLASH_MODEL, IdeClient, - ideContext, + ideContextStore, getErrorMessage, getAllGeminiMdFilenames, AuthType, @@ -707,8 +707,8 @@ Logging in with Google... Please restart Gemini CLI to continue. }, [terminalWidth, refreshStatic]); useEffect(() => { - const unsubscribe = ideContext.subscribeToIdeContext(setIdeContextState); - setIdeContextState(ideContext.getIdeContext()); + const unsubscribe = ideContextStore.subscribe(setIdeContextState); + setIdeContextState(ideContextStore.get()); return unsubscribe; }, []); diff --git a/packages/cli/src/ui/commands/ideCommand.ts b/packages/cli/src/ui/commands/ideCommand.ts index bf1f1e4231..b5f002b641 100644 --- a/packages/cli/src/ui/commands/ideCommand.ts +++ b/packages/cli/src/ui/commands/ideCommand.ts @@ -15,7 +15,7 @@ import { import { getIdeInstaller, IDEConnectionStatus, - ideContext, + ideContextStore, GEMINI_CLI_COMPANION_EXTENSION_NAME, } from '@google/gemini-cli-core'; import path from 'node:path'; @@ -90,7 +90,7 @@ async function getIdeStatusMessageWithFiles(ideClient: IdeClient): Promise<{ switch (connection.status) { case IDEConnectionStatus.Connected: { let content = `🟢 Connected to ${ideClient.getDetectedIdeDisplayName()}`; - const context = ideContext.getIdeContext(); + const context = ideContextStore.get(); const openFiles = context?.workspaceState?.openFiles; if (openFiles && openFiles.length > 0) { content += formatFileList(openFiles); diff --git a/packages/cli/src/ui/hooks/useIdeTrustListener.ts b/packages/cli/src/ui/hooks/useIdeTrustListener.ts index 88b3d5ff97..c585100747 100644 --- a/packages/cli/src/ui/hooks/useIdeTrustListener.ts +++ b/packages/cli/src/ui/hooks/useIdeTrustListener.ts @@ -5,7 +5,7 @@ */ import { useCallback, useEffect, useState, useSyncExternalStore } from 'react'; -import { IdeClient, ideContext } from '@google/gemini-cli-core'; +import { IdeClient, ideContextStore } from '@google/gemini-cli-core'; /** * This hook listens for trust status updates from the IDE companion extension. @@ -26,8 +26,7 @@ export function useIdeTrustListener() { }; }, []); - const getSnapshot = () => - ideContext.getIdeContext()?.workspaceState?.isTrusted; + const getSnapshot = () => ideContextStore.get()?.workspaceState?.isTrusted; const isIdeTrusted = useSyncExternalStore(subscribe, getSnapshot); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 3dcf72bd0c..7321d112a6 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -49,7 +49,7 @@ import { import { shouldAttemptBrowserLaunch } from '../utils/browser.js'; import type { MCPOAuthConfig } from '../mcp/oauth-provider.js'; import { IdeClient } from '../ide/ide-client.js'; -import { ideContext } from '../ide/ideContext.js'; +import { ideContextStore } from '../ide/ideContext.js'; import type { FileSystemService } from '../services/fileSystemService.js'; import { StandardFileSystemService } from '../services/fileSystemService.js'; import { @@ -839,7 +839,7 @@ export class Config { // restarts in the more common path. If the user chooses to mark the folder // as untrusted, the CLI will restart and we will have the trust value // reloaded. - const context = ideContext.getIdeContext(); + const context = ideContextStore.get(); if (context?.workspaceState?.isTrusted !== undefined) { return context.workspaceState.isTrusted; } diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 21f7c88989..3dd4b7b053 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -39,7 +39,7 @@ import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js'; import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; import { setSimulate429 } from '../utils/testUtils.js'; import { tokenLimit } from './tokenLimits.js'; -import { ideContext } from '../ide/ideContext.js'; +import { ideContextStore } from '../ide/ideContext.js'; import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js'; import type { ModelRouterService } from '../routing/modelRouterService.js'; @@ -1007,7 +1007,7 @@ describe('Gemini Client (client.ts)', () => { it('should include editor context when ideMode is enabled', async () => { // Arrange - vi.mocked(ideContext.getIdeContext).mockReturnValue({ + vi.mocked(ideContextStore.get).mockReturnValue({ workspaceState: { openFiles: [ { @@ -1056,7 +1056,7 @@ describe('Gemini Client (client.ts)', () => { } // Assert - expect(ideContext.getIdeContext).toHaveBeenCalled(); + expect(ideContextStore.get).toHaveBeenCalled(); const expectedContext = ` Here is the user's editor context as a JSON object. This is for your information only. \`\`\`json @@ -1086,7 +1086,7 @@ ${JSON.stringify( it('should not add context if ideMode is enabled but no open files', async () => { // Arrange - vi.mocked(ideContext.getIdeContext).mockReturnValue({ + vi.mocked(ideContextStore.get).mockReturnValue({ workspaceState: { openFiles: [], }, @@ -1118,7 +1118,7 @@ ${JSON.stringify( } // Assert - expect(ideContext.getIdeContext).toHaveBeenCalled(); + expect(ideContextStore.get).toHaveBeenCalled(); // The `turn.run` method is now called with the model name as the first // argument. We use `expect.any(String)` because this test is // concerned with the IDE context logic, not the model routing, @@ -1132,7 +1132,7 @@ ${JSON.stringify( it('should add context if ideMode is enabled and there is one active file', async () => { // Arrange - vi.mocked(ideContext.getIdeContext).mockReturnValue({ + vi.mocked(ideContextStore.get).mockReturnValue({ workspaceState: { openFiles: [ { @@ -1172,7 +1172,7 @@ ${JSON.stringify( } // Assert - expect(ideContext.getIdeContext).toHaveBeenCalled(); + expect(ideContextStore.get).toHaveBeenCalled(); const expectedContext = ` Here is the user's editor context as a JSON object. This is for your information only. \`\`\`json @@ -1201,7 +1201,7 @@ ${JSON.stringify( it('should add context if ideMode is enabled and there are open files but no active file', async () => { // Arrange - vi.mocked(ideContext.getIdeContext).mockReturnValue({ + vi.mocked(ideContextStore.get).mockReturnValue({ workspaceState: { openFiles: [ { @@ -1242,7 +1242,7 @@ ${JSON.stringify( } // Assert - expect(ideContext.getIdeContext).toHaveBeenCalled(); + expect(ideContextStore.get).toHaveBeenCalled(); const expectedContext = ` Here is the user's editor context as a JSON object. This is for your information only. \`\`\`json @@ -1836,7 +1836,7 @@ ${JSON.stringify( }; // Setup current context - vi.mocked(ideContext.getIdeContext).mockReturnValue({ + vi.mocked(ideContextStore.get).mockReturnValue({ workspaceState: { openFiles: [ { ...currentActiveFile, isActive: true, timestamp: Date.now() }, @@ -1898,7 +1898,7 @@ ${JSON.stringify( }; // Setup current context (same as previous) - vi.mocked(ideContext.getIdeContext).mockReturnValue({ + vi.mocked(ideContextStore.get).mockReturnValue({ workspaceState: { openFiles: [ { ...activeFile, isActive: true, timestamp: Date.now() }, @@ -1968,7 +1968,7 @@ ${JSON.stringify( client['chat'] = mockChat as GeminiChat; vi.spyOn(client['config'], 'getIdeMode').mockReturnValue(true); - vi.mocked(ideContext.getIdeContext).mockReturnValue({ + vi.mocked(ideContextStore.get).mockReturnValue({ workspaceState: { openFiles: [{ path: '/path/to/file.ts', timestamp: Date.now() }], }, @@ -2065,7 +2065,7 @@ ${JSON.stringify( openFiles: [{ path: '/path/to/fileA.ts', timestamp: Date.now() }], }, }; - vi.mocked(ideContext.getIdeContext).mockReturnValue(initialIdeContext); + vi.mocked(ideContextStore.get).mockReturnValue(initialIdeContext); // Act: Send the tool response let stream = client.sendMessageStream( @@ -2124,7 +2124,7 @@ ${JSON.stringify( openFiles: [{ path: '/path/to/fileB.ts', timestamp: Date.now() }], }, }; - vi.mocked(ideContext.getIdeContext).mockReturnValue(newIdeContext); + vi.mocked(ideContextStore.get).mockReturnValue(newIdeContext); // Act: Send a new, regular user message stream = client.sendMessageStream( @@ -2165,7 +2165,7 @@ ${JSON.stringify( ], }, }; - vi.mocked(ideContext.getIdeContext).mockReturnValue(contextA); + vi.mocked(ideContextStore.get).mockReturnValue(contextA); // Act: Send a regular message to establish the initial context let stream = client.sendMessageStream( @@ -2208,7 +2208,7 @@ ${JSON.stringify( ], }, }; - vi.mocked(ideContext.getIdeContext).mockReturnValue(contextB); + vi.mocked(ideContextStore.get).mockReturnValue(contextB); // Act: Send the tool response stream = client.sendMessageStream( @@ -2263,7 +2263,7 @@ ${JSON.stringify( ], }, }; - vi.mocked(ideContext.getIdeContext).mockReturnValue(contextC); + vi.mocked(ideContextStore.get).mockReturnValue(contextC); // Act: Send a new, regular user message stream = client.sendMessageStream( diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index e1436234e2..21392ed12b 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -36,7 +36,7 @@ import { DEFAULT_THINKING_MODE, } from '../config/models.js'; import { LoopDetectionService } from '../services/loopDetectionService.js'; -import { ideContext } from '../ide/ideContext.js'; +import { ideContextStore } from '../ide/ideContext.js'; import { logChatCompression, logNextSpeakerCheck, @@ -263,7 +263,7 @@ export class GeminiClient { contextParts: string[]; newIdeContext: IdeContext | undefined; } { - const currentIdeContext = ideContext.getIdeContext(); + const currentIdeContext = ideContextStore.get(); if (!currentIdeContext) { return { contextParts: [], newIdeContext: undefined }; } diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts index 6c6454b781..bcba4b81d5 100644 --- a/packages/core/src/ide/ide-client.ts +++ b/packages/core/src/ide/ide-client.ts @@ -8,7 +8,7 @@ import * as fs from 'node:fs'; import { isSubpath } from '../utils/paths.js'; import { detectIde, type DetectedIde, getIdeInfo } from '../ide/detect-ide.js'; import { - ideContext, + ideContextStore, IdeDiffAcceptedNotificationSchema, IdeDiffClosedNotificationSchema, CloseDiffResponseSchema, @@ -362,7 +362,7 @@ export class IdeClient { } if (status === IDEConnectionStatus.Disconnected) { - ideContext.clearIdeContext(); + ideContextStore.clear(); } } @@ -562,7 +562,7 @@ export class IdeClient { this.client.setNotificationHandler( IdeContextNotificationSchema, (notification) => { - ideContext.setIdeContext(notification.params); + ideContextStore.set(notification.params); const isTrusted = notification.params.workspaceState?.isTrusted; if (isTrusted !== undefined) { for (const listener of this.trustChangeListeners) { diff --git a/packages/core/src/ide/ideContext.test.ts b/packages/core/src/ide/ideContext.test.ts index 7970d06aec..0e3c7ec53b 100644 --- a/packages/core/src/ide/ideContext.test.ts +++ b/packages/core/src/ide/ideContext.test.ts @@ -9,7 +9,7 @@ import { IDE_MAX_SELECTED_TEXT_LENGTH, } from './constants.js'; import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest'; -import { createIdeContextStore } from './ideContext.js'; +import { IdeContextStore } from './ideContext.js'; import { type IdeContext, FileSchema, @@ -19,11 +19,11 @@ import { describe('ideContext', () => { describe('createIdeContextStore', () => { - let ideContextStore: ReturnType; + let ideContextStore: IdeContextStore; beforeEach(() => { // Create a fresh, isolated instance for each test - ideContextStore = createIdeContextStore(); + ideContextStore = new IdeContextStore(); }); afterEach(() => { @@ -31,7 +31,7 @@ describe('ideContext', () => { }); it('should return undefined initially for ide context', () => { - expect(ideContextStore.getIdeContext()).toBeUndefined(); + expect(ideContextStore.get()).toBeUndefined(); }); it('should set and retrieve the ide context', () => { @@ -48,9 +48,9 @@ describe('ideContext', () => { }, }; - ideContextStore.setIdeContext(testFile); + ideContextStore.set(testFile); - const activeFile = ideContextStore.getIdeContext(); + const activeFile = ideContextStore.get(); expect(activeFile).toEqual(testFile); }); @@ -67,7 +67,7 @@ describe('ideContext', () => { ], }, }; - ideContextStore.setIdeContext(firstFile); + ideContextStore.set(firstFile); const secondFile = { workspaceState: { @@ -81,9 +81,9 @@ describe('ideContext', () => { ], }, }; - ideContextStore.setIdeContext(secondFile); + ideContextStore.set(secondFile); - const activeFile = ideContextStore.getIdeContext(); + const activeFile = ideContextStore.get(); expect(activeFile).toEqual(secondFile); }); @@ -100,16 +100,16 @@ describe('ideContext', () => { ], }, }; - ideContextStore.setIdeContext(testFile); - expect(ideContextStore.getIdeContext()).toEqual(testFile); + ideContextStore.set(testFile); + expect(ideContextStore.get()).toEqual(testFile); }); it('should notify subscribers when ide context changes', () => { const subscriber1 = vi.fn(); const subscriber2 = vi.fn(); - ideContextStore.subscribeToIdeContext(subscriber1); - ideContextStore.subscribeToIdeContext(subscriber2); + ideContextStore.subscribe(subscriber1); + ideContextStore.subscribe(subscriber2); const testFile = { workspaceState: { @@ -123,7 +123,7 @@ describe('ideContext', () => { ], }, }; - ideContextStore.setIdeContext(testFile); + ideContextStore.set(testFile); expect(subscriber1).toHaveBeenCalledTimes(1); expect(subscriber1).toHaveBeenCalledWith(testFile); @@ -143,7 +143,7 @@ describe('ideContext', () => { ], }, }; - ideContextStore.setIdeContext(newFile); + ideContextStore.set(newFile); expect(subscriber1).toHaveBeenCalledTimes(2); expect(subscriber1).toHaveBeenCalledWith(newFile); @@ -155,10 +155,10 @@ describe('ideContext', () => { const subscriber1 = vi.fn(); const subscriber2 = vi.fn(); - const unsubscribe1 = ideContextStore.subscribeToIdeContext(subscriber1); - ideContextStore.subscribeToIdeContext(subscriber2); + const unsubscribe1 = ideContextStore.subscribe(subscriber1); + ideContextStore.subscribe(subscriber2); - ideContextStore.setIdeContext({ + ideContextStore.set({ workspaceState: { openFiles: [ { @@ -175,7 +175,7 @@ describe('ideContext', () => { unsubscribe1(); - ideContextStore.setIdeContext({ + ideContextStore.set({ workspaceState: { openFiles: [ { @@ -205,21 +205,21 @@ describe('ideContext', () => { }, }; - ideContextStore.setIdeContext(testFile); + ideContextStore.set(testFile); - expect(ideContextStore.getIdeContext()).toEqual(testFile); + expect(ideContextStore.get()).toEqual(testFile); - ideContextStore.clearIdeContext(); + ideContextStore.clear(); - expect(ideContextStore.getIdeContext()).toBeUndefined(); + expect(ideContextStore.get()).toBeUndefined(); }); it('should set the context and notify subscribers when no workspaceState is present', () => { const subscriber = vi.fn(); - ideContextStore.subscribeToIdeContext(subscriber); + ideContextStore.subscribe(subscriber); const context: IdeContext = {}; - ideContextStore.setIdeContext(context); - expect(ideContextStore.getIdeContext()).toBe(context); + ideContextStore.set(context); + expect(ideContextStore.get()).toBe(context); expect(subscriber).toHaveBeenCalledWith(context); }); @@ -229,10 +229,8 @@ describe('ideContext', () => { openFiles: [], }, }; - ideContextStore.setIdeContext(context); - expect( - ideContextStore.getIdeContext()?.workspaceState?.openFiles, - ).toEqual([]); + ideContextStore.set(context); + expect(ideContextStore.get()?.workspaceState?.openFiles).toEqual([]); }); it('should sort openFiles by timestamp in descending order', () => { @@ -245,9 +243,8 @@ describe('ideContext', () => { ], }, }; - ideContextStore.setIdeContext(context); - const openFiles = - ideContextStore.getIdeContext()?.workspaceState?.openFiles; + ideContextStore.set(context); + const openFiles = ideContextStore.get()?.workspaceState?.openFiles; expect(openFiles?.[0]?.path).toBe('file2.ts'); expect(openFiles?.[1]?.path).toBe('file3.ts'); expect(openFiles?.[2]?.path).toBe('file1.ts'); @@ -279,9 +276,8 @@ describe('ideContext', () => { ], }, }; - ideContextStore.setIdeContext(context); - const openFiles = - ideContextStore.getIdeContext()?.workspaceState?.openFiles; + ideContextStore.set(context); + const openFiles = ideContextStore.get()?.workspaceState?.openFiles; expect(openFiles?.[0]?.isActive).toBe(true); expect(openFiles?.[0]?.cursor).toBeDefined(); expect(openFiles?.[0]?.selectedText).toBeDefined(); @@ -309,10 +305,9 @@ describe('ideContext', () => { ], }, }; - ideContextStore.setIdeContext(context); + ideContextStore.set(context); const selectedText = - ideContextStore.getIdeContext()?.workspaceState?.openFiles?.[0] - ?.selectedText; + ideContextStore.get()?.workspaceState?.openFiles?.[0]?.selectedText; expect(selectedText).toHaveLength( IDE_MAX_SELECTED_TEXT_LENGTH + '... [TRUNCATED]'.length, ); @@ -333,10 +328,9 @@ describe('ideContext', () => { ], }, }; - ideContextStore.setIdeContext(context); + ideContextStore.set(context); const selectedText = - ideContextStore.getIdeContext()?.workspaceState?.openFiles?.[0] - ?.selectedText; + ideContextStore.get()?.workspaceState?.openFiles?.[0]?.selectedText; expect(selectedText).toBe(shortText); }); @@ -354,9 +348,8 @@ describe('ideContext', () => { openFiles: files, }, }; - ideContextStore.setIdeContext(context); - const openFiles = - ideContextStore.getIdeContext()?.workspaceState?.openFiles; + ideContextStore.set(context); + const openFiles = ideContextStore.get()?.workspaceState?.openFiles; expect(openFiles).toHaveLength(IDE_MAX_OPEN_FILES); }); }); diff --git a/packages/core/src/ide/ideContext.ts b/packages/core/src/ide/ideContext.ts index 57d021f66a..cb57b99f4b 100644 --- a/packages/core/src/ide/ideContext.ts +++ b/packages/core/src/ide/ideContext.ts @@ -69,25 +69,18 @@ export type DiffUpdateResult = content: undefined; }; -type IdeContextSubscriber = (ideContext: IdeContext | undefined) => void; +type IdeContextSubscriber = (ideContext?: IdeContext) => void; -/** - * Creates a new store for managing the IDE's context. - * This factory function encapsulates the state and logic, allowing for the creation - * of isolated instances, which is particularly useful for testing. - * - * @returns An object with methods to interact with the IDE context. - */ -export function createIdeContextStore() { - let ideContextState: IdeContext | undefined = undefined; - const subscribers = new Set(); +export class IdeContextStore { + private ideContextState?: IdeContext; + private readonly subscribers = new Set(); /** * Notifies all registered subscribers about the current IDE context. */ - function notifySubscribers(): void { - for (const subscriber of subscribers) { - subscriber(ideContextState); + private notifySubscribers(): void { + for (const subscriber of this.subscribers) { + subscriber(this.ideContextState); } } @@ -95,11 +88,11 @@ export function createIdeContextStore() { * Sets the IDE context and notifies all registered subscribers of the change. * @param newIdeContext The new IDE context from the IDE. */ - function setIdeContext(newIdeContext: IdeContext): void { + set(newIdeContext: IdeContext): void { const { workspaceState } = newIdeContext; if (!workspaceState) { - ideContextState = newIdeContext; - notifySubscribers(); + this.ideContextState = newIdeContext; + this.notifySubscribers(); return; } @@ -147,24 +140,24 @@ export function createIdeContextStore() { workspaceState.openFiles = openFiles.slice(0, IDE_MAX_OPEN_FILES); } } - ideContextState = newIdeContext; - notifySubscribers(); + this.ideContextState = newIdeContext; + this.notifySubscribers(); } /** * Clears the IDE context and notifies all registered subscribers of the change. */ - function clearIdeContext(): void { - ideContextState = undefined; - notifySubscribers(); + clear(): void { + this.ideContextState = undefined; + this.notifySubscribers(); } /** * Retrieves the current IDE context. * @returns The `IdeContext` object if a file is active; otherwise, `undefined`. */ - function getIdeContext(): IdeContext | undefined { - return ideContextState; + get(): IdeContext | undefined { + return this.ideContextState; } /** @@ -176,22 +169,15 @@ export function createIdeContextStore() { * @param subscriber The function to be called when the IDE context changes. * @returns A function that, when called, will unsubscribe the provided subscriber. */ - function subscribeToIdeContext(subscriber: IdeContextSubscriber): () => void { - subscribers.add(subscriber); + subscribe(subscriber: IdeContextSubscriber): () => void { + this.subscribers.add(subscriber); return () => { - subscribers.delete(subscriber); + this.subscribers.delete(subscriber); }; } - - return { - setIdeContext, - getIdeContext, - subscribeToIdeContext, - clearIdeContext, - }; } /** * The default, shared instance of the IDE context store for the application. */ -export const ideContext = createIdeContextStore(); +export const ideContextStore = new IdeContextStore(); diff --git a/packages/core/src/utils/ide-trust.ts b/packages/core/src/utils/ide-trust.ts index 1cfaa88b47..c1ac801651 100644 --- a/packages/core/src/utils/ide-trust.ts +++ b/packages/core/src/utils/ide-trust.ts @@ -4,12 +4,12 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { ideContext } from '../ide/ideContext.js'; +import { ideContextStore } from '../ide/ideContext.js'; /** * Gets the workspace trust from the IDE if available. * @returns A boolean if the IDE provides a trust value, otherwise undefined. */ export function getIdeTrust(): boolean | undefined { - return ideContext.getIdeContext()?.workspaceState?.isTrusted; + return ideContextStore.get()?.workspaceState?.isTrusted; }