MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-07-10 01:50:20 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat(core): Land ContextCompressionService (#24483)

Browse Source
This commit is contained in:
joshualitt
2026-04-02 09:22:04 -07:00
committed by GitHub
parent 826bc3359a
commit aee2dd8577
32 changed files with 1160 additions and 229 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/config/config.test.ts
+14 -14
View File
@@ -221,8 +221,8 @@ vi.mock('../utils/fetch.js', () => ({
setGlobalProxy: mockSetGlobalProxy,
}));
vi.mock('../context/contextManager.js', () => ({
ContextManager: vi.fn().mockImplementation(() => ({
vi.mock('../context/memoryContextManager.js', () => ({
MemoryContextManager: vi.fn().mockImplementation(() => ({
refresh: vi.fn(),
getGlobalMemory: vi.fn().mockReturnValue(''),
getExtensionMemory: vi.fn().mockReturnValue(''),
@@ -237,7 +237,7 @@ import { tokenLimit } from '../core/tokenLimits.js';
import { getCodeAssistServer } from '../code_assist/codeAssist.js';
import { getExperiments } from '../code_assist/experiments/experiments.js';
import type { CodeAssistServer } from '../code_assist/server.js';
import { ContextManager } from '../context/contextManager.js';
import { MemoryContextManager } from '../context/memoryContextManager.js';
import { UserTierId } from '../code_assist/types.js';
import type {
ModelConfigService,
@@ -3022,11 +3022,11 @@ describe('Config Quota & Preview Model Access', () => {
describe('Config JIT Initialization', () => {
let config: Config;
let mockContextManager: ContextManager;
let mockMemoryContextManager: MemoryContextManager;
beforeEach(() => {
vi.clearAllMocks();
mockContextManager = {
mockMemoryContextManager = {
refresh: vi.fn(),
getGlobalMemory: vi.fn().mockReturnValue('Global Memory'),
getExtensionMemory: vi.fn().mockReturnValue('Extension Memory'),
@@ -3035,13 +3035,13 @@ describe('Config JIT Initialization', () => {
.mockReturnValue('Environment Memory\n\nMCP Instructions'),
getUserProjectMemory: vi.fn().mockReturnValue(''),
getLoadedPaths: vi.fn().mockReturnValue(new Set(['/path/to/GEMINI.md'])),
} as unknown as ContextManager;
(ContextManager as unknown as Mock).mockImplementation(
() => mockContextManager,
} as unknown as MemoryContextManager;
(MemoryContextManager as unknown as Mock).mockImplementation(
() => mockMemoryContextManager,
);
});
it('should initialize ContextManager, load memory, and delegate to it when experimentalJitContext is enabled', async () => {
it('should initialize MemoryContextManager, load memory, and delegate to it when experimentalJitContext is enabled', async () => {
const params: ConfigParameters = {
sessionId: 'test-session',
targetDir: '/tmp/test',
@@ -3055,8 +3055,8 @@ describe('Config JIT Initialization', () => {
config = new Config(params);
await config.initialize();
expect(ContextManager).toHaveBeenCalledWith(config);
expect(mockContextManager.refresh).toHaveBeenCalled();
expect(MemoryContextManager).toHaveBeenCalledWith(config);
expect(mockMemoryContextManager.refresh).toHaveBeenCalled();
expect(config.getUserMemory()).toEqual({
global: 'Global Memory',
extension: 'Extension Memory',
@@ -3079,12 +3079,12 @@ describe('Config JIT Initialization', () => {
expect(sessionMemory).toContain('</project_context>');
expect(sessionMemory).toContain('</loaded_context>');
// Verify state update (delegated to ContextManager)
// Verify state update (delegated to MemoryContextManager)
expect(config.getGeminiMdFileCount()).toBe(1);
expect(config.getGeminiMdFilePaths()).toEqual(['/path/to/GEMINI.md']);
});
it('should NOT initialize ContextManager when experimentalJitContext is disabled', async () => {
it('should NOT initialize MemoryContextManager when experimentalJitContext is disabled', async () => {
const params: ConfigParameters = {
sessionId: 'test-session',
targetDir: '/tmp/test',
@@ -3098,7 +3098,7 @@ describe('Config JIT Initialization', () => {
config = new Config(params);
await config.initialize();
expect(ContextManager).not.toHaveBeenCalled();
expect(MemoryContextManager).not.toHaveBeenCalled();
expect(config.getUserMemory()).toBe('Initial Memory');
});
packages/core/src/config/config.ts
+32 -55
View File
@@ -11,7 +11,11 @@ import { inspect } from 'node:util';
import process from 'node:process';
import { z } from 'zod';
import type { ConversationRecord } from '../services/chatRecordingService.js';
import type { AgentHistoryProviderConfig } from '../services/types.js';
import type {
AgentHistoryProviderConfig,
ContextManagementConfig,
ToolOutputMaskingConfig,
} from '../context/types.js';
export type { ConversationRecord };
import {
AuthType,
@@ -120,7 +124,7 @@ import {
type ModelConfigServiceConfig,
} from '../services/modelConfigService.js';
import { DEFAULT_MODEL_CONFIGS } from './defaultModelConfigs.js';
import { ContextManager } from '../context/contextManager.js';
import { MemoryContextManager } from '../context/memoryContextManager.js';
import { TrackerService } from '../services/trackerService.js';
import type { GenerateContentParameters } from '@google/genai';
@@ -210,32 +214,6 @@ export interface OutputSettings {
format?: OutputFormat;
}
export interface ToolOutputMaskingConfig {
protectionThresholdTokens: number;
minPrunableThresholdTokens: number;
protectLatestTurn: boolean;
}
export interface ContextManagementConfig {
enabled: boolean;
historyWindow: {
maxTokens: number;
retainedTokens: number;
};
messageLimits: {
normalMaxTokens: number;
retainedMaxTokens: number;
normalizationHeadRatio: number;
};
tools: {
distillation: {
maxOutputTokens: number;
summarizationThresholdTokens: number;
};
outputMasking: ToolOutputMaskingConfig;
};
}
export interface GemmaModelRouterSettings {
enabled?: boolean;
classifier?: {
@@ -962,7 +940,7 @@ export class Config implements McpContext, AgentLoopContext {
private readonly trackerEnabled: boolean;
private readonly planModeRoutingEnabled: boolean;
private readonly modelSteering: boolean;
private contextManager?: ContextManager;
private memoryContextManager?: MemoryContextManager;
private readonly contextManagement: ContextManagementConfig;
private terminalBackground: string | undefined = undefined;
private remoteAdminSettings: AdminControlsSettings | undefined;
@@ -1493,8 +1471,8 @@ export class Config implements McpContext, AgentLoopContext {
}
if (this.experimentalJitContext) {
this.contextManager = new ContextManager(this);
await this.contextManager.refresh();
this.memoryContextManager = new MemoryContextManager(this);
await this.memoryContextManager.refresh();
}
await this._geminiClient.initialize();
@@ -2302,12 +2280,12 @@ export class Config implements McpContext, AgentLoopContext {
}
getUserMemory(): string | HierarchicalMemory {
if (this.experimentalJitContext && this.contextManager) {
if (this.experimentalJitContext && this.memoryContextManager) {
return {
global: this.contextManager.getGlobalMemory(),
extension: this.contextManager.getExtensionMemory(),
project: this.contextManager.getEnvironmentMemory(),
userProjectMemory: this.contextManager.getUserProjectMemory(),
global: this.memoryContextManager.getGlobalMemory(),
extension: this.memoryContextManager.getExtensionMemory(),
project: this.memoryContextManager.getEnvironmentMemory(),
userProjectMemory: this.memoryContextManager.getUserProjectMemory(),
};
}
return this.userMemory;
@@ -2317,8 +2295,8 @@ export class Config implements McpContext, AgentLoopContext {
* Refreshes the MCP context, including memory, tools, and system instructions.
*/
async refreshMcpContext(): Promise<void> {
if (this.experimentalJitContext && this.contextManager) {
await this.contextManager.refresh();
if (this.experimentalJitContext && this.memoryContextManager) {
await this.memoryContextManager.refresh();
} else {
const { refreshServerHierarchicalMemory } = await import(
'../utils/memoryDiscovery.js'
@@ -2344,9 +2322,10 @@ export class Config implements McpContext, AgentLoopContext {
* via system instruction updates.
*/
getSystemInstructionMemory(): string | HierarchicalMemory {
if (this.experimentalJitContext && this.contextManager) {
const global = this.contextManager.getGlobalMemory();
const userProjectMemory = this.contextManager.getUserProjectMemory();
if (this.experimentalJitContext && this.memoryContextManager) {
const global = this.memoryContextManager.getGlobalMemory();
const userProjectMemory =
this.memoryContextManager.getUserProjectMemory();
if (userProjectMemory?.trim()) {
return { global, userProjectMemory };
}
@@ -2361,12 +2340,12 @@ export class Config implements McpContext, AgentLoopContext {
* disabled (Tier 2 memory is already in the system instruction).
*/
getSessionMemory(): string {
if (!this.experimentalJitContext || !this.contextManager) {
if (!this.experimentalJitContext || !this.memoryContextManager) {
return '';
}
const sections: string[] = [];
const extension = this.contextManager.getExtensionMemory();
const project = this.contextManager.getEnvironmentMemory();
const extension = this.memoryContextManager.getExtensionMemory();
const project = this.memoryContextManager.getEnvironmentMemory();
if (extension?.trim()) {
sections.push(
`<extension_context>\n${extension.trim()}\n</extension_context>`,
@@ -2380,22 +2359,22 @@ export class Config implements McpContext, AgentLoopContext {
}
getGlobalMemory(): string {
return this.contextManager?.getGlobalMemory() ?? '';
return this.memoryContextManager?.getGlobalMemory() ?? '';
}
getEnvironmentMemory(): string {
return this.contextManager?.getEnvironmentMemory() ?? '';
return this.memoryContextManager?.getEnvironmentMemory() ?? '';
}
getContextManager(): ContextManager | undefined {
return this.contextManager;
getMemoryContextManager(): MemoryContextManager | undefined {
return this.memoryContextManager;
}
isJitContextEnabled(): boolean {
return this.experimentalJitContext;
}
isAutoDistillationEnabled(): boolean {
isContextManagementEnabled(): boolean {
return this.contextManagement.enabled;
}
@@ -2413,8 +2392,6 @@ export class Config implements McpContext, AgentLoopContext {
get agentHistoryProviderConfig(): AgentHistoryProviderConfig {
return {
isTruncationEnabled: this.contextManagement.enabled,
isSummarizationEnabled: this.contextManagement.enabled,
maxTokens: this.contextManagement.historyWindow.maxTokens,
retainedTokens: this.contextManagement.historyWindow.retainedTokens,
normalMessageTokens: this.contextManagement.messageLimits.normalMaxTokens,
@@ -2471,8 +2448,8 @@ export class Config implements McpContext, AgentLoopContext {
}
getGeminiMdFileCount(): number {
if (this.experimentalJitContext && this.contextManager) {
return this.contextManager.getLoadedPaths().size;
if (this.experimentalJitContext && this.memoryContextManager) {
return this.memoryContextManager.getLoadedPaths().size;
}
return this.geminiMdFileCount;
}
@@ -2482,8 +2459,8 @@ export class Config implements McpContext, AgentLoopContext {
}
getGeminiMdFilePaths(): string[] {
if (this.experimentalJitContext && this.contextManager) {
return Array.from(this.contextManager.getLoadedPaths());
if (this.experimentalJitContext && this.memoryContextManager) {
return Array.from(this.memoryContextManager.getLoadedPaths());
}
return this.geminiMdFilePaths;
}