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

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

packages/core/src/config/config.ts

@@ -65,6 +65,7 @@ import { OutputFormat } from '../output/types.js';
 import type { ModelConfigServiceConfig } from '../services/modelConfigService.js';
 import { ModelConfigService } from '../services/modelConfigService.js';
 import { DEFAULT_MODEL_CONFIGS } from './defaultModelConfigs.js';
+import { ContextManager } from '../services/contextManager.js';
 
 // Re-export OAuth config type
 export type { MCPOAuthConfig, AnyToolInvocation };
@@ -309,6 +310,7 @@ export interface ConfigParameters {
   };
   previewFeatures?: boolean;
   enableModelAvailabilityService?: boolean;
+  experimentalJitContext?: boolean;
 }
 
 export class Config {
@@ -428,6 +430,9 @@ export class Config {
   private previewModelBypassMode = false;
   private readonly enableModelAvailabilityService: boolean;
 
+  private readonly experimentalJitContext: boolean;
+  private contextManager?: ContextManager;
+
   constructor(params: ConfigParameters) {
     this.sessionId = params.sessionId;
     this.embeddingModel =
@@ -486,6 +491,7 @@ export class Config {
     this.model = params.model;
     this.enableModelAvailabilityService =
       params.enableModelAvailabilityService ?? false;
+    this.experimentalJitContext = params.experimentalJitContext ?? false;
     this.modelAvailabilityService = new ModelAvailabilityService();
     this.previewFeatures = params.previewFeatures ?? undefined;
     this.maxSessionTurns = params.maxSessionTurns ?? -1;
@@ -651,6 +657,10 @@ export class Config {
       await this.hookSystem.initialize();
     }
 
+    if (this.experimentalJitContext) {
+      this.contextManager = new ContextManager(this);
+    }
+
     await this.geminiClient.initialize();
   }
 
@@ -958,6 +968,22 @@ export class Config {
     this.userMemory = newUserMemory;
   }
 
+  getGlobalMemory(): string {
+    return this.contextManager?.getGlobalMemory() ?? '';
+  }
+
+  getEnvironmentMemory(): string {
+    return this.contextManager?.getEnvironmentMemory() ?? '';
+  }
+
+  getContextManager(): ContextManager | undefined {
+    return this.contextManager;
+  }
+
+  isJitContextEnabled(): boolean {
+    return this.experimentalJitContext;
+  }
+
   getGeminiMdFileCount(): number {
     return this.geminiMdFileCount;
   }

packages/core/src/index.ts

@@ -82,6 +82,7 @@ export * from './services/fileDiscoveryService.js';
 export * from './services/gitService.js';
 export * from './services/chatRecordingService.js';
 export * from './services/fileSystemService.js';
+export * from './services/contextManager.js';
 
 // Export IDE specific logic
 export * from './ide/ide-client.js';

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;
+  }
+}

packages/core/src/utils/memoryDiscovery.ts

@@ -34,7 +34,7 @@ const logger = {
     console.error('[ERROR] [MemoryDiscovery]', ...args),
 };
 
-interface GeminiFileContent {
+export interface GeminiFileContent {
   filePath: string;
   content: string | null;
 }
@@ -304,7 +304,7 @@ async function readGeminiMdFiles(
   return results;
 }
 
-function concatenateInstructions(
+export function concatenateInstructions(
   instructionContents: GeminiFileContent[],
   // CWD is needed to resolve relative paths for display markers
   currentWorkingDirectoryForDisplay: string,