feat(core): Render memory hierarchically in context. (#18350)

This commit is contained in:
joshualitt
2026-02-09 18:01:59 -08:00
committed by GitHub
parent 5d0570b113
commit 89d4556c45
25 changed files with 1189 additions and 530 deletions
@@ -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');
});
+76 -28
View File
@@ -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 {