From d7a304bcffc7d2340f3de762f14286596954ce4e Mon Sep 17 00:00:00 2001 From: Brandon Keiji Date: Wed, 23 Jul 2025 14:48:35 -0700 Subject: [PATCH] feat(memory): make directory search limit on memory discovery configurable with settings.json (#4460) --- docs/cli/configuration.md | 2 +- packages/cli/src/config/config.test.ts | 3 +- packages/cli/src/config/config.ts | 3 + packages/cli/src/config/settings.ts | 1 + packages/cli/src/ui/App.tsx | 3 +- .../cli/src/ui/commands/memoryCommand.test.ts | 64 +++++++++++++++---- packages/cli/src/ui/commands/memoryCommand.ts | 21 ++++-- packages/core/src/config/config.test.ts | 36 ----------- packages/core/src/config/config.ts | 16 ----- .../core/src/utils/memoryDiscovery.test.ts | 31 +++++++-- packages/core/src/utils/memoryDiscovery.ts | 7 +- 11 files changed, 104 insertions(+), 83 deletions(-) diff --git a/docs/cli/configuration.md b/docs/cli/configuration.md index 1ff178b8..ad0a0df0 100644 --- a/docs/cli/configuration.md +++ b/docs/cli/configuration.md @@ -429,7 +429,7 @@ This example demonstrates how you can provide general project context, specific - Location: The CLI searches for the configured context file in the current working directory and then in each parent directory up to either the project root (identified by a `.git` folder) or your home directory. - Scope: Provides context relevant to the entire project or a significant portion of it. 3. **Sub-directory Context Files (Contextual/Local):** - - Location: The CLI also scans for the configured context file in subdirectories _below_ the current working directory (respecting common ignore patterns like `node_modules`, `.git`, etc.). + - Location: The CLI also scans for the configured context file in subdirectories _below_ the current working directory (respecting common ignore patterns like `node_modules`, `.git`, etc.). The breadth of this search is limited to 200 directories by default, but can be configured with a `memoryDiscoveryMaxDirs` field in your `settings.json` file. - Scope: Allows for highly specific instructions relevant to a particular component, module, or subsection of your project. - **Concatenation & UI Indication:** The contents of all found context files are concatenated (with separators indicating their origin and path) and provided as part of the system prompt to the Gemini model. The CLI footer displays the count of loaded context files, giving you a quick visual cue about the active instructional context. - **Commands for Memory Management:** diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index 0c0761cc..c0e9c215 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -37,7 +37,7 @@ vi.mock('@google/gemini-cli-core', async () => { ...actualServer, loadEnvironment: vi.fn(), loadServerHierarchicalMemory: vi.fn( - (cwd, debug, fileService, extensionPaths) => + (cwd, debug, fileService, extensionPaths, _maxDirs) => Promise.resolve({ memoryContent: extensionPaths?.join(',') || '', fileCount: extensionPaths?.length || 0, @@ -491,6 +491,7 @@ describe('Hierarchical Memory Loading (config.ts) - Placeholder Suite', () => { respectGitIgnore: false, respectGeminiIgnore: true, }, + undefined, // maxDirs ); }); diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index ec84db52..650f3aa2 100644 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -225,6 +225,7 @@ export async function loadHierarchicalGeminiMemory( currentWorkingDirectory: string, debugMode: boolean, fileService: FileDiscoveryService, + settings: Settings, extensionContextFilePaths: string[] = [], fileFilteringOptions?: FileFilteringOptions, ): Promise<{ memoryContent: string; fileCount: number }> { @@ -242,6 +243,7 @@ export async function loadHierarchicalGeminiMemory( fileService, extensionContextFilePaths, fileFilteringOptions, + settings.memoryDiscoveryMaxDirs, ); } @@ -298,6 +300,7 @@ export async function loadCliConfig( process.cwd(), debugMode, fileService, + settings, extensionContextFilePaths, fileFiltering, ); diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts index bc2206a7..c8885d48 100644 --- a/packages/cli/src/config/settings.ts +++ b/packages/cli/src/config/settings.ts @@ -100,6 +100,7 @@ export interface Settings { // Add other settings here. ideMode?: boolean; + memoryDiscoveryMaxDirs?: number; } export interface SettingsError { diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index baf21395..566d6bd5 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -236,6 +236,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { process.cwd(), config.getDebugMode(), config.getFileService(), + settings.merged, config.getExtensionContextFilePaths(), config.getFileFilteringOptions(), ); @@ -267,7 +268,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { ); console.error('Error refreshing memory:', error); } - }, [config, addItem]); + }, [config, addItem, settings.merged]); // Watch for model changes (e.g., from Flash fallback) useEffect(() => { diff --git a/packages/cli/src/ui/commands/memoryCommand.test.ts b/packages/cli/src/ui/commands/memoryCommand.test.ts index 47d098b1..74614fa7 100644 --- a/packages/cli/src/ui/commands/memoryCommand.test.ts +++ b/packages/cli/src/ui/commands/memoryCommand.test.ts @@ -9,7 +9,12 @@ import { memoryCommand } from './memoryCommand.js'; import { type CommandContext, SlashCommand } from './types.js'; import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; import { MessageType } from '../types.js'; -import { getErrorMessage } from '@google/gemini-cli-core'; +import { LoadedSettings } from '../../config/settings.js'; +import { + getErrorMessage, + loadServerHierarchicalMemory, + type FileDiscoveryService, +} from '@google/gemini-cli-core'; vi.mock('@google/gemini-cli-core', async (importOriginal) => { const original = @@ -20,9 +25,12 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { if (error instanceof Error) return error.message; return String(error); }), + loadServerHierarchicalMemory: vi.fn(), }; }); +const mockLoadServerHierarchicalMemory = loadServerHierarchicalMemory as Mock; + describe('memoryCommand', () => { let mockContext: CommandContext; @@ -139,19 +147,37 @@ describe('memoryCommand', () => { describe('/memory refresh', () => { let refreshCommand: SlashCommand; - let mockRefreshMemory: Mock; + let mockSetUserMemory: Mock; + let mockSetGeminiMdFileCount: Mock; beforeEach(() => { refreshCommand = getSubCommand('refresh'); - mockRefreshMemory = vi.fn(); + mockSetUserMemory = vi.fn(); + mockSetGeminiMdFileCount = vi.fn(); + const mockConfig = { + setUserMemory: mockSetUserMemory, + setGeminiMdFileCount: mockSetGeminiMdFileCount, + getWorkingDir: () => '/test/dir', + getDebugMode: () => false, + getFileService: () => ({}) as FileDiscoveryService, + getExtensionContextFilePaths: () => [], + getFileFilteringOptions: () => ({ + ignore: [], + include: [], + }), + }; + mockContext = createMockCommandContext({ services: { - config: { - refreshMemory: mockRefreshMemory, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - } as any, + config: Promise.resolve(mockConfig), + settings: { + merged: { + memoryDiscoveryMaxDirs: 1000, + }, + } as LoadedSettings, }, }); + mockLoadServerHierarchicalMemory.mockClear(); }); it('should display success message when memory is refreshed with content', async () => { @@ -161,7 +187,7 @@ describe('memoryCommand', () => { memoryContent: 'new memory content', fileCount: 2, }; - mockRefreshMemory.mockResolvedValue(refreshResult); + mockLoadServerHierarchicalMemory.mockResolvedValue(refreshResult); await refreshCommand.action(mockContext, ''); @@ -173,7 +199,13 @@ describe('memoryCommand', () => { expect.any(Number), ); - expect(mockRefreshMemory).toHaveBeenCalledOnce(); + expect(loadServerHierarchicalMemory).toHaveBeenCalledOnce(); + expect(mockSetUserMemory).toHaveBeenCalledWith( + refreshResult.memoryContent, + ); + expect(mockSetGeminiMdFileCount).toHaveBeenCalledWith( + refreshResult.fileCount, + ); expect(mockContext.ui.addItem).toHaveBeenCalledWith( { @@ -188,11 +220,13 @@ describe('memoryCommand', () => { if (!refreshCommand.action) throw new Error('Command has no action'); const refreshResult = { memoryContent: '', fileCount: 0 }; - mockRefreshMemory.mockResolvedValue(refreshResult); + mockLoadServerHierarchicalMemory.mockResolvedValue(refreshResult); await refreshCommand.action(mockContext, ''); - expect(mockRefreshMemory).toHaveBeenCalledOnce(); + expect(loadServerHierarchicalMemory).toHaveBeenCalledOnce(); + expect(mockSetUserMemory).toHaveBeenCalledWith(''); + expect(mockSetGeminiMdFileCount).toHaveBeenCalledWith(0); expect(mockContext.ui.addItem).toHaveBeenCalledWith( { @@ -207,11 +241,13 @@ describe('memoryCommand', () => { if (!refreshCommand.action) throw new Error('Command has no action'); const error = new Error('Failed to read memory files.'); - mockRefreshMemory.mockRejectedValue(error); + mockLoadServerHierarchicalMemory.mockRejectedValue(error); await refreshCommand.action(mockContext, ''); - expect(mockRefreshMemory).toHaveBeenCalledOnce(); + expect(loadServerHierarchicalMemory).toHaveBeenCalledOnce(); + expect(mockSetUserMemory).not.toHaveBeenCalled(); + expect(mockSetGeminiMdFileCount).not.toHaveBeenCalled(); expect(mockContext.ui.addItem).toHaveBeenCalledWith( { @@ -243,7 +279,7 @@ describe('memoryCommand', () => { expect.any(Number), ); - expect(mockRefreshMemory).not.toHaveBeenCalled(); + expect(loadServerHierarchicalMemory).not.toHaveBeenCalled(); }); }); }); diff --git a/packages/cli/src/ui/commands/memoryCommand.ts b/packages/cli/src/ui/commands/memoryCommand.ts index afa43031..fe698c0f 100644 --- a/packages/cli/src/ui/commands/memoryCommand.ts +++ b/packages/cli/src/ui/commands/memoryCommand.ts @@ -4,7 +4,10 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { getErrorMessage } from '@google/gemini-cli-core'; +import { + getErrorMessage, + loadServerHierarchicalMemory, +} from '@google/gemini-cli-core'; import { MessageType } from '../types.js'; import { CommandKind, @@ -81,10 +84,20 @@ export const memoryCommand: SlashCommand = { ); try { - const result = await context.services.config?.refreshMemory(); + const config = await context.services.config; + if (config) { + const { memoryContent, fileCount } = + await loadServerHierarchicalMemory( + config.getWorkingDir(), + config.getDebugMode(), + config.getFileService(), + config.getExtensionContextFilePaths(), + config.getFileFilteringOptions(), + context.services.settings.merged.memoryDiscoveryMaxDirs, + ); + config.setUserMemory(memoryContent); + config.setGeminiMdFileCount(fileCount); - if (result) { - const { memoryContent, fileCount } = result; const successMessage = memoryContent.length > 0 ? `Memory refreshed successfully. Loaded ${memoryContent.length} characters from ${fileCount} file(s).` diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 44300a83..9fec505f 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -18,7 +18,6 @@ import { } from '../core/contentGenerator.js'; import { GeminiClient } from '../core/client.js'; import { GitService } from '../services/gitService.js'; -import { loadServerHierarchicalMemory } from '../utils/memoryDiscovery.js'; // Mock dependencies that might be called during Config construction or createServerConfig vi.mock('../tools/tool-registry', () => { @@ -313,39 +312,4 @@ describe('Server Config (config.ts)', () => { expect(config.getTelemetryOtlpEndpoint()).toBe(DEFAULT_OTLP_ENDPOINT); }); }); - - describe('refreshMemory', () => { - it('should update memory and file count on successful refresh', async () => { - const config = new Config(baseParams); - const mockMemoryData = { - memoryContent: 'new memory content', - fileCount: 5, - }; - - (loadServerHierarchicalMemory as Mock).mockResolvedValue(mockMemoryData); - - const result = await config.refreshMemory(); - - expect(loadServerHierarchicalMemory).toHaveBeenCalledWith( - config.getWorkingDir(), - config.getDebugMode(), - config.getFileService(), - config.getExtensionContextFilePaths(), - config.getFileFilteringOptions(), - ); - - expect(config.getUserMemory()).toBe(mockMemoryData.memoryContent); - expect(config.getGeminiMdFileCount()).toBe(mockMemoryData.fileCount); - expect(result).toEqual(mockMemoryData); - }); - - it('should propagate errors from loadServerHierarchicalMemory', async () => { - const config = new Config(baseParams); - const testError = new Error('Failed to load memory'); - - (loadServerHierarchicalMemory as Mock).mockRejectedValue(testError); - - await expect(config.refreshMemory()).rejects.toThrow(testError); - }); - }); }); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index d4427093..4a9abfdc 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -30,7 +30,6 @@ import { WebSearchTool } from '../tools/web-search.js'; import { GeminiClient } from '../core/client.js'; import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; import { GitService } from '../services/gitService.js'; -import { loadServerHierarchicalMemory } from '../utils/memoryDiscovery.js'; import { getProjectTempDir } from '../utils/paths.js'; import { initializeTelemetry, @@ -577,21 +576,6 @@ export class Config { return this.gitService; } - async refreshMemory(): Promise<{ memoryContent: string; fileCount: number }> { - const { memoryContent, fileCount } = await loadServerHierarchicalMemory( - this.getWorkingDir(), - this.getDebugMode(), - this.getFileService(), - this.getExtensionContextFilePaths(), - this.getFileFilteringOptions(), - ); - - this.setUserMemory(memoryContent); - this.setGeminiMdFileCount(fileCount); - - return { memoryContent, fileCount }; - } - async createToolRegistry(): Promise { const registry = new ToolRegistry(this); diff --git a/packages/core/src/utils/memoryDiscovery.test.ts b/packages/core/src/utils/memoryDiscovery.test.ts index 3051da0e..2fb2fcb1 100644 --- a/packages/core/src/utils/memoryDiscovery.test.ts +++ b/packages/core/src/utils/memoryDiscovery.test.ts @@ -319,18 +319,35 @@ My code memory }); }); - it('should respect MAX_DIRECTORIES_TO_SCAN_FOR_MEMORY during downward scan', async () => { - // the max depth is 200 so it will give up before searching all these. - for (let i = 0; i < 250; i++) { + it('should respect the maxDirs parameter during downward scan', async () => { + const consoleDebugSpy = vi + .spyOn(console, 'debug') + .mockImplementation(() => {}); + + for (let i = 0; i < 100; i++) { await createEmptyDir(path.join(cwd, `deep_dir_${i}`)); } - // "much_deeper" is alphabetically after "deep_dir_*" so it won't be loaded - await createTestFile( - path.join(cwd, 'much_deeper', DEFAULT_CONTEXT_FILENAME), - 'Ignored memory', + // Pass the custom limit directly to the function + await loadServerHierarchicalMemory( + cwd, + true, + new FileDiscoveryService(projectRoot), + [], + { + respectGitIgnore: true, + respectGeminiIgnore: true, + }, + 50, // maxDirs ); + expect(consoleDebugSpy).toHaveBeenCalledWith( + expect.stringContaining('[DEBUG] [BfsFileSearch]'), + expect.stringContaining('Scanning [50/50]:'), + ); + + vi.mocked(console.debug).mockRestore(); + const result = await loadServerHierarchicalMemory( cwd, false, diff --git a/packages/core/src/utils/memoryDiscovery.ts b/packages/core/src/utils/memoryDiscovery.ts index 33231823..88c82373 100644 --- a/packages/core/src/utils/memoryDiscovery.ts +++ b/packages/core/src/utils/memoryDiscovery.ts @@ -33,8 +33,6 @@ const logger = { console.error('[ERROR] [MemoryDiscovery]', ...args), }; -const MAX_DIRECTORIES_TO_SCAN_FOR_MEMORY = 200; - interface GeminiFileContent { filePath: string; content: string | null; @@ -90,6 +88,7 @@ async function getGeminiMdFilePathsInternal( fileService: FileDiscoveryService, extensionContextFilePaths: string[] = [], fileFilteringOptions: FileFilteringOptions, + maxDirs: number, ): Promise { const allPaths = new Set(); const geminiMdFilenames = getAllGeminiMdFilenames(); @@ -194,7 +193,7 @@ async function getGeminiMdFilePathsInternal( const downwardPaths = await bfsFileSearch(resolvedCwd, { fileName: geminiMdFilename, - maxDirs: MAX_DIRECTORIES_TO_SCAN_FOR_MEMORY, + maxDirs, debug: debugMode, fileService, fileFilteringOptions: mergedOptions, // Pass merged options as fileFilter @@ -295,6 +294,7 @@ export async function loadServerHierarchicalMemory( fileService: FileDiscoveryService, extensionContextFilePaths: string[] = [], fileFilteringOptions?: FileFilteringOptions, + maxDirs: number = 200, ): Promise<{ memoryContent: string; fileCount: number }> { if (debugMode) logger.debug( @@ -311,6 +311,7 @@ export async function loadServerHierarchicalMemory( fileService, extensionContextFilePaths, fileFilteringOptions || DEFAULT_MEMORY_FILE_FILTERING_OPTIONS, + maxDirs, ); if (filePaths.length === 0) { if (debugMode) logger.debug('No GEMINI.md files found in hierarchy.');