From 82d8680dccbe35a4308be3348bbb3c11835904de Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Tue, 17 Mar 2026 13:20:32 -0700 Subject: [PATCH] refactor(core): align JIT memory placement with tiered context model (#22766) --- packages/cli/src/test-utils/mockConfig.ts | 5 +-- packages/core/src/config/config.test.ts | 15 ++++++++ packages/core/src/config/config.ts | 37 +++++++++++++++++++ packages/core/src/core/client.test.ts | 15 +++++--- packages/core/src/core/client.ts | 6 +-- .../core/src/utils/environmentContext.test.ts | 13 ++++++- packages/core/src/utils/environmentContext.ts | 12 ++++-- 7 files changed, 86 insertions(+), 17 deletions(-) diff --git a/packages/cli/src/test-utils/mockConfig.ts b/packages/cli/src/test-utils/mockConfig.ts index 59d19b3412..d4f11212e3 100644 --- a/packages/cli/src/test-utils/mockConfig.ts +++ b/packages/cli/src/test-utils/mockConfig.ts @@ -17,7 +17,6 @@ import { * Creates a mocked Config object with default values and allows overrides. */ export const createMockConfig = (overrides: Partial = {}): Config => - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion ({ getSandbox: vi.fn(() => undefined), getQuestion: vi.fn(() => ''), @@ -79,6 +78,8 @@ export const createMockConfig = (overrides: Partial = {}): Config => getFileService: vi.fn().mockReturnValue({}), getGitService: vi.fn().mockResolvedValue({}), getUserMemory: vi.fn().mockReturnValue(''), + getSystemInstructionMemory: vi.fn().mockReturnValue(''), + getSessionMemory: vi.fn().mockReturnValue(''), getGeminiMdFilePaths: vi.fn().mockReturnValue([]), getShowMemoryUsage: vi.fn().mockReturnValue(false), getAccessibility: vi.fn().mockReturnValue({}), @@ -182,11 +183,9 @@ export function createMockSettings( overrides: Record = {}, ): LoadedSettings { const merged = createTestMergedSettings( - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion (overrides['merged'] as Partial) || {}, ); - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion return { system: { settings: {} }, systemDefaults: { settings: {} }, diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 573a6bedde..a4ef0cbaac 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -3063,6 +3063,21 @@ describe('Config JIT Initialization', () => { project: 'Environment Memory\n\nMCP Instructions', }); + // Tier 1: system instruction gets only global memory + expect(config.getSystemInstructionMemory()).toBe('Global Memory'); + + // Tier 2: session memory gets extension + project formatted with XML tags + const sessionMemory = config.getSessionMemory(); + expect(sessionMemory).toContain(''); + expect(sessionMemory).toContain(''); + expect(sessionMemory).toContain('Extension Memory'); + expect(sessionMemory).toContain(''); + expect(sessionMemory).toContain(''); + expect(sessionMemory).toContain('Environment Memory'); + expect(sessionMemory).toContain('MCP Instructions'); + expect(sessionMemory).toContain(''); + expect(sessionMemory).toContain(''); + // Verify state update (delegated to ContextManager) expect(config.getGeminiMdFileCount()).toBe(1); expect(config.getGeminiMdFilePaths()).toEqual(['/path/to/GEMINI.md']); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 2e9102250c..64e78c1776 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -2056,6 +2056,43 @@ export class Config implements McpContext, AgentLoopContext { this.userMemory = newUserMemory; } + /** + * Returns memory for the system instruction. + * When JIT is enabled, only global memory (Tier 1) goes in the system + * instruction. Extension and project memory (Tier 2) are placed in the + * first user message instead, per the tiered context model. + */ + getSystemInstructionMemory(): string | HierarchicalMemory { + if (this.experimentalJitContext && this.contextManager) { + return this.contextManager.getGlobalMemory(); + } + return this.userMemory; + } + + /** + * Returns Tier 2 memory (extension + project) for injection into the first + * user message when JIT is enabled. Returns empty string when JIT is + * disabled (Tier 2 memory is already in the system instruction). + */ + getSessionMemory(): string { + if (!this.experimentalJitContext || !this.contextManager) { + return ''; + } + const sections: string[] = []; + const extension = this.contextManager.getExtensionMemory(); + const project = this.contextManager.getEnvironmentMemory(); + if (extension?.trim()) { + sections.push( + `\n${extension.trim()}\n`, + ); + } + if (project?.trim()) { + sections.push(`\n${project.trim()}\n`); + } + if (sections.length === 0) return ''; + return `\n\n${sections.join('\n')}\n`; + } + getGlobalMemory(): string { return this.contextManager?.getGlobalMemory() ?? ''; } diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 984ab2c199..77c4a5a498 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -216,6 +216,8 @@ describe('Gemini Client (client.ts)', () => { getUserMemory: vi.fn().mockReturnValue(''), getGlobalMemory: vi.fn().mockReturnValue(''), getEnvironmentMemory: vi.fn().mockReturnValue(''), + getSystemInstructionMemory: vi.fn().mockReturnValue(''), + getSessionMemory: vi.fn().mockReturnValue(''), isJitContextEnabled: vi.fn().mockReturnValue(false), getContextManager: vi.fn().mockReturnValue(undefined), getToolOutputMaskingEnabled: vi.fn().mockReturnValue(false), @@ -1961,12 +1963,11 @@ ${JSON.stringify( }); }); - it('should use getGlobalMemory for system instruction when JIT is enabled', async () => { + it('should use getSystemInstructionMemory for system instruction when JIT is enabled', async () => { vi.mocked(mockConfig.isJitContextEnabled).mockReturnValue(true); - vi.mocked(mockConfig.getGlobalMemory).mockReturnValue( + vi.mocked(mockConfig.getSystemInstructionMemory).mockReturnValue( 'Global JIT Memory', ); - vi.mocked(mockConfig.getUserMemory).mockReturnValue('Full JIT Memory'); const { getCoreSystemPrompt } = await import('./prompts.js'); const mockGetCoreSystemPrompt = vi.mocked(getCoreSystemPrompt); @@ -1975,13 +1976,15 @@ ${JSON.stringify( expect(mockGetCoreSystemPrompt).toHaveBeenCalledWith( mockConfig, - 'Full JIT Memory', + 'Global JIT Memory', ); }); - it('should use getUserMemory for system instruction when JIT is disabled', async () => { + it('should use getSystemInstructionMemory for system instruction when JIT is disabled', async () => { vi.mocked(mockConfig.isJitContextEnabled).mockReturnValue(false); - vi.mocked(mockConfig.getUserMemory).mockReturnValue('Legacy Memory'); + vi.mocked(mockConfig.getSystemInstructionMemory).mockReturnValue( + 'Legacy Memory', + ); const { getCoreSystemPrompt } = await import('./prompts.js'); const mockGetCoreSystemPrompt = vi.mocked(getCoreSystemPrompt); diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 985670c7da..c398a356ff 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -344,7 +344,7 @@ export class GeminiClient { return; } - const systemMemory = this.config.getUserMemory(); + const systemMemory = this.config.getSystemInstructionMemory(); const systemInstruction = getCoreSystemPrompt(this.config, systemMemory); this.getChat().setSystemInstruction(systemInstruction); } @@ -364,7 +364,7 @@ export class GeminiClient { const history = await getInitialChatHistory(this.config, extraHistory); try { - const systemMemory = this.config.getUserMemory(); + const systemMemory = this.config.getSystemInstructionMemory(); const systemInstruction = getCoreSystemPrompt(this.config, systemMemory); return new GeminiChat( this.config, @@ -1027,7 +1027,7 @@ export class GeminiClient { } = desiredModelConfig; try { - const userMemory = this.config.getUserMemory(); + const userMemory = this.config.getSystemInstructionMemory(); const systemInstruction = getCoreSystemPrompt(this.config, userMemory); const { model, diff --git a/packages/core/src/utils/environmentContext.test.ts b/packages/core/src/utils/environmentContext.test.ts index 42b2316955..51be00b61b 100644 --- a/packages/core/src/utils/environmentContext.test.ts +++ b/packages/core/src/utils/environmentContext.test.ts @@ -165,16 +165,27 @@ describe('getEnvironmentContext', () => { expect(getFolderStructure).not.toHaveBeenCalled(); }); - it('should exclude environment memory when JIT context is enabled', async () => { + it('should use session memory instead of environment memory when JIT context is enabled', async () => { (mockConfig as Record)['isJitContextEnabled'] = vi .fn() .mockReturnValue(true); + (mockConfig as Record)['getSessionMemory'] = vi + .fn() + .mockReturnValue( + '\n\n\nExt Memory\n\n\nProj Memory\n\n', + ); const parts = await getEnvironmentContext(mockConfig as Config); const context = parts[0].text; expect(context).not.toContain('Mock Environment Memory'); expect(mockConfig.getEnvironmentMemory).not.toHaveBeenCalled(); + expect(context).toContain(''); + expect(context).toContain(''); + expect(context).toContain('Ext Memory'); + expect(context).toContain(''); + expect(context).toContain('Proj Memory'); + expect(context).toContain(''); }); it('should include environment memory when JIT context is disabled', async () => { diff --git a/packages/core/src/utils/environmentContext.ts b/packages/core/src/utils/environmentContext.ts index d5bdd2d75b..abdf6faae9 100644 --- a/packages/core/src/utils/environmentContext.ts +++ b/packages/core/src/utils/environmentContext.ts @@ -57,11 +57,15 @@ export async function getEnvironmentContext(config: Config): Promise { ? await getDirectoryContextString(config) : ''; const tempDir = config.storage.getProjectTempDir(); - // When JIT context is enabled, project memory is already included in the - // system instruction via renderUserMemory(). Skip it here to avoid sending - // the same GEMINI.md content twice. + // Tiered context model (see issue #11488): + // - Tier 1 (global): system instruction only + // - Tier 2 (extension + project): first user message (here) + // - Tier 3 (subdirectory): tool output (JIT) + // When JIT is enabled, Tier 2 memory is provided by getSessionMemory(). + // When JIT is disabled, all memory is in the system instruction and + // getEnvironmentMemory() provides the project memory for this message. const environmentMemory = config.isJitContextEnabled?.() - ? '' + ? config.getSessionMemory() : config.getEnvironmentMemory(); const context = `