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

packages/core/src/config/config.test.ts

@@ -186,7 +186,15 @@ vi.mock('../utils/fetch.js', () => ({
   setGlobalProxy: mockSetGlobalProxy,
 }));
 
-vi.mock('../services/contextManager.js');
+vi.mock('../services/contextManager.js', () => ({
+  ContextManager: vi.fn().mockImplementation(() => ({
+    refresh: vi.fn(),
+    getGlobalMemory: vi.fn().mockReturnValue(''),
+    getExtensionMemory: vi.fn().mockReturnValue(''),
+    getEnvironmentMemory: vi.fn().mockReturnValue(''),
+    getLoadedPaths: vi.fn().mockReturnValue(new Set()),
+  })),
+}));
 
 import { BaseLlmClient } from '../core/baseLlmClient.js';
 import { tokenLimit } from '../core/tokenLimits.js';
@@ -2059,23 +2067,19 @@ describe('Config Quota & Preview Model Access', () => {
 
 describe('Config JIT Initialization', () => {
   let config: Config;
-  let mockContextManager: {
-    refresh: Mock;
-    getGlobalMemory: Mock;
-    getEnvironmentMemory: Mock;
-    getLoadedPaths: Mock;
-  };
+  let mockContextManager: ContextManager;
 
   beforeEach(() => {
     vi.clearAllMocks();
     mockContextManager = {
       refresh: vi.fn(),
       getGlobalMemory: vi.fn().mockReturnValue('Global Memory'),
+      getExtensionMemory: vi.fn().mockReturnValue('Extension Memory'),
       getEnvironmentMemory: vi
         .fn()
         .mockReturnValue('Environment Memory\n\nMCP Instructions'),
       getLoadedPaths: vi.fn().mockReturnValue(new Set(['/path/to/GEMINI.md'])),
-    };
+    } as unknown as ContextManager;
     (ContextManager as unknown as Mock).mockImplementation(
       () => mockContextManager,
     );
@@ -2097,9 +2101,11 @@ describe('Config JIT Initialization', () => {
 
     expect(ContextManager).toHaveBeenCalledWith(config);
     expect(mockContextManager.refresh).toHaveBeenCalled();
-    expect(config.getUserMemory()).toBe(
-      'Global Memory\n\nEnvironment Memory\n\nMCP Instructions',
-    );
+    expect(config.getUserMemory()).toEqual({
+      global: 'Global Memory',
+      extension: 'Extension Memory',
+      project: 'Environment Memory\n\nMCP Instructions',
+    });
 
     // Verify state update (delegated to ContextManager)
     expect(config.getGeminiMdFileCount()).toBe(1);

packages/core/src/config/config.ts

@@ -101,6 +101,7 @@ import { HookSystem } from '../hooks/index.js';
 import type { UserTierId } from '../code_assist/types.js';
 import type { RetrieveUserQuotaResponse } from '../code_assist/types.js';
 import type { AdminControlsSettings } from '../code_assist/types.js';
+import type { HierarchicalMemory } from './memory.js';
 import { getCodeAssistServer } from '../code_assist/codeAssist.js';
 import type { Experiments } from '../code_assist/experiments/experiments.js';
 import { AgentRegistry } from '../agents/registry.js';
@@ -384,7 +385,7 @@ export interface ConfigParameters {
   mcpServerCommand?: string;
   mcpServers?: Record<string, MCPServerConfig>;
   mcpEnablementCallbacks?: McpEnablementCallbacks;
-  userMemory?: string;
+  userMemory?: string | HierarchicalMemory;
   geminiMdFileCount?: number;
   geminiMdFilePaths?: string[];
   approvalMode?: ApprovalMode;
@@ -519,7 +520,7 @@ export class Config {
   private readonly extensionsEnabled: boolean;
   private mcpServers: Record<string, MCPServerConfig> | undefined;
   private readonly mcpEnablementCallbacks?: McpEnablementCallbacks;
-  private userMemory: string;
+  private userMemory: string | HierarchicalMemory;
   private geminiMdFileCount: number;
   private geminiMdFilePaths: string[];
   private readonly showMemoryUsage: boolean;
@@ -1379,14 +1380,13 @@ export class Config {
     this.mcpServers = mcpServers;
   }
 
-  getUserMemory(): string {
+  getUserMemory(): string | HierarchicalMemory {
     if (this.experimentalJitContext && this.contextManager) {
-      return [
-        this.contextManager.getGlobalMemory(),
-        this.contextManager.getEnvironmentMemory(),
-      ]
-        .filter(Boolean)
-        .join('\n\n');
+      return {
+        global: this.contextManager.getGlobalMemory(),
+        extension: this.contextManager.getExtensionMemory(),
+        project: this.contextManager.getEnvironmentMemory(),
+      };
     }
     return this.userMemory;
   }
@@ -1409,7 +1409,7 @@ export class Config {
     }
   }
 
-  setUserMemory(newUserMemory: string): void {
+  setUserMemory(newUserMemory: string | HierarchicalMemory): void {
     this.userMemory = newUserMemory;
   }
 

packages/core/src/config/memory.test.ts

@@ -0,0 +1,104 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect } from 'vitest';
+import { flattenMemory } from './memory.js';
+
+describe('memory', () => {
+  describe('flattenMemory', () => {
+    it('should return empty string for null or undefined', () => {
+      expect(flattenMemory(undefined)).toBe('');
+      expect(flattenMemory(null as unknown as undefined)).toBe('');
+    });
+
+    it('should return the string itself if a string is provided', () => {
+      expect(flattenMemory('raw string')).toBe('raw string');
+    });
+
+    it('should return empty string for an empty object', () => {
+      expect(flattenMemory({})).toBe('');
+    });
+
+    it('should return content with headers even if only global memory is present', () => {
+      expect(flattenMemory({ global: 'global content' })).toBe(
+        `--- Global ---
+global content`,
+      );
+    });
+
+    it('should return content with headers even if only extension memory is present', () => {
+      expect(flattenMemory({ extension: 'extension content' })).toBe(
+        `--- Extension ---
+extension content`,
+      );
+    });
+
+    it('should return content with headers even if only project memory is present', () => {
+      expect(flattenMemory({ project: 'project content' })).toBe(
+        `--- Project ---
+project content`,
+      );
+    });
+
+    it('should include headers if multiple levels are present (global + project)', () => {
+      const result = flattenMemory({
+        global: 'global content',
+        project: 'project content',
+      });
+      expect(result).toContain('--- Global ---');
+      expect(result).toContain('global content');
+      expect(result).toContain('--- Project ---');
+      expect(result).toContain('project content');
+      expect(result).not.toContain('--- Extension ---');
+    });
+
+    it('should include headers if all levels are present', () => {
+      const result = flattenMemory({
+        global: 'global content',
+        extension: 'extension content',
+        project: 'project content',
+      });
+      expect(result).toContain('--- Global ---');
+      expect(result).toContain('--- Extension ---');
+      expect(result).toContain('--- Project ---');
+      expect(result).toBe(
+        `--- Global ---
+global content
+
+--- Extension ---
+extension content
+
+--- Project ---
+project content`,
+      );
+    });
+
+    it('should trim content and ignore empty strings', () => {
+      const result = flattenMemory({
+        global: '  trimmed global  ',
+        extension: '   ',
+        project: 'project\n',
+      });
+      expect(result).toBe(
+        `--- Global ---
+trimmed global
+
+--- Project ---
+project`,
+      );
+    });
+
+    it('should return empty string if all levels are only whitespace', () => {
+      expect(
+        flattenMemory({
+          global: '  ',
+          extension: '\n',
+          project: ' 	 ',
+        }),
+      ).toBe('');
+    });
+  });
+});

packages/core/src/config/memory.ts

@@ -0,0 +1,34 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export interface HierarchicalMemory {
+  global?: string;
+  extension?: string;
+  project?: string;
+}
+
+/**
+ * Flattens hierarchical memory into a single string for display or legacy use.
+ */
+export function flattenMemory(memory?: string | HierarchicalMemory): string {
+  if (!memory) return '';
+  if (typeof memory === 'string') return memory;
+
+  const sections: Array<{ name: string; content: string }> = [];
+  if (memory.global?.trim()) {
+    sections.push({ name: 'Global', content: memory.global.trim() });
+  }
+  if (memory.extension?.trim()) {
+    sections.push({ name: 'Extension', content: memory.extension.trim() });
+  }
+  if (memory.project?.trim()) {
+    sections.push({ name: 'Project', content: memory.project.trim() });
+  }
+
+  if (sections.length === 0) return '';
+
+  return sections.map((s) => `--- ${s.name} ---\n${s.content}`).join('\n\n');
+}