mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-28 14:04:41 -07:00
feat(core): Land ContextCompressionService (#24483)
This commit is contained in:
joshualitt
@@ -0,0 +1,288 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||
import { ContextCompressionService } from './contextCompressionService.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { Content } from '@google/genai';
|
||||
import * as fsSync from 'node:fs';
|
||||
|
||||
vi.mock('node:fs/promises', () => ({
|
||||
readFile: vi.fn(),
|
||||
writeFile: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('node:fs', () => ({
|
||||
existsSync: vi.fn(),
|
||||
}));
|
||||
|
||||
describe('ContextCompressionService', () => {
|
||||
let mockConfig: Partial<Config>;
|
||||
let service: ContextCompressionService;
|
||||
const generateContentMock: ReturnType<typeof vi.fn> = vi.fn();
|
||||
const generateJsonMock: ReturnType<typeof vi.fn> = vi.fn();
|
||||
|
||||
beforeEach(() => {
|
||||
mockConfig = {
|
||||
storage: {
|
||||
getProjectTempDir: vi.fn().mockReturnValue('/mock/temp/dir'),
|
||||
},
|
||||
isContextManagementEnabled: vi.fn().mockResolvedValue(true),
|
||||
getBaseLlmClient: vi.fn().mockReturnValue({
|
||||
generateContent: generateContentMock,
|
||||
generateJson: generateJsonMock,
|
||||
}),
|
||||
} as unknown as Config;
|
||||
|
||||
vi.mocked(fsSync.existsSync).mockReturnValue(false);
|
||||
|
||||
service = new ContextCompressionService(mockConfig as Config);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('compressHistory', () => {
|
||||
it('bypasses compression if feature flag is false', async () => {
|
||||
mockConfig.isContextManagementEnabled = vi.fn().mockResolvedValue(false);
|
||||
const history: Content[] = [{ role: 'user', parts: [{ text: 'hello' }] }];
|
||||
|
||||
const res = await service.compressHistory(history, 'test prompt');
|
||||
expect(res).toStrictEqual(history);
|
||||
});
|
||||
|
||||
it('protects files that were read within the RECENT_TURNS_PROTECTED window', async () => {
|
||||
const history: Content[] = [
|
||||
// Turn 0 & 1 (Old)
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{
|
||||
functionCall: {
|
||||
name: 'read_file',
|
||||
args: { filepath: 'src/app.ts' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'read_file',
|
||||
response: {
|
||||
output: '--- src/app.ts ---\nLine 1\nLine 2\nLine 3',
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
|
||||
// Padding (Turns 2 & 3)
|
||||
{ role: 'model', parts: [{ text: 'res 1' }] },
|
||||
{ role: 'user', parts: [{ text: 'res 2' }] },
|
||||
|
||||
// Padding (Turns 4 & 5)
|
||||
{ role: 'model', parts: [{ text: 'res 3' }] },
|
||||
{ role: 'user', parts: [{ text: 'res 4' }] },
|
||||
|
||||
// Recent Turn (Turn 6 & 7, inside window, cutoff is Math.max(0, 8 - 4) = 4)
|
||||
// Here the model explicitly reads the file again
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{
|
||||
functionCall: {
|
||||
name: 'read_file',
|
||||
args: { filepath: 'src/app.ts' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'read_file',
|
||||
response: {
|
||||
output: '--- src/app.ts ---\nLine 1\nLine 2\nLine 3',
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
const res = await service.compressHistory(history, 'test prompt');
|
||||
|
||||
// Because src/app.ts was re-read recently (index 6 is >= 4), the OLD response at index 1 is PROTECTED.
|
||||
// It should NOT be compressed.
|
||||
const compressedOutput =
|
||||
res[1].parts![0].functionResponse!.response!['output'];
|
||||
expect(compressedOutput).toBe(
|
||||
'--- src/app.ts ---\nLine 1\nLine 2\nLine 3',
|
||||
);
|
||||
// Verify generateContentMock wasn't called because it bypassed the LLM routing
|
||||
expect(generateContentMock).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('compresses files read outside the protected window', async () => {
|
||||
const history: Content[] = [
|
||||
// Turn 0: The original function call to read the file
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{
|
||||
functionCall: {
|
||||
name: 'read_file',
|
||||
args: { filepath: 'src/old.ts' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
// Turn 1: The tool output response
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'read_file',
|
||||
response: {
|
||||
output: '--- src/old.ts ---\nLine 1\nLine 2\nLine 3\nLine 4',
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
// Padding turns to push it out of the recent window
|
||||
{ role: 'model', parts: [{ text: 'msg 2' }] },
|
||||
{ role: 'user', parts: [{ text: 'res 2' }] },
|
||||
{ role: 'model', parts: [{ text: 'msg 3' }] },
|
||||
{ role: 'user', parts: [{ text: 'res 3' }] },
|
||||
{ role: 'model', parts: [{ text: 'msg 4' }] },
|
||||
{ role: 'user', parts: [{ text: 'res 4' }] },
|
||||
];
|
||||
|
||||
// Mock the routing request to return PARTIAL
|
||||
generateJsonMock.mockResolvedValueOnce({
|
||||
'src/old.ts': {
|
||||
level: 'PARTIAL',
|
||||
start_line: 2,
|
||||
end_line: 3,
|
||||
},
|
||||
});
|
||||
|
||||
const res = await service.compressHistory(history, 'test prompt');
|
||||
const compressedOutput =
|
||||
res[1].parts![0].functionResponse!.response!['output'];
|
||||
|
||||
expect(compressedOutput).toContain('[Showing lines 2–3 of 4 in old.ts.');
|
||||
expect(compressedOutput).toContain('2 | Line 2');
|
||||
expect(compressedOutput).toContain('3 | Line 3');
|
||||
});
|
||||
|
||||
it('returns SUMMARY and hits cache on subsequent requests', async () => {
|
||||
const history1: Content[] = [
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{
|
||||
functionCall: {
|
||||
name: 'read_file',
|
||||
args: { filepath: 'src/index.ts' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'read_file',
|
||||
response: {
|
||||
output: `--- src/index.ts ---\nVery long content here...`,
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{ role: 'model', parts: [{ text: 'p1' }] },
|
||||
{ role: 'user', parts: [{ text: 'p2' }] },
|
||||
{ role: 'model', parts: [{ text: 'p3' }] },
|
||||
{ role: 'user', parts: [{ text: 'p4' }] },
|
||||
{ role: 'model', parts: [{ text: 'p5' }] },
|
||||
{ role: 'user', parts: [{ text: 'p6' }] },
|
||||
];
|
||||
|
||||
// 1st request: routing says SUMMARY
|
||||
generateJsonMock.mockResolvedValueOnce({
|
||||
'src/index.ts': { level: 'SUMMARY' },
|
||||
});
|
||||
// 2nd request: the actual summarization call
|
||||
generateContentMock.mockResolvedValueOnce({
|
||||
candidates: [
|
||||
{ content: { parts: [{ text: 'This is a cached summary.' }] } },
|
||||
],
|
||||
});
|
||||
|
||||
await service.compressHistory(history1, 'test query');
|
||||
expect(generateJsonMock).toHaveBeenCalledTimes(1);
|
||||
expect(generateContentMock).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Time passes, we get a new query. The file is still old.
|
||||
const history2: Content[] = [
|
||||
...history1,
|
||||
{ role: 'model', parts: [{ text: 'p7' }] },
|
||||
{ role: 'user', parts: [{ text: 'p8' }] },
|
||||
];
|
||||
|
||||
// 3rd request: routing says SUMMARY again.
|
||||
generateJsonMock.mockResolvedValueOnce({
|
||||
'src/index.ts': { level: 'SUMMARY' },
|
||||
});
|
||||
|
||||
const res = await service.compressHistory(history2, 'new query');
|
||||
|
||||
// It should NOT make a 3rd fetch call for routing, since content has not changed and state is cached.
|
||||
expect(generateJsonMock).toHaveBeenCalledTimes(1);
|
||||
expect(generateContentMock).toHaveBeenCalledTimes(1);
|
||||
|
||||
const compressedOutput =
|
||||
res[1].parts![0].functionResponse!.response!['output'];
|
||||
expect(compressedOutput).toContain('This is a cached summary.');
|
||||
});
|
||||
it('returns unmodified history if structural validation fails', async () => {
|
||||
// Creating a broken history where functionCall is NOT followed by user functionResponse
|
||||
const brokenHistory: Content[] = [
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{
|
||||
functionCall: {
|
||||
name: 'read_file',
|
||||
args: { filepath: 'src/index.ts' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
// Missing user functionResponse!
|
||||
{ role: 'model', parts: [{ text: 'Wait, I am a model again.' }] },
|
||||
{ role: 'user', parts: [{ text: 'This is invalid.' }] },
|
||||
{ role: 'model', parts: [{ text: 'Yep.' }] },
|
||||
{ role: 'user', parts: [{ text: 'Padding.' }] },
|
||||
{ role: 'model', parts: [{ text: 'Padding.' }] },
|
||||
];
|
||||
|
||||
const res = await service.compressHistory(brokenHistory, 'test query');
|
||||
|
||||
// Because it's broken, it should return the exact same array by reference.
|
||||
expect(res).toBe(brokenHistory);
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user