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

packages/core/src/services/contextManager.ts

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