From 6ae75c9f32a968efa50857a8f24b958a58a84fd6 Mon Sep 17 00:00:00 2001 From: Jacob MacDonald Date: Thu, 7 Aug 2025 07:34:40 -0700 Subject: [PATCH] Add a context percentage threshold setting for auto compression (#5721) --- docs/cli/configuration.md | 12 +++ packages/cli/src/config/config.test.ts | 39 ++++++++ packages/cli/src/config/config.ts | 1 + packages/cli/src/config/settings.test.ts | 118 +++++++++++++++++++++++ packages/cli/src/config/settings.ts | 21 ++++ packages/core/src/config/config.ts | 11 +++ packages/core/src/core/client.test.ts | 10 +- packages/core/src/core/client.ts | 14 ++- 8 files changed, 219 insertions(+), 7 deletions(-) diff --git a/docs/cli/configuration.md b/docs/cli/configuration.md index 5c917a3f..9fc74adb 100644 --- a/docs/cli/configuration.md +++ b/docs/cli/configuration.md @@ -268,6 +268,18 @@ In addition to a project settings file, a project's `.gemini` directory can cont "loadMemoryFromIncludeDirectories": true ``` +- **`chatCompression`** (object): + - **Description:** Controls the settings for chat history compression, both automatic and + when manually invoked through the /compress command. + - **Properties:** + - **`contextPercentageThreshold`** (number): A value between 0 and 1 that specifies the token threshold for compression as a percentage of the model's total token limit. For example, a value of `0.6` will trigger compression when the chat history exceeds 60% of the token limit. + - **Example:** + ```json + "chatCompression": { + "contextPercentageThreshold": 0.6 + } + ``` + ### Example `settings.json`: ```json diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index 6a7e3b57..b670fbc8 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -1123,3 +1123,42 @@ describe('loadCliConfig with includeDirectories', () => { ); }); }); + +describe('loadCliConfig chatCompression', () => { + const originalArgv = process.argv; + const originalEnv = { ...process.env }; + + beforeEach(() => { + vi.resetAllMocks(); + vi.mocked(os.homedir).mockReturnValue('/mock/home/user'); + process.env.GEMINI_API_KEY = 'test-api-key'; + }); + + afterEach(() => { + process.argv = originalArgv; + process.env = originalEnv; + vi.restoreAllMocks(); + }); + + it('should pass chatCompression settings to the core config', async () => { + process.argv = ['node', 'script.js']; + const argv = await parseArguments(); + const settings: Settings = { + chatCompression: { + contextPercentageThreshold: 0.5, + }, + }; + const config = await loadCliConfig(settings, [], 'test-session', argv); + expect(config.getChatCompression()).toEqual({ + contextPercentageThreshold: 0.5, + }); + }); + + it('should have undefined chatCompression if not in settings', async () => { + process.argv = ['node', 'script.js']; + const argv = await parseArguments(); + const settings: Settings = {}; + const config = await loadCliConfig(settings, [], 'test-session', argv); + expect(config.getChatCompression()).toBeUndefined(); + }); +}); diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 2c942c08..a47d8301 100644 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -482,6 +482,7 @@ export async function loadCliConfig( summarizeToolOutput: settings.summarizeToolOutput, ideMode, ideModeFeature, + chatCompression: settings.chatCompression, folderTrustFeature, }); } diff --git a/packages/cli/src/config/settings.test.ts b/packages/cli/src/config/settings.test.ts index d0266720..f68b13e3 100644 --- a/packages/cli/src/config/settings.test.ts +++ b/packages/cli/src/config/settings.test.ts @@ -113,6 +113,7 @@ describe('Settings Loading and Merging', () => { customThemes: {}, mcpServers: {}, includeDirectories: [], + chatCompression: {}, }); expect(settings.errors.length).toBe(0); }); @@ -147,6 +148,7 @@ describe('Settings Loading and Merging', () => { customThemes: {}, mcpServers: {}, includeDirectories: [], + chatCompression: {}, }); }); @@ -181,6 +183,7 @@ describe('Settings Loading and Merging', () => { customThemes: {}, mcpServers: {}, includeDirectories: [], + chatCompression: {}, }); }); @@ -213,6 +216,7 @@ describe('Settings Loading and Merging', () => { customThemes: {}, mcpServers: {}, includeDirectories: [], + chatCompression: {}, }); }); @@ -251,6 +255,7 @@ describe('Settings Loading and Merging', () => { customThemes: {}, mcpServers: {}, includeDirectories: [], + chatCompression: {}, }); }); @@ -301,6 +306,7 @@ describe('Settings Loading and Merging', () => { customThemes: {}, mcpServers: {}, includeDirectories: [], + chatCompression: {}, }); }); @@ -622,6 +628,116 @@ describe('Settings Loading and Merging', () => { expect(settings.merged.mcpServers).toEqual({}); }); + it('should merge chatCompression settings, with workspace taking precedence', () => { + (mockFsExistsSync as Mock).mockReturnValue(true); + const userSettingsContent = { + chatCompression: { contextPercentageThreshold: 0.5 }, + }; + const workspaceSettingsContent = { + chatCompression: { contextPercentageThreshold: 0.8 }, + }; + + (fs.readFileSync as Mock).mockImplementation( + (p: fs.PathOrFileDescriptor) => { + if (p === USER_SETTINGS_PATH) + return JSON.stringify(userSettingsContent); + if (p === MOCK_WORKSPACE_SETTINGS_PATH) + return JSON.stringify(workspaceSettingsContent); + return '{}'; + }, + ); + + const settings = loadSettings(MOCK_WORKSPACE_DIR); + + expect(settings.user.settings.chatCompression).toEqual({ + contextPercentageThreshold: 0.5, + }); + expect(settings.workspace.settings.chatCompression).toEqual({ + contextPercentageThreshold: 0.8, + }); + expect(settings.merged.chatCompression).toEqual({ + contextPercentageThreshold: 0.8, + }); + }); + + it('should handle chatCompression when only in user settings', () => { + (mockFsExistsSync as Mock).mockImplementation( + (p: fs.PathLike) => p === USER_SETTINGS_PATH, + ); + const userSettingsContent = { + chatCompression: { contextPercentageThreshold: 0.5 }, + }; + (fs.readFileSync as Mock).mockImplementation( + (p: fs.PathOrFileDescriptor) => { + if (p === USER_SETTINGS_PATH) + return JSON.stringify(userSettingsContent); + return '{}'; + }, + ); + + const settings = loadSettings(MOCK_WORKSPACE_DIR); + expect(settings.merged.chatCompression).toEqual({ + contextPercentageThreshold: 0.5, + }); + }); + + it('should have chatCompression as an empty object if not in any settings file', () => { + (mockFsExistsSync as Mock).mockReturnValue(false); // No settings files exist + (fs.readFileSync as Mock).mockReturnValue('{}'); + const settings = loadSettings(MOCK_WORKSPACE_DIR); + expect(settings.merged.chatCompression).toEqual({}); + }); + + it('should ignore chatCompression if contextPercentageThreshold is invalid', () => { + const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); + (mockFsExistsSync as Mock).mockImplementation( + (p: fs.PathLike) => p === USER_SETTINGS_PATH, + ); + const userSettingsContent = { + chatCompression: { contextPercentageThreshold: 1.5 }, + }; + (fs.readFileSync as Mock).mockImplementation( + (p: fs.PathOrFileDescriptor) => { + if (p === USER_SETTINGS_PATH) + return JSON.stringify(userSettingsContent); + return '{}'; + }, + ); + + const settings = loadSettings(MOCK_WORKSPACE_DIR); + expect(settings.merged.chatCompression).toBeUndefined(); + expect(warnSpy).toHaveBeenCalledWith( + 'Invalid value for chatCompression.contextPercentageThreshold: "1.5". Please use a value between 0 and 1. Using default compression settings.', + ); + warnSpy.mockRestore(); + }); + + it('should deep merge chatCompression settings', () => { + (mockFsExistsSync as Mock).mockReturnValue(true); + const userSettingsContent = { + chatCompression: { contextPercentageThreshold: 0.5 }, + }; + const workspaceSettingsContent = { + chatCompression: {}, + }; + + (fs.readFileSync as Mock).mockImplementation( + (p: fs.PathOrFileDescriptor) => { + if (p === USER_SETTINGS_PATH) + return JSON.stringify(userSettingsContent); + if (p === MOCK_WORKSPACE_SETTINGS_PATH) + return JSON.stringify(workspaceSettingsContent); + return '{}'; + }, + ); + + const settings = loadSettings(MOCK_WORKSPACE_DIR); + + expect(settings.merged.chatCompression).toEqual({ + contextPercentageThreshold: 0.5, + }); + }); + it('should merge includeDirectories from all scopes', () => { (mockFsExistsSync as Mock).mockReturnValue(true); const systemSettingsContent = { @@ -695,6 +811,7 @@ describe('Settings Loading and Merging', () => { customThemes: {}, mcpServers: {}, includeDirectories: [], + chatCompression: {}, }); // Check that error objects are populated in settings.errors @@ -1132,6 +1249,7 @@ describe('Settings Loading and Merging', () => { customThemes: {}, mcpServers: {}, includeDirectories: [], + chatCompression: {}, }); }); }); diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts index 64500845..8005ad65 100644 --- a/packages/cli/src/config/settings.ts +++ b/packages/cli/src/config/settings.ts @@ -13,6 +13,7 @@ import { GEMINI_CONFIG_DIR as GEMINI_DIR, getErrorMessage, BugCommandSettings, + ChatCompressionSettings, TelemetrySettings, AuthType, } from '@google/gemini-cli-core'; @@ -134,6 +135,8 @@ export interface Settings { includeDirectories?: string[]; loadMemoryFromIncludeDirectories?: boolean; + + chatCompression?: ChatCompressionSettings; } export interface SettingsError { @@ -194,6 +197,11 @@ export class LoadedSettings { ...(user.includeDirectories || []), ...(workspace.includeDirectories || []), ], + chatCompression: { + ...(system.chatCompression || {}), + ...(user.chatCompression || {}), + ...(workspace.chatCompression || {}), + }, }; } @@ -482,6 +490,19 @@ export function loadSettings(workspaceDir: string): LoadedSettings { settingsErrors, ); + // Validate chatCompression settings + const chatCompression = loadedSettings.merged.chatCompression; + const threshold = chatCompression?.contextPercentageThreshold; + if ( + threshold != null && + (typeof threshold !== 'number' || threshold < 0 || threshold > 1) + ) { + console.warn( + `Invalid value for chatCompression.contextPercentageThreshold: "${threshold}". Please use a value between 0 and 1. Using default compression settings.`, + ); + delete loadedSettings.merged.chatCompression; + } + // Load environment with merged settings loadEnvironment(loadedSettings.merged); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 005573da..4848bfb6 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -69,6 +69,10 @@ export interface BugCommandSettings { urlTemplate: string; } +export interface ChatCompressionSettings { + contextPercentageThreshold?: number; +} + export interface SummarizeToolOutputSettings { tokenBudget?: number; } @@ -191,6 +195,7 @@ export interface ConfigParameters { folderTrustFeature?: boolean; ideMode?: boolean; loadMemoryFromIncludeDirectories?: boolean; + chatCompression?: ChatCompressionSettings; } export class Config { @@ -252,6 +257,7 @@ export class Config { | undefined; private readonly experimentalAcp: boolean = false; private readonly loadMemoryFromIncludeDirectories: boolean = false; + private readonly chatCompression: ChatCompressionSettings | undefined; constructor(params: ConfigParameters) { this.sessionId = params.sessionId; @@ -316,6 +322,7 @@ export class Config { } this.loadMemoryFromIncludeDirectories = params.loadMemoryFromIncludeDirectories ?? false; + this.chatCompression = params.chatCompression; if (params.contextFileName) { setGeminiMdFilename(params.contextFileName); @@ -667,6 +674,10 @@ export class Config { return this.ideClient; } + getChatCompression(): ChatCompressionSettings | undefined { + return this.chatCompression; + } + async getGitService(): Promise { if (!this.gitService) { this.gitService = new GitService(this.targetDir); diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 1e39758a..ff901a8b 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -206,6 +206,7 @@ describe('Gemini Client (client.ts)', () => { }), getGeminiClient: vi.fn(), setFallbackMode: vi.fn(), + getChatCompression: vi.fn().mockReturnValue(undefined), }; const MockedConfig = vi.mocked(Config, true); MockedConfig.mockImplementation( @@ -531,14 +532,19 @@ describe('Gemini Client (client.ts)', () => { expect(newChat).toBe(initialChat); }); - it('should trigger summarization if token count is at threshold', async () => { + it('should trigger summarization if token count is at threshold with contextPercentageThreshold setting', async () => { const MOCKED_TOKEN_LIMIT = 1000; + const MOCKED_CONTEXT_PERCENTAGE_THRESHOLD = 0.5; vi.mocked(tokenLimit).mockReturnValue(MOCKED_TOKEN_LIMIT); + vi.spyOn(client['config'], 'getChatCompression').mockReturnValue({ + contextPercentageThreshold: MOCKED_CONTEXT_PERCENTAGE_THRESHOLD, + }); mockGetHistory.mockReturnValue([ { role: 'user', parts: [{ text: '...history...' }] }, ]); - const originalTokenCount = 1000 * 0.7; + const originalTokenCount = + MOCKED_TOKEN_LIMIT * MOCKED_CONTEXT_PERCENTAGE_THRESHOLD; const newTokenCount = 100; mockCountTokens diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index a16a72cc..13e60039 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -596,12 +596,16 @@ export class GeminiClient { return null; } + const contextPercentageThreshold = + this.config.getChatCompression()?.contextPercentageThreshold; + // Don't compress if not forced and we are under the limit. - if ( - !force && - originalTokenCount < this.COMPRESSION_TOKEN_THRESHOLD * tokenLimit(model) - ) { - return null; + if (!force) { + const threshold = + contextPercentageThreshold ?? this.COMPRESSION_TOKEN_THRESHOLD; + if (originalTokenCount < threshold * tokenLimit(model)) { + return null; + } } let compressBeforeIndex = findIndexAfterFraction(