refactor(core): align JIT memory placement with tiered context model (#22766)

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('<loaded_context>');
+    expect(sessionMemory).toContain('<extension_context>');
+    expect(sessionMemory).toContain('Extension Memory');
+    expect(sessionMemory).toContain('</extension_context>');
+    expect(sessionMemory).toContain('<project_context>');
+    expect(sessionMemory).toContain('Environment Memory');
+    expect(sessionMemory).toContain('MCP Instructions');
+    expect(sessionMemory).toContain('</project_context>');
+    expect(sessionMemory).toContain('</loaded_context>');
+
     // Verify state update (delegated to ContextManager)
     expect(config.getGeminiMdFileCount()).toBe(1);
     expect(config.getGeminiMdFilePaths()).toEqual(['/path/to/GEMINI.md']);

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(
+        `<extension_context>\n${extension.trim()}\n</extension_context>`,
+      );
+    }
+    if (project?.trim()) {
+      sections.push(`<project_context>\n${project.trim()}\n</project_context>`);
+    }
+    if (sections.length === 0) return '';
+    return `\n<loaded_context>\n${sections.join('\n')}\n</loaded_context>`;
+  }
+
   getGlobalMemory(): string {
     return this.contextManager?.getGlobalMemory() ?? '';
   }

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);

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,

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<string, unknown>)['isJitContextEnabled'] = vi
       .fn()
       .mockReturnValue(true);
+    (mockConfig as Record<string, unknown>)['getSessionMemory'] = vi
+      .fn()
+      .mockReturnValue(
+        '\n<loaded_context>\n<extension_context>\nExt Memory\n</extension_context>\n<project_context>\nProj Memory\n</project_context>\n</loaded_context>',
+      );
 
     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('<loaded_context>');
+    expect(context).toContain('<extension_context>');
+    expect(context).toContain('Ext Memory');
+    expect(context).toContain('<project_context>');
+    expect(context).toContain('Proj Memory');
+    expect(context).toContain('</loaded_context>');
   });
 
   it('should include environment memory when JIT context is disabled', async () => {

packages/core/src/utils/environmentContext.ts

@@ -57,11 +57,15 @@ export async function getEnvironmentContext(config: Config): Promise<Part[]> {
     ? 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 = `