mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-24 03:54:43 -07:00
feat(core): Render memory hierarchically in context. (#18350)
This commit is contained in:
@@ -16,8 +16,10 @@ vi.mock('../utils/memoryDiscovery.js', async (importOriginal) => {
|
||||
await importOriginal<typeof import('../utils/memoryDiscovery.js')>();
|
||||
return {
|
||||
...actual,
|
||||
loadGlobalMemory: vi.fn(),
|
||||
loadEnvironmentMemory: vi.fn(),
|
||||
getGlobalMemoryPaths: vi.fn(),
|
||||
getExtensionMemoryPaths: vi.fn(),
|
||||
getEnvironmentMemoryPaths: vi.fn(),
|
||||
readGeminiMdFiles: vi.fn(),
|
||||
loadJitSubdirectoryMemory: vi.fn(),
|
||||
concatenateInstructions: vi
|
||||
.fn()
|
||||
@@ -33,10 +35,13 @@ describe('ContextManager', () => {
|
||||
mockConfig = {
|
||||
getDebugMode: vi.fn().mockReturnValue(false),
|
||||
getWorkingDir: vi.fn().mockReturnValue('/app'),
|
||||
getImportFormat: vi.fn().mockReturnValue('tree'),
|
||||
getWorkspaceContext: vi.fn().mockReturnValue({
|
||||
getDirectories: vi.fn().mockReturnValue(['/app']),
|
||||
}),
|
||||
getExtensionLoader: vi.fn().mockReturnValue({}),
|
||||
getExtensionLoader: vi.fn().mockReturnValue({
|
||||
getExtensions: vi.fn().mockReturnValue([]),
|
||||
}),
|
||||
getMcpClientManager: vi.fn().mockReturnValue({
|
||||
getMcpInstructions: vi.fn().mockReturnValue('MCP Instructions'),
|
||||
}),
|
||||
@@ -46,66 +51,60 @@ describe('ContextManager', () => {
|
||||
contextManager = new ContextManager(mockConfig);
|
||||
vi.clearAllMocks();
|
||||
vi.spyOn(coreEvents, 'emit');
|
||||
vi.mocked(memoryDiscovery.getExtensionMemoryPaths).mockReturnValue([]);
|
||||
});
|
||||
|
||||
describe('refresh', () => {
|
||||
it('should load and format global and environment memory', async () => {
|
||||
const mockGlobalResult: memoryDiscovery.MemoryLoadResult = {
|
||||
files: [
|
||||
{ path: '/home/user/.gemini/GEMINI.md', content: 'Global Content' },
|
||||
],
|
||||
};
|
||||
vi.mocked(memoryDiscovery.loadGlobalMemory).mockResolvedValue(
|
||||
mockGlobalResult,
|
||||
const globalPaths = ['/home/user/.gemini/GEMINI.md'];
|
||||
const envPaths = ['/app/GEMINI.md'];
|
||||
|
||||
vi.mocked(memoryDiscovery.getGlobalMemoryPaths).mockResolvedValue(
|
||||
globalPaths,
|
||||
);
|
||||
vi.mocked(memoryDiscovery.getEnvironmentMemoryPaths).mockResolvedValue(
|
||||
envPaths,
|
||||
);
|
||||
|
||||
const mockEnvResult: memoryDiscovery.MemoryLoadResult = {
|
||||
files: [{ path: '/app/GEMINI.md', content: 'Env Content' }],
|
||||
};
|
||||
vi.mocked(memoryDiscovery.loadEnvironmentMemory).mockResolvedValue(
|
||||
mockEnvResult,
|
||||
);
|
||||
vi.mocked(memoryDiscovery.readGeminiMdFiles).mockResolvedValue([
|
||||
{ filePath: globalPaths[0], content: 'Global Content' },
|
||||
{ filePath: envPaths[0], content: 'Env Content' },
|
||||
]);
|
||||
|
||||
await contextManager.refresh();
|
||||
|
||||
expect(memoryDiscovery.loadGlobalMemory).toHaveBeenCalledWith(false);
|
||||
expect(contextManager.getGlobalMemory()).toMatch(
|
||||
/--- Context from: .*GEMINI.md ---/,
|
||||
);
|
||||
expect(contextManager.getGlobalMemory()).toContain('Global Content');
|
||||
|
||||
expect(memoryDiscovery.loadEnvironmentMemory).toHaveBeenCalledWith(
|
||||
expect(memoryDiscovery.getGlobalMemoryPaths).toHaveBeenCalled();
|
||||
expect(memoryDiscovery.getEnvironmentMemoryPaths).toHaveBeenCalledWith(
|
||||
['/app'],
|
||||
expect.anything(),
|
||||
false,
|
||||
);
|
||||
expect(contextManager.getEnvironmentMemory()).toContain(
|
||||
'--- Context from: GEMINI.md ---',
|
||||
expect(memoryDiscovery.readGeminiMdFiles).toHaveBeenCalledWith(
|
||||
expect.arrayContaining([...globalPaths, ...envPaths]),
|
||||
false,
|
||||
'tree',
|
||||
);
|
||||
|
||||
expect(contextManager.getGlobalMemory()).toContain('Global Content');
|
||||
expect(contextManager.getEnvironmentMemory()).toContain('Env Content');
|
||||
expect(contextManager.getEnvironmentMemory()).toContain(
|
||||
'MCP Instructions',
|
||||
);
|
||||
|
||||
expect(contextManager.getLoadedPaths()).toContain(
|
||||
'/home/user/.gemini/GEMINI.md',
|
||||
);
|
||||
expect(contextManager.getLoadedPaths()).toContain('/app/GEMINI.md');
|
||||
expect(contextManager.getLoadedPaths()).toContain(globalPaths[0]);
|
||||
expect(contextManager.getLoadedPaths()).toContain(envPaths[0]);
|
||||
});
|
||||
|
||||
it('should emit MemoryChanged event when memory is refreshed', async () => {
|
||||
const mockGlobalResult = {
|
||||
files: [{ path: '/app/GEMINI.md', content: 'content' }],
|
||||
};
|
||||
const mockEnvResult = {
|
||||
files: [{ path: '/app/src/GEMINI.md', content: 'env content' }],
|
||||
};
|
||||
vi.mocked(memoryDiscovery.loadGlobalMemory).mockResolvedValue(
|
||||
mockGlobalResult,
|
||||
);
|
||||
vi.mocked(memoryDiscovery.loadEnvironmentMemory).mockResolvedValue(
|
||||
mockEnvResult,
|
||||
);
|
||||
vi.mocked(memoryDiscovery.getGlobalMemoryPaths).mockResolvedValue([
|
||||
'/app/GEMINI.md',
|
||||
]);
|
||||
vi.mocked(memoryDiscovery.getEnvironmentMemoryPaths).mockResolvedValue([
|
||||
'/app/src/GEMINI.md',
|
||||
]);
|
||||
vi.mocked(memoryDiscovery.readGeminiMdFiles).mockResolvedValue([
|
||||
{ filePath: '/app/GEMINI.md', content: 'content' },
|
||||
{ filePath: '/app/src/GEMINI.md', content: 'env content' },
|
||||
]);
|
||||
|
||||
await contextManager.refresh();
|
||||
|
||||
@@ -116,18 +115,16 @@ describe('ContextManager', () => {
|
||||
|
||||
it('should not load environment memory if folder is not trusted', async () => {
|
||||
vi.mocked(mockConfig.isTrustedFolder).mockReturnValue(false);
|
||||
const mockGlobalResult = {
|
||||
files: [
|
||||
{ path: '/home/user/.gemini/GEMINI.md', content: 'Global Content' },
|
||||
],
|
||||
};
|
||||
vi.mocked(memoryDiscovery.loadGlobalMemory).mockResolvedValue(
|
||||
mockGlobalResult,
|
||||
);
|
||||
vi.mocked(memoryDiscovery.getGlobalMemoryPaths).mockResolvedValue([
|
||||
'/home/user/.gemini/GEMINI.md',
|
||||
]);
|
||||
vi.mocked(memoryDiscovery.readGeminiMdFiles).mockResolvedValue([
|
||||
{ filePath: '/home/user/.gemini/GEMINI.md', content: 'Global Content' },
|
||||
]);
|
||||
|
||||
await contextManager.refresh();
|
||||
|
||||
expect(memoryDiscovery.loadEnvironmentMemory).not.toHaveBeenCalled();
|
||||
expect(memoryDiscovery.getEnvironmentMemoryPaths).not.toHaveBeenCalled();
|
||||
expect(contextManager.getEnvironmentMemory()).toBe('');
|
||||
expect(contextManager.getGlobalMemory()).toContain('Global Content');
|
||||
});
|
||||
|
||||
@@ -5,10 +5,14 @@
|
||||
*/
|
||||
|
||||
import {
|
||||
loadGlobalMemory,
|
||||
loadEnvironmentMemory,
|
||||
loadJitSubdirectoryMemory,
|
||||
concatenateInstructions,
|
||||
getGlobalMemoryPaths,
|
||||
getExtensionMemoryPaths,
|
||||
getEnvironmentMemoryPaths,
|
||||
readGeminiMdFiles,
|
||||
categorizeAndConcatenate,
|
||||
type GeminiFileContent,
|
||||
} from '../utils/memoryDiscovery.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { coreEvents, CoreEvent } from '../utils/events.js';
|
||||
@@ -17,51 +21,91 @@ export class ContextManager {
|
||||
private readonly loadedPaths: Set<string> = new Set();
|
||||
private readonly config: Config;
|
||||
private globalMemory: string = '';
|
||||
private environmentMemory: string = '';
|
||||
private extensionMemory: string = '';
|
||||
private projectMemory: string = '';
|
||||
|
||||
constructor(config: Config) {
|
||||
this.config = config;
|
||||
}
|
||||
|
||||
/**
|
||||
* Refreshes the memory by reloading global and environment memory.
|
||||
* Refreshes the memory by reloading global, extension, and project memory.
|
||||
*/
|
||||
async refresh(): Promise<void> {
|
||||
this.loadedPaths.clear();
|
||||
await this.loadGlobalMemory();
|
||||
await this.loadEnvironmentMemory();
|
||||
const debugMode = this.config.getDebugMode();
|
||||
|
||||
const paths = await this.discoverMemoryPaths(debugMode);
|
||||
const contentsMap = await this.loadMemoryContents(paths, debugMode);
|
||||
|
||||
this.categorizeMemoryContents(paths, contentsMap);
|
||||
this.emitMemoryChanged();
|
||||
}
|
||||
|
||||
private async loadGlobalMemory(): Promise<void> {
|
||||
const result = await loadGlobalMemory(this.config.getDebugMode());
|
||||
this.markAsLoaded(result.files.map((f) => f.path));
|
||||
this.globalMemory = concatenateInstructions(
|
||||
result.files.map((f) => ({ filePath: f.path, content: f.content })),
|
||||
this.config.getWorkingDir(),
|
||||
);
|
||||
private async discoverMemoryPaths(debugMode: boolean) {
|
||||
const [global, extension, project] = await Promise.all([
|
||||
getGlobalMemoryPaths(debugMode),
|
||||
Promise.resolve(
|
||||
getExtensionMemoryPaths(this.config.getExtensionLoader()),
|
||||
),
|
||||
this.config.isTrustedFolder()
|
||||
? getEnvironmentMemoryPaths(
|
||||
[...this.config.getWorkspaceContext().getDirectories()],
|
||||
debugMode,
|
||||
)
|
||||
: Promise.resolve([]),
|
||||
]);
|
||||
|
||||
return { global, extension, project };
|
||||
}
|
||||
|
||||
private async loadEnvironmentMemory(): Promise<void> {
|
||||
if (!this.config.isTrustedFolder()) {
|
||||
this.environmentMemory = '';
|
||||
return;
|
||||
}
|
||||
const result = await loadEnvironmentMemory(
|
||||
[...this.config.getWorkspaceContext().getDirectories()],
|
||||
this.config.getExtensionLoader(),
|
||||
this.config.getDebugMode(),
|
||||
private async loadMemoryContents(
|
||||
paths: { global: string[]; extension: string[]; project: string[] },
|
||||
debugMode: boolean,
|
||||
) {
|
||||
const allPaths = Array.from(
|
||||
new Set([...paths.global, ...paths.extension, ...paths.project]),
|
||||
);
|
||||
this.markAsLoaded(result.files.map((f) => f.path));
|
||||
const envMemory = concatenateInstructions(
|
||||
result.files.map((f) => ({ filePath: f.path, content: f.content })),
|
||||
this.config.getWorkingDir(),
|
||||
|
||||
const allContents = await readGeminiMdFiles(
|
||||
allPaths,
|
||||
debugMode,
|
||||
this.config.getImportFormat(),
|
||||
);
|
||||
|
||||
this.markAsLoaded(
|
||||
allContents.filter((c) => c.content !== null).map((c) => c.filePath),
|
||||
);
|
||||
|
||||
return new Map(allContents.map((c) => [c.filePath, c]));
|
||||
}
|
||||
|
||||
private categorizeMemoryContents(
|
||||
paths: { global: string[]; extension: string[]; project: string[] },
|
||||
contentsMap: Map<string, GeminiFileContent>,
|
||||
) {
|
||||
const workingDir = this.config.getWorkingDir();
|
||||
const hierarchicalMemory = categorizeAndConcatenate(
|
||||
paths,
|
||||
contentsMap,
|
||||
workingDir,
|
||||
);
|
||||
|
||||
this.globalMemory = hierarchicalMemory.global || '';
|
||||
this.extensionMemory = hierarchicalMemory.extension || '';
|
||||
|
||||
const mcpInstructions =
|
||||
this.config.getMcpClientManager()?.getMcpInstructions() || '';
|
||||
this.environmentMemory = [envMemory, mcpInstructions.trimStart()]
|
||||
const projectMemoryWithMcp = [
|
||||
hierarchicalMemory.project,
|
||||
mcpInstructions.trimStart(),
|
||||
]
|
||||
.filter(Boolean)
|
||||
.join('\n\n');
|
||||
|
||||
this.projectMemory = this.config.isTrustedFolder()
|
||||
? projectMemoryWithMcp
|
||||
: '';
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -103,8 +147,12 @@ export class ContextManager {
|
||||
return this.globalMemory;
|
||||
}
|
||||
|
||||
getExtensionMemory(): string {
|
||||
return this.extensionMemory;
|
||||
}
|
||||
|
||||
getEnvironmentMemory(): string {
|
||||
return this.environmentMemory;
|
||||
return this.projectMemory;
|
||||
}
|
||||
|
||||
private markAsLoaded(paths: string[]): void {
|
||||
|
||||
Reference in New Issue
Block a user