feat(core): Implement JIT context memory loading and UI sync (#14469)

packages/core/src/services/contextManager.test.ts

@@ -8,7 +8,7 @@ import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { ContextManager } from './contextManager.js';
 import * as memoryDiscovery from '../utils/memoryDiscovery.js';
 import type { Config } from '../config/config.js';
-import type { ExtensionLoader } from '../utils/extensionLoader.js';
+import { coreEvents, CoreEvent } from '../utils/events.js';
 
 // Mock memoryDiscovery module
 vi.mock('../utils/memoryDiscovery.js', async (importOriginal) => {
@@ -19,6 +19,9 @@ vi.mock('../utils/memoryDiscovery.js', async (importOriginal) => {
     loadGlobalMemory: vi.fn(),
     loadEnvironmentMemory: vi.fn(),
     loadJitSubdirectoryMemory: vi.fn(),
+    concatenateInstructions: vi
+      .fn()
+      .mockImplementation(actual.concatenateInstructions),
   };
 });
 
@@ -30,58 +33,84 @@ describe('ContextManager', () => {
     mockConfig = {
       getDebugMode: vi.fn().mockReturnValue(false),
       getWorkingDir: vi.fn().mockReturnValue('/app'),
+      getWorkspaceContext: vi.fn().mockReturnValue({
+        getDirectories: vi.fn().mockReturnValue(['/app']),
+      }),
+      getExtensionLoader: vi.fn().mockReturnValue({}),
+      getMcpClientManager: vi.fn().mockReturnValue({
+        getMcpInstructions: vi.fn().mockReturnValue('MCP Instructions'),
+      }),
     } as unknown as Config;
 
     contextManager = new ContextManager(mockConfig);
     vi.clearAllMocks();
+    vi.spyOn(coreEvents, 'emit');
   });
 
-  describe('loadGlobalMemory', () => {
-    it('should load and format global memory', async () => {
-      const mockResult: memoryDiscovery.MemoryLoadResult = {
+  describe('refresh', () => {
+    it('should load and format global and environment memory', async () => {
+      const mockGlobalResult: memoryDiscovery.MemoryLoadResult = {
         files: [
           { path: '/home/user/.gemini/GEMINI.md', content: 'Global Content' },
         ],
       };
-      vi.mocked(memoryDiscovery.loadGlobalMemory).mockResolvedValue(mockResult);
-
-      const result = await contextManager.loadGlobalMemory();
-
-      expect(memoryDiscovery.loadGlobalMemory).toHaveBeenCalledWith(false);
-      // The path will be relative to CWD (/app), so it might contain ../
-      expect(result).toMatch(/--- Context from: .*GEMINI.md ---/);
-      expect(result).toContain('Global Content');
-      expect(contextManager.getLoadedPaths()).toContain(
-        '/home/user/.gemini/GEMINI.md',
+      vi.mocked(memoryDiscovery.loadGlobalMemory).mockResolvedValue(
+        mockGlobalResult,
       );
-      expect(contextManager.getGlobalMemory()).toBe(result);
-    });
-  });
 
-  describe('loadEnvironmentMemory', () => {
-    it('should load and format environment memory', async () => {
-      const mockResult: memoryDiscovery.MemoryLoadResult = {
+      const mockEnvResult: memoryDiscovery.MemoryLoadResult = {
         files: [{ path: '/app/GEMINI.md', content: 'Env Content' }],
       };
       vi.mocked(memoryDiscovery.loadEnvironmentMemory).mockResolvedValue(
-        mockResult,
+        mockEnvResult,
       );
-      const mockExtensionLoader = {} as unknown as ExtensionLoader;
 
-      const result = await contextManager.loadEnvironmentMemory(
-        ['/app'],
-        mockExtensionLoader,
+      await contextManager.refresh();
+
+      expect(memoryDiscovery.loadGlobalMemory).toHaveBeenCalledWith(false);
+      expect(contextManager.getGlobalMemory()).toMatch(
+        /--- Context from: .*GEMINI.md ---/,
       );
+      expect(contextManager.getGlobalMemory()).toContain('Global Content');
 
       expect(memoryDiscovery.loadEnvironmentMemory).toHaveBeenCalledWith(
         ['/app'],
-        mockExtensionLoader,
+        expect.anything(),
         false,
       );
-      expect(result).toContain('--- Context from: GEMINI.md ---');
-      expect(result).toContain('Env Content');
+      expect(contextManager.getEnvironmentMemory()).toContain(
+        '--- Context from: GEMINI.md ---',
+      );
+      expect(contextManager.getEnvironmentMemory()).toContain('Env Content');
+      expect(contextManager.getEnvironmentMemory()).toContain(
+        'MCP Instructions',
+      );
+
+      expect(contextManager.getLoadedPaths()).toContain(
+        '/home/user/.gemini/GEMINI.md',
+      );
       expect(contextManager.getLoadedPaths()).toContain('/app/GEMINI.md');
-      expect(contextManager.getEnvironmentMemory()).toBe(result);
+    });
+
+    it('should emit MemoryChanged event when memory is refreshed', async () => {
+      const mockGlobalResult = {
+        files: [{ path: '/app/GEMINI.md', content: 'content' }],
+      };
+      const mockEnvResult = {
+        files: [{ path: '/app/src/GEMINI.md', content: 'env content' }],
+      };
+      vi.mocked(memoryDiscovery.loadGlobalMemory).mockResolvedValue(
+        mockGlobalResult,
+      );
+      vi.mocked(memoryDiscovery.loadEnvironmentMemory).mockResolvedValue(
+        mockEnvResult,
+      );
+
+      await contextManager.refresh();
+
+      expect(coreEvents.emit).toHaveBeenCalledWith(CoreEvent.MemoryChanged, {
+        fileCount: 2,
+      });
     });
   });
 
@@ -122,27 +151,4 @@ describe('ContextManager', () => {
       expect(result).toBe('');
     });
   });
-
-  describe('reset', () => {
-    it('should clear loaded paths and memory', async () => {
-      // Setup some state
-      const mockResult: memoryDiscovery.MemoryLoadResult = {
-        files: [
-          { path: '/home/user/.gemini/GEMINI.md', content: 'Global Content' },
-        ],
-      };
-      vi.mocked(memoryDiscovery.loadGlobalMemory).mockResolvedValue(mockResult);
-      await contextManager.loadGlobalMemory();
-
-      expect(contextManager.getLoadedPaths().size).toBeGreaterThan(0);
-      expect(contextManager.getGlobalMemory()).toBeTruthy();
-
-      // Reset
-      contextManager.reset();
-
-      expect(contextManager.getLoadedPaths().size).toBe(0);
-      expect(contextManager.getGlobalMemory()).toBe('');
-      expect(contextManager.getEnvironmentMemory()).toBe('');
-    });
-  });
 });

packages/core/src/services/contextManager.ts

@@ -10,8 +10,8 @@ import {
   loadJitSubdirectoryMemory,
   concatenateInstructions,
 } from '../utils/memoryDiscovery.js';
-import type { ExtensionLoader } from '../utils/extensionLoader.js';
 import type { Config } from '../config/config.js';
+import { coreEvents, CoreEvent } from '../utils/events.js';
 
 export class ContextManager {
   private readonly loadedPaths: Set<string> = new Set();
@@ -24,36 +24,40 @@ export class ContextManager {
   }
 
   /**
-   * Loads the global memory (Tier 1) and returns the formatted content.
+   * Refreshes the memory by reloading global and environment memory.
    */
-  async loadGlobalMemory(): Promise<string> {
+  async refresh(): Promise<void> {
+    this.loadedPaths.clear();
+    await this.loadGlobalMemory();
+    await this.loadEnvironmentMemory();
+    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(),
     );
-    return this.globalMemory;
   }
 
-  /**
-   * Loads the environment memory (Tier 2) and returns the formatted content.
-   */
-  async loadEnvironmentMemory(
-    trustedRoots: string[],
-    extensionLoader: ExtensionLoader,
-  ): Promise<string> {
+  private async loadEnvironmentMemory(): Promise<void> {
     const result = await loadEnvironmentMemory(
-      trustedRoots,
-      extensionLoader,
+      [...this.config.getWorkspaceContext().getDirectories()],
+      this.config.getExtensionLoader(),
       this.config.getDebugMode(),
     );
     this.markAsLoaded(result.files.map((f) => f.path));
-    this.environmentMemory = concatenateInstructions(
+    const envMemory = concatenateInstructions(
       result.files.map((f) => ({ filePath: f.path, content: f.content })),
       this.config.getWorkingDir(),
     );
-    return this.environmentMemory;
+    const mcpInstructions =
+      this.config.getMcpClientManager()?.getMcpInstructions() || '';
+    this.environmentMemory = [envMemory, mcpInstructions.trimStart()]
+      .filter(Boolean)
+      .join('\n\n');
   }
 
   /**
@@ -82,6 +86,12 @@ export class ContextManager {
     );
   }
 
+  private emitMemoryChanged(): void {
+    coreEvents.emit(CoreEvent.MemoryChanged, {
+      fileCount: this.loadedPaths.size,
+    });
+  }
+
   getGlobalMemory(): string {
     return this.globalMemory;
   }
@@ -96,15 +106,6 @@ export class ContextManager {
     }
   }
 
-  /**
-   * Resets the loaded paths tracking and memory. Useful for testing or full reloads.
-   */
-  reset(): void {
-    this.loadedPaths.clear();
-    this.globalMemory = '';
-    this.environmentMemory = '';
-  }
-
   getLoadedPaths(): ReadonlySet<string> {
     return this.loadedPaths;
   }