MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-03-11 06:31:01 -07:00
Code Issues Packages Projects Releases Wiki Activity

Refactor IdeContextStore (#8278)

Browse Source
This commit is contained in:
Tommaso Sciortino
2025-09-11 11:22:20 -07:00
committed by GitHub
parent 538e6cd19a
commit d892cde0b0
10 changed files with 92 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
packages/cli/src/ui/AppContainer.tsx
View File

@@ -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;
}, []);

4
packages/cli/src/ui/commands/ideCommand.ts
View File

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

5
packages/cli/src/ui/hooks/useIdeTrustListener.ts
View File

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

4
packages/core/src/config/config.ts
View File

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

34
packages/core/src/core/client.test.ts
View File

@@ -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(

4
packages/core/src/core/client.ts
View File

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

6
packages/core/src/ide/ide-client.ts
View File

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

83
packages/core/src/ide/ideContext.test.ts
View File

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

56
packages/core/src/ide/ideContext.ts
View File

@@ -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<IdeContextSubscriber>();
export class IdeContextStore {
private ideContextState?: IdeContext;
private readonly subscribers = new Set<IdeContextSubscriber>();
/**
* 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();

4
packages/core/src/utils/ide-trust.ts
View File

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