From 47603ef8e1a031c5d326769a71f6148fa6fca16a Mon Sep 17 00:00:00 2001 From: Jacob MacDonald Date: Fri, 7 Nov 2025 09:07:25 -0800 Subject: [PATCH] Reload gemini memory on extension load/unload + memory refresh refactor (#12651) --- packages/cli/src/config/config.ts | 62 ++------------ packages/cli/src/ui/AppContainer.tsx | 44 ++++------ .../cli/src/ui/commands/directoryCommand.tsx | 22 +---- .../cli/src/ui/commands/memoryCommand.test.ts | 49 +++-------- packages/cli/src/ui/commands/memoryCommand.ts | 27 ++---- packages/cli/src/ui/commands/types.ts | 1 - .../ui/hooks/slashCommandProcessor.test.tsx | 1 - .../cli/src/ui/hooks/slashCommandProcessor.ts | 3 - .../src/ui/noninteractive/nonInteractiveUi.ts | 1 - packages/core/src/config/config.ts | 14 ++++ packages/core/src/utils/events.ts | 73 ++++------------ .../core/src/utils/extensionLoader.test.ts | 84 ++++++++++++++----- packages/core/src/utils/extensionLoader.ts | 57 ++++++++++--- .../core/src/utils/memoryDiscovery.test.ts | 59 ++++++++++++- packages/core/src/utils/memoryDiscovery.ts | 38 +++++++++ 15 files changed, 276 insertions(+), 259 deletions(-) diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index a241a27d51..73c2398d6b 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -4,21 +4,14 @@ * SPDX-License-Identifier: Apache-2.0 */ -import * as fs from 'node:fs'; -import * as path from 'node:path'; -import { homedir } from 'node:os'; import yargs from 'yargs/yargs'; import { hideBin } from 'yargs/helpers'; import process from 'node:process'; import { mcpCommand } from '../commands/mcp.js'; -import type { - FileFilteringOptions, - OutputFormat, -} from '@google/gemini-cli-core'; +import type { OutputFormat } from '@google/gemini-cli-core'; import { extensionsCommand } from '../commands/extensions.js'; import { Config, - loadServerHierarchicalMemory, setGeminiMdFilename as setServerGeminiMdFilename, getCurrentGeminiMdFilename, ApprovalMode, @@ -36,6 +29,7 @@ import { getPty, EDIT_TOOL_NAME, debugLogger, + loadServerHierarchicalMemory, } from '@google/gemini-cli-core'; import type { Settings } from './settings.js'; @@ -47,10 +41,7 @@ import { appEvents } from '../utils/events.js'; import { isWorkspaceTrusted } from './trustedFolders.js'; import { createPolicyEngineConfig } from './policy.js'; import { ExtensionManager } from './extension-manager.js'; -import type { - ExtensionEvents, - ExtensionLoader, -} from '@google/gemini-cli-core/src/utils/extensionLoader.js'; +import type { ExtensionEvents } from '@google/gemini-cli-core/src/utils/extensionLoader.js'; import { requestConsentNonInteractive } from './extensions/consent.js'; import { promptForSetting } from './extensions/extensionSettings.js'; import type { EventEmitter } from 'node:stream'; @@ -297,49 +288,6 @@ export async function parseArguments(settings: Settings): Promise { return result as unknown as CliArgs; } -// This function is now a thin wrapper around the server's implementation. -// It's kept in the CLI for now as App.tsx directly calls it for memory refresh. -// TODO: Consider if App.tsx should get memory via a server call or if Config should refresh itself. -export async function loadHierarchicalGeminiMemory( - currentWorkingDirectory: string, - includeDirectoriesToReadGemini: readonly string[] = [], - debugMode: boolean, - fileService: FileDiscoveryService, - settings: Settings, - extensionLoader: ExtensionLoader, - folderTrust: boolean, - memoryImportFormat: 'flat' | 'tree' = 'tree', - fileFilteringOptions?: FileFilteringOptions, -): Promise<{ memoryContent: string; fileCount: number; filePaths: string[] }> { - // FIX: Use real, canonical paths for a reliable comparison to handle symlinks. - const realCwd = fs.realpathSync(path.resolve(currentWorkingDirectory)); - const realHome = fs.realpathSync(path.resolve(homedir())); - const isHomeDirectory = realCwd === realHome; - - // If it is the home directory, pass an empty string to the core memory - // function to signal that it should skip the workspace search. - const effectiveCwd = isHomeDirectory ? '' : currentWorkingDirectory; - - if (debugMode) { - debugLogger.debug( - `CLI: Delegating hierarchical memory load to server for CWD: ${currentWorkingDirectory} (memoryImportFormat: ${memoryImportFormat})`, - ); - } - - // Directly call the server function with the corrected path. - return loadServerHierarchicalMemory( - effectiveCwd, - includeDirectoriesToReadGemini, - debugMode, - fileService, - extensionLoader, - folderTrust, - memoryImportFormat, - fileFilteringOptions, - settings.context?.discoveryMaxDirs, - ); -} - /** * Creates a filter function to determine if a tool should be excluded. * @@ -437,18 +385,18 @@ export async function loadCliConfig( // Call the (now wrapper) loadHierarchicalGeminiMemory which calls the server's version const { memoryContent, fileCount, filePaths } = - await loadHierarchicalGeminiMemory( + await loadServerHierarchicalMemory( cwd, settings.context?.loadMemoryFromIncludeDirectories ? includeDirectories : [], debugMode, fileService, - settings, extensionManager, trustedFolder, memoryImportFormat, memoryFileFiltering, + settings.context?.discoveryMaxDirs, ); const question = argv.promptInteractive || argv.prompt || ''; diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 79e1bcdaa9..0eee49d44e 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -48,10 +48,11 @@ import { debugLogger, coreEvents, CoreEvent, + refreshServerHierarchicalMemory, type ModelChangedPayload, + type MemoryChangedPayload, } from '@google/gemini-cli-core'; import { validateAuthMethod } from '../config/auth.js'; -import { loadHierarchicalGeminiMemory } from '../config/config.js'; import process from 'node:process'; import { useHistory } from './hooks/useHistoryManager.js'; import { useMemoryMonitor } from './hooks/useMemoryMonitor.js'; @@ -160,9 +161,6 @@ export const AppContainer = (props: AppContainerProps) => { const [showDebugProfiler, setShowDebugProfiler] = useState(false); const [copyModeEnabled, setCopyModeEnabled] = useState(false); - const [geminiMdFileCount, setGeminiMdFileCount] = useState( - initializationResult.geminiMdFileCount, - ); const [shellModeActive, setShellModeActive] = useState(false); const [modelSwitchedFromQuotaError, setModelSwitchedFromQuotaError] = useState(false); @@ -569,7 +567,6 @@ Logging in with Google... Please restart Gemini CLI to continue. refreshStatic, toggleVimEnabled, setIsProcessing, - setGeminiMdFileCount, slashCommandActions, extensionsUpdateStateInternal, isConfigInitialized, @@ -584,26 +581,8 @@ Logging in with Google... Please restart Gemini CLI to continue. Date.now(), ); try { - const { memoryContent, fileCount, filePaths } = - await loadHierarchicalGeminiMemory( - process.cwd(), - settings.merged.context?.loadMemoryFromIncludeDirectories - ? config.getWorkspaceContext().getDirectories() - : [], - config.getDebugMode(), - config.getFileService(), - settings.merged, - config.getExtensionLoader(), - config.isTrustedFolder(), - settings.merged.context?.importFormat || 'tree', // Use setting or default to 'tree' - config.getFileFilteringOptions(), - ); - - config.setUserMemory(memoryContent); - config.setGeminiMdFileCount(fileCount); - config.setGeminiMdFilePaths(filePaths); - - setGeminiMdFileCount(fileCount); + const { memoryContent, fileCount } = + await refreshServerHierarchicalMemory(config); historyManager.addItem( { @@ -635,7 +614,7 @@ Logging in with Google... Please restart Gemini CLI to continue. ); debugLogger.warn('Error refreshing memory:', error); } - }, [config, historyManager, settings.merged]); + }, [config, historyManager]); const cancelHandlerRef = useRef<() => void>(() => {}); @@ -1248,6 +1227,19 @@ Logging in with Google... Please restart Gemini CLI to continue. [pendingSlashCommandHistoryItems, pendingGeminiHistoryItems], ); + const [geminiMdFileCount, setGeminiMdFileCount] = useState( + config.getGeminiMdFileCount(), + ); + useEffect(() => { + const handleMemoryChanged = (result: MemoryChangedPayload) => { + setGeminiMdFileCount(result.fileCount); + }; + coreEvents.on(CoreEvent.MemoryChanged, handleMemoryChanged); + return () => { + coreEvents.off(CoreEvent.MemoryChanged, handleMemoryChanged); + }; + }, []); + const uiState: UIState = useMemo( () => ({ history: historyManager.history, diff --git a/packages/cli/src/ui/commands/directoryCommand.tsx b/packages/cli/src/ui/commands/directoryCommand.tsx index ee078356c5..fe7076af0d 100644 --- a/packages/cli/src/ui/commands/directoryCommand.tsx +++ b/packages/cli/src/ui/commands/directoryCommand.tsx @@ -9,7 +9,7 @@ import { CommandKind } from './types.js'; import { MessageType } from '../types.js'; import * as os from 'node:os'; import * as path from 'node:path'; -import { loadServerHierarchicalMemory } from '@google/gemini-cli-core'; +import { refreshServerHierarchicalMemory } from '@google/gemini-cli-core'; export function expandHomeDir(p: string): string { if (!p) { @@ -94,25 +94,7 @@ export const directoryCommand: SlashCommand = { try { if (config.shouldLoadMemoryFromIncludeDirectories()) { - const { memoryContent, fileCount } = - await loadServerHierarchicalMemory( - config.getWorkingDir(), - [ - ...config.getWorkspaceContext().getDirectories(), - ...pathsToAdd, - ], - config.getDebugMode(), - config.getFileService(), - config.getExtensionLoader(), - config.getFolderTrust(), - context.services.settings.merged.context?.importFormat || - 'tree', // Use setting or default to 'tree' - config.getFileFilteringOptions(), - context.services.settings.merged.context?.discoveryMaxDirs, - ); - config.setUserMemory(memoryContent); - config.setGeminiMdFileCount(fileCount); - context.ui.setGeminiMdFileCount(fileCount); + await refreshServerHierarchicalMemory(config); } addItem( { diff --git a/packages/cli/src/ui/commands/memoryCommand.test.ts b/packages/cli/src/ui/commands/memoryCommand.test.ts index 523e0be0f1..f555a9c925 100644 --- a/packages/cli/src/ui/commands/memoryCommand.test.ts +++ b/packages/cli/src/ui/commands/memoryCommand.test.ts @@ -13,11 +13,11 @@ import { MessageType } from '../types.js'; import type { LoadedSettings } from '../../config/settings.js'; import { getErrorMessage, + refreshServerHierarchicalMemory, SimpleExtensionLoader, type FileDiscoveryService, } from '@google/gemini-cli-core'; import type { LoadServerHierarchicalMemoryResponse } from '@google/gemini-cli-core/index.js'; -import { loadHierarchicalGeminiMemory } from '../../config/config.js'; vi.mock('@google/gemini-cli-core', async (importOriginal) => { const original = @@ -28,19 +28,12 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { if (error instanceof Error) return error.message; return String(error); }), + refreshServerHierarchicalMemory: vi.fn(), }; }); -vi.mock('../../config/config.js', async (importOriginal) => { - const original = - await importOriginal(); - return { - ...original, - loadHierarchicalGeminiMemory: vi.fn(), - }; -}); - -const mockLoadHierarchicalGeminiMemory = loadHierarchicalGeminiMemory as Mock; +const mockRefreshServerHierarchicalMemory = + refreshServerHierarchicalMemory as Mock; describe('memoryCommand', () => { let mockContext: CommandContext; @@ -203,11 +196,8 @@ describe('memoryCommand', () => { }, } as unknown as LoadedSettings, }, - ui: { - setGeminiMdFileCount: vi.fn(), - }, }); - mockLoadHierarchicalGeminiMemory.mockClear(); + mockRefreshServerHierarchicalMemory.mockClear(); }); it('should display success message when memory is refreshed with content', async () => { @@ -218,7 +208,7 @@ describe('memoryCommand', () => { fileCount: 2, filePaths: ['/path/one/GEMINI.md', '/path/two/GEMINI.md'], }; - mockLoadHierarchicalGeminiMemory.mockResolvedValue(refreshResult); + mockRefreshServerHierarchicalMemory.mockResolvedValue(refreshResult); await refreshCommand.action(mockContext, ''); @@ -230,19 +220,7 @@ describe('memoryCommand', () => { expect.any(Number), ); - expect(mockLoadHierarchicalGeminiMemory).toHaveBeenCalledOnce(); - expect(mockSetUserMemory).toHaveBeenCalledWith( - refreshResult.memoryContent, - ); - expect(mockSetGeminiMdFileCount).toHaveBeenCalledWith( - refreshResult.fileCount, - ); - expect(mockSetGeminiMdFilePaths).toHaveBeenCalledWith( - refreshResult.filePaths, - ); - expect(mockContext.ui.setGeminiMdFileCount).toHaveBeenCalledWith( - refreshResult.fileCount, - ); + expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce(); expect(mockContext.ui.addItem).toHaveBeenCalledWith( { @@ -257,14 +235,11 @@ describe('memoryCommand', () => { if (!refreshCommand.action) throw new Error('Command has no action'); const refreshResult = { memoryContent: '', fileCount: 0, filePaths: [] }; - mockLoadHierarchicalGeminiMemory.mockResolvedValue(refreshResult); + mockRefreshServerHierarchicalMemory.mockResolvedValue(refreshResult); await refreshCommand.action(mockContext, ''); - expect(mockLoadHierarchicalGeminiMemory).toHaveBeenCalledOnce(); - expect(mockSetUserMemory).toHaveBeenCalledWith(''); - expect(mockSetGeminiMdFileCount).toHaveBeenCalledWith(0); - expect(mockSetGeminiMdFilePaths).toHaveBeenCalledWith([]); + expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce(); expect(mockContext.ui.addItem).toHaveBeenCalledWith( { @@ -279,11 +254,11 @@ describe('memoryCommand', () => { if (!refreshCommand.action) throw new Error('Command has no action'); const error = new Error('Failed to read memory files.'); - mockLoadHierarchicalGeminiMemory.mockRejectedValue(error); + mockRefreshServerHierarchicalMemory.mockRejectedValue(error); await refreshCommand.action(mockContext, ''); - expect(mockLoadHierarchicalGeminiMemory).toHaveBeenCalledOnce(); + expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce(); expect(mockSetUserMemory).not.toHaveBeenCalled(); expect(mockSetGeminiMdFileCount).not.toHaveBeenCalled(); expect(mockSetGeminiMdFilePaths).not.toHaveBeenCalled(); @@ -318,7 +293,7 @@ describe('memoryCommand', () => { expect.any(Number), ); - expect(mockLoadHierarchicalGeminiMemory).not.toHaveBeenCalled(); + expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled(); }); }); diff --git a/packages/cli/src/ui/commands/memoryCommand.ts b/packages/cli/src/ui/commands/memoryCommand.ts index ffe04fbe08..61d0ec4467 100644 --- a/packages/cli/src/ui/commands/memoryCommand.ts +++ b/packages/cli/src/ui/commands/memoryCommand.ts @@ -4,9 +4,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { getErrorMessage } from '@google/gemini-cli-core'; +import { + getErrorMessage, + refreshServerHierarchicalMemory, +} from '@google/gemini-cli-core'; import { MessageType } from '../types.js'; -import { loadHierarchicalGeminiMemory } from '../../config/config.js'; import type { SlashCommand, SlashCommandActionReturn } from './types.js'; import { CommandKind } from './types.js'; @@ -80,26 +82,9 @@ export const memoryCommand: SlashCommand = { try { const config = await context.services.config; - const settings = context.services.settings; if (config) { - const { memoryContent, fileCount, filePaths } = - await loadHierarchicalGeminiMemory( - config.getWorkingDir(), - config.shouldLoadMemoryFromIncludeDirectories() - ? config.getWorkspaceContext().getDirectories() - : [], - config.getDebugMode(), - config.getFileService(), - settings.merged, - config.getExtensionLoader(), - config.isTrustedFolder(), - settings.merged.context?.importFormat || 'tree', - config.getFileFilteringOptions(), - ); - config.setUserMemory(memoryContent); - config.setGeminiMdFileCount(fileCount); - config.setGeminiMdFilePaths(filePaths); - context.ui.setGeminiMdFileCount(fileCount); + const { memoryContent, fileCount } = + await refreshServerHierarchicalMemory(config); const successMessage = memoryContent.length > 0 diff --git a/packages/cli/src/ui/commands/types.ts b/packages/cli/src/ui/commands/types.ts index 99e514fbba..5bd77d2ddf 100644 --- a/packages/cli/src/ui/commands/types.ts +++ b/packages/cli/src/ui/commands/types.ts @@ -68,7 +68,6 @@ export interface CommandContext { toggleCorgiMode: () => void; toggleDebugProfiler: () => void; toggleVimEnabled: () => Promise; - setGeminiMdFileCount: (count: number) => void; reloadCommands: () => void; extensionsUpdateState: Map; dispatchExtensionStateUpdate: (action: ExtensionUpdateAction) => void; diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.tsx b/packages/cli/src/ui/hooks/slashCommandProcessor.test.tsx index 7fa0db1852..4b247c9361 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.tsx +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.tsx @@ -164,7 +164,6 @@ describe('useSlashCommandProcessor', () => { vi.fn(), // refreshStatic vi.fn(), // toggleVimEnabled setIsProcessing, - vi.fn(), // setGeminiMdFileCount { openAuthDialog: mockOpenAuthDialog, openThemeDialog: mockOpenThemeDialog, diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index fe8be20012..0fc0570188 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -73,7 +73,6 @@ export const useSlashCommandProcessor = ( refreshStatic: () => void, toggleVimEnabled: () => Promise, setIsProcessing: (isProcessing: boolean) => void, - setGeminiMdFileCount: (count: number) => void, actions: SlashCommandProcessorActions, extensionsUpdateState: Map, isConfigInitialized: boolean, @@ -207,7 +206,6 @@ export const useSlashCommandProcessor = ( toggleCorgiMode: actions.toggleCorgiMode, toggleDebugProfiler: actions.toggleDebugProfiler, toggleVimEnabled, - setGeminiMdFileCount, reloadCommands, extensionsUpdateState, dispatchExtensionStateUpdate: actions.dispatchExtensionStateUpdate, @@ -234,7 +232,6 @@ export const useSlashCommandProcessor = ( setPendingItem, toggleVimEnabled, sessionShellAllowlist, - setGeminiMdFileCount, reloadCommands, extensionsUpdateState, ], diff --git a/packages/cli/src/ui/noninteractive/nonInteractiveUi.ts b/packages/cli/src/ui/noninteractive/nonInteractiveUi.ts index 0b3de8bc49..c8bfcf4276 100644 --- a/packages/cli/src/ui/noninteractive/nonInteractiveUi.ts +++ b/packages/cli/src/ui/noninteractive/nonInteractiveUi.ts @@ -23,7 +23,6 @@ export function createNonInteractiveUI(): CommandContext['ui'] { toggleCorgiMode: () => {}, toggleDebugProfiler: () => {}, toggleVimEnabled: async () => false, - setGeminiMdFileCount: (_count) => {}, reloadCommands: () => {}, extensionsUpdateState: new Map(), dispatchExtensionStateUpdate: (_action: ExtensionUpdateAction) => {}, diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index cbe26bf7d3..b53530c654 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -266,6 +266,8 @@ export interface ConfigParameters { folderTrust?: boolean; ideMode?: boolean; loadMemoryFromIncludeDirectories?: boolean; + importFormat?: 'tree' | 'flat'; + discoveryMaxDirs?: number; compressionThreshold?: number; interactive?: boolean; trustedFolder?: boolean; @@ -369,6 +371,8 @@ export class Config { | undefined; private readonly experimentalZedIntegration: boolean = false; private readonly loadMemoryFromIncludeDirectories: boolean = false; + private readonly importFormat: 'tree' | 'flat'; + private readonly discoveryMaxDirs: number; private readonly compressionThreshold: number | undefined; private readonly interactive: boolean; private readonly ptyInfo: string; @@ -479,6 +483,8 @@ export class Config { this.ideMode = params.ideMode ?? false; this.loadMemoryFromIncludeDirectories = params.loadMemoryFromIncludeDirectories ?? false; + this.importFormat = params.importFormat ?? 'tree'; + this.discoveryMaxDirs = params.discoveryMaxDirs ?? 200; this.compressionThreshold = params.compressionThreshold; this.interactive = params.interactive ?? false; this.ptyInfo = params.ptyInfo ?? 'child_process'; @@ -707,6 +713,14 @@ export class Config { return this.loadMemoryFromIncludeDirectories; } + getImportFormat(): 'tree' | 'flat' { + return this.importFormat; + } + + getDiscoveryMaxDirs(): number { + return this.discoveryMaxDirs; + } + getContentGeneratorConfig(): ContentGeneratorConfig { return this.contentGeneratorConfig; } diff --git a/packages/core/src/utils/events.ts b/packages/core/src/utils/events.ts index 386200fad7..3df9c39a99 100644 --- a/packages/core/src/utils/events.ts +++ b/packages/core/src/utils/events.ts @@ -5,6 +5,7 @@ */ import { EventEmitter } from 'node:events'; +import type { LoadServerHierarchicalMemoryResponse } from './memoryDiscovery.js'; /** * Defines the severity level for user-facing feedback. @@ -53,13 +54,26 @@ export interface ModelChangedPayload { model: string; } +/** + * Payload for the 'memory-changed' event. + */ +export type MemoryChangedPayload = LoadServerHierarchicalMemoryResponse; + export enum CoreEvent { UserFeedback = 'user-feedback', FallbackModeChanged = 'fallback-mode-changed', ModelChanged = 'model-changed', + MemoryChanged = 'memory-changed', } -export class CoreEventEmitter extends EventEmitter { +export interface CoreEvents { + [CoreEvent.UserFeedback]: [UserFeedbackPayload]; + [CoreEvent.FallbackModeChanged]: [FallbackModeChangedPayload]; + [CoreEvent.ModelChanged]: [ModelChangedPayload]; + [CoreEvent.MemoryChanged]: [MemoryChangedPayload]; +} + +export class CoreEventEmitter extends EventEmitter { private _feedbackBacklog: UserFeedbackPayload[] = []; private static readonly MAX_BACKLOG_SIZE = 10000; @@ -116,63 +130,6 @@ export class CoreEventEmitter extends EventEmitter { this.emit(CoreEvent.UserFeedback, payload); } } - - override on( - event: CoreEvent.UserFeedback, - listener: (payload: UserFeedbackPayload) => void, - ): this; - override on( - event: CoreEvent.FallbackModeChanged, - listener: (payload: FallbackModeChangedPayload) => void, - ): this; - override on( - event: CoreEvent.ModelChanged, - listener: (payload: ModelChangedPayload) => void, - ): this; - override on( - event: string | symbol, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - listener: (...args: any[]) => void, - ): this { - return super.on(event, listener); - } - - override off( - event: CoreEvent.UserFeedback, - listener: (payload: UserFeedbackPayload) => void, - ): this; - override off( - event: CoreEvent.FallbackModeChanged, - listener: (payload: FallbackModeChangedPayload) => void, - ): this; - override off( - event: CoreEvent.ModelChanged, - listener: (payload: ModelChangedPayload) => void, - ): this; - override off( - event: string | symbol, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - listener: (...args: any[]) => void, - ): this { - return super.off(event, listener); - } - - override emit( - event: CoreEvent.UserFeedback, - payload: UserFeedbackPayload, - ): boolean; - override emit( - event: CoreEvent.FallbackModeChanged, - payload: FallbackModeChangedPayload, - ): boolean; - override emit( - event: CoreEvent.ModelChanged, - payload: ModelChangedPayload, - ): boolean; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - override emit(event: string | symbol, ...args: any[]): boolean { - return super.emit(event, ...args); - } } export const coreEvents = new CoreEventEmitter(); diff --git a/packages/core/src/utils/extensionLoader.test.ts b/packages/core/src/utils/extensionLoader.test.ts index 3329dc36f7..cb175b2f64 100644 --- a/packages/core/src/utils/extensionLoader.test.ts +++ b/packages/core/src/utils/extensionLoader.test.ts @@ -9,6 +9,16 @@ import { SimpleExtensionLoader } from './extensionLoader.js'; import type { Config } from '../config/config.js'; import { type McpClientManager } from '../tools/mcp-client-manager.js'; +const mockRefreshServerHierarchicalMemory = vi.hoisted(() => vi.fn()); + +vi.mock('./memoryDiscovery.js', async (importActual) => { + const actual = await importActual(); + return { + ...actual, + refreshServerHierarchicalMemory: mockRefreshServerHierarchicalMemory, + }; +}); + describe('SimpleExtensionLoader', () => { let mockConfig: Config; let extensionReloadingEnabled: boolean; @@ -79,29 +89,59 @@ describe('SimpleExtensionLoader', () => { ).toHaveBeenCalledExactlyOnceWith(activeExtension); }); - it.each([true, false])( - 'should only call `start` and `stop` if extension reloading is enabled ($i)', - async (reloadingEnabled) => { - extensionReloadingEnabled = reloadingEnabled; - const loader = new SimpleExtensionLoader([]); - await loader.start(mockConfig); - expect(mockMcpClientManager.startExtension).not.toHaveBeenCalled(); - await loader.loadExtension(activeExtension); - if (reloadingEnabled) { - expect( - mockMcpClientManager.startExtension, - ).toHaveBeenCalledExactlyOnceWith(activeExtension); - } else { + describe.each([true, false])( + 'when enableExtensionReloading === $i', + (reloadingEnabled) => { + beforeEach(() => { + extensionReloadingEnabled = reloadingEnabled; + }); + + it(`should ${reloadingEnabled ? '' : 'not '}reload extension features`, async () => { + const loader = new SimpleExtensionLoader([]); + await loader.start(mockConfig); expect(mockMcpClientManager.startExtension).not.toHaveBeenCalled(); - } - await loader.unloadExtension(activeExtension); - if (reloadingEnabled) { - expect( - mockMcpClientManager.stopExtension, - ).toHaveBeenCalledExactlyOnceWith(activeExtension); - } else { - expect(mockMcpClientManager.stopExtension).not.toHaveBeenCalled(); - } + await loader.loadExtension(activeExtension); + if (reloadingEnabled) { + expect( + mockMcpClientManager.startExtension, + ).toHaveBeenCalledExactlyOnceWith(activeExtension); + expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce(); + } else { + expect(mockMcpClientManager.startExtension).not.toHaveBeenCalled(); + expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled(); + } + mockRefreshServerHierarchicalMemory.mockClear(); + + await loader.unloadExtension(activeExtension); + if (reloadingEnabled) { + expect( + mockMcpClientManager.stopExtension, + ).toHaveBeenCalledExactlyOnceWith(activeExtension); + expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce(); + } else { + expect(mockMcpClientManager.stopExtension).not.toHaveBeenCalled(); + expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled(); + } + }); + + it.runIf(reloadingEnabled)( + 'Should only reload memory once all extensions are done', + async () => { + const anotherExtension = { + ...activeExtension, + name: 'another-extension', + }; + const loader = new SimpleExtensionLoader([]); + await loader.loadExtension(activeExtension); + await loader.start(mockConfig); + expect(mockRefreshServerHierarchicalMemory).not.toHaveBeenCalled(); + await Promise.all([ + loader.unloadExtension(activeExtension), + loader.loadExtension(anotherExtension), + ]); + expect(mockRefreshServerHierarchicalMemory).toHaveBeenCalledOnce(); + }, + ); }, ); }); diff --git a/packages/core/src/utils/extensionLoader.ts b/packages/core/src/utils/extensionLoader.ts index f47ff95015..411ff6a7d1 100644 --- a/packages/core/src/utils/extensionLoader.ts +++ b/packages/core/src/utils/extensionLoader.ts @@ -6,6 +6,7 @@ import type { EventEmitter } from 'node:events'; import type { Config, GeminiCLIExtension } from '../config/config.js'; +import { refreshServerHierarchicalMemory } from './memoryDiscovery.js'; export abstract class ExtensionLoader { // Assigned in `start`. @@ -18,6 +19,9 @@ export abstract class ExtensionLoader { protected stoppingCount: number = 0; protected stopCompletedCount: number = 0; + // Whether or not we are currently executing `start` + private isStarting: boolean = false; + constructor(private readonly eventEmitter?: EventEmitter) {} /** @@ -32,16 +36,21 @@ export abstract class ExtensionLoader { * McpClientManager, PromptRegistry, and GeminiChat set up. */ async start(config: Config): Promise { - if (!this.config) { - this.config = config; - } else { - throw new Error('Already started, you may only call `start` once.'); + this.isStarting = true; + try { + if (!this.config) { + this.config = config; + } else { + throw new Error('Already started, you may only call `start` once.'); + } + await Promise.all( + this.getExtensions() + .filter((e) => e.isActive) + .map(this.startExtension.bind(this)), + ); + } finally { + this.isStarting = false; } - await Promise.all( - this.getExtensions() - .filter((e) => e.isActive) - .map(this.startExtension.bind(this)), - ); } /** @@ -64,12 +73,15 @@ export abstract class ExtensionLoader { }); try { await this.config.getMcpClientManager()!.startExtension(extension); + // Note: Context files are loaded only once all extensions are done + // loading/unloading to reduce churn, see the `maybeRefreshMemories` call + // below. + // TODO: Update custom command updating away from the event based system // and call directly into a custom command manager here. See the // useSlashCommandProcessor hook which responds to events fired here today. // TODO: Move all enablement of extension features here, including at least: - // - context file loading // - excluded tool configuration } finally { this.startCompletedCount++; @@ -81,6 +93,25 @@ export abstract class ExtensionLoader { this.startingCount = 0; this.startCompletedCount = 0; } + await this.maybeRefreshMemories(); + } + } + + private async maybeRefreshMemories(): Promise { + if (!this.config) { + throw new Error( + 'Cannot refresh gemini memories prior to calling `start`.', + ); + } + if ( + !this.isStarting && // Don't refresh memories on the first call to `start`. + this.startingCount === this.startCompletedCount && + this.stoppingCount === this.stopCompletedCount + ) { + // Wait until all extensions are done starting and stopping before we + // reload memory, this is somewhat expensive and also busts the context + // cache, we want to only do it once. + await refreshServerHierarchicalMemory(this.config); } } @@ -119,12 +150,15 @@ export abstract class ExtensionLoader { try { await this.config.getMcpClientManager()!.stopExtension(extension); + // Note: Context files are loaded only once all extensions are done + // loading/unloading to reduce churn, see the `maybeRefreshMemories` call + // below. + // TODO: Update custom command updating away from the event based system // and call directly into a custom command manager here. See the // useSlashCommandProcessor hook which responds to events fired here today. // TODO: Remove all extension features here, including at least: - // - context files // - excluded tools } finally { this.stopCompletedCount++; @@ -136,6 +170,7 @@ export abstract class ExtensionLoader { this.stoppingCount = 0; this.stopCompletedCount = 0; } + await this.maybeRefreshMemories(); } } diff --git a/packages/core/src/utils/memoryDiscovery.test.ts b/packages/core/src/utils/memoryDiscovery.test.ts index fba1901f43..fd34310afe 100644 --- a/packages/core/src/utils/memoryDiscovery.test.ts +++ b/packages/core/src/utils/memoryDiscovery.test.ts @@ -13,6 +13,7 @@ import { loadGlobalMemory, loadEnvironmentMemory, loadJitSubdirectoryMemory, + refreshServerHierarchicalMemory, } from './memoryDiscovery.js'; import { setGeminiMdFilename, @@ -20,8 +21,10 @@ import { } from '../tools/memoryTool.js'; import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; import { GEMINI_DIR } from './paths.js'; -import type { GeminiCLIExtension } from '../config/config.js'; +import { Config, type GeminiCLIExtension } from '../config/config.js'; +import { Storage } from '../config/storage.js'; import { SimpleExtensionLoader } from './extensionLoader.js'; +import { CoreEvent, coreEvents } from './events.js'; vi.mock('os', async (importOriginal) => { const actualOs = await importOriginal(); @@ -876,4 +879,58 @@ included directory memory expect(result.files.find((f) => f.path === outerMemory)).toBeUndefined(); }); }); + + it('refreshServerHierarchicalMemory should refresh memory and update config', async () => { + const extensionLoader = new SimpleExtensionLoader([]); + const config = new Config({ + sessionId: '1', + targetDir: cwd, + cwd, + debugMode: false, + model: 'fake-model', + extensionLoader, + }); + const result = await loadServerHierarchicalMemory( + config.getWorkingDir(), + config.shouldLoadMemoryFromIncludeDirectories() + ? config.getWorkspaceContext().getDirectories() + : [], + config.getDebugMode(), + config.getFileService(), + config.getExtensionLoader(), + config.isTrustedFolder(), + config.getImportFormat(), + ); + expect(result.fileCount).equals(0); + + // Now add an extension with a memory file + const extensionsDir = new Storage(homedir).getExtensionsDir(); + const extensionPath = path.join(extensionsDir, 'new-extension'); + const contextFilePath = path.join(extensionPath, 'CustomContext.md'); + await fsPromises.mkdir(extensionPath, { recursive: true }); + await fsPromises.writeFile(contextFilePath, 'Really cool custom context!'); + await extensionLoader.loadExtension({ + name: 'new-extension', + isActive: true, + contextFiles: [contextFilePath], + version: '1.0.0', + id: '1234', + path: extensionPath, + }); + + const mockEventListener = vi.fn(); + coreEvents.on(CoreEvent.MemoryChanged, mockEventListener); + const refreshResult = await refreshServerHierarchicalMemory(config); + expect(refreshResult.fileCount).equals(1); + expect(config.getGeminiMdFileCount()).equals(refreshResult.fileCount); + expect(refreshResult.memoryContent).toContain( + 'Really cool custom context!', + ); + expect(config.getUserMemory()).equals(refreshResult.memoryContent); + expect(refreshResult.filePaths[0]).toContain( + path.join(extensionPath, 'CustomContext.md'), + ); + expect(config.getGeminiMdFilePaths()).equals(refreshResult.filePaths); + expect(mockEventListener).toHaveBeenCalledExactlyOnceWith(refreshResult); + }); }); diff --git a/packages/core/src/utils/memoryDiscovery.ts b/packages/core/src/utils/memoryDiscovery.ts index bd98f0ab01..6a74abc8a6 100644 --- a/packages/core/src/utils/memoryDiscovery.ts +++ b/packages/core/src/utils/memoryDiscovery.ts @@ -17,6 +17,8 @@ import { DEFAULT_MEMORY_FILE_FILTERING_OPTIONS } from '../config/constants.js'; import { GEMINI_DIR } from './paths.js'; import type { ExtensionLoader } from './extensionLoader.js'; import { debugLogger } from './debugLogger.js'; +import type { Config } from '../config/config.js'; +import { CoreEvent, coreEvents } from './events.js'; // Simple console logger, similar to the one previously in CLI's config.ts // TODO: Integrate with a more robust server-side logger if available/appropriate. @@ -481,6 +483,15 @@ export async function loadServerHierarchicalMemory( fileFilteringOptions?: FileFilteringOptions, maxDirs: number = 200, ): Promise { + // FIX: Use real, canonical paths for a reliable comparison to handle symlinks. + const realCwd = await fs.realpath(path.resolve(currentWorkingDirectory)); + const realHome = await fs.realpath(path.resolve(homedir())); + const isHomeDirectory = realCwd === realHome; + + // If it is the home directory, pass an empty string to the core memory + // function to signal that it should skip the workspace search. + currentWorkingDirectory = isHomeDirectory ? '' : currentWorkingDirectory; + if (debugMode) logger.debug( `Loading server hierarchical memory for CWD: ${currentWorkingDirectory} (importFormat: ${importFormat})`, @@ -538,6 +549,33 @@ export async function loadServerHierarchicalMemory( }; } +/** + * Loads the hierarchical memory and resets the state of `config` as needed such + * that it reflects the new memory. + * + * Returns the result of the call to `loadHierarchicalGeminiMemory`. + */ +export async function refreshServerHierarchicalMemory(config: Config) { + const result = await loadServerHierarchicalMemory( + config.getWorkingDir(), + config.shouldLoadMemoryFromIncludeDirectories() + ? config.getWorkspaceContext().getDirectories() + : [], + config.getDebugMode(), + config.getFileService(), + config.getExtensionLoader(), + config.isTrustedFolder(), + config.getImportFormat(), + config.getFileFilteringOptions(), + config.getDiscoveryMaxDirs(), + ); + config.setUserMemory(result.memoryContent); + config.setGeminiMdFileCount(result.fileCount); + config.setGeminiMdFilePaths(result.filePaths); + coreEvents.emit(CoreEvent.MemoryChanged, result); + return result; +} + export async function loadJitSubdirectoryMemory( targetPath: string, trustedRoots: string[],