diff --git a/core/src/common.ts b/core/src/common.ts index 697fd90f..1dffac1d 100644 --- a/core/src/common.ts +++ b/core/src/common.ts @@ -106,6 +106,8 @@ export { type CodeExecutionResult, type File, } from './code_executors/code_execution_utils.js'; +export {AnchoredContextCompactor} from './context/anchored_context_compactor.js'; +export type {AnchoredContextCompactorOptions} from './context/anchored_context_compactor.js'; export type {BaseContextCompactor} from './context/base_context_compactor.js'; export type {BaseSummarizer} from './context/summarizers/base_summarizer.js'; export {LlmSummarizer} from './context/summarizers/llm_summarizer.js'; @@ -116,7 +118,7 @@ export {TrajectoryThoughtPruningCompactor} from './context/trajectory_thought_pr export type {TrajectoryThoughtPruningCompactorOptions} from './context/trajectory_thought_pruning_compactor.js'; export {TruncatingContextCompactor} from './context/truncating_context_compactor.js'; export type {TruncatingContextCompactorOptions} from './context/truncating_context_compactor.js'; -export {isCompactedEvent} from './events/compacted_event.js'; +export {isCompactedEvent, isScratchpadEvent} from './events/compacted_event.js'; export type {CompactedEvent} from './events/compacted_event.js'; export { createEvent, diff --git a/core/src/context/anchored_context_compactor.ts b/core/src/context/anchored_context_compactor.ts new file mode 100644 index 00000000..ab5a6171 --- /dev/null +++ b/core/src/context/anchored_context_compactor.ts @@ -0,0 +1,165 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {InvocationContext} from '../agents/invocation_context.js'; +import {CompactedEvent, isScratchpadEvent} from '../events/compacted_event.js'; +import { + Event, + getEventTokens, + getFunctionCalls, + getFunctionResponses, +} from '../events/event.js'; +import {BaseContextCompactor} from './base_context_compactor.js'; +import {BaseSummarizer} from './summarizers/base_summarizer.js'; + +export interface AnchoredContextCompactorOptions { + /** The maximum number of tokens to retain in the session history before compaction. */ + tokenThreshold: number; + /** + * The minimum number of raw events to keep at the end of the session. + * Compaction will not affect these tail events (unless needed for tool splits). + */ + eventRetentionSize: number; + /** The summarizer used to create the compacted event content. */ + summarizer: BaseSummarizer; +} + +/** + * A context compactor that maintains a single persistent 'Scratchpad' or + * 'State Tracker' event at the top of the context history. + * + * When compaction is triggered, it merges new raw events into the existing + * Scratchpad event and discards them from the active history view. + */ +export class AnchoredContextCompactor implements BaseContextCompactor { + private readonly tokenThreshold: number; + private readonly eventRetentionSize: number; + private readonly summarizer: BaseSummarizer; + + constructor(options: AnchoredContextCompactorOptions) { + this.tokenThreshold = options.tokenThreshold; + this.eventRetentionSize = options.eventRetentionSize; + this.summarizer = options.summarizer; + } + + private getActiveEvents(events: Event[]): Event[] { + let latestScratchpad: CompactedEvent | undefined = undefined; + + for (let i = events.length - 1; i >= 0; i--) { + const e = events[i]; + if (isScratchpadEvent(e)) { + latestScratchpad = e; + break; + } + } + + if (!latestScratchpad) { + return events; + } + + const activeRawEvents = events.filter( + (e) => e.timestamp > latestScratchpad!.endTime && !isScratchpadEvent(e), + ); + + return [latestScratchpad, ...activeRawEvents]; + } + + shouldCompact( + invocationContext: InvocationContext, + ): boolean | Promise { + const events = invocationContext.session.events; + const activeEvents = this.getActiveEvents(events); + const hasScratchpad = + activeEvents.length > 0 && isScratchpadEvent(activeEvents[0]); + const rawEvents = hasScratchpad ? activeEvents.slice(1) : activeEvents; + + if (rawEvents.length <= this.eventRetentionSize) { + return false; + } + + const totalTokens = activeEvents.reduce( + (sum, event) => sum + getEventTokens(event), + 0, + ); + + return totalTokens > this.tokenThreshold; + } + + async compact(invocationContext: InvocationContext): Promise { + const events = invocationContext.session.events; + const activeEvents = this.getActiveEvents(events); + const hasScratchpad = + activeEvents.length > 0 && isScratchpadEvent(activeEvents[0]); + const rawEvents = hasScratchpad ? activeEvents.slice(1) : activeEvents; + + if (rawEvents.length <= this.eventRetentionSize) { + return; + } + + // Determine the baseline index to retain from the active raw events. + let retainStartIndex = Math.max( + 0, + rawEvents.length - this.eventRetentionSize, + ); + + // Prevent splitting between a tool call and its response. + while (retainStartIndex > 0) { + const eventToRetain = rawEvents[retainStartIndex]; + const previousEvent = rawEvents[retainStartIndex - 1]; + + if ( + getFunctionResponses(eventToRetain).length > 0 && + getFunctionCalls(previousEvent).length > 0 + ) { + retainStartIndex--; + } else { + // No conflict, safe to split here. + break; + } + } + + if (retainStartIndex === 0) { + // Cannot compact if we have to retain everything + return; + } + + // Extract raw events to compact. + const rawEventsToCompact = rawEvents.slice(0, retainStartIndex); + + let scratchpadEvent: CompactedEvent; + + if (hasScratchpad) { + const existingScratchpad = activeEvents[0] as CompactedEvent; + scratchpadEvent = await this.summarizer.summarize([ + existingScratchpad, + ...rawEventsToCompact, + ]); + } else { + scratchpadEvent = await this.summarizer.summarize(rawEventsToCompact); + } + + // Ensure the event is marked as scratchpad and has system author. + const updatedScratchpad = { + ...scratchpadEvent, + isScratchpad: true, + author: 'system', + } as CompactedEvent; + + // Reconstruct the events list: inactive events + new scratchpad + active retained events + const inactiveEvents = events.slice(0, events.indexOf(activeEvents[0])); + const retainedRawEvents = rawEvents.slice(retainStartIndex); + + const newEventsList = [ + ...inactiveEvents, + updatedScratchpad, + ...retainedRawEvents, + ]; + + // Mutate the original session events array. + events.length = 0; + events.push(...newEventsList); + } +} diff --git a/core/src/context/token_based_context_compactor.ts b/core/src/context/token_based_context_compactor.ts index 68ca529a..a713c649 100644 --- a/core/src/context/token_based_context_compactor.ts +++ b/core/src/context/token_based_context_compactor.ts @@ -6,7 +6,7 @@ import {InvocationContext} from '../agents/invocation_context.js'; import {CompactedEvent, isCompactedEvent} from '../events/compacted_event.js'; -import {Event, stringifyContent} from '../events/event.js'; +import {Event, getEventTokens} from '../events/event.js'; import {BaseContextCompactor} from './base_context_compactor.js'; import {BaseSummarizer} from './summarizers/base_summarizer.js'; @@ -142,15 +142,6 @@ export class TokenBasedContextCompactor implements BaseContextCompactor { } } -function getEventTokens(event: Event): number { - if (event.usageMetadata?.promptTokenCount !== undefined) { - return event.usageMetadata.promptTokenCount; - } - // Estimate: 4 chars per token. - const contentStr = stringifyContent(event); - return Math.ceil(contentStr.length / 4); -} - function hasFunctionCall(event: Event): boolean { return !!event.content?.parts?.some( (part) => part.functionCall !== undefined, diff --git a/core/src/events/compacted_event.ts b/core/src/events/compacted_event.ts index fa303b6d..572afbbc 100644 --- a/core/src/events/compacted_event.ts +++ b/core/src/events/compacted_event.ts @@ -30,6 +30,11 @@ export interface CompactedEvent extends Event { * The summarized content of the compacted events. */ compactedContent: string; + + /** + * Identifies this compacted event as the persistent context scratchpad. + */ + isScratchpad?: boolean; } /** @@ -39,6 +44,17 @@ export function isCompactedEvent(event: Event): event is CompactedEvent { return 'isCompacted' in event && event.isCompacted === true; } +/** + * Type guard to check if an event is a scratchpad CompactedEvent. + */ +export function isScratchpadEvent( + event: Event, +): event is CompactedEvent & {isScratchpad: true} { + return ( + isCompactedEvent(event) && (event as CompactedEvent).isScratchpad === true + ); +} + export function createCompactedEvent( params: Partial = {}, ): CompactedEvent { @@ -48,5 +64,6 @@ export function createCompactedEvent( startTime: params.startTime!, endTime: params.endTime!, compactedContent: params.compactedContent!, + isScratchpad: params.isScratchpad, }; } diff --git a/core/src/events/event.ts b/core/src/events/event.ts index 399be199..4b039732 100644 --- a/core/src/events/event.ts +++ b/core/src/events/event.ts @@ -166,6 +166,18 @@ export function stringifyContent(event: Event): string { .join(''); } +/** + * Estimates the number of tokens in the event based on usage metadata or characters. + */ +export function getEventTokens(event: Event): number { + if (event.usageMetadata?.promptTokenCount !== undefined) { + return event.usageMetadata.promptTokenCount; + } + // Estimate: 4 chars per token. + const contentStr = stringifyContent(event); + return Math.ceil(contentStr.length / 4); +} + /** * Returns whether the event contains any thought parts. */ diff --git a/core/test/context/anchored_context_compactor_test.ts b/core/test/context/anchored_context_compactor_test.ts new file mode 100644 index 00000000..ae1a26d0 --- /dev/null +++ b/core/test/context/anchored_context_compactor_test.ts @@ -0,0 +1,268 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + AnchoredContextCompactor, + BaseAgent, + BaseSummarizer, + CompactedEvent, + Event, + InvocationContext, + PluginManager, + Session, + isScratchpadEvent, +} from '@google/adk'; +import {describe, expect, it} from 'vitest'; + +class MockSummarizer implements BaseSummarizer { + async summarize(events: Event[]): Promise { + return { + id: 'mock-id', + invocationId: '', + author: 'system', + actions: { + stateDelta: {}, + artifactDelta: {}, + requestedAuthConfigs: [], + requestedToolConfirmations: {}, + }, + timestamp: Date.now(), + isCompacted: true, + startTime: events[0].timestamp, + endTime: events[events.length - 1].timestamp, + compactedContent: `Mock summary of ${events.length} events`, + content: { + role: 'model', + parts: [{text: `Mock summary of ${events.length} events`}], + }, + } as CompactedEvent; + } +} + +class FailingSummarizer implements BaseSummarizer { + async summarize(): Promise { + throw new Error('Summarization failed'); + } +} + +function createMockEvent( + id: string, + tokenCount?: number, + isFuncCall?: boolean, + isFuncResp?: boolean, +): Event { + const event: Event = { + id, + timestamp: Date.now(), + content: {parts: []}, + } as unknown as Event; + if (tokenCount !== undefined) { + event.usageMetadata = {promptTokenCount: tokenCount}; + } + if (isFuncCall) { + event.content!.parts!.push({functionCall: {name: 'mock', args: {}}}); + } + if (isFuncResp) { + event.content!.parts!.push({ + functionResponse: {name: 'mock', response: {}}, + }); + } + return event; +} + +function createMockScratchpadEvent( + id: string, + tokenCount?: number, + contentStr?: string, +): CompactedEvent { + const event = createMockEvent(id, tokenCount) as CompactedEvent; + event.isCompacted = true; + event.isScratchpad = true; + event.author = 'system'; + event.startTime = Date.now() - 10000; + event.endTime = Date.now() - 5000; + event.compactedContent = contentStr || 'Existing scratchpad content'; + return event; +} + +function createMockInvocationContext(events: Event[]): InvocationContext { + const session = { + id: 'test-session', + events, + appName: 'test-app', + userId: 'test-user', + } as unknown as Session; + return new InvocationContext({ + invocationId: 'test-invocation', + agent: {} as BaseAgent, + session, + pluginManager: new PluginManager([]), + }); +} + +describe('AnchoredContextCompactor', () => { + it('should not compact if event count is within retention size', async () => { + const compactor = new AnchoredContextCompactor({ + tokenThreshold: 10, + eventRetentionSize: 3, + summarizer: new MockSummarizer(), + }); + + const context = createMockInvocationContext([ + createMockEvent('1', 5), + createMockEvent('2', 5), + createMockEvent('3', 5), + ]); + + expect(await compactor.shouldCompact(context)).toBe(false); + + await compactor.compact(context); + expect(context.session.events.length).toBe(3); + }); + + it('should not compact if event count including scratchpad is within retention size plus scratchpad', async () => { + const compactor = new AnchoredContextCompactor({ + tokenThreshold: 10, + eventRetentionSize: 2, + summarizer: new MockSummarizer(), + }); + + const context = createMockInvocationContext([ + createMockScratchpadEvent('scratchpad', 5), + createMockEvent('1', 5), + createMockEvent('2', 5), + ]); + + // Raw events count = 2 <= retention size 2. So should not compact. + expect(await compactor.shouldCompact(context)).toBe(false); + + await compactor.compact(context); + expect(context.session.events.length).toBe(3); + }); + + it('should compact if token threshold exceeded (first compaction - no scratchpad)', async () => { + const compactor = new AnchoredContextCompactor({ + tokenThreshold: 10, + eventRetentionSize: 2, + summarizer: new MockSummarizer(), + }); + + // Total tokens: 5 * 4 = 20 > 10. Raw events: 4 > 2. + const context = createMockInvocationContext([ + createMockEvent('1', 5), + createMockEvent('2', 5), + createMockEvent('3', 5), + createMockEvent('4', 5), + ]); + + expect(await compactor.shouldCompact(context)).toBe(true); + + await compactor.compact(context); + + // Should compact '1' and '2' (index 0 and 1) into scratchpad. + // Kept events should be '3' and '4' (index 2 and 3). + // Result events array should be [scratchpad, '3', '4']. + expect(context.session.events.length).toBe(3); + const firstEvent = context.session.events[0]; + expect(isScratchpadEvent(firstEvent)).toBe(true); + expect((firstEvent as CompactedEvent).compactedContent).toBe( + 'Mock summary of 2 events', + ); + expect(context.session.events[1].id).toBe('3'); + expect(context.session.events[2].id).toBe('4'); + }); + + it('should compact if token threshold exceeded (subsequent compaction - merges into existing scratchpad)', async () => { + const compactor = new AnchoredContextCompactor({ + tokenThreshold: 10, + eventRetentionSize: 2, + summarizer: new MockSummarizer(), + }); + + // We have a scratchpad, and 4 raw events. + // Total raw events = 4 > retention size 2. + // Total tokens = 5 (scratchpad) + 5 * 4 (raw) = 25 > 10. + const context = createMockInvocationContext([ + createMockScratchpadEvent('scratchpad', 5, 'existing summary'), + createMockEvent('1', 5), + createMockEvent('2', 5), + createMockEvent('3', 5), + createMockEvent('4', 5), + ]); + + expect(await compactor.shouldCompact(context)).toBe(true); + + await compactor.compact(context); + + // Should merge scratchpad, '1', and '2' (3 events in total) into a new scratchpad. + // Kept events should be '3' and '4'. + // Result events array: [new_scratchpad, '3', '4']. + expect(context.session.events.length).toBe(3); + const firstEvent = context.session.events[0]; + expect(isScratchpadEvent(firstEvent)).toBe(true); + expect((firstEvent as CompactedEvent).compactedContent).toBe( + 'Mock summary of 3 events', + ); + expect(context.session.events[1].id).toBe('3'); + expect(context.session.events[2].id).toBe('4'); + }); + + it('should not split tool call and responses', async () => { + const compactor = new AnchoredContextCompactor({ + tokenThreshold: 10, + eventRetentionSize: 2, + summarizer: new MockSummarizer(), + }); + + // 4 raw events. Retention = 2 implies events[2] ('3') and events[3] ('4') retained. + // But '3' is a response and '2' is a call. They must not be split. + // So retention start index drops to 1, meaning we retain '2', '3', '4'. + // Compaction should only compress event '1' (index 0). + const context = createMockInvocationContext([ + createMockEvent('1', 5), + createMockEvent('2', 5, true, false), // call + createMockEvent('3', 5, false, true), // response + createMockEvent('4', 5), + ]); + + await compactor.compact(context); + + expect(context.session.events.length).toBe(4); // scratchpad, '2', '3', '4' + expect(isScratchpadEvent(context.session.events[0])).toBe(true); + expect((context.session.events[0] as CompactedEvent).compactedContent).toBe( + 'Mock summary of 1 events', + ); + expect(context.session.events[1].id).toBe('2'); + expect(context.session.events[2].id).toBe('3'); + expect(context.session.events[3].id).toBe('4'); + }); + + it('should not mutate history if summarizer fails', async () => { + const compactor = new AnchoredContextCompactor({ + tokenThreshold: 10, + eventRetentionSize: 2, + summarizer: new FailingSummarizer(), + }); + + const context = createMockInvocationContext([ + createMockEvent('1', 5), + createMockEvent('2', 5), + createMockEvent('3', 5), + createMockEvent('4', 5), + ]); + + await expect(compactor.compact(context)).rejects.toThrow( + 'Summarization failed', + ); + + // History should remain exactly as before. + expect(context.session.events.length).toBe(4); + expect(context.session.events[0].id).toBe('1'); + expect(context.session.events[1].id).toBe('2'); + expect(context.session.events[2].id).toBe('3'); + expect(context.session.events[3].id).toBe('4'); + }); +}); diff --git a/tests/e2e/context_compaction/e2e_anchored_compaction_test.ts b/tests/e2e/context_compaction/e2e_anchored_compaction_test.ts new file mode 100644 index 00000000..779cf128 --- /dev/null +++ b/tests/e2e/context_compaction/e2e_anchored_compaction_test.ts @@ -0,0 +1,146 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + AnchoredContextCompactor, + BasePlugin, + CompactedEvent, + ContextCompactionTrigger, + Gemini, + InMemoryRunner, + InvocationContext, + isCompactedEvent, + isScratchpadEvent, + LlmAgent, + LlmSummarizer, +} from '@google/adk'; +import {createUserContent} from '@google/genai'; +import * as dotenv from 'dotenv'; +import * as fs from 'fs'; +import * as path from 'path'; +import {describe, expect, it} from 'vitest'; + +class TestCompactionPlugin extends BasePlugin { + beforeCalled = false; + afterCalled = false; + + constructor() { + super('TestCompactionPlugin'); + } + + // eslint-disable-next-line @typescript-eslint/no-unused-vars + override async beforeContextCompaction(params: { + invocationContext: InvocationContext; + trigger: ContextCompactionTrigger; + }) { + this.beforeCalled = true; + } + + // eslint-disable-next-line @typescript-eslint/no-unused-vars + override async afterContextCompaction(params: { + invocationContext: InvocationContext; + trigger: ContextCompactionTrigger; + }) { + this.afterCalled = true; + } +} + +function createAnchoredCompactionAgent(): LlmAgent { + const compactor = new AnchoredContextCompactor({ + tokenThreshold: 200, // Artificially low token limit. + eventRetentionSize: 2, // Keep the last 2 events uncompacted out of those triggered. + summarizer: new LlmSummarizer({ + llm: new Gemini({model: 'gemini-2.5-flash'}), + }), + }); + + return new LlmAgent({ + name: 'anchored_compaction_agent', + description: + 'An agent configured to test live anchored context compaction.', + instruction: + 'You are a helpful conversational AI. Please provide short, single-sentence answers.', + model: 'gemini-2.5-flash', + contextCompactors: [compactor], + }); +} + +describe('E2e Anchored Context Compaction', () => { + const envPath = path.resolve(__dirname, '.env'); + const envExists = fs.existsSync(envPath); + + if (envExists) { + dotenv.config({path: envPath}); + } + + const hasAKey = + !!process.env.GEMINI_API_KEY || + !!process.env.GOOGLE_GENAI_API_KEY || + !!process.env.GOOGLE_CLOUD_PROJECT; + + it.skipIf(!hasAKey)( + 'should hit token threshold and maintain a persistent scratchpad at index 0', + async () => { + const agent = createAnchoredCompactionAgent(); + const plugin = new TestCompactionPlugin(); + const runner = new InMemoryRunner({ + agent, + appName: 'e2e_test', + plugins: [plugin], + }); + const session = await runner.sessionService.createSession({ + appName: 'e2e_test', + userId: 'test_user', + }); + + const turns = [ + 'Tell me a long story about a brave knight named Sir Galahad exploring a dragon-infested cave.', + 'What happens after he finds the treasure?', + 'Can you summarize his entire adventure in 3 sentences?', + ]; + + for (const prompt of turns) { + const responseGen = runner.runAsync({ + userId: 'test_user', + sessionId: session.id, + newMessage: createUserContent(prompt), + }); + + for await (const _ of responseGen) { + // Drain the generator to let the agent run and append events + } + } + + // Now retrieve the session and check its events + const updatedSession = await runner.sessionService.getSession({ + appName: 'e2e_test', + userId: 'test_user', + sessionId: session.id, + }); + + const events = updatedSession!.events; + + // Find if there is a CompactedEvent + const compactedEvents = events.filter(isCompactedEvent); + expect(compactedEvents.length).toBeGreaterThan(0); + + // In AnchoredContextCompactor, the scratchpad should be at index 0 + const firstEvent = events[0]; + expect(isScratchpadEvent(firstEvent)).toBe(true); + expect(firstEvent.author).toBe('system'); + expect((firstEvent as CompactedEvent).compactedContent).toBeTruthy(); + + // Verify that there is at most one scratchpad event + const scratchpads = events.filter(isScratchpadEvent); + expect(scratchpads.length).toBe(1); + + // Verify that the plugin callbacks were called + expect(plugin.beforeCalled).toBe(true); + expect(plugin.afterCalled).toBe(true); + }, + 30000, + ); +}); diff --git a/tests/integration/context_compaction/anchored/agent.ts b/tests/integration/context_compaction/anchored/agent.ts new file mode 100644 index 00000000..579f82eb --- /dev/null +++ b/tests/integration/context_compaction/anchored/agent.ts @@ -0,0 +1,55 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {AnchoredContextCompactor, LlmAgent, LlmSummarizer} from '@google/adk'; +import {GeminiWithMockResponses} from '../../test_case_utils.js'; + +export const compactor = new AnchoredContextCompactor({ + tokenThreshold: 40, + eventRetentionSize: 2, + summarizer: new LlmSummarizer({ + llm: new GeminiWithMockResponses([ + // First compaction summary response + { + candidates: [ + { + content: { + role: 'model', + parts: [{text: 'Compacted summary of turn 1 and 2.'}], + }, + }, + ], + usageMetadata: { + promptTokenCount: 10, + totalTokenCount: 15, + }, + }, + // Second compaction summary response (if triggered again) + { + candidates: [ + { + content: { + role: 'model', + parts: [{text: 'Compacted summary including turn 3 and 4.'}], + }, + }, + ], + usageMetadata: { + promptTokenCount: 10, + totalTokenCount: 15, + }, + }, + ]), + }), +}); + +export const rootAgent = new LlmAgent({ + name: 'anchored_compaction_agent', + model: 'gemini-2.5-flash', + description: 'Agent to demonstrate anchored context compaction.', + instruction: 'You are a helpful assistant that answers concisely.', + contextCompactors: [compactor], +}); diff --git a/tests/integration/context_compaction/anchored/agent_test.ts b/tests/integration/context_compaction/anchored/agent_test.ts new file mode 100644 index 00000000..fefbd2ec --- /dev/null +++ b/tests/integration/context_compaction/anchored/agent_test.ts @@ -0,0 +1,207 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + CompactedEvent, + Event, + InMemoryRunner, + isScratchpadEvent, +} from '@google/adk'; +import {createUserContent} from '@google/genai'; +import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest'; +import {GeminiWithMockResponses} from '../../test_case_utils.js'; +import {rootAgent} from './agent.js'; + +function getActiveEvents(events: Event[]): Event[] { + let scratchpad: CompactedEvent | undefined = undefined; + for (let i = events.length - 1; i >= 0; i--) { + if (isScratchpadEvent(events[i])) { + scratchpad = events[i] as CompactedEvent; + break; + } + } + if (!scratchpad) return events; + return [ + scratchpad, + ...events.filter( + (e) => !isScratchpadEvent(e) && e.timestamp > scratchpad!.endTime, + ), + ]; +} + +describe('Anchored Context Compaction', () => { + let mockTime = 1000; + beforeEach(() => { + mockTime = 1000; + vi.spyOn(Date, 'now').mockImplementation(() => { + return mockTime++; + }); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it('should iteratively merge into a persistent scratchpad at index 0', async () => { + // 4 turns of agent mock responses. + // Each turn responds with some content and has promptTokenCount of 25 to trigger compaction. + rootAgent.model = new GeminiWithMockResponses([ + // Turn 1 response + { + candidates: [ + { + content: { + role: 'model', + parts: [{text: 'Response to Message 1.'}], + }, + }, + ], + usageMetadata: { + promptTokenCount: 25, + totalTokenCount: 35, + }, + }, + // Turn 2 response + { + candidates: [ + { + content: { + role: 'model', + parts: [{text: 'Response to Message 2.'}], + }, + }, + ], + usageMetadata: { + promptTokenCount: 25, + totalTokenCount: 35, + }, + }, + // Turn 3 response + { + candidates: [ + { + content: { + role: 'model', + parts: [{text: 'Response to Message 3.'}], + }, + }, + ], + usageMetadata: { + promptTokenCount: 25, + totalTokenCount: 35, + }, + }, + // Turn 4 response + { + candidates: [ + { + content: { + role: 'model', + parts: [{text: 'Response to Message 4.'}], + }, + }, + ], + usageMetadata: { + promptTokenCount: 25, + totalTokenCount: 35, + }, + }, + ]); + + const runner = new InMemoryRunner({ + agent: rootAgent, + appName: 'anchored_compaction_agent', + }); + const session = await runner.sessionService.createSession({ + appName: 'anchored_compaction_agent', + userId: 'test_user', + }); + + // Turn 1 + for await (const _ of runner.runAsync({ + userId: 'test_user', + sessionId: session.id, + newMessage: createUserContent('Message 1'), + })) { + // intentionally empty + } + + // Turn 2 + for await (const _ of runner.runAsync({ + userId: 'test_user', + sessionId: session.id, + newMessage: createUserContent('Message 2'), + })) { + // intentionally empty + } + + // Turn 3: Token threshold (40) exceeded (Turn 1 and Turn 2 total tokens > 40). + // Compaction should run and merge Turn 1 + Turn 2 into Scratchpad 1. + // Events list should become: [Scratchpad 1, Turn 3 User, Turn 3 Model] + for await (const _ of runner.runAsync({ + userId: 'test_user', + sessionId: session.id, + newMessage: createUserContent('Message 3'), + })) { + // intentionally empty + } + + // Assert first compaction results + let updatedSession = await runner.sessionService.getSession({ + sessionId: session.id, + userId: 'test_user', + appName: 'anchored_compaction_agent', + }); + let activeEvents = getActiveEvents(updatedSession!.events); + + expect(activeEvents.length).toBe(4); // Scratchpad + Retained Model 2 + Turn 3 User + Turn 3 Model + expect(isScratchpadEvent(activeEvents[0])).toBe(true); + expect(activeEvents[0].author).toBe('system'); + expect((activeEvents[0] as CompactedEvent).compactedContent).toBe( + 'Compacted summary of turn 1 and 2.', + ); + expect(activeEvents[1].content?.parts?.[0].text).toBe( + 'Response to Message 2.', + ); + expect(activeEvents[2].content?.parts?.[0].text).toBe('Message 3'); + expect(activeEvents[3].content?.parts?.[0].text).toBe( + 'Response to Message 3.', + ); + + // Turn 4: Token threshold (40) exceeded again (Scratchpad 1 (15) + Turn 3 (25+25) = 65 > 40). + // Compaction should run and merge Scratchpad 1 + Turn 3 into Scratchpad 2. + // Events list should become: [Scratchpad 2, Turn 4 User, Turn 4 Model] + for await (const _ of runner.runAsync({ + userId: 'test_user', + sessionId: session.id, + newMessage: createUserContent('Message 4'), + })) { + // intentionally empty + } + + // Assert second compaction results + updatedSession = await runner.sessionService.getSession({ + sessionId: session.id, + userId: 'test_user', + appName: 'anchored_compaction_agent', + }); + activeEvents = getActiveEvents(updatedSession!.events); + + expect(activeEvents.length).toBe(4); // Scratchpad 2 + Retained Model 3 + Turn 4 User + Turn 4 Model + expect(isScratchpadEvent(activeEvents[0])).toBe(true); + expect(activeEvents[0].author).toBe('system'); + expect((activeEvents[0] as CompactedEvent).compactedContent).toBe( + 'Compacted summary including turn 3 and 4.', + ); + expect(activeEvents[1].content?.parts?.[0].text).toBe( + 'Response to Message 3.', + ); + expect(activeEvents[2].content?.parts?.[0].text).toBe('Message 4'); + expect(activeEvents[3].content?.parts?.[0].text).toBe( + 'Response to Message 4.', + ); + }); +});