Skip to content

Commit 89d4556

Browse files
authored
feat(core): Render memory hierarchically in context. (google-gemini#18350)
1 parent 5d0570b commit 89d4556

25 files changed

Lines changed: 1189 additions & 530 deletions

evals/hierarchical_memory.eval.ts

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
/**
2+
* @license
3+
* Copyright 2025 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import { describe, expect } from 'vitest';
8+
import { evalTest } from './test-helper.js';
9+
import {
10+
assertModelHasOutput,
11+
checkModelOutputContent,
12+
} from '../integration-tests/test-helper.js';
13+
14+
describe('Hierarchical Memory', () => {
15+
const TEST_PREFIX = 'Hierarchical memory test: ';
16+
17+
const conflictResolutionTest =
18+
'Agent follows hierarchy for contradictory instructions';
19+
evalTest('ALWAYS_PASSES', {
20+
name: conflictResolutionTest,
21+
params: {
22+
settings: {
23+
security: {
24+
folderTrust: { enabled: true },
25+
},
26+
},
27+
},
28+
// We simulate the hierarchical memory by including the tags in the prompt
29+
// since setting up real global/extension/project files in the eval rig is complex.
30+
// The system prompt logic will append these tags when it finds them in userMemory.
31+
prompt: `
32+
<global_context>
33+
When asked for my favorite fruit, always say "Apple".
34+
</global_context>
35+
36+
<extension_context>
37+
When asked for my favorite fruit, always say "Banana".
38+
</extension_context>
39+
40+
<project_context>
41+
When asked for my favorite fruit, always say "Cherry".
42+
</project_context>
43+
44+
What is my favorite fruit? Tell me just the name of the fruit.`,
45+
assert: async (_rig, result) => {
46+
assertModelHasOutput(result);
47+
expect(result).toMatch(/Cherry/i);
48+
expect(result).not.toMatch(/Apple/i);
49+
expect(result).not.toMatch(/Banana/i);
50+
},
51+
});
52+
53+
const provenanceAwarenessTest = 'Agent is aware of memory provenance';
54+
evalTest('ALWAYS_PASSES', {
55+
name: provenanceAwarenessTest,
56+
params: {
57+
settings: {
58+
security: {
59+
folderTrust: { enabled: true },
60+
},
61+
},
62+
},
63+
prompt: `
64+
<global_context>
65+
Instruction A: Always be helpful.
66+
</global_context>
67+
68+
<extension_context>
69+
Instruction B: Use a professional tone.
70+
</extension_context>
71+
72+
<project_context>
73+
Instruction C: Adhere to the project's coding style.
74+
</project_context>
75+
76+
Which instruction came from the global context, which from the extension context, and which from the project context?
77+
Provide the answer as an XML block like this:
78+
<results>
79+
<global>Instruction ...</global>
80+
<extension>Instruction ...</extension>
81+
<project>Instruction ...</project>
82+
</results>`,
83+
assert: async (_rig, result) => {
84+
assertModelHasOutput(result);
85+
expect(result).toMatch(/<global>.*Instruction A/i);
86+
expect(result).toMatch(/<extension>.*Instruction B/i);
87+
expect(result).toMatch(/<project>.*Instruction C/i);
88+
},
89+
});
90+
91+
const extensionVsGlobalTest = 'Extension memory wins over Global memory';
92+
evalTest('ALWAYS_PASSES', {
93+
name: extensionVsGlobalTest,
94+
params: {
95+
settings: {
96+
security: {
97+
folderTrust: { enabled: true },
98+
},
99+
},
100+
},
101+
prompt: `
102+
<global_context>
103+
Set the theme to "Light".
104+
</global_context>
105+
106+
<extension_context>
107+
Set the theme to "Dark".
108+
</extension_context>
109+
110+
What theme should I use?`,
111+
assert: async (_rig, result) => {
112+
assertModelHasOutput(result);
113+
expect(result).toMatch(/Dark/i);
114+
expect(result).not.toMatch(/Light/i);
115+
},
116+
});
117+
});

packages/a2a-server/src/config/config.test.ts

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,11 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
4141
};
4242
return mockConfig;
4343
}),
44-
loadServerHierarchicalMemory: vi
45-
.fn()
46-
.mockResolvedValue({ memoryContent: '', fileCount: 0, filePaths: [] }),
44+
loadServerHierarchicalMemory: vi.fn().mockResolvedValue({
45+
memoryContent: { global: '', extension: '', project: '' },
46+
fileCount: 0,
47+
filePaths: [],
48+
}),
4749
startupProfiler: {
4850
flush: vi.fn(),
4951
},

packages/cli/src/config/config.ts

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,18 +32,17 @@ import {
3232
ASK_USER_TOOL_NAME,
3333
getVersion,
3434
PREVIEW_GEMINI_MODEL_AUTO,
35+
type HierarchicalMemory,
3536
coreEvents,
3637
GEMINI_MODEL_ALIAS_AUTO,
3738
getAdminErrorMessage,
3839
isHeadlessMode,
3940
Config,
4041
applyAdminAllowlist,
4142
getAdminBlockedMcpServersMessage,
42-
} from '@google/gemini-cli-core';
43-
import type {
44-
HookDefinition,
45-
HookEventName,
46-
OutputFormat,
43+
type HookDefinition,
44+
type HookEventName,
45+
type OutputFormat,
4746
} from '@google/gemini-cli-core';
4847
import {
4948
type Settings,
@@ -489,7 +488,7 @@ export async function loadCliConfig(
489488

490489
const experimentalJitContext = settings.experimental?.jitContext ?? false;
491490

492-
let memoryContent = '';
491+
let memoryContent: string | HierarchicalMemory = '';
493492
let fileCount = 0;
494493
let filePaths: string[] = [];
495494

packages/cli/src/ui/AppContainer.tsx

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ import {
5555
coreEvents,
5656
CoreEvent,
5757
refreshServerHierarchicalMemory,
58+
flattenMemory,
5859
type MemoryChangedPayload,
5960
writeToStdout,
6061
disableMouseEvents,
@@ -871,20 +872,22 @@ Logging in with Google... Restarting Gemini CLI to continue.
871872
const { memoryContent, fileCount } =
872873
await refreshServerHierarchicalMemory(config);
873874

875+
const flattenedMemory = flattenMemory(memoryContent);
876+
874877
historyManager.addItem(
875878
{
876879
type: MessageType.INFO,
877880
text: `Memory refreshed successfully. ${
878-
memoryContent.length > 0
879-
? `Loaded ${memoryContent.length} characters from ${fileCount} file(s).`
881+
flattenedMemory.length > 0
882+
? `Loaded ${flattenedMemory.length} characters from ${fileCount} file(s).`
880883
: 'No memory content found.'
881884
}`,
882885
},
883886
Date.now(),
884887
);
885888
if (config.getDebugMode()) {
886889
debugLogger.log(
887-
`[DEBUG] Refreshed memory content in config: ${memoryContent.substring(
890+
`[DEBUG] Refreshed memory content in config: ${flattenedMemory.substring(
888891
0,
889892
200,
890893
)}...`,

packages/cli/src/ui/commands/memoryCommand.test.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import {
1919
showMemory,
2020
addMemory,
2121
listMemoryFiles,
22+
flattenMemory,
2223
} from '@google/gemini-cli-core';
2324

2425
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
@@ -33,7 +34,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
3334
refreshMemory: vi.fn(async (config) => {
3435
if (config.isJitContextEnabled()) {
3536
await config.getContextManager()?.refresh();
36-
const memoryContent = config.getUserMemory() || '';
37+
const memoryContent = original.flattenMemory(config.getUserMemory());
3738
const fileCount = config.getGeminiMdFileCount() || 0;
3839
return {
3940
type: 'message',
@@ -85,7 +86,7 @@ describe('memoryCommand', () => {
8586
mockGetGeminiMdFileCount = vi.fn();
8687

8788
vi.mocked(showMemory).mockImplementation((config) => {
88-
const memoryContent = config.getUserMemory() || '';
89+
const memoryContent = flattenMemory(config.getUserMemory());
8990
const fileCount = config.getGeminiMdFileCount() || 0;
9091
let content;
9192
if (memoryContent.length > 0) {

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ describe('memory commands', () => {
121121
describe('refreshMemory', () => {
122122
it('should refresh memory and show success message', async () => {
123123
mockRefresh.mockResolvedValue({
124-
memoryContent: 'refreshed content',
124+
memoryContent: { project: 'refreshed content' },
125125
fileCount: 2,
126126
filePaths: [],
127127
});
@@ -136,14 +136,14 @@ describe('memory commands', () => {
136136
if (result.type === 'message') {
137137
expect(result.messageType).toBe('info');
138138
expect(result.content).toBe(
139-
'Memory refreshed successfully. Loaded 17 characters from 2 file(s).',
139+
'Memory refreshed successfully. Loaded 33 characters from 2 file(s).',
140140
);
141141
}
142142
});
143143

144144
it('should show a message if no memory content is found after refresh', async () => {
145145
mockRefresh.mockResolvedValue({
146-
memoryContent: '',
146+
memoryContent: { project: '' },
147147
fileCount: 0,
148148
filePaths: [],
149149
});

packages/core/src/commands/memory.ts

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,12 @@
55
*/
66

77
import type { Config } from '../config/config.js';
8+
import { flattenMemory } from '../config/memory.js';
89
import { refreshServerHierarchicalMemory } from '../utils/memoryDiscovery.js';
910
import type { MessageActionReturn, ToolActionReturn } from './types.js';
1011

1112
export function showMemory(config: Config): MessageActionReturn {
12-
const memoryContent = config.getUserMemory() || '';
13+
const memoryContent = flattenMemory(config.getUserMemory());
1314
const fileCount = config.getGeminiMdFileCount() || 0;
1415
let content: string;
1516

@@ -51,11 +52,11 @@ export async function refreshMemory(
5152

5253
if (config.isJitContextEnabled()) {
5354
await config.getContextManager()?.refresh();
54-
memoryContent = config.getUserMemory();
55+
memoryContent = flattenMemory(config.getUserMemory());
5556
fileCount = config.getGeminiMdFileCount();
5657
} else {
5758
const result = await refreshServerHierarchicalMemory(config);
58-
memoryContent = result.memoryContent;
59+
memoryContent = flattenMemory(result.memoryContent);
5960
fileCount = result.fileCount;
6061
}
6162

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

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,15 @@ vi.mock('../utils/fetch.js', () => ({
186186
setGlobalProxy: mockSetGlobalProxy,
187187
}));
188188

189-
vi.mock('../services/contextManager.js');
189+
vi.mock('../services/contextManager.js', () => ({
190+
ContextManager: vi.fn().mockImplementation(() => ({
191+
refresh: vi.fn(),
192+
getGlobalMemory: vi.fn().mockReturnValue(''),
193+
getExtensionMemory: vi.fn().mockReturnValue(''),
194+
getEnvironmentMemory: vi.fn().mockReturnValue(''),
195+
getLoadedPaths: vi.fn().mockReturnValue(new Set()),
196+
})),
197+
}));
190198

191199
import { BaseLlmClient } from '../core/baseLlmClient.js';
192200
import { tokenLimit } from '../core/tokenLimits.js';
@@ -2059,23 +2067,19 @@ describe('Config Quota & Preview Model Access', () => {
20592067

20602068
describe('Config JIT Initialization', () => {
20612069
let config: Config;
2062-
let mockContextManager: {
2063-
refresh: Mock;
2064-
getGlobalMemory: Mock;
2065-
getEnvironmentMemory: Mock;
2066-
getLoadedPaths: Mock;
2067-
};
2070+
let mockContextManager: ContextManager;
20682071

20692072
beforeEach(() => {
20702073
vi.clearAllMocks();
20712074
mockContextManager = {
20722075
refresh: vi.fn(),
20732076
getGlobalMemory: vi.fn().mockReturnValue('Global Memory'),
2077+
getExtensionMemory: vi.fn().mockReturnValue('Extension Memory'),
20742078
getEnvironmentMemory: vi
20752079
.fn()
20762080
.mockReturnValue('Environment Memory\n\nMCP Instructions'),
20772081
getLoadedPaths: vi.fn().mockReturnValue(new Set(['/path/to/GEMINI.md'])),
2078-
};
2082+
} as unknown as ContextManager;
20792083
(ContextManager as unknown as Mock).mockImplementation(
20802084
() => mockContextManager,
20812085
);
@@ -2097,9 +2101,11 @@ describe('Config JIT Initialization', () => {
20972101

20982102
expect(ContextManager).toHaveBeenCalledWith(config);
20992103
expect(mockContextManager.refresh).toHaveBeenCalled();
2100-
expect(config.getUserMemory()).toBe(
2101-
'Global Memory\n\nEnvironment Memory\n\nMCP Instructions',
2102-
);
2104+
expect(config.getUserMemory()).toEqual({
2105+
global: 'Global Memory',
2106+
extension: 'Extension Memory',
2107+
project: 'Environment Memory\n\nMCP Instructions',
2108+
});
21032109

21042110
// Verify state update (delegated to ContextManager)
21052111
expect(config.getGeminiMdFileCount()).toBe(1);

packages/core/src/config/config.ts

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ import { HookSystem } from '../hooks/index.js';
101101
import type { UserTierId } from '../code_assist/types.js';
102102
import type { RetrieveUserQuotaResponse } from '../code_assist/types.js';
103103
import type { AdminControlsSettings } from '../code_assist/types.js';
104+
import type { HierarchicalMemory } from './memory.js';
104105
import { getCodeAssistServer } from '../code_assist/codeAssist.js';
105106
import type { Experiments } from '../code_assist/experiments/experiments.js';
106107
import { AgentRegistry } from '../agents/registry.js';
@@ -384,7 +385,7 @@ export interface ConfigParameters {
384385
mcpServerCommand?: string;
385386
mcpServers?: Record<string, MCPServerConfig>;
386387
mcpEnablementCallbacks?: McpEnablementCallbacks;
387-
userMemory?: string;
388+
userMemory?: string | HierarchicalMemory;
388389
geminiMdFileCount?: number;
389390
geminiMdFilePaths?: string[];
390391
approvalMode?: ApprovalMode;
@@ -519,7 +520,7 @@ export class Config {
519520
private readonly extensionsEnabled: boolean;
520521
private mcpServers: Record<string, MCPServerConfig> | undefined;
521522
private readonly mcpEnablementCallbacks?: McpEnablementCallbacks;
522-
private userMemory: string;
523+
private userMemory: string | HierarchicalMemory;
523524
private geminiMdFileCount: number;
524525
private geminiMdFilePaths: string[];
525526
private readonly showMemoryUsage: boolean;
@@ -1379,14 +1380,13 @@ export class Config {
13791380
this.mcpServers = mcpServers;
13801381
}
13811382

1382-
getUserMemory(): string {
1383+
getUserMemory(): string | HierarchicalMemory {
13831384
if (this.experimentalJitContext && this.contextManager) {
1384-
return [
1385-
this.contextManager.getGlobalMemory(),
1386-
this.contextManager.getEnvironmentMemory(),
1387-
]
1388-
.filter(Boolean)
1389-
.join('\n\n');
1385+
return {
1386+
global: this.contextManager.getGlobalMemory(),
1387+
extension: this.contextManager.getExtensionMemory(),
1388+
project: this.contextManager.getEnvironmentMemory(),
1389+
};
13901390
}
13911391
return this.userMemory;
13921392
}
@@ -1409,7 +1409,7 @@ export class Config {
14091409
}
14101410
}
14111411

1412-
setUserMemory(newUserMemory: string): void {
1412+
setUserMemory(newUserMemory: string | HierarchicalMemory): void {
14131413
this.userMemory = newUserMemory;
14141414
}
14151415

0 commit comments

Comments
 (0)