feat(core): Implement JIT context manager and setting (#14324)

Co-authored-by: Abhi <43648792+abhipatel12@users.noreply.github.com>

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

@@ -0,0 +1,148 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+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';
+
+// Mock memoryDiscovery module
+vi.mock('../utils/memoryDiscovery.js', async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import('../utils/memoryDiscovery.js')>();
+  return {
+    ...actual,
+    loadGlobalMemory: vi.fn(),
+    loadEnvironmentMemory: vi.fn(),
+    loadJitSubdirectoryMemory: vi.fn(),
+  };
+});
+
+describe('ContextManager', () => {
+  let contextManager: ContextManager;
+  let mockConfig: Config;
+
+  beforeEach(() => {
+    mockConfig = {
+      getDebugMode: vi.fn().mockReturnValue(false),
+      getWorkingDir: vi.fn().mockReturnValue('/app'),
+    } as unknown as Config;
+
+    contextManager = new ContextManager(mockConfig);
+    vi.clearAllMocks();
+  });
+
+  describe('loadGlobalMemory', () => {
+    it('should load and format global memory', async () => {
+      const mockResult: 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',
+      );
+      expect(contextManager.getGlobalMemory()).toBe(result);
+    });
+  });
+
+  describe('loadEnvironmentMemory', () => {
+    it('should load and format environment memory', async () => {
+      const mockResult: memoryDiscovery.MemoryLoadResult = {
+        files: [{ path: '/app/GEMINI.md', content: 'Env Content' }],
+      };
+      vi.mocked(memoryDiscovery.loadEnvironmentMemory).mockResolvedValue(
+        mockResult,
+      );
+      const mockExtensionLoader = {} as unknown as ExtensionLoader;
+
+      const result = await contextManager.loadEnvironmentMemory(
+        ['/app'],
+        mockExtensionLoader,
+      );
+
+      expect(memoryDiscovery.loadEnvironmentMemory).toHaveBeenCalledWith(
+        ['/app'],
+        mockExtensionLoader,
+        false,
+      );
+      expect(result).toContain('--- Context from: GEMINI.md ---');
+      expect(result).toContain('Env Content');
+      expect(contextManager.getLoadedPaths()).toContain('/app/GEMINI.md');
+      expect(contextManager.getEnvironmentMemory()).toBe(result);
+    });
+  });
+
+  describe('discoverContext', () => {
+    it('should discover and load new context', async () => {
+      const mockResult: memoryDiscovery.MemoryLoadResult = {
+        files: [{ path: '/app/src/GEMINI.md', content: 'Src Content' }],
+      };
+      vi.mocked(memoryDiscovery.loadJitSubdirectoryMemory).mockResolvedValue(
+        mockResult,
+      );
+
+      const result = await contextManager.discoverContext('/app/src/file.ts', [
+        '/app',
+      ]);
+
+      expect(memoryDiscovery.loadJitSubdirectoryMemory).toHaveBeenCalledWith(
+        '/app/src/file.ts',
+        ['/app'],
+        expect.any(Set),
+        false,
+      );
+      expect(result).toMatch(/--- Context from: src[\\/]GEMINI\.md ---/);
+      expect(result).toContain('Src Content');
+      expect(contextManager.getLoadedPaths()).toContain('/app/src/GEMINI.md');
+    });
+
+    it('should return empty string if no new files found', async () => {
+      const mockResult: memoryDiscovery.MemoryLoadResult = { files: [] };
+      vi.mocked(memoryDiscovery.loadJitSubdirectoryMemory).mockResolvedValue(
+        mockResult,
+      );
+
+      const result = await contextManager.discoverContext('/app/src/file.ts', [
+        '/app',
+      ]);
+
+      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

@@ -0,0 +1,111 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  loadGlobalMemory,
+  loadEnvironmentMemory,
+  loadJitSubdirectoryMemory,
+  concatenateInstructions,
+} from '../utils/memoryDiscovery.js';
+import type { ExtensionLoader } from '../utils/extensionLoader.js';
+import type { Config } from '../config/config.js';
+
+export class ContextManager {
+  private readonly loadedPaths: Set<string> = new Set();
+  private readonly config: Config;
+  private globalMemory: string = '';
+  private environmentMemory: string = '';
+
+  constructor(config: Config) {
+    this.config = config;
+  }
+
+  /**
+   * Loads the global memory (Tier 1) and returns the formatted content.
+   */
+  async loadGlobalMemory(): Promise<string> {
+    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> {
+    const result = await loadEnvironmentMemory(
+      trustedRoots,
+      extensionLoader,
+      this.config.getDebugMode(),
+    );
+    this.markAsLoaded(result.files.map((f) => f.path));
+    this.environmentMemory = concatenateInstructions(
+      result.files.map((f) => ({ filePath: f.path, content: f.content })),
+      this.config.getWorkingDir(),
+    );
+    return this.environmentMemory;
+  }
+
+  /**
+   * Discovers and loads context for a specific accessed path (Tier 3 - JIT).
+   * Traverses upwards from the accessed path to the project root.
+   */
+  async discoverContext(
+    accessedPath: string,
+    trustedRoots: string[],
+  ): Promise<string> {
+    const result = await loadJitSubdirectoryMemory(
+      accessedPath,
+      trustedRoots,
+      this.loadedPaths,
+      this.config.getDebugMode(),
+    );
+
+    if (result.files.length === 0) {
+      return '';
+    }
+
+    this.markAsLoaded(result.files.map((f) => f.path));
+    return concatenateInstructions(
+      result.files.map((f) => ({ filePath: f.path, content: f.content })),
+      this.config.getWorkingDir(),
+    );
+  }
+
+  getGlobalMemory(): string {
+    return this.globalMemory;
+  }
+
+  getEnvironmentMemory(): string {
+    return this.environmentMemory;
+  }
+
+  private markAsLoaded(paths: string[]): void {
+    for (const p of paths) {
+      this.loadedPaths.add(p);
+    }
+  }
+
+  /**
+   * 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;
+  }
+}