From 752a521423630589e49f9b5c1aed3b05173f686f Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Wed, 3 Dec 2025 04:09:46 +0800 Subject: [PATCH] feat(core): Implement JIT context manager and setting (#14324) Co-authored-by: Abhi <43648792+abhipatel12@users.noreply.github.com> --- docs/get-started/configuration.md | 5 + packages/cli/src/config/config.ts | 1 + packages/cli/src/config/settingsSchema.ts | 9 ++ packages/core/src/config/config.ts | 26 +++ packages/core/src/index.ts | 1 + .../core/src/services/contextManager.test.ts | 148 ++++++++++++++++++ packages/core/src/services/contextManager.ts | 111 +++++++++++++ packages/core/src/utils/memoryDiscovery.ts | 4 +- schemas/settings.schema.json | 7 + 9 files changed, 310 insertions(+), 2 deletions(-) create mode 100644 packages/core/src/services/contextManager.test.ts create mode 100644 packages/core/src/services/contextManager.ts diff --git a/docs/get-started/configuration.md b/docs/get-started/configuration.md index 6f1e46a18c..255e828571 100644 --- a/docs/get-started/configuration.md +++ b/docs/get-started/configuration.md @@ -782,6 +782,11 @@ their corresponding top-level category object in your `settings.json` file. - **Default:** `false` - **Requires restart:** Yes +- **`experimental.jitContext`** (boolean): + - **Description:** Enable Just-In-Time (JIT) context loading. + - **Default:** `false` + - **Requires restart:** Yes + - **`experimental.codebaseInvestigatorSettings.enabled`** (boolean): - **Description:** Enable the Codebase Investigator agent. - **Default:** `true` diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 3a85a0760a..7625a68d5d 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -635,6 +635,7 @@ export async function loadCliConfig( enableExtensionReloading: settings.experimental?.extensionReloading, enableModelAvailabilityService: settings.experimental?.isModelAvailabilityServiceEnabled, + experimentalJitContext: settings.experimental?.jitContext, noBrowser: !!process.env['NO_BROWSER'], summarizeToolOutput: settings.model?.summarizeToolOutput, ideMode, diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index d4ea853727..8821deab7c 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -1310,6 +1310,15 @@ const SETTINGS_SCHEMA = { description: 'Enable model routing using new availability service.', showInDialog: false, }, + jitContext: { + type: 'boolean', + label: 'JIT Context Loading', + category: 'Experimental', + requiresRestart: true, + default: false, + description: 'Enable Just-In-Time (JIT) context loading.', + showInDialog: false, + }, codebaseInvestigatorSettings: { type: 'object', label: 'Codebase Investigator Settings', diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 4693810f0f..11674a5d7e 100644 --- a/packages/core/src/config/config.ts +++ b/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; } diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index d3c711473a..d69f4c1bd7 100644 --- a/packages/core/src/index.ts +++ b/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'; diff --git a/packages/core/src/services/contextManager.test.ts b/packages/core/src/services/contextManager.test.ts new file mode 100644 index 0000000000..4f33c9f62d --- /dev/null +++ b/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(); + 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(''); + }); + }); +}); diff --git a/packages/core/src/services/contextManager.ts b/packages/core/src/services/contextManager.ts new file mode 100644 index 0000000000..dd094a10de --- /dev/null +++ b/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 = 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 { + 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 { + 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 { + 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 { + return this.loadedPaths; + } +} diff --git a/packages/core/src/utils/memoryDiscovery.ts b/packages/core/src/utils/memoryDiscovery.ts index e91cfd2ff1..cbbbad880f 100644 --- a/packages/core/src/utils/memoryDiscovery.ts +++ b/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, diff --git a/schemas/settings.schema.json b/schemas/settings.schema.json index 0d5c317912..dbac18cf55 100644 --- a/schemas/settings.schema.json +++ b/schemas/settings.schema.json @@ -1279,6 +1279,13 @@ "default": false, "type": "boolean" }, + "jitContext": { + "title": "JIT Context Loading", + "description": "Enable Just-In-Time (JIT) context loading.", + "markdownDescription": "Enable Just-In-Time (JIT) context loading.\n\n- Category: `Experimental`\n- Requires restart: `yes`\n- Default: `false`", + "default": false, + "type": "boolean" + }, "codebaseInvestigatorSettings": { "title": "Codebase Investigator Settings", "description": "Configuration for Codebase Investigator.",