mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-20 18:14:29 -07:00
feat(core): Land AgentHistoryProvider. (#23978)
This commit is contained in:
@@ -0,0 +1,17 @@
|
||||
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
|
||||
|
||||
exports[`AgentHistoryProvider > should handle summarizer failures gracefully 1`] = `
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"text": "[System Note: Prior conversation history was truncated. The most recent user message before truncation was:]
|
||||
|
||||
Message 18",
|
||||
},
|
||||
{
|
||||
"text": "Message 20",
|
||||
},
|
||||
],
|
||||
"role": "user",
|
||||
}
|
||||
`;
|
||||
@@ -0,0 +1,138 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { AgentHistoryProvider } from './agentHistoryProvider.js';
|
||||
import type { Content, GenerateContentResponse } from '@google/genai';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { BaseLlmClient } from '../core/baseLlmClient.js';
|
||||
|
||||
describe('AgentHistoryProvider', () => {
|
||||
let config: Config;
|
||||
let provider: AgentHistoryProvider;
|
||||
let generateContentMock: ReturnType<typeof vi.fn>;
|
||||
|
||||
beforeEach(() => {
|
||||
config = {
|
||||
isExperimentalAgentHistoryTruncationEnabled: vi
|
||||
.fn()
|
||||
.mockReturnValue(false),
|
||||
isExperimentalAgentHistorySummarizationEnabled: vi
|
||||
.fn()
|
||||
.mockReturnValue(false),
|
||||
getBaseLlmClient: vi.fn(),
|
||||
} as unknown as Config;
|
||||
|
||||
generateContentMock = vi.fn().mockResolvedValue({
|
||||
candidates: [{ content: { parts: [{ text: 'Mock intent summary' }] } }],
|
||||
} as unknown as GenerateContentResponse);
|
||||
|
||||
config.getBaseLlmClient = vi.fn().mockReturnValue({
|
||||
generateContent: generateContentMock,
|
||||
} as unknown as BaseLlmClient);
|
||||
|
||||
provider = new AgentHistoryProvider(config, {
|
||||
truncationThreshold: 30,
|
||||
retainedMessages: 15,
|
||||
});
|
||||
});
|
||||
|
||||
const createMockHistory = (count: number): Content[] =>
|
||||
Array.from({ length: count }).map((_, i) => ({
|
||||
role: i % 2 === 0 ? 'user' : 'model',
|
||||
parts: [{ text: `Message ${i}` }],
|
||||
}));
|
||||
|
||||
it('should return history unchanged if truncation is disabled', async () => {
|
||||
vi.spyOn(
|
||||
config,
|
||||
'isExperimentalAgentHistoryTruncationEnabled',
|
||||
).mockReturnValue(false);
|
||||
|
||||
const history = createMockHistory(40);
|
||||
const result = await provider.manageHistory(history);
|
||||
|
||||
expect(result).toBe(history);
|
||||
expect(result.length).toBe(40);
|
||||
});
|
||||
|
||||
it('should return history unchanged if length is under threshold', async () => {
|
||||
vi.spyOn(
|
||||
config,
|
||||
'isExperimentalAgentHistoryTruncationEnabled',
|
||||
).mockReturnValue(true);
|
||||
|
||||
const history = createMockHistory(20); // Threshold is 30
|
||||
const result = await provider.manageHistory(history);
|
||||
|
||||
expect(result).toBe(history);
|
||||
expect(result.length).toBe(20);
|
||||
});
|
||||
|
||||
it('should truncate mechanically to RETAINED_MESSAGES without summarization when sum flag is off', async () => {
|
||||
vi.spyOn(
|
||||
config,
|
||||
'isExperimentalAgentHistoryTruncationEnabled',
|
||||
).mockReturnValue(true);
|
||||
vi.spyOn(
|
||||
config,
|
||||
'isExperimentalAgentHistorySummarizationEnabled',
|
||||
).mockReturnValue(false);
|
||||
|
||||
const history = createMockHistory(35); // Above 30 threshold, should truncate to 15
|
||||
const result = await provider.manageHistory(history);
|
||||
|
||||
expect(result.length).toBe(15);
|
||||
expect(generateContentMock).not.toHaveBeenCalled();
|
||||
|
||||
// Check fallback message logic
|
||||
// Messages 20 to 34 are retained. Message 20 is 'user'.
|
||||
expect(result[0].role).toBe('user');
|
||||
expect(result[0].parts![0].text).toContain(
|
||||
'System Note: Prior conversation history was truncated',
|
||||
);
|
||||
});
|
||||
|
||||
it('should call summarizer and prepend summary when summarization is enabled', async () => {
|
||||
vi.spyOn(
|
||||
config,
|
||||
'isExperimentalAgentHistoryTruncationEnabled',
|
||||
).mockReturnValue(true);
|
||||
vi.spyOn(
|
||||
config,
|
||||
'isExperimentalAgentHistorySummarizationEnabled',
|
||||
).mockReturnValue(true);
|
||||
|
||||
const history = createMockHistory(35);
|
||||
const result = await provider.manageHistory(history);
|
||||
|
||||
expect(generateContentMock).toHaveBeenCalled();
|
||||
expect(result.length).toBe(15); // retained messages
|
||||
expect(result[0].role).toBe('user');
|
||||
expect(result[0].parts![0].text).toContain('<intent_summary>');
|
||||
expect(result[0].parts![0].text).toContain('Mock intent summary');
|
||||
});
|
||||
|
||||
it('should handle summarizer failures gracefully', async () => {
|
||||
vi.spyOn(
|
||||
config,
|
||||
'isExperimentalAgentHistoryTruncationEnabled',
|
||||
).mockReturnValue(true);
|
||||
vi.spyOn(
|
||||
config,
|
||||
'isExperimentalAgentHistorySummarizationEnabled',
|
||||
).mockReturnValue(true);
|
||||
|
||||
generateContentMock.mockRejectedValue(new Error('API Error'));
|
||||
|
||||
const history = createMockHistory(35);
|
||||
const result = await provider.manageHistory(history);
|
||||
|
||||
expect(generateContentMock).toHaveBeenCalled();
|
||||
expect(result.length).toBe(15);
|
||||
expect(result[0]).toMatchSnapshot();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,185 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Content } from '@google/genai';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { getResponseText } from '../utils/partUtils.js';
|
||||
import { LlmRole } from '../telemetry/llmRole.js';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
|
||||
export interface AgentHistoryProviderConfig {
|
||||
truncationThreshold: number;
|
||||
retainedMessages: number;
|
||||
}
|
||||
|
||||
export class AgentHistoryProvider {
|
||||
constructor(
|
||||
private readonly config: Config,
|
||||
private readonly providerConfig: AgentHistoryProviderConfig,
|
||||
) {}
|
||||
|
||||
/**
|
||||
* Evaluates the chat history and performs truncation and summarization if necessary.
|
||||
* Returns a new array of Content if truncation occurred, otherwise returns the original array.
|
||||
*/
|
||||
async manageHistory(
|
||||
history: readonly Content[],
|
||||
abortSignal?: AbortSignal,
|
||||
): Promise<readonly Content[]> {
|
||||
if (!this.shouldTruncate(history)) {
|
||||
return history;
|
||||
}
|
||||
|
||||
const { messagesToKeep, messagesToTruncate } =
|
||||
this.splitHistoryForTruncation(history);
|
||||
|
||||
debugLogger.log(
|
||||
`AgentHistoryProvider: Truncating ${messagesToTruncate.length} messages, retaining ${messagesToKeep.length} messages.`,
|
||||
);
|
||||
|
||||
const summaryText = await this.getSummaryText(
|
||||
messagesToTruncate,
|
||||
abortSignal,
|
||||
);
|
||||
|
||||
return this.mergeSummaryWithHistory(summaryText, messagesToKeep);
|
||||
}
|
||||
|
||||
private shouldTruncate(history: readonly Content[]): boolean {
|
||||
if (!this.config.isExperimentalAgentHistoryTruncationEnabled()) {
|
||||
return false;
|
||||
}
|
||||
return history.length > this.providerConfig.truncationThreshold;
|
||||
}
|
||||
|
||||
private splitHistoryForTruncation(history: readonly Content[]): {
|
||||
messagesToKeep: readonly Content[];
|
||||
messagesToTruncate: readonly Content[];
|
||||
} {
|
||||
return {
|
||||
messagesToKeep: history.slice(-this.providerConfig.retainedMessages),
|
||||
messagesToTruncate: history.slice(
|
||||
0,
|
||||
history.length - this.providerConfig.retainedMessages,
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
private getFallbackSummaryText(
|
||||
messagesToTruncate: readonly Content[],
|
||||
): string {
|
||||
const defaultNote =
|
||||
'System Note: Prior conversation history was truncated to maintain performance and focus. Important context should have been saved to memory.';
|
||||
|
||||
let lastUserText = '';
|
||||
for (let i = messagesToTruncate.length - 1; i >= 0; i--) {
|
||||
const msg = messagesToTruncate[i];
|
||||
if (msg.role === 'user') {
|
||||
lastUserText =
|
||||
msg.parts
|
||||
?.map((p) => p.text || '')
|
||||
.join('')
|
||||
.trim() || '';
|
||||
if (lastUserText) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (lastUserText) {
|
||||
return `[System Note: Prior conversation history was truncated. The most recent user message before truncation was:]\n\n${lastUserText}`;
|
||||
}
|
||||
|
||||
return defaultNote;
|
||||
}
|
||||
|
||||
private async getSummaryText(
|
||||
messagesToTruncate: readonly Content[],
|
||||
abortSignal?: AbortSignal,
|
||||
): Promise<string> {
|
||||
if (!this.config.isExperimentalAgentHistorySummarizationEnabled()) {
|
||||
debugLogger.log(
|
||||
'AgentHistoryProvider: Summarization disabled, using fallback note.',
|
||||
);
|
||||
return this.getFallbackSummaryText(messagesToTruncate);
|
||||
}
|
||||
|
||||
try {
|
||||
const summary = await this.generateIntentSummary(
|
||||
messagesToTruncate,
|
||||
abortSignal,
|
||||
);
|
||||
debugLogger.log('AgentHistoryProvider: Summarization successful.');
|
||||
return summary;
|
||||
} catch (error) {
|
||||
debugLogger.log('AgentHistoryProvider: Summarization failed.', error);
|
||||
return this.getFallbackSummaryText(messagesToTruncate);
|
||||
}
|
||||
}
|
||||
|
||||
private mergeSummaryWithHistory(
|
||||
summaryText: string,
|
||||
messagesToKeep: readonly Content[],
|
||||
): readonly Content[] {
|
||||
if (messagesToKeep.length === 0) {
|
||||
return [{ role: 'user', parts: [{ text: summaryText }] }];
|
||||
}
|
||||
|
||||
// To ensure strict user/model alternating roles required by the Gemini API,
|
||||
// we merge the summary into the first retained message if it's from the 'user'.
|
||||
const firstRetainedMessage = messagesToKeep[0];
|
||||
if (firstRetainedMessage.role === 'user') {
|
||||
const mergedParts = [
|
||||
{ text: summaryText },
|
||||
...(firstRetainedMessage.parts || []),
|
||||
];
|
||||
const mergedMessage: Content = {
|
||||
role: 'user',
|
||||
parts: mergedParts,
|
||||
};
|
||||
return [mergedMessage, ...messagesToKeep.slice(1)];
|
||||
} else {
|
||||
const summaryMessage: Content = {
|
||||
role: 'user',
|
||||
parts: [{ text: summaryText }],
|
||||
};
|
||||
return [summaryMessage, ...messagesToKeep];
|
||||
}
|
||||
}
|
||||
|
||||
private async generateIntentSummary(
|
||||
messagesToTruncate: readonly Content[],
|
||||
abortSignal?: AbortSignal,
|
||||
): Promise<string> {
|
||||
const prompt = `Create a succinct, agent-continuity focused intent summary of the truncated conversation history.
|
||||
Distill the essence of the ongoing work by capturing:
|
||||
- The Original Mandate: What the user (or calling agent) originally requested and why.
|
||||
- The Agent's Strategy: How you (the agent) are approaching the task and where the work is taking place (e.g., specific files, directories, or architectural layers).
|
||||
- Evolving Context: Any significant shifts in the user's intent or the agent's technical approach over the course of the truncated history.
|
||||
|
||||
Write this summary to orient the active agent. Do NOT predict next steps or summarize the current task state, as those are covered by the active history. Focus purely on foundational context and strategic continuity.`;
|
||||
|
||||
const summaryResponse = await this.config
|
||||
.getBaseLlmClient()
|
||||
.generateContent({
|
||||
modelConfigKey: { model: 'agent-history-provider-summarizer' },
|
||||
contents: [
|
||||
...messagesToTruncate,
|
||||
{
|
||||
role: 'user',
|
||||
parts: [{ text: prompt }],
|
||||
},
|
||||
],
|
||||
promptId: 'agent-history-provider',
|
||||
abortSignal: abortSignal ?? new AbortController().signal,
|
||||
role: LlmRole.UTILITY_COMPRESSOR,
|
||||
});
|
||||
|
||||
let summary = getResponseText(summaryResponse) ?? '';
|
||||
summary = summary.replace(/<\/?intent_summary>/g, '').trim();
|
||||
return `<intent_summary>\n${summary}\n</intent_summary>`;
|
||||
}
|
||||
}
|
||||
@@ -256,5 +256,9 @@
|
||||
"chat-compression-default": {
|
||||
"model": "gemini-3-pro-preview",
|
||||
"generateContentConfig": {}
|
||||
},
|
||||
"agent-history-provider-summarizer": {
|
||||
"model": "gemini-3-flash-preview",
|
||||
"generateContentConfig": {}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -256,5 +256,9 @@
|
||||
"chat-compression-default": {
|
||||
"model": "gemini-3-pro-preview",
|
||||
"generateContentConfig": {}
|
||||
},
|
||||
"agent-history-provider-summarizer": {
|
||||
"model": "gemini-3-flash-preview",
|
||||
"generateContentConfig": {}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user