Refactor IdeContextStore (#8278)

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);
+    const unsubscribe = ideContextStore.subscribe(setIdeContextState);
-    setIdeContextState(ideContext.getIdeContext());
+    setIdeContextState(ideContextStore.get());
     return unsubscribe;
   }, []);
 

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);

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 = () =>
+  const getSnapshot = () => ideContextStore.get()?.workspaceState?.isTrusted;
-    ideContext.getIdeContext()?.workspaceState?.isTrusted;
 
   const isIdeTrusted = useSyncExternalStore(subscribe, getSnapshot);
 

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;
     }

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(

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 };
     }

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) {

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<typeof createIdeContextStore>;
+    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);
+      ideContextStore.set(testFile);
-      expect(ideContextStore.getIdeContext()).toEqual(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.subscribe(subscriber1);
-      ideContextStore.subscribeToIdeContext(subscriber2);
+      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);
+      const unsubscribe1 = ideContextStore.subscribe(subscriber1);
-      ideContextStore.subscribeToIdeContext(subscriber2);
+      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);
+      ideContextStore.set(context);
-      expect(ideContextStore.getIdeContext()).toBe(context);
+      expect(ideContextStore.get()).toBe(context);
       expect(subscriber).toHaveBeenCalledWith(context);
     });
 
@@ -229,10 +229,8 @@ describe('ideContext', () => {
           openFiles: [],
         },
       };
-      ideContextStore.setIdeContext(context);
+      ideContextStore.set(context);
-      expect(
+      expect(ideContextStore.get()?.workspaceState?.openFiles).toEqual([]);
-        ideContextStore.getIdeContext()?.workspaceState?.openFiles,
-      ).toEqual([]);
     });
 
     it('should sort openFiles by timestamp in descending order', () => {
@@ -245,9 +243,8 @@ describe('ideContext', () => {
           ],
         },
       };
-      ideContextStore.setIdeContext(context);
+      ideContextStore.set(context);
-      const openFiles =
+      const openFiles = ideContextStore.get()?.workspaceState?.openFiles;
-        ideContextStore.getIdeContext()?.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);
+      ideContextStore.set(context);
-      const openFiles =
+      const openFiles = ideContextStore.get()?.workspaceState?.openFiles;
-        ideContextStore.getIdeContext()?.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]
+        ideContextStore.get()?.workspaceState?.openFiles?.[0]?.selectedText;
-          ?.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]
+        ideContextStore.get()?.workspaceState?.openFiles?.[0]?.selectedText;
-          ?.selectedText;
       expect(selectedText).toBe(shortText);
     });
 
@@ -354,9 +348,8 @@ describe('ideContext', () => {
           openFiles: files,
         },
       };
-      ideContextStore.setIdeContext(context);
+      ideContextStore.set(context);
-      const openFiles =
+      const openFiles = ideContextStore.get()?.workspaceState?.openFiles;
-        ideContextStore.getIdeContext()?.workspaceState?.openFiles;
       expect(openFiles).toHaveLength(IDE_MAX_OPEN_FILES);
     });
   });

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;
 
-/**
+export class IdeContextStore {
- * Creates a new store for managing the IDE's context.
+  private ideContextState?: IdeContext;
- * This factory function encapsulates the state and logic, allowing for the creation
+  private readonly subscribers = new Set<IdeContextSubscriber>();
- * 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<IdeContextSubscriber>();
 
   /**
    * Notifies all registered subscribers about the current IDE context.
    */
-  function notifySubscribers(): void {
+  private notifySubscribers(): void {
-    for (const subscriber of subscribers) {
+    for (const subscriber of this.subscribers) {
-      subscriber(ideContextState);
+      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;
+      this.ideContextState = newIdeContext;
-      notifySubscribers();
+      this.notifySubscribers();
       return;
     }
 
@@ -147,24 +140,24 @@ export function createIdeContextStore() {
         workspaceState.openFiles = openFiles.slice(0, IDE_MAX_OPEN_FILES);
       }
     }
-    ideContextState = newIdeContext;
+    this.ideContextState = newIdeContext;
-    notifySubscribers();
+    this.notifySubscribers();
   }
 
   /**
    * Clears the IDE context and notifies all registered subscribers of the change.
    */
-  function clearIdeContext(): void {
+  clear(): void {
-    ideContextState = undefined;
+    this.ideContextState = undefined;
-    notifySubscribers();
+    this.notifySubscribers();
   }
 
   /**
    * Retrieves the current IDE context.
    * @returns The `IdeContext` object if a file is active; otherwise, `undefined`.
    */
-  function getIdeContext(): IdeContext | undefined {
+  get(): IdeContext | undefined {
-    return ideContextState;
+    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 {
+  subscribe(subscriber: IdeContextSubscriber): () => void {
-    subscribers.add(subscriber);
+    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();

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;
 }