mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-07-17 13:30:53 -07:00
feat(core): introduce decoupled ContextManager and Sidecar architecture
This commit is contained in:
Your Name
@@ -699,6 +699,7 @@ export interface ConfigParameters {
|
||||
experimentalJitContext?: boolean;
|
||||
autoDistillation?: boolean;
|
||||
experimentalMemoryManager?: boolean;
|
||||
experimentalContextSidecarConfig?: string;
|
||||
experimentalAgentHistoryTruncation?: boolean;
|
||||
experimentalAgentHistoryTruncationThreshold?: number;
|
||||
experimentalAgentHistoryRetainedMessages?: number;
|
||||
@@ -940,6 +941,7 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
private readonly adminSkillsEnabled: boolean;
|
||||
private readonly experimentalJitContext: boolean;
|
||||
private readonly experimentalMemoryManager: boolean;
|
||||
private readonly experimentalContextSidecarConfig?: string;
|
||||
private readonly memoryBoundaryMarkers: readonly string[];
|
||||
private readonly topicUpdateNarration: boolean;
|
||||
private readonly disableLLMCorrection: boolean;
|
||||
@@ -1151,6 +1153,8 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
|
||||
this.experimentalJitContext = params.experimentalJitContext ?? false;
|
||||
this.experimentalMemoryManager = params.experimentalMemoryManager ?? false;
|
||||
this.experimentalContextSidecarConfig =
|
||||
params.experimentalContextSidecarConfig;
|
||||
this.memoryBoundaryMarkers = params.memoryBoundaryMarkers ?? ['.git'];
|
||||
this.contextManagement = {
|
||||
enabled: params.contextManagement?.enabled ?? false,
|
||||
@@ -2413,6 +2417,10 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
return this.experimentalMemoryManager;
|
||||
}
|
||||
|
||||
getExperimentalContextSidecarConfig(): string | undefined {
|
||||
return this.experimentalContextSidecarConfig;
|
||||
}
|
||||
|
||||
getContextManagementConfig(): ContextManagementConfig {
|
||||
return this.contextManagement;
|
||||
}
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,113 @@
|
||||
import { IrMapper } from './ir/mapper.js';
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import {
|
||||
createMockContextConfig,
|
||||
setupContextComponentTest,
|
||||
} from './testing/contextTestUtils.js';
|
||||
|
||||
describe('ContextManager Barrier Tests', () => {
|
||||
it('Soft Barrier (retainedTokens): should inject ready variants and shrink projection', async () => {
|
||||
const config = createMockContextConfig();
|
||||
const { chatHistory, contextManager } = setupContextComponentTest(config);
|
||||
|
||||
// 1. Shrink limits: 1 char = 1 token. RetainedTokens = 10. MaxTokens = 100.
|
||||
IrMapper.setConfig({ charsPerToken: 1 });
|
||||
|
||||
contextManager['sidecar'].budget.retainedTokens = 5;
|
||||
contextManager['sidecar'].budget.maxTokens = 100;
|
||||
|
||||
// 2. Build tiny history: 5 turns (10 messages). 2 tokens per turn.
|
||||
const tinyHistory = [];
|
||||
for (let i = 0; i < 5; i++) {
|
||||
tinyHistory.push({ role: 'user', parts: [{ text: `U${i}` }] });
|
||||
tinyHistory.push({ role: 'model', parts: [{ text: `M${i}` }] });
|
||||
}
|
||||
|
||||
// Set history directly to avoid event races
|
||||
await chatHistory.set(tinyHistory);
|
||||
|
||||
// 3. Pre-verify baseline length.
|
||||
const baseline = await contextManager.projectCompressedHistory();
|
||||
expect(baseline.length).toBe(10);
|
||||
|
||||
// 4. Emit a fake snapshot covering the first 3 pairs (6 messages)
|
||||
const targetEp = contextManager['pristineEpisodes'][2];
|
||||
const replacedIds = contextManager['pristineEpisodes']
|
||||
.slice(0, 3)
|
||||
.map((ep) => ep.id);
|
||||
|
||||
contextManager['eventBus'].emitVariantReady({
|
||||
targetId: targetEp.id,
|
||||
variantId: 'snapshot',
|
||||
variant: {
|
||||
status: 'ready',
|
||||
type: 'snapshot',
|
||||
replacedEpisodeIds: replacedIds,
|
||||
episode: {
|
||||
id: 'snapshot-ep',
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: 't1',
|
||||
type: 'USER_PROMPT',
|
||||
semanticParts: [],
|
||||
metadata: {
|
||||
originalTokens: 0,
|
||||
currentTokens: 0,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
yield: {
|
||||
id: 'y1',
|
||||
type: 'AGENT_YIELD',
|
||||
text: '<SNAP>',
|
||||
metadata: {
|
||||
originalTokens: 5,
|
||||
currentTokens: 5,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
steps: [],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// 5. Verify Projection shrinks: 6 original messages replaced by 1 snapshot episode (1 text part) -> length 5.
|
||||
const projection = await contextManager.projectCompressedHistory();
|
||||
expect(projection.length).toBe(5);
|
||||
// console.dir(projection, {depth: null});
|
||||
// projection[0] should be the snapshot yield
|
||||
expect(projection[0].parts![0].text).toBe('<SNAP>');
|
||||
});
|
||||
|
||||
it('Hard Barrier (maxTokens): should ruthlessly truncate unprotected episodes', async () => {
|
||||
const config = createMockContextConfig();
|
||||
const { chatHistory, contextManager } = setupContextComponentTest(config);
|
||||
|
||||
// 1. Shrink limits: maxTokens = 15.
|
||||
IrMapper.setConfig({ charsPerToken: 1 });
|
||||
contextManager['sidecar'].budget.maxTokens = 15;
|
||||
|
||||
// 2. Build history: 2 turns. Total = 24 tokens.
|
||||
const history = [
|
||||
{ role: 'user', parts: [{ text: 'U0' }] },
|
||||
{ role: 'model', parts: [{ text: 'M0_LARGE!!' }] },
|
||||
{ role: 'user', parts: [{ text: 'U1' }] },
|
||||
{ role: 'model', parts: [{ text: 'M1_LARGE!!' }] },
|
||||
];
|
||||
await chatHistory.set(history);
|
||||
|
||||
const projection = await contextManager.projectCompressedHistory();
|
||||
|
||||
// Because Turn 0 is architecturally protected (system prompt/initialization), it SURVIVES!
|
||||
// Turn 1 is dropped to satisfy the maxTokens constraint.
|
||||
expect(projection.length).toBe(2);
|
||||
expect(projection[0].parts![0].text).toBe('U0');
|
||||
expect(projection[1].parts![0].text).toBe('M0_LARGE!!');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,74 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||
import { IrMapper } from './ir/mapper.js';
|
||||
import {
|
||||
createSyntheticHistory,
|
||||
createMockContextConfig,
|
||||
setupContextComponentTest,
|
||||
} from './testing/contextTestUtils.js';
|
||||
|
||||
describe('ContextManager Sync Pressure Barrier Tests', () => {
|
||||
beforeEach(() => {
|
||||
vi.useFakeTimers();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it('should instantly truncate history when maxTokens is exceeded using truncate strategy', async () => {
|
||||
// 1. Setup
|
||||
const config = createMockContextConfig();
|
||||
const { chatHistory, contextManager } = setupContextComponentTest(config);
|
||||
|
||||
// 2. Add System Prompt (Episode 0 - Protected)
|
||||
chatHistory.set([
|
||||
{ role: 'user', parts: [{ text: 'System prompt' }] },
|
||||
{ role: 'model', parts: [{ text: 'Understood.' }] },
|
||||
]);
|
||||
|
||||
// 3. Add massive history that blows past the 150k maxTokens limit
|
||||
// 20 turns * 10,000 tokens/turn = ~200,000 tokens
|
||||
const massiveHistory = createSyntheticHistory(20, 35000);
|
||||
chatHistory.set([...chatHistory.get(), ...massiveHistory]);
|
||||
|
||||
// 4. Add the Latest Turn (Protected)
|
||||
chatHistory.set([
|
||||
...chatHistory.get(),
|
||||
{ role: 'user', parts: [{ text: 'Final question.' }] },
|
||||
{ role: 'model', parts: [{ text: 'Final answer.' }] },
|
||||
]);
|
||||
|
||||
const rawHistoryLength = chatHistory.get().length;
|
||||
IrMapper.setConfig({ charsPerToken: 1 });
|
||||
|
||||
// 5. Project History (Triggers Sync Barrier)
|
||||
const projection = await contextManager.projectCompressedHistory();
|
||||
|
||||
// 6. Assertions
|
||||
// The barrier should have dropped several older episodes to get under 150k.
|
||||
|
||||
expect(projection.length).toBeLessThan(rawHistoryLength);
|
||||
|
||||
// Verify Episode 0 (System) is perfectly preserved at the front
|
||||
|
||||
expect(projection[0].role).toBe('user');
|
||||
expect(projection[0].parts![0].text).toBe('System prompt');
|
||||
|
||||
// Verify the latest turn is perfectly preserved at the back
|
||||
const lastUser = projection[projection.length - 2];
|
||||
const lastModel = projection[projection.length - 1];
|
||||
|
||||
expect(lastUser.role).toBe('user');
|
||||
expect(lastUser.parts![0].text).toBe('Final question.');
|
||||
|
||||
expect(lastModel.role).toBe('model');
|
||||
expect(lastModel.parts![0].text).toBe('Final answer.');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,149 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
describe,
|
||||
it,
|
||||
expect,
|
||||
vi,
|
||||
beforeEach,
|
||||
beforeAll,
|
||||
afterAll,
|
||||
} from 'vitest';
|
||||
import { ContextManager } from './contextManager.js';
|
||||
import { ContextEnvironmentImpl } from './sidecar/environmentImpl.js';
|
||||
import { SidecarLoader } from './sidecar/SidecarLoader.js';
|
||||
import { ContextTracer } from './tracer.js';
|
||||
|
||||
import type { Content } from '@google/genai';
|
||||
|
||||
expect.addSnapshotSerializer({
|
||||
test: (val) =>
|
||||
typeof val === 'string' &&
|
||||
/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i.test(val),
|
||||
print: () => '"<UUID>"',
|
||||
});
|
||||
|
||||
describe('ContextManager Golden Tests', () => {
|
||||
beforeAll(() => {
|
||||
vi.useFakeTimers();
|
||||
vi.setSystemTime(new Date(2026, 3, 2).getTime());
|
||||
vi.spyOn(Math, 'random').mockReturnValue(0.5);
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
vi.useRealTimers();
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
let mockConfig: any; // eslint-disable-line @typescript-eslint/no-explicit-any
|
||||
let contextManager: ContextManager;
|
||||
|
||||
beforeEach(() => {
|
||||
mockConfig = {
|
||||
isContextManagementEnabled: vi.fn().mockReturnValue(true),
|
||||
getTargetDir: vi.fn().mockReturnValue('/tmp'),
|
||||
getSessionId: vi.fn().mockReturnValue('test-session'),
|
||||
getToolOutputMaskingConfig: vi.fn().mockResolvedValue({
|
||||
enabled: true,
|
||||
minPrunableThresholdTokens: 50,
|
||||
protectLatestTurn: false,
|
||||
protectionThresholdTokens: 100,
|
||||
}),
|
||||
storage: { getProjectTempDir: vi.fn().mockReturnValue('/tmp') },
|
||||
getUsageStatisticsEnabled: vi.fn().mockReturnValue(false),
|
||||
getBaseLlmClient: vi.fn().mockReturnValue({
|
||||
generateJson: vi.fn().mockResolvedValue({
|
||||
'test_file.txt': { level: 'SUMMARY' },
|
||||
}),
|
||||
generateContent: vi.fn().mockResolvedValue({
|
||||
candidates: [
|
||||
{ content: { parts: [{ text: 'This is a summary.' }] } },
|
||||
],
|
||||
}),
|
||||
}),
|
||||
};
|
||||
|
||||
const sidecar = SidecarLoader.fromLegacyConfig(mockConfig as any);
|
||||
const tracer = new ContextTracer('/tmp', 'test-session');
|
||||
const env = new ContextEnvironmentImpl(
|
||||
{} as any,
|
||||
'test',
|
||||
'/tmp',
|
||||
'/tmp',
|
||||
tracer,
|
||||
4,
|
||||
);
|
||||
contextManager = new ContextManager(sidecar, env, tracer);
|
||||
});
|
||||
|
||||
const createLargeHistory = (): Content[] => [
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{ text: 'A long long time ago, '.repeat(500) }, // Squashing target
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'model',
|
||||
parts: [{ text: 'in a galaxy far far away...' }],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'some_tool',
|
||||
response: { output: 'TOOL OUTPUT DATA '.repeat(500) }, // Masking target
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{ text: '--- test_file.txt ---\n' + 'FILE DATA '.repeat(1000) }, // Semantic target
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
it('should process history and match golden snapshot', async () => {
|
||||
const history = createLargeHistory();
|
||||
(contextManager as any).pristineEpisodes = (
|
||||
await import('./ir/mapper.js')
|
||||
).IrMapper.toIr(history);
|
||||
const result = await contextManager.projectCompressedHistory();
|
||||
expect(result).toMatchSnapshot();
|
||||
});
|
||||
|
||||
it('should not modify history when under budget', async () => {
|
||||
const history = createLargeHistory();
|
||||
(contextManager as any).pristineEpisodes = (
|
||||
await import('./ir/mapper.js')
|
||||
).IrMapper.toIr(history);
|
||||
// In Golden Tests, we just want to ensure the logic doesn't throw or alter unprotected history in weird ways.
|
||||
// Since we're skipping processors due to being under budget, it should equal history.
|
||||
const tracer2 = new ContextTracer('/tmp', 'test2');
|
||||
contextManager = new ContextManager(
|
||||
{
|
||||
pipelines: {
|
||||
eagerBackground: [],
|
||||
normalProcessingGraph: [],
|
||||
retainedProcessingGraph: [],
|
||||
},
|
||||
} as any,
|
||||
{} as any,
|
||||
tracer2,
|
||||
);
|
||||
|
||||
(contextManager as any).pristineEpisodes = (
|
||||
await import('./ir/mapper.js')
|
||||
).IrMapper.toIr(history);
|
||||
const result = await contextManager.projectCompressedHistory();
|
||||
|
||||
expect(result.length).toEqual(history.length);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,586 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type { Content } from '@google/genai';
|
||||
|
||||
import type { AgentChatHistory } from '../core/agentChatHistory.js';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
import { IrMapper } from './ir/mapper.js';
|
||||
import type { Episode } from './ir/types.js';
|
||||
|
||||
import { ContextEventBus } from './eventBus.js';
|
||||
import { ContextTracer } from './tracer.js';
|
||||
|
||||
import { StateSnapshotWorker } from './workers/stateSnapshotWorker.js';
|
||||
|
||||
import type { ContextEnvironment } from './sidecar/environment.js';
|
||||
|
||||
import type { SidecarConfig } from './sidecar/types.js';
|
||||
import { ProcessorRegistry } from './sidecar/registry.js';
|
||||
import type { ContextProcessor } from './pipeline.js';
|
||||
import type { AsyncContextWorker } from './workers/asyncContextWorker.js';
|
||||
|
||||
import { ToolMaskingProcessor } from './processors/toolMaskingProcessor.js';
|
||||
import { BlobDegradationProcessor } from './processors/blobDegradationProcessor.js';
|
||||
import { SemanticCompressionProcessor } from './processors/semanticCompressionProcessor.js';
|
||||
import { HistorySquashingProcessor } from './processors/historySquashingProcessor.js';
|
||||
|
||||
export class ContextManager {
|
||||
// The stateful, pristine Episodic Intermediate Representation graph.
|
||||
// This allows the agent to remember and summarize continuously without losing data across turns.
|
||||
private pristineEpisodes: Episode[] = [];
|
||||
private unsubscribeHistory?: () => void;
|
||||
private readonly eventBus: ContextEventBus;
|
||||
|
||||
// Internal sub-components
|
||||
// Synchronous processors are instantiated but effectively used as singletons within this class
|
||||
private workers: AsyncContextWorker[] = [];
|
||||
|
||||
constructor(
|
||||
private sidecar: SidecarConfig,
|
||||
private env: ContextEnvironment,
|
||||
private readonly tracer: ContextTracer,
|
||||
) {
|
||||
this.eventBus = new ContextEventBus();
|
||||
|
||||
// Register built-ins
|
||||
ProcessorRegistry.register({
|
||||
id: 'ToolMaskingProcessor',
|
||||
create: (env, opts) => new ToolMaskingProcessor(env, opts as any),
|
||||
});
|
||||
ProcessorRegistry.register({
|
||||
id: 'BlobDegradationProcessor',
|
||||
create: (env, opts) => new BlobDegradationProcessor(env),
|
||||
});
|
||||
ProcessorRegistry.register({
|
||||
id: 'SemanticCompressionProcessor',
|
||||
create: (env, opts) => new SemanticCompressionProcessor(env, opts as any),
|
||||
});
|
||||
ProcessorRegistry.register({
|
||||
id: 'HistorySquashingProcessor',
|
||||
create: (env, opts) => new HistorySquashingProcessor(env, opts as any),
|
||||
});
|
||||
ProcessorRegistry.register({
|
||||
id: 'StateSnapshotWorker',
|
||||
create: (env, opts) => new StateSnapshotWorker(env),
|
||||
});
|
||||
|
||||
this.eventBus.onVariantReady((event) => {
|
||||
// Find the target episode in the pristine graph
|
||||
const targetEp = this.pristineEpisodes.find(
|
||||
(ep) => ep.id === event.targetId,
|
||||
);
|
||||
if (targetEp) {
|
||||
if (!targetEp.variants) {
|
||||
targetEp.variants = {};
|
||||
}
|
||||
targetEp.variants[event.variantId] = event.variant;
|
||||
this.tracer.logEvent(
|
||||
'ContextManager',
|
||||
`Received async variant [${event.variantId}] for Episode ${event.targetId}`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`ContextManager: Received async variant [${event.variantId}] for Episode ${event.targetId}.`,
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
// Initialize synchronous fallback processors
|
||||
// Order matters: Fast, lossless masking -> Intelligent degradation -> Brutal truncation fallback
|
||||
|
||||
// Initialize and start background subconscious workers
|
||||
for (const bgDef of this.sidecar.pipelines.eagerBackground) {
|
||||
const worker = ProcessorRegistry.get(bgDef.processorId).create(
|
||||
this.env,
|
||||
bgDef.options,
|
||||
) as AsyncContextWorker;
|
||||
worker.start(this.eventBus);
|
||||
this.workers.push(worker);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Safely stops background workers and clears event listeners.
|
||||
*/
|
||||
shutdown() {
|
||||
for (const worker of this.workers) {
|
||||
worker.stop();
|
||||
}
|
||||
if (this.unsubscribeHistory) {
|
||||
this.unsubscribeHistory();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Subscribes to the core AgentChatHistory to natively track all message events,
|
||||
* converting them seamlessly into pristine Episodes.
|
||||
*/
|
||||
subscribeToHistory(chatHistory: AgentChatHistory) {
|
||||
if (this.unsubscribeHistory) {
|
||||
this.unsubscribeHistory();
|
||||
}
|
||||
|
||||
this.unsubscribeHistory = chatHistory.subscribe((event) => {
|
||||
// Rebuild the pristine IR graph from the full source history on every change.
|
||||
// We must map the FULL array at once because IrMapper groups adjacent
|
||||
// function calls and responses into unified Episodes. Pushing messages
|
||||
// individually would shatter these episodic boundaries.
|
||||
this.pristineEpisodes = IrMapper.toIr(chatHistory.get());
|
||||
this.tracer.logEvent(
|
||||
'ContextManager',
|
||||
'Rebuilt pristine graph from chat history update',
|
||||
{ episodeCount: this.pristineEpisodes.length },
|
||||
);
|
||||
this.checkTriggers();
|
||||
});
|
||||
}
|
||||
|
||||
private checkTriggers() {
|
||||
if (!this.sidecar.budget) return;
|
||||
|
||||
const mngConfig = this.sidecar;
|
||||
|
||||
// Calculate tokens based on the *Working Buffer View*, not the raw pristine log.
|
||||
// This solves Bug 2: The View shrinks when variants are applied, preventing infinite GC loops.
|
||||
const workingBuffer = this.getWorkingBufferView();
|
||||
const currentTokens = this.calculateIrTokens(workingBuffer);
|
||||
|
||||
this.tracer.logEvent('ContextManager', 'Evaluated triggers', {
|
||||
currentTokens,
|
||||
retainedTokens: mngConfig.budget.retainedTokens,
|
||||
});
|
||||
|
||||
// 1. Eager Compute Trigger (Continuous Streaming)
|
||||
// Broadcast the full pristine log to the async workers so they can proactively summarize partial massive files.
|
||||
this.eventBus.emitChunkReceived({ episodes: this.pristineEpisodes });
|
||||
|
||||
// 2. The Ship of Theseus Trigger (retainedTokens crossed)
|
||||
// If we exceed 65k, tell the background processors to opportunistically synthesize the oldest nodes.
|
||||
if (currentTokens > mngConfig.budget.retainedTokens) {
|
||||
const deficit = currentTokens - mngConfig.budget.retainedTokens;
|
||||
this.tracer.logEvent(
|
||||
'ContextManager',
|
||||
'Budget crossed. Emitting ConsolidationNeeded',
|
||||
{ deficit },
|
||||
);
|
||||
console.log(
|
||||
'EMITTING CONSOLIDATION. Buffer:',
|
||||
workingBuffer.length,
|
||||
'Deficit:',
|
||||
deficit,
|
||||
);
|
||||
this.eventBus.emitConsolidationNeeded({
|
||||
episodes: workingBuffer, // Pass the working buffer so they know what still needs compression
|
||||
targetDeficit: deficit,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a computed view of the pristine log.
|
||||
* Sweeps backwards (newest to oldest), tracking rolling tokens.
|
||||
* When rollingTokens > retainedTokens, it injects the "best" available ready variant
|
||||
* (snapshot > summary > masked) instead of the raw text.
|
||||
* Handles N-to-1 variant skipping automatically.
|
||||
*/
|
||||
/**
|
||||
* Applies the data-driven Sidecar configuration graphs.
|
||||
* Splits the episodes into the 'retained' and 'normal' ranges,
|
||||
* runs their respective processor pipelines sequentially, and recombines them.
|
||||
*/
|
||||
private async applyProcessorGraphs(episodes: Episode[]): Promise<Episode[]> {
|
||||
const mngConfig = this.sidecar;
|
||||
const retainedLimit = mngConfig.budget.retainedTokens;
|
||||
|
||||
// If we're incredibly small, maybe we just run the retained graph on everything?
|
||||
// Let's divide the episodes exactly at the retained boundary.
|
||||
const retainedWindow: Episode[] = [];
|
||||
const normalWindow: Episode[] = [];
|
||||
let rollingTokens = 0;
|
||||
|
||||
// Scan backwards to fill the retained window
|
||||
for (let i = episodes.length - 1; i >= 0; i--) {
|
||||
const ep = episodes[i];
|
||||
const epTokens = this.calculateIrTokens([ep]);
|
||||
if (
|
||||
(rollingTokens + epTokens <= retainedLimit &&
|
||||
normalWindow.length === 0) ||
|
||||
retainedWindow.length === 0
|
||||
) {
|
||||
// We always put at least the latest episode in the retained window.
|
||||
// We only add to retainedWindow if we haven't already started the normalWindow (contiguous block).
|
||||
retainedWindow.unshift(ep);
|
||||
rollingTokens += epTokens;
|
||||
} else {
|
||||
normalWindow.unshift(ep);
|
||||
}
|
||||
}
|
||||
|
||||
const protectedIds = new Set<string>();
|
||||
// We must protect the System Episode, which is always index 0 of pristineEpisodes.
|
||||
if (this.pristineEpisodes.length > 0) {
|
||||
protectedIds.add(this.pristineEpisodes[0].id); // Structural invariant
|
||||
}
|
||||
|
||||
const createAccountingState = (currentTotal: number) => ({
|
||||
currentTokens: currentTotal,
|
||||
maxTokens: mngConfig.budget.maxTokens,
|
||||
retainedTokens: mngConfig.budget.retainedTokens,
|
||||
deficitTokens: Math.max(0, currentTotal - mngConfig.budget.maxTokens),
|
||||
protectedEpisodeIds: protectedIds,
|
||||
isBudgetSatisfied: currentTotal <= mngConfig.budget.maxTokens, // We use maxTokens here so processors don't prematurely short-circuit if they are trying to prevent a barrier hit
|
||||
});
|
||||
|
||||
// Run Retained Graph
|
||||
let processedRetained = [...retainedWindow];
|
||||
for (const def of mngConfig.pipelines.retainedProcessingGraph) {
|
||||
const processor = ProcessorRegistry.get(def.processorId).create(
|
||||
this.env,
|
||||
def.options,
|
||||
) as ContextProcessor;
|
||||
this.tracer.logEvent(
|
||||
'ContextManager',
|
||||
`Running ${processor.name} on retained window.`,
|
||||
);
|
||||
const state = createAccountingState(
|
||||
this.calculateIrTokens([...normalWindow, ...processedRetained]),
|
||||
);
|
||||
processedRetained = await processor.process(processedRetained, state);
|
||||
}
|
||||
|
||||
// Run Normal Graph
|
||||
let processedNormal = [...normalWindow];
|
||||
for (const def of mngConfig.pipelines.normalProcessingGraph) {
|
||||
const processor = ProcessorRegistry.get(def.processorId).create(
|
||||
this.env,
|
||||
def.options,
|
||||
) as ContextProcessor;
|
||||
this.tracer.logEvent(
|
||||
'ContextManager',
|
||||
`Running ${processor.name} on normal window.`,
|
||||
);
|
||||
const state = createAccountingState(
|
||||
this.calculateIrTokens([...processedNormal, ...processedRetained]),
|
||||
);
|
||||
processedNormal = await processor.process(processedNormal, state);
|
||||
}
|
||||
|
||||
return [...processedNormal, ...processedRetained];
|
||||
}
|
||||
|
||||
public getWorkingBufferView(): Episode[] {
|
||||
const mngConfig = this.sidecar;
|
||||
const retainedTokens = mngConfig.budget.retainedTokens;
|
||||
|
||||
let currentEpisodes: Episode[] = [];
|
||||
let rollingTokens = 0;
|
||||
const skippedIds = new Set<string>();
|
||||
this.tracer.logEvent('ViewGenerator', 'Generating Working Buffer View');
|
||||
|
||||
for (let i = this.pristineEpisodes.length - 1; i >= 0; i--) {
|
||||
const ep = this.pristineEpisodes[i];
|
||||
|
||||
// If this episode was already replaced by an N-to-1 Snapshot injected earlier in the sweep, skip it entirely!
|
||||
// This solves Bug 1 (Duplicate Projection).
|
||||
if (skippedIds.has(ep.id)) {
|
||||
this.tracer.logEvent(
|
||||
'ViewGenerator',
|
||||
`Skipping episode [${ep.id}] due to N-to-1 replacement.`,
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
let projectedEp = {
|
||||
...ep,
|
||||
trigger: {
|
||||
...ep.trigger,
|
||||
metadata: {
|
||||
...ep.trigger.metadata,
|
||||
transformations: [...ep.trigger.metadata.transformations],
|
||||
},
|
||||
semanticParts:
|
||||
ep.trigger.type === 'USER_PROMPT'
|
||||
? [...ep.trigger.semanticParts.map((sp) => ({ ...sp }))]
|
||||
: undefined,
|
||||
} as any,
|
||||
steps: ep.steps.map(
|
||||
(step) =>
|
||||
({
|
||||
...step,
|
||||
metadata: {
|
||||
...step.metadata,
|
||||
transformations: [...step.metadata.transformations],
|
||||
},
|
||||
}) as any,
|
||||
),
|
||||
yield: ep.yield
|
||||
? {
|
||||
...ep.yield,
|
||||
metadata: {
|
||||
...ep.yield.metadata,
|
||||
transformations: [...ep.yield.metadata.transformations],
|
||||
},
|
||||
}
|
||||
: undefined,
|
||||
};
|
||||
|
||||
const epTokens = this.calculateIrTokens([projectedEp]);
|
||||
|
||||
if (ep.variants) {
|
||||
console.log(
|
||||
'Checking variants for',
|
||||
ep.id,
|
||||
'rollingTokens:',
|
||||
rollingTokens,
|
||||
'retained:',
|
||||
retainedTokens,
|
||||
);
|
||||
}
|
||||
if (rollingTokens > retainedTokens && ep.variants) {
|
||||
console.log('EVALUATING VARIANTS FOR', ep.id);
|
||||
const snapshot = ep.variants['snapshot'];
|
||||
const summary = ep.variants['summary'];
|
||||
const masked = ep.variants['masked'];
|
||||
|
||||
if (
|
||||
snapshot &&
|
||||
snapshot.status === 'ready' &&
|
||||
snapshot.type === 'snapshot'
|
||||
) {
|
||||
projectedEp = snapshot.episode as any;
|
||||
// Mark all the episodes this snapshot covers to be skipped by the backwards sweep.
|
||||
for (const id of snapshot.replacedEpisodeIds) {
|
||||
skippedIds.add(id);
|
||||
}
|
||||
this.tracer.logEvent(
|
||||
'ViewGenerator',
|
||||
`Episode [${ep.id}] has SnapshotVariant. Selecting variant over raw text. Added [${snapshot.replacedEpisodeIds.join(',')}] to skippedIds.`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`Opportunistically swapped Episodes [${snapshot.replacedEpisodeIds.join(', ')}] for pre-computed Snapshot variant.`,
|
||||
);
|
||||
} else if (
|
||||
summary &&
|
||||
summary.status === 'ready' &&
|
||||
summary.type === 'summary'
|
||||
) {
|
||||
projectedEp.steps = [
|
||||
{
|
||||
id: ep.id + '-summary',
|
||||
type: 'AGENT_THOUGHT',
|
||||
text: summary.text,
|
||||
metadata: {
|
||||
originalTokens: epTokens,
|
||||
currentTokens: summary.recoveredTokens || 50,
|
||||
transformations: [
|
||||
{
|
||||
processorName: 'AsyncSemanticCompressor',
|
||||
action: 'SUMMARIZED',
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
] as any;
|
||||
projectedEp.yield = undefined;
|
||||
this.tracer.logEvent(
|
||||
'ViewGenerator',
|
||||
`Episode [${ep.id}] has SummaryVariant. Selecting variant over raw text.`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`Opportunistically swapped Episode ${ep.id} for pre-computed Summary variant.`,
|
||||
);
|
||||
} else if (
|
||||
masked &&
|
||||
masked.status === 'ready' &&
|
||||
masked.type === 'masked'
|
||||
) {
|
||||
if (
|
||||
projectedEp.trigger.type === 'USER_PROMPT' &&
|
||||
projectedEp.trigger.semanticParts.length > 0
|
||||
) {
|
||||
projectedEp.trigger.semanticParts[0].presentation = {
|
||||
text: masked.text,
|
||||
tokens: masked.recoveredTokens || 10,
|
||||
};
|
||||
}
|
||||
this.tracer.logEvent(
|
||||
'ViewGenerator',
|
||||
`Episode [${ep.id}] has MaskedVariant. Selecting variant over raw text.`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`Opportunistically swapped Episode ${ep.id} for pre-computed Masked variant.`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
currentEpisodes.unshift(projectedEp);
|
||||
rollingTokens += this.calculateIrTokens([projectedEp]);
|
||||
}
|
||||
|
||||
return currentEpisodes;
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a temporary, compressed Content[] array to be used exclusively for the LLM request.
|
||||
* This does NOT mutate the pristine episodic graph.
|
||||
*/
|
||||
async projectCompressedHistory(): Promise<Content[]> {
|
||||
if (!this.sidecar.budget) {
|
||||
return this._projectAndDump(IrMapper.fromIr(this.pristineEpisodes));
|
||||
}
|
||||
|
||||
const mngConfig = this.sidecar;
|
||||
const maxTokens = mngConfig.budget.maxTokens;
|
||||
this.tracer.logEvent('ContextManager', 'Projection requested.');
|
||||
|
||||
// Get the dynamically computed Working Buffer View
|
||||
let currentEpisodes = this.getWorkingBufferView();
|
||||
|
||||
currentEpisodes = await this.applyProcessorGraphs(currentEpisodes);
|
||||
|
||||
let currentTokens = this.calculateIrTokens(currentEpisodes);
|
||||
|
||||
if (currentTokens <= maxTokens) {
|
||||
this.tracer.logEvent(
|
||||
'ContextManager',
|
||||
`View is within maxTokens (${currentTokens} <= ${maxTokens}). Returning view.`,
|
||||
);
|
||||
return this._projectAndDump(IrMapper.fromIr(currentEpisodes));
|
||||
}
|
||||
|
||||
this.tracer.logEvent(
|
||||
'ContextManager',
|
||||
`View exceeds maxTokens (${currentTokens} > ${maxTokens}). Hitting Synchronous Pressure Barrier. Strategy: ${mngConfig.gcBackstop.strategy}`,
|
||||
);
|
||||
// --- The Synchronous Pressure Barrier ---
|
||||
// The background eager workers couldn't keep up, or a massive file was pasted.
|
||||
// The Working Buffer View is still over the absolute hard limit (maxTokens).
|
||||
// We MUST reduce tokens before returning, or the API request will 400.
|
||||
|
||||
debugLogger.log(
|
||||
`Context Manager Synchronous Barrier triggered: View at ${currentTokens} tokens (limit: ${maxTokens}). Strategy: ${mngConfig.gcBackstop.strategy}`,
|
||||
);
|
||||
|
||||
// Calculate target based on gcTarget
|
||||
let targetTokens = maxTokens;
|
||||
|
||||
if (mngConfig.gcBackstop.target === 'max') {
|
||||
targetTokens = mngConfig.budget.retainedTokens;
|
||||
} else if (mngConfig.gcBackstop.target === 'freeNTokens') {
|
||||
targetTokens =
|
||||
maxTokens - (mngConfig.gcBackstop.freeTokensTarget ?? 10000);
|
||||
}
|
||||
|
||||
// Structural invariant: We ALWAYS protect the architectural initialization turn (Turn 0)
|
||||
// We do NOT arbitrarily protect recent episodes (like currentEpisodes.length - 1)
|
||||
// because an episode can be unboundedly large, and protecting it would crash the LLM.
|
||||
const protectedEpisodeId =
|
||||
this.pristineEpisodes.length > 0 ? this.pristineEpisodes[0].id : null;
|
||||
|
||||
let remainingTokens = currentTokens;
|
||||
|
||||
const truncated: Episode[] = [];
|
||||
|
||||
const strategy = mngConfig.gcBackstop.strategy;
|
||||
|
||||
for (const ep of currentEpisodes) {
|
||||
const epTokens = this.calculateIrTokens([ep]);
|
||||
if (remainingTokens > targetTokens && ep.id !== protectedEpisodeId) {
|
||||
console.log(
|
||||
'DROPPING EPISODE:',
|
||||
ep.id,
|
||||
'rem:',
|
||||
remainingTokens,
|
||||
'tgt:',
|
||||
targetTokens,
|
||||
);
|
||||
|
||||
remainingTokens -= epTokens;
|
||||
if (strategy === 'truncate') {
|
||||
this.tracer.logEvent('Barrier', `Truncating episode [${ep.id}].`);
|
||||
|
||||
debugLogger.log(`Barrier (truncate): Dropped Episode ${ep.id}`);
|
||||
} else if (strategy === 'compress') {
|
||||
this.tracer.logEvent(
|
||||
'Barrier',
|
||||
`Compress fallback to truncate for [${ep.id}].`,
|
||||
);
|
||||
debugLogger.warn(
|
||||
`Synchronous compress barrier not fully implemented, truncating Episode ${ep.id}.`,
|
||||
);
|
||||
} else if (strategy === 'rollingSummarizer') {
|
||||
this.tracer.logEvent(
|
||||
'Barrier',
|
||||
`RollingSummarizer fallback to truncate for [${ep.id}].`,
|
||||
);
|
||||
debugLogger.warn(
|
||||
`Synchronous rollingSummarizer barrier not fully implemented, truncating Episode ${ep.id}.`,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
console.log(
|
||||
'KEEPING EPISODE:',
|
||||
ep.id,
|
||||
'rem:',
|
||||
remainingTokens,
|
||||
'tgt:',
|
||||
targetTokens,
|
||||
);
|
||||
truncated.push(ep);
|
||||
}
|
||||
}
|
||||
currentEpisodes = truncated;
|
||||
|
||||
const finalTokens = this.calculateIrTokens(currentEpisodes);
|
||||
this.tracer.logEvent(
|
||||
'ContextManager',
|
||||
`Finished projection. Final token count: ${finalTokens}.`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`Context Manager finished. Final actual token count: ${finalTokens}.`,
|
||||
);
|
||||
|
||||
return this._projectAndDump(IrMapper.fromIr(currentEpisodes));
|
||||
}
|
||||
|
||||
private async _projectAndDump(contents: Content[]): Promise<Content[]> {
|
||||
if (process.env['GEMINI_DUMP_CONTEXT'] === 'true') {
|
||||
try {
|
||||
const fs = await import('node:fs/promises');
|
||||
const path = await import('node:path');
|
||||
const dumpPath = path.join(
|
||||
this.env.getTraceDir(),
|
||||
'.gemini',
|
||||
'projected_context.json',
|
||||
);
|
||||
await fs.mkdir(path.dirname(dumpPath), { recursive: true });
|
||||
await fs.writeFile(
|
||||
dumpPath,
|
||||
JSON.stringify(contents, null, 2),
|
||||
'utf-8',
|
||||
);
|
||||
debugLogger.log(
|
||||
`[Observability] Context successfully dumped to ${dumpPath}`,
|
||||
);
|
||||
} catch (e) {
|
||||
debugLogger.error(`Failed to dump context: ${e}`);
|
||||
}
|
||||
}
|
||||
return contents;
|
||||
}
|
||||
|
||||
private calculateIrTokens(episodes: Episode[]): number {
|
||||
let tokens = 0;
|
||||
for (const ep of episodes) {
|
||||
if (ep.trigger) tokens += ep.trigger.metadata.currentTokens;
|
||||
for (const step of ep.steps) {
|
||||
tokens += step.metadata.currentTokens;
|
||||
}
|
||||
if (ep.yield) tokens += ep.yield.metadata.currentTokens;
|
||||
}
|
||||
return tokens;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { EventEmitter } from 'node:events';
|
||||
import type { Episode, Variant } from './ir/types.js';
|
||||
|
||||
export interface ContextConsolidationEvent {
|
||||
episodes: Episode[];
|
||||
targetDeficit: number;
|
||||
}
|
||||
|
||||
export interface IrChunkReceivedEvent {
|
||||
episodes: Episode[];
|
||||
}
|
||||
|
||||
export interface VariantReadyEvent {
|
||||
targetId: string; // The Episode or Step ID this variant attaches to
|
||||
variantId: string; // A unique ID for the variant itself
|
||||
variant: Variant;
|
||||
}
|
||||
|
||||
export class ContextEventBus extends EventEmitter {
|
||||
emitChunkReceived(event: IrChunkReceivedEvent) {
|
||||
this.emit('IR_CHUNK_RECEIVED', event);
|
||||
}
|
||||
|
||||
onChunkReceived(listener: (event: IrChunkReceivedEvent) => void) {
|
||||
this.on('IR_CHUNK_RECEIVED', listener);
|
||||
}
|
||||
|
||||
emitConsolidationNeeded(event: ContextConsolidationEvent) {
|
||||
this.emit('BUDGET_RETAINED_CROSSED', event);
|
||||
}
|
||||
|
||||
onConsolidationNeeded(listener: (event: ContextConsolidationEvent) => void) {
|
||||
this.on('BUDGET_RETAINED_CROSSED', listener);
|
||||
}
|
||||
|
||||
emitVariantReady(event: VariantReadyEvent) {
|
||||
this.emit('VARIANT_READY', event);
|
||||
}
|
||||
|
||||
onVariantReady(listener: (event: VariantReadyEvent) => void) {
|
||||
this.on('VARIANT_READY', listener);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { IrMapper } from './mapper.js';
|
||||
import type { Content } from '@google/genai';
|
||||
import type { UserPrompt, ToolExecution } from './types.js';
|
||||
|
||||
describe('IrMapper', () => {
|
||||
it('should correctly map a complex conversation into Episodes and back', () => {
|
||||
const rawHistory: Content[] = [
|
||||
{ role: 'user', parts: [{ text: 'Can you read file A and B?' }] },
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{ text: 'Let me check those files.' },
|
||||
{
|
||||
functionCall: {
|
||||
id: 'call_1',
|
||||
name: 'read_file',
|
||||
args: { filepath: 'A.txt' },
|
||||
},
|
||||
},
|
||||
{
|
||||
functionCall: {
|
||||
id: 'call_2',
|
||||
name: 'read_file',
|
||||
args: { filepath: 'B.txt' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
id: 'call_1',
|
||||
name: 'read_file',
|
||||
response: { output: 'Contents of A' },
|
||||
},
|
||||
},
|
||||
{
|
||||
functionResponse: {
|
||||
id: 'call_2',
|
||||
name: 'read_file',
|
||||
response: { output: 'Contents of B' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{ text: 'Thanks. Now I will compile.' },
|
||||
{
|
||||
functionCall: {
|
||||
id: 'call_3',
|
||||
name: 'shell',
|
||||
args: { cmd: 'make' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
id: 'call_3',
|
||||
name: 'shell',
|
||||
response: { output: 'success' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{ role: 'model', parts: [{ text: 'Everything is done!' }] },
|
||||
];
|
||||
|
||||
const episodes = IrMapper.toIr(rawHistory);
|
||||
|
||||
expect(episodes).toHaveLength(1);
|
||||
const ep = episodes[0];
|
||||
|
||||
expect(ep.trigger.type).toBe('USER_PROMPT');
|
||||
expect(
|
||||
((ep.trigger as UserPrompt).semanticParts[0] as { text: string }).text,
|
||||
).toBe('Can you read file A and B?');
|
||||
|
||||
// Steps should be: Thought, ToolExecution(A), ToolExecution(B), Thought, ToolExecution(make)
|
||||
expect(ep.steps).toHaveLength(5);
|
||||
expect(ep.steps[0].type).toBe('AGENT_THOUGHT');
|
||||
expect(ep.steps[1].type).toBe('TOOL_EXECUTION');
|
||||
expect((ep.steps[1] as ToolExecution).toolName).toBe('read_file');
|
||||
expect((ep.steps[1] as ToolExecution).intent).toEqual({
|
||||
filepath: 'A.txt',
|
||||
});
|
||||
expect((ep.steps[1] as ToolExecution).observation).toEqual({
|
||||
output: 'Contents of A',
|
||||
});
|
||||
|
||||
expect(ep.steps[2].type).toBe('TOOL_EXECUTION');
|
||||
expect((ep.steps[2] as ToolExecution).intent).toEqual({
|
||||
filepath: 'B.txt',
|
||||
});
|
||||
|
||||
expect(ep.steps[3].type).toBe('AGENT_THOUGHT');
|
||||
|
||||
expect(ep.steps[4].type).toBe('TOOL_EXECUTION');
|
||||
expect((ep.steps[4] as ToolExecution).toolName).toBe('shell');
|
||||
|
||||
expect(ep.yield?.type).toBe('AGENT_YIELD');
|
||||
expect(ep.yield?.text).toBe('Everything is done!');
|
||||
|
||||
// Test Re-serialization
|
||||
const reconstituted = IrMapper.fromIr(episodes);
|
||||
|
||||
// Compare basic structure (the reconstituted version might have slightly different grouping of calls/responses
|
||||
// based on flush logic, but semantically equivalent)
|
||||
expect(reconstituted[0]).toEqual(rawHistory[0]);
|
||||
// Reconstituted history is identical except tool IDs will be reassigned because IrMapper discards string IDs in favor of deterministic object hash IDs
|
||||
expect(reconstituted[1].parts![0]).toEqual(rawHistory[1].parts![0]);
|
||||
|
||||
// The exact structural equivalence isn't mathematically perfect because Gemini allows mixing text and calls
|
||||
// in one Content block, but the flat representation is semantically identical.
|
||||
});
|
||||
|
||||
it('should guarantee WeakMap ID stability across continuous mapping', () => {
|
||||
// 1. Initial history
|
||||
const history: Content[] = [
|
||||
{ role: 'user', parts: [{ text: 'Hello' }] },
|
||||
{ role: 'model', parts: [{ text: 'Hi there' }] },
|
||||
];
|
||||
|
||||
const initialIr = IrMapper.toIr(history);
|
||||
expect(initialIr).toHaveLength(1);
|
||||
|
||||
// Save the uniquely generated deterministic ID for the first episode
|
||||
const episodeId = initialIr[0].id;
|
||||
const triggerId = initialIr[0].trigger.id;
|
||||
|
||||
// 2. Push new history (simulating a continuing conversation)
|
||||
history.push({ role: 'user', parts: [{ text: 'How are you?' }] });
|
||||
history.push({ role: 'model', parts: [{ text: 'I am an AI.' }] });
|
||||
|
||||
const updatedIr = IrMapper.toIr(history);
|
||||
expect(updatedIr).toHaveLength(2);
|
||||
|
||||
// 3. Verify ID Stability
|
||||
// The exact same ID must be generated for the first episode because the underlying Content object reference hasn't changed.
|
||||
// This proves the WeakMap successfully pinned the reference!
|
||||
expect(updatedIr[0].id).toBe(episodeId);
|
||||
expect(updatedIr[0].trigger.id).toBe(triggerId);
|
||||
|
||||
// Ensure the new episode has a different ID
|
||||
expect(updatedIr[1].id).not.toBe(episodeId);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,307 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Content, Part } from '@google/genai';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import type {
|
||||
Episode,
|
||||
IrMetadata,
|
||||
SemanticPart,
|
||||
ToolExecution,
|
||||
AgentThought,
|
||||
AgentYield,
|
||||
UserPrompt,
|
||||
} from './types.js';
|
||||
import { estimateContextTokenCountSync as estimateTokenCountSync } from '../utils/contextTokenCalculator.js';
|
||||
|
||||
// WeakMap to provide stable, deterministic identity across parses for the exact same Content/Part references
|
||||
const nodeIdentityMap = new WeakMap<object, string>();
|
||||
|
||||
function getStableId(obj: object): string {
|
||||
let id = nodeIdentityMap.get(obj);
|
||||
if (!id) {
|
||||
id = randomUUID();
|
||||
nodeIdentityMap.set(obj, id);
|
||||
}
|
||||
return id;
|
||||
}
|
||||
|
||||
export class IrMapper {
|
||||
static setConfig(cfg: { charsPerToken?: number }) {
|
||||
this.config = cfg;
|
||||
}
|
||||
private static config: { charsPerToken?: number } | undefined;
|
||||
|
||||
/**
|
||||
* Translates a flat Gemini Content[] array into our rich Episodic Intermediate Representation.
|
||||
* Groups adjacent function calls and responses into unified ToolExecution nodes.
|
||||
*/
|
||||
static toIr(history: readonly Content[]): Episode[] {
|
||||
const episodes: Episode[] = [];
|
||||
let currentEpisode: Partial<Episode> | null = null;
|
||||
const pendingCallParts: Map<string, Part> = new Map();
|
||||
|
||||
const createMetadata = (parts: Part[]): IrMetadata => {
|
||||
const tokens = estimateTokenCountSync(parts, 0, IrMapper.config);
|
||||
return {
|
||||
originalTokens: tokens,
|
||||
currentTokens: tokens,
|
||||
transformations: [],
|
||||
};
|
||||
};
|
||||
|
||||
const finalizeEpisode = () => {
|
||||
if (currentEpisode && currentEpisode.trigger) {
|
||||
episodes.push(currentEpisode as unknown as Episode); // eslint-disable-line @typescript-eslint/no-unsafe-type-assertion
|
||||
}
|
||||
currentEpisode = null;
|
||||
};
|
||||
|
||||
for (const msg of history) {
|
||||
if (!msg.parts) continue;
|
||||
|
||||
if (msg.role === 'user') {
|
||||
const hasToolResponses = msg.parts.some((p) => !!p.functionResponse);
|
||||
const hasUserParts = msg.parts.some(
|
||||
(p) => !!p.text || !!p.inlineData || !!p.fileData,
|
||||
);
|
||||
|
||||
if (hasToolResponses) {
|
||||
if (!currentEpisode) {
|
||||
currentEpisode = {
|
||||
id: getStableId(msg),
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: getStableId(msg.parts[0] || msg),
|
||||
type: 'SYSTEM_EVENT',
|
||||
name: 'history_resume',
|
||||
payload: {},
|
||||
metadata: createMetadata([]),
|
||||
},
|
||||
steps: [],
|
||||
};
|
||||
}
|
||||
|
||||
for (const part of msg.parts) {
|
||||
if (part.functionResponse) {
|
||||
const callId = part.functionResponse.id || '';
|
||||
const matchingCall = pendingCallParts.get(callId);
|
||||
|
||||
const intentTokens = matchingCall
|
||||
? estimateTokenCountSync([matchingCall])
|
||||
: 0;
|
||||
const obsTokens = estimateTokenCountSync([part]);
|
||||
|
||||
const step: ToolExecution = {
|
||||
id: getStableId(part),
|
||||
type: 'TOOL_EXECUTION',
|
||||
toolName: part.functionResponse.name || 'unknown',
|
||||
intent:
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
(matchingCall?.functionCall?.args as unknown as Record<
|
||||
string,
|
||||
unknown
|
||||
>) || {},
|
||||
observation:
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
(part.functionResponse.response as unknown as Record<
|
||||
string,
|
||||
unknown
|
||||
>) || {},
|
||||
tokens: {
|
||||
intent: intentTokens,
|
||||
observation: obsTokens,
|
||||
},
|
||||
metadata: {
|
||||
originalTokens: intentTokens + obsTokens,
|
||||
currentTokens: intentTokens + obsTokens,
|
||||
transformations: [],
|
||||
},
|
||||
};
|
||||
currentEpisode.steps!.push(step);
|
||||
if (callId) pendingCallParts.delete(callId);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (hasUserParts) {
|
||||
finalizeEpisode();
|
||||
|
||||
const semanticParts: SemanticPart[] = [];
|
||||
for (const p of msg.parts) {
|
||||
if (p.text !== undefined)
|
||||
semanticParts.push({ type: 'text', text: p.text });
|
||||
else if (p.inlineData)
|
||||
semanticParts.push({
|
||||
type: 'inline_data',
|
||||
mimeType: p.inlineData.mimeType || '',
|
||||
data: p.inlineData.data || '',
|
||||
});
|
||||
else if (p.fileData)
|
||||
semanticParts.push({
|
||||
type: 'file_data',
|
||||
mimeType: p.fileData.mimeType || '',
|
||||
fileUri: p.fileData.fileUri || '',
|
||||
});
|
||||
else if (!p.functionResponse)
|
||||
semanticParts.push({ type: 'raw_part', part: p }); // Preserve unknowns
|
||||
}
|
||||
|
||||
const trigger: UserPrompt = {
|
||||
id: getStableId(msg.parts[0] || msg),
|
||||
type: 'USER_PROMPT',
|
||||
semanticParts,
|
||||
metadata: createMetadata(
|
||||
msg.parts.filter((p) => !p.functionResponse),
|
||||
),
|
||||
};
|
||||
|
||||
currentEpisode = {
|
||||
id: getStableId(msg),
|
||||
timestamp: Date.now(),
|
||||
trigger,
|
||||
steps: [],
|
||||
};
|
||||
}
|
||||
} else if (msg.role === 'model') {
|
||||
if (!currentEpisode) {
|
||||
currentEpisode = {
|
||||
id: getStableId(msg),
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: getStableId(msg.parts[0] || msg),
|
||||
type: 'SYSTEM_EVENT',
|
||||
name: 'model_init',
|
||||
payload: {},
|
||||
metadata: createMetadata([]),
|
||||
},
|
||||
steps: [],
|
||||
};
|
||||
}
|
||||
|
||||
for (const part of msg.parts) {
|
||||
if (part.functionCall) {
|
||||
const callId = part.functionCall.id || '';
|
||||
if (callId) pendingCallParts.set(callId, part);
|
||||
} else if (part.text) {
|
||||
const thought: AgentThought = {
|
||||
id: getStableId(part),
|
||||
type: 'AGENT_THOUGHT',
|
||||
text: part.text,
|
||||
metadata: createMetadata([part]),
|
||||
};
|
||||
currentEpisode.steps!.push(thought);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (currentEpisode) {
|
||||
if (currentEpisode.steps && currentEpisode.steps.length > 0) {
|
||||
const lastStep = currentEpisode.steps[currentEpisode.steps.length - 1];
|
||||
if (lastStep.type === 'AGENT_THOUGHT') {
|
||||
const yieldNode: AgentYield = {
|
||||
id: lastStep.id,
|
||||
type: 'AGENT_YIELD',
|
||||
text: lastStep.text,
|
||||
metadata: lastStep.metadata,
|
||||
};
|
||||
currentEpisode.steps.pop();
|
||||
currentEpisode.yield = yieldNode;
|
||||
}
|
||||
}
|
||||
finalizeEpisode();
|
||||
}
|
||||
|
||||
return episodes;
|
||||
}
|
||||
|
||||
/**
|
||||
* Re-serializes the Episodic IR back into a flat Gemini Content[] array.
|
||||
*/
|
||||
static fromIr(episodes: Episode[]): Content[] {
|
||||
const history: Content[] = [];
|
||||
|
||||
for (const ep of episodes) {
|
||||
// 1. Serialize Trigger
|
||||
if (ep.trigger.type === 'USER_PROMPT') {
|
||||
const parts: Part[] = [];
|
||||
for (const sp of ep.trigger.semanticParts) {
|
||||
if (sp.presentation) {
|
||||
parts.push({ text: sp.presentation.text });
|
||||
} else if (sp.type === 'text') {
|
||||
parts.push({ text: sp.text });
|
||||
} else if (sp.type === 'inline_data') {
|
||||
parts.push({
|
||||
inlineData: { mimeType: sp.mimeType, data: sp.data },
|
||||
});
|
||||
} else if (sp.type === 'file_data') {
|
||||
parts.push({
|
||||
fileData: { mimeType: sp.mimeType, fileUri: sp.fileUri },
|
||||
});
|
||||
} else if (sp.type === 'raw_part') {
|
||||
// eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion, @typescript-eslint/no-unsafe-type-assertion
|
||||
parts.push(sp.part as unknown as Part);
|
||||
}
|
||||
}
|
||||
if (parts.length > 0) history.push({ role: 'user', parts });
|
||||
}
|
||||
|
||||
// 2. Serialize Steps
|
||||
let pendingModelParts: Part[] = [];
|
||||
let pendingUserParts: Part[] = [];
|
||||
|
||||
const flushPending = () => {
|
||||
if (pendingModelParts.length > 0) {
|
||||
history.push({ role: 'model', parts: [...pendingModelParts] });
|
||||
pendingModelParts = [];
|
||||
}
|
||||
if (pendingUserParts.length > 0) {
|
||||
history.push({ role: 'user', parts: [...pendingUserParts] });
|
||||
pendingUserParts = [];
|
||||
}
|
||||
};
|
||||
|
||||
for (const step of ep.steps) {
|
||||
if (step.type === 'AGENT_THOUGHT') {
|
||||
if (pendingUserParts.length > 0) flushPending();
|
||||
pendingModelParts.push({
|
||||
text: step.presentation?.text ?? step.text,
|
||||
});
|
||||
} else if (step.type === 'TOOL_EXECUTION') {
|
||||
pendingModelParts.push({
|
||||
functionCall: {
|
||||
name: step.toolName,
|
||||
args: step.intent as unknown as Record<string, unknown>, // eslint-disable-line @typescript-eslint/no-unsafe-type-assertion
|
||||
id: step.id,
|
||||
},
|
||||
});
|
||||
const observation = step.presentation
|
||||
? step.presentation.observation
|
||||
: step.observation;
|
||||
pendingUserParts.push({
|
||||
functionResponse: {
|
||||
name: step.toolName,
|
||||
response: observation as unknown as Record<string, unknown>, // eslint-disable-line @typescript-eslint/no-unsafe-type-assertion
|
||||
id: step.id,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
flushPending();
|
||||
|
||||
// 3. Serialize Yield
|
||||
if (ep.yield) {
|
||||
history.push({
|
||||
role: 'model',
|
||||
parts: [{ text: ep.yield.presentation?.text ?? ep.yield.text }],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return history;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
/**
|
||||
* Universal Audit Metadata
|
||||
* Tracks the lifecycle and transformations of a node or part within the IR.
|
||||
* This guarantees perfect reversibility and enables long-term memory offloading.
|
||||
*/
|
||||
export interface IrMetadata {
|
||||
/** The estimated number of tokens this entity originally consumed. */
|
||||
originalTokens: number;
|
||||
/** The current estimated number of tokens this entity consumes in its degraded state. */
|
||||
currentTokens: number;
|
||||
/** An audit trail of all transformations applied by ContextProcessors. */
|
||||
transformations: Array<{
|
||||
processorName: string;
|
||||
action:
|
||||
| 'MASKED'
|
||||
| 'TRUNCATED'
|
||||
| 'DEGRADED'
|
||||
| 'SUMMARIZED'
|
||||
| 'EVICTED'
|
||||
| 'SYNTHESIZED';
|
||||
timestamp: number;
|
||||
/** Pointer to where the original uncompressed payload was saved (if applicable) */
|
||||
diskPointer?: string;
|
||||
}>;
|
||||
}
|
||||
|
||||
export type IrNodeType =
|
||||
| 'USER_PROMPT'
|
||||
| 'SYSTEM_EVENT'
|
||||
| 'AGENT_THOUGHT'
|
||||
| 'TOOL_EXECUTION'
|
||||
| 'AGENT_YIELD';
|
||||
|
||||
/** Base interface for all nodes in the Episodic IR */
|
||||
export type VariantStatus = 'computing' | 'ready' | 'failed';
|
||||
|
||||
export interface BaseVariant {
|
||||
status: VariantStatus;
|
||||
recoveredTokens?: number;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
export interface SummaryVariant extends BaseVariant {
|
||||
type: 'summary';
|
||||
text: string;
|
||||
}
|
||||
|
||||
export interface MaskedVariant extends BaseVariant {
|
||||
type: 'masked';
|
||||
text: string;
|
||||
}
|
||||
|
||||
export interface SnapshotVariant extends BaseVariant {
|
||||
type: 'snapshot';
|
||||
episode: Episode;
|
||||
replacedEpisodeIds: string[];
|
||||
}
|
||||
|
||||
export type Variant = SummaryVariant | MaskedVariant | SnapshotVariant;
|
||||
|
||||
/** Base interface for all nodes in the Episodic IR */
|
||||
export interface IrNode {
|
||||
readonly id: string;
|
||||
readonly type: IrNodeType;
|
||||
metadata: IrMetadata;
|
||||
variants?: Record<string, Variant>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Semantic Parts for User Prompts
|
||||
* Ensures we can safely truncate text without deleting multi-modal parts (like images).
|
||||
*/
|
||||
export type SemanticPart =
|
||||
| {
|
||||
type: 'text';
|
||||
text: string;
|
||||
presentation?: { text: string; tokens: number };
|
||||
}
|
||||
| {
|
||||
type: 'inline_data';
|
||||
mimeType: string;
|
||||
data: string;
|
||||
presentation?: { text: string; tokens: number };
|
||||
}
|
||||
| {
|
||||
type: 'file_data';
|
||||
mimeType: string;
|
||||
fileUri: string;
|
||||
presentation?: { text: string; tokens: number };
|
||||
}
|
||||
| {
|
||||
type: 'raw_part';
|
||||
part: unknown;
|
||||
presentation?: { text: string; tokens: number };
|
||||
};
|
||||
|
||||
/**
|
||||
* Trigger Nodes
|
||||
* Events that wake the agent up and initiate an Episode.
|
||||
*/
|
||||
export interface UserPrompt extends IrNode {
|
||||
readonly type: 'USER_PROMPT';
|
||||
/** The semantic breakdown of the user's multi-modal input */
|
||||
semanticParts: SemanticPart[];
|
||||
}
|
||||
|
||||
export interface SystemEvent extends IrNode {
|
||||
readonly type: 'SYSTEM_EVENT';
|
||||
name: string;
|
||||
payload: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export type EpisodeTrigger = UserPrompt | SystemEvent;
|
||||
|
||||
/**
|
||||
* Step Nodes
|
||||
* The internal autonomous actions taken by the agent during its loop.
|
||||
*/
|
||||
export interface AgentThought extends IrNode {
|
||||
readonly type: 'AGENT_THOUGHT';
|
||||
text: string;
|
||||
/** Overrides the rendered output for this thought */
|
||||
presentation?: {
|
||||
text: string;
|
||||
tokens: number;
|
||||
};
|
||||
}
|
||||
|
||||
export interface ToolExecution extends IrNode {
|
||||
readonly type: 'TOOL_EXECUTION';
|
||||
/** The name of the tool invoked */
|
||||
toolName: string;
|
||||
|
||||
/** The arguments passed to the tool (The 'FunctionCall') */
|
||||
intent: Record<string, unknown>;
|
||||
|
||||
/** The result returned by the tool (The 'FunctionResponse') */
|
||||
observation: string | Record<string, unknown>;
|
||||
|
||||
/** Granular token tracking for the different lifecycle phases of the tool */
|
||||
tokens: {
|
||||
intent: number;
|
||||
observation: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* The presentation layer. If defined, the IrMapper uses this instead of the
|
||||
* raw observation to build the functionResponse.
|
||||
* This preserves the immutable raw data for semantic queries while modifying the rendered output.
|
||||
*/
|
||||
presentation?: {
|
||||
intent?: Record<string, unknown>;
|
||||
observation?: string | Record<string, unknown>;
|
||||
tokens: {
|
||||
intent: number;
|
||||
observation: number;
|
||||
};
|
||||
};
|
||||
}
|
||||
|
||||
export type EpisodeStep = AgentThought | ToolExecution;
|
||||
|
||||
/**
|
||||
* Resolution Node
|
||||
* The final message where the agent yields control back to the user.
|
||||
*/
|
||||
export interface AgentYield extends IrNode {
|
||||
readonly type: 'AGENT_YIELD';
|
||||
text: string;
|
||||
presentation?: {
|
||||
text: string;
|
||||
tokens: number;
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* The Episode
|
||||
* A discrete, continuous run of the agent. Represents the full cycle from
|
||||
* taking control (Trigger) to returning control (Yield), encompassing all
|
||||
* internal reasoning and observations (Steps).
|
||||
*/
|
||||
export interface Episode {
|
||||
readonly id: string;
|
||||
/** When the episode began */
|
||||
readonly timestamp: number;
|
||||
variants?: Record<string, Variant>;
|
||||
|
||||
/** The event that initiated this run */
|
||||
trigger: EpisodeTrigger;
|
||||
|
||||
/** The sequence of autonomous actions and observations */
|
||||
steps: EpisodeStep[];
|
||||
|
||||
/** The final handover back to the user (can be undefined if the episode was aborted/errored) */
|
||||
yield?: AgentYield;
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Episode } from './ir/types.js';
|
||||
|
||||
/**
|
||||
* State object passed through the processing pipeline.
|
||||
* Contains global accounting logic and semantic protection rules.
|
||||
*/
|
||||
export interface ContextAccountingState {
|
||||
readonly currentTokens: number;
|
||||
readonly maxTokens: number;
|
||||
readonly retainedTokens: number;
|
||||
|
||||
/** The exact number of tokens that need to be trimmed to reach the retainedTokens goal */
|
||||
readonly deficitTokens: number;
|
||||
|
||||
/**
|
||||
* Set of Episode IDs that the orchestrator has deemed highly protected.
|
||||
* Processors should generally skip mutating these episodes unless doing proactive/required transforms.
|
||||
*/
|
||||
readonly protectedEpisodeIds: Set<string>;
|
||||
|
||||
/**
|
||||
* True if currentTokens <= retainedTokens.
|
||||
*/
|
||||
readonly isBudgetSatisfied: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Interface for all context degradation strategies.
|
||||
*/
|
||||
export interface ContextProcessor {
|
||||
/** Unique name for telemetry and logging. */
|
||||
readonly name: string;
|
||||
|
||||
/**
|
||||
* Processes the episodic history payload based on the current accounting state.
|
||||
* Processors should return a new or mutated array of episodes.
|
||||
*/
|
||||
process(
|
||||
episodes: Episode[],
|
||||
state: ContextAccountingState,
|
||||
): Promise<Episode[]>;
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import { createMockEnvironment } from '../testing/contextTestUtils.js';
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { BlobDegradationProcessor } from './blobDegradationProcessor.js';
|
||||
import type { Episode, UserPrompt } from '../ir/types.js';
|
||||
import type { ContextAccountingState } from '../pipeline.js';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import * as fsPromises from 'node:fs/promises';
|
||||
|
||||
vi.mock('node:fs/promises');
|
||||
|
||||
describe('BlobDegradationProcessor', () => {
|
||||
let processor: BlobDegradationProcessor;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
|
||||
processor = new BlobDegradationProcessor(createMockEnvironment());
|
||||
});
|
||||
|
||||
const getDummyState = (
|
||||
isSatisfied = false,
|
||||
deficit = 0,
|
||||
protectedIds = new Set<string>(),
|
||||
): ContextAccountingState => ({
|
||||
currentTokens: 5000,
|
||||
maxTokens: 10000,
|
||||
retainedTokens: 4000,
|
||||
deficitTokens: deficit,
|
||||
protectedEpisodeIds: protectedIds,
|
||||
isBudgetSatisfied: isSatisfied,
|
||||
});
|
||||
|
||||
it('degrades inline_data into a text reference and saves to disk', async () => {
|
||||
const dummyImageBase64 = Buffer.from('fake-image-data').toString('base64');
|
||||
|
||||
const ep: Episode = {
|
||||
id: 'ep-1',
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: randomUUID(),
|
||||
type: 'USER_PROMPT',
|
||||
semanticParts: [
|
||||
{ type: 'text', text: 'Look at this image:' },
|
||||
{
|
||||
type: 'inline_data',
|
||||
mimeType: 'image/png',
|
||||
data: dummyImageBase64,
|
||||
},
|
||||
],
|
||||
metadata: {
|
||||
originalTokens: 300,
|
||||
currentTokens: 300,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
steps: [],
|
||||
};
|
||||
|
||||
// Fake token calculator says inlineData costs 258 tokens, text costs 10
|
||||
const state = getDummyState(false, 500, new Set());
|
||||
const result = await processor.process([ep], state);
|
||||
|
||||
const parts = (result[0].trigger as UserPrompt).semanticParts;
|
||||
|
||||
// Text part should be untouched
|
||||
expect(parts[0].presentation).toBeUndefined();
|
||||
|
||||
// Inline data should be degraded
|
||||
expect(parts[1].presentation).toBeDefined();
|
||||
expect(parts[1].presentation!.text).toContain(
|
||||
'[Multi-Modal Blob (image/png',
|
||||
);
|
||||
expect(parts[1].presentation!.text).toContain(
|
||||
'degraded to text to preserve context window',
|
||||
);
|
||||
|
||||
expect(fsPromises.writeFile).toHaveBeenCalledTimes(1);
|
||||
expect(result[0].trigger.metadata.transformations.length).toBe(1);
|
||||
});
|
||||
|
||||
it('degrades file_data into a text reference without disk write', async () => {
|
||||
const ep: Episode = {
|
||||
id: 'ep-2',
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: randomUUID(),
|
||||
type: 'USER_PROMPT',
|
||||
semanticParts: [
|
||||
{
|
||||
type: 'file_data',
|
||||
mimeType: 'application/pdf',
|
||||
fileUri: 'gs://fake-bucket/doc.pdf',
|
||||
},
|
||||
],
|
||||
metadata: {
|
||||
originalTokens: 300,
|
||||
currentTokens: 300,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
steps: [],
|
||||
};
|
||||
|
||||
const state = getDummyState(false, 500, new Set());
|
||||
const result = await processor.process([ep], state);
|
||||
|
||||
const parts = (result[0].trigger as UserPrompt).semanticParts;
|
||||
expect(parts[0].presentation).toBeDefined();
|
||||
expect(parts[0].presentation!.text).toContain(
|
||||
'[File Reference (application/pdf)',
|
||||
);
|
||||
expect(parts[0].presentation!.text).toContain(
|
||||
'Original URI: gs://fake-bucket/doc.pdf',
|
||||
);
|
||||
|
||||
expect(fsPromises.writeFile).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,123 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type { Episode } from '../ir/types.js';
|
||||
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { estimateContextTokenCountSync as estimateTokenCountSync } from '../utils/contextTokenCalculator.js';
|
||||
import { sanitizeFilenamePart } from '../../utils/fileUtils.js';
|
||||
import * as fsPromises from 'node:fs/promises';
|
||||
import path from 'node:path';
|
||||
import type { Part } from '@google/genai';
|
||||
|
||||
export class BlobDegradationProcessor implements ContextProcessor {
|
||||
readonly name = 'BlobDegradation';
|
||||
private env: ContextEnvironment;
|
||||
|
||||
constructor(env: ContextEnvironment, options: Record<string, unknown> = {}) {
|
||||
this.env = env;
|
||||
}
|
||||
|
||||
async process(
|
||||
episodes: Episode[],
|
||||
state: ContextAccountingState,
|
||||
): Promise<Episode[]> {
|
||||
if (state.isBudgetSatisfied) {
|
||||
return episodes;
|
||||
}
|
||||
|
||||
let currentDeficit = state.deficitTokens;
|
||||
const newEpisodes = [...episodes];
|
||||
let directoryCreated = false;
|
||||
|
||||
let blobOutputsDir = path.join(
|
||||
this.env.getProjectTempDir(),
|
||||
'degraded-blobs',
|
||||
);
|
||||
const sessionId = this.env.getSessionId();
|
||||
if (sessionId) {
|
||||
blobOutputsDir = path.join(
|
||||
blobOutputsDir,
|
||||
`session-${sanitizeFilenamePart(sessionId)}`,
|
||||
);
|
||||
}
|
||||
|
||||
const ensureDir = async () => {
|
||||
if (!directoryCreated) {
|
||||
await fsPromises.mkdir(blobOutputsDir, { recursive: true });
|
||||
directoryCreated = true;
|
||||
}
|
||||
};
|
||||
|
||||
// Forward scan, looking for bloated non-text parts to degrade
|
||||
for (let i = 0; i < newEpisodes.length; i++) {
|
||||
if (currentDeficit <= 0) break;
|
||||
const ep = newEpisodes[i];
|
||||
if (state.protectedEpisodeIds.has(ep.id)) continue;
|
||||
|
||||
if (ep.trigger.type === 'USER_PROMPT') {
|
||||
for (const part of ep.trigger.semanticParts) {
|
||||
if (currentDeficit <= 0) break;
|
||||
// We only target non-text parts that haven't already been masked
|
||||
if (part.type === 'text' || part.presentation) continue;
|
||||
|
||||
let newText = '';
|
||||
let tokensSaved = 0;
|
||||
|
||||
if (part.type === 'inline_data') {
|
||||
await ensureDir();
|
||||
const ext = part.mimeType.split('/')[1] || 'bin';
|
||||
const fileName = `blob_${Date.now()}_${Math.random().toString(36).substring(7)}.${ext}`;
|
||||
const filePath = path.join(blobOutputsDir, fileName);
|
||||
|
||||
// Base64 to buffer
|
||||
const buffer = Buffer.from(part.data, 'base64');
|
||||
await fsPromises.writeFile(filePath, buffer);
|
||||
|
||||
const mb = (buffer.byteLength / 1024 / 1024).toFixed(2);
|
||||
newText = `[Multi-Modal Blob (${part.mimeType}, ${mb}MB) degraded to text to preserve context window. Saved to: ${filePath}]`;
|
||||
|
||||
// Re-calculate tokens. Images are expensive (~258 tokens). The text is cheap (~20 tokens).
|
||||
const oldTokens = estimateTokenCountSync([
|
||||
{ inlineData: { mimeType: part.mimeType, data: part.data } },
|
||||
]);
|
||||
const newTokens = estimateTokenCountSync([{ text: newText }]);
|
||||
tokensSaved = oldTokens - newTokens;
|
||||
} else if (part.type === 'file_data') {
|
||||
newText = `[File Reference (${part.mimeType}) degraded to text to preserve context window. Original URI: ${part.fileUri}]`;
|
||||
const oldTokens = estimateTokenCountSync([
|
||||
{ fileData: { mimeType: part.mimeType, fileUri: part.fileUri } },
|
||||
]);
|
||||
const newTokens = estimateTokenCountSync([{ text: newText }]);
|
||||
tokensSaved = oldTokens - newTokens;
|
||||
} else if (part.type === 'raw_part') {
|
||||
newText = `[Unknown Part degraded to text to preserve context window.]`;
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
const oldTokens = estimateTokenCountSync([part.part as Part]);
|
||||
const newTokens = estimateTokenCountSync([{ text: newText }]);
|
||||
tokensSaved = oldTokens - newTokens;
|
||||
}
|
||||
|
||||
if (newText && tokensSaved > 0) {
|
||||
const newTokens = estimateTokenCountSync([{ text: newText }], 0, {
|
||||
charsPerToken: this.env.getCharsPerToken(),
|
||||
});
|
||||
part.presentation = { text: newText, tokens: newTokens };
|
||||
|
||||
ep.trigger.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'DEGRADED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
|
||||
currentDeficit -= tokensSaved;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return newEpisodes;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import { createMockEnvironment } from '../testing/contextTestUtils.js';
|
||||
import { describe, it, expect, beforeEach } from 'vitest';
|
||||
import { HistorySquashingProcessor } from './historySquashingProcessor.js';
|
||||
import type {
|
||||
Episode,
|
||||
UserPrompt,
|
||||
AgentThought,
|
||||
AgentYield,
|
||||
} from '../ir/types.js';
|
||||
import type { ContextAccountingState } from '../pipeline.js';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
|
||||
describe('HistorySquashingProcessor', () => {
|
||||
let processor: HistorySquashingProcessor;
|
||||
|
||||
beforeEach(() => {
|
||||
processor = new HistorySquashingProcessor(createMockEnvironment(), {
|
||||
maxTokensPerNode: 100,
|
||||
});
|
||||
});
|
||||
|
||||
const getDummyState = (
|
||||
isSatisfied = false,
|
||||
deficit = 0,
|
||||
protectedIds = new Set<string>(),
|
||||
): ContextAccountingState => ({
|
||||
currentTokens: 5000,
|
||||
maxTokens: 10000,
|
||||
retainedTokens: 4000,
|
||||
deficitTokens: deficit,
|
||||
protectedEpisodeIds: protectedIds,
|
||||
isBudgetSatisfied: isSatisfied,
|
||||
});
|
||||
|
||||
const createDummyEpisode = (
|
||||
id: string,
|
||||
userText: string,
|
||||
modelThought: string,
|
||||
): Episode => ({
|
||||
id,
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: randomUUID(),
|
||||
type: 'USER_PROMPT',
|
||||
semanticParts: [{ type: 'text', text: userText }],
|
||||
metadata: {
|
||||
originalTokens: 1000,
|
||||
currentTokens: 1000,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
steps: [
|
||||
{
|
||||
id: randomUUID(),
|
||||
type: 'AGENT_THOUGHT',
|
||||
text: modelThought,
|
||||
metadata: {
|
||||
originalTokens: 1000,
|
||||
currentTokens: 1000,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
it('bypasses processing if budget is satisfied', async () => {
|
||||
const episodes = [createDummyEpisode('1', 'short text', 'short thought')];
|
||||
const state = getDummyState(true);
|
||||
|
||||
const result = await processor.process(episodes, state);
|
||||
|
||||
expect(result).toStrictEqual(episodes);
|
||||
expect(
|
||||
(result[0].trigger as UserPrompt).semanticParts[0].presentation,
|
||||
).toBeUndefined();
|
||||
});
|
||||
|
||||
it('skips protected episodes', async () => {
|
||||
// 500 chars = ~125 tokens. Limit is 100 tokens, so it WOULD truncate if not protected.
|
||||
const longText = 'A'.repeat(500);
|
||||
const episodes = [createDummyEpisode('ep-1', longText, 'short thought')];
|
||||
const state = getDummyState(false, 100, new Set(['ep-1']));
|
||||
|
||||
const result = await processor.process(episodes, state);
|
||||
|
||||
expect(
|
||||
(result[0].trigger as UserPrompt).semanticParts[0].presentation,
|
||||
).toBeUndefined();
|
||||
});
|
||||
|
||||
it('truncates both UserPrompts and AgentThoughts', async () => {
|
||||
const longUser = 'U'.repeat(1000); // ~250 tokens
|
||||
const longModel = 'M'.repeat(1000); // ~250 tokens
|
||||
const episodes = [createDummyEpisode('ep-2', longUser, longModel)];
|
||||
const state = getDummyState(false, 500, new Set()); // High deficit, force truncation
|
||||
|
||||
const result = await processor.process(episodes, state);
|
||||
|
||||
const userPart = (result[0].trigger as UserPrompt).semanticParts[0];
|
||||
const thoughtPart = result[0].steps[0] as AgentThought;
|
||||
|
||||
expect(userPart.presentation).toBeDefined();
|
||||
expect(userPart.presentation!.text).toContain(
|
||||
'[... OMITTED 600 chars ...]',
|
||||
);
|
||||
|
||||
expect(thoughtPart.presentation).toBeDefined();
|
||||
expect(thoughtPart.presentation!.text).toContain(
|
||||
'[... OMITTED 600 chars ...]',
|
||||
);
|
||||
|
||||
// Check audit trails
|
||||
expect(result[0].trigger.metadata.transformations.length).toBe(1);
|
||||
expect(thoughtPart.metadata.transformations.length).toBe(1);
|
||||
});
|
||||
|
||||
it('stops processing once deficit is resolved', async () => {
|
||||
const longUser1 = 'A'.repeat(1000);
|
||||
const longUser2 = 'B'.repeat(1000);
|
||||
const episodes = [
|
||||
createDummyEpisode('ep-3', longUser1, 'short'),
|
||||
createDummyEpisode('ep-4', longUser2, 'short'),
|
||||
];
|
||||
|
||||
// Set deficit to exactly what ONE truncation will save
|
||||
// Original = ~250 tokens. Limit = 100. Truncation saves ~150 tokens.
|
||||
const state = getDummyState(false, 150, new Set());
|
||||
|
||||
const result = await processor.process(episodes, state);
|
||||
|
||||
// First episode should be truncated
|
||||
const ep1Part = (result[0].trigger as UserPrompt).semanticParts[0];
|
||||
expect(ep1Part.presentation).toBeDefined();
|
||||
|
||||
// Second episode should be untouched because the deficit hit 0
|
||||
const ep2Part = (result[1].trigger as UserPrompt).semanticParts[0];
|
||||
expect(ep2Part.presentation).toBeUndefined();
|
||||
});
|
||||
|
||||
it('truncates IrNodes', async () => {
|
||||
const longYield = 'Y'.repeat(1000); // ~250 tokens
|
||||
const ep = createDummyEpisode('ep-5', 'short', 'short');
|
||||
ep.yield = {
|
||||
id: randomUUID(),
|
||||
type: 'AGENT_YIELD',
|
||||
text: longYield,
|
||||
metadata: {
|
||||
originalTokens: 250,
|
||||
currentTokens: 250,
|
||||
transformations: [],
|
||||
},
|
||||
};
|
||||
|
||||
const state = getDummyState(false, 500, new Set());
|
||||
const result = await processor.process([ep], state);
|
||||
|
||||
const yieldPart = result[0].yield as AgentYield;
|
||||
const yieldPresentation = yieldPart.presentation as { text: string };
|
||||
expect(yieldPresentation).toBeDefined();
|
||||
expect(yieldPresentation.text).toContain('[... OMITTED 600 chars ...]');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,132 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Episode } from '../ir/types.js';
|
||||
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { truncateProportionally } from '../truncation.js';
|
||||
|
||||
export class HistorySquashingProcessor implements ContextProcessor {
|
||||
readonly name = 'HistorySquashing';
|
||||
private options: { maxTokensPerNode: number };
|
||||
|
||||
constructor(env: ContextEnvironment, options: { maxTokensPerNode: number }) {
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
private tryApplySquash(
|
||||
text: string,
|
||||
limitChars: number,
|
||||
currentDeficit: number,
|
||||
setPresentation: (p: { text: string; tokens: number }) => void,
|
||||
recordAudit: () => void,
|
||||
): number {
|
||||
if (currentDeficit <= 0) return 0;
|
||||
const originalLength = text.length;
|
||||
if (originalLength <= limitChars) return 0;
|
||||
|
||||
const newText = truncateProportionally(
|
||||
text,
|
||||
limitChars,
|
||||
`\n\n[... OMITTED ${originalLength - limitChars} chars ...]\n\n`,
|
||||
);
|
||||
|
||||
if (newText !== text) {
|
||||
const newTokens = Math.floor(newText.length / 4);
|
||||
const oldTokens = Math.floor(originalLength / 4);
|
||||
const tokensSaved = oldTokens - newTokens;
|
||||
|
||||
setPresentation({ text: newText, tokens: newTokens });
|
||||
recordAudit();
|
||||
return tokensSaved;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
async process(
|
||||
episodes: Episode[],
|
||||
state: ContextAccountingState,
|
||||
): Promise<Episode[]> {
|
||||
if (state.isBudgetSatisfied) {
|
||||
return episodes;
|
||||
}
|
||||
|
||||
const { maxTokensPerNode } = this.options;
|
||||
// We estimate 4 chars per token for truncation logic
|
||||
const limitChars = maxTokensPerNode * 4;
|
||||
|
||||
// We track how many tokens we still need to cut. If we hit 0, we can stop early!
|
||||
let currentDeficit = state.deficitTokens;
|
||||
const newEpisodes = [...episodes];
|
||||
|
||||
for (let i = 0; i < newEpisodes.length; i++) {
|
||||
if (currentDeficit <= 0) break;
|
||||
if (state.protectedEpisodeIds.has(newEpisodes[i].id)) continue;
|
||||
|
||||
const ep = newEpisodes[i];
|
||||
|
||||
// 1. Squash User Prompts
|
||||
if (ep.trigger.type === 'USER_PROMPT') {
|
||||
for (const part of ep.trigger.semanticParts) {
|
||||
if (part.type === 'text') {
|
||||
const saved = this.tryApplySquash(
|
||||
part.text,
|
||||
limitChars,
|
||||
currentDeficit,
|
||||
(p) => (part.presentation = p),
|
||||
() =>
|
||||
ep.trigger.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'TRUNCATED',
|
||||
timestamp: Date.now(),
|
||||
}),
|
||||
);
|
||||
currentDeficit -= saved;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Squash Model Thoughts
|
||||
for (const step of ep.steps) {
|
||||
if (currentDeficit <= 0) break;
|
||||
if (step.type === 'AGENT_THOUGHT') {
|
||||
const saved = this.tryApplySquash(
|
||||
step.text,
|
||||
limitChars,
|
||||
currentDeficit,
|
||||
(p) => (step.presentation = p),
|
||||
() =>
|
||||
step.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'TRUNCATED',
|
||||
timestamp: Date.now(),
|
||||
}),
|
||||
);
|
||||
currentDeficit -= saved;
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Squash Agent Yields
|
||||
if (currentDeficit > 0 && ep.yield) {
|
||||
const saved = this.tryApplySquash(
|
||||
ep.yield.text,
|
||||
limitChars,
|
||||
currentDeficit,
|
||||
(p) => (ep.yield!.presentation = p),
|
||||
() =>
|
||||
ep.yield!.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'TRUNCATED',
|
||||
timestamp: Date.now(),
|
||||
}),
|
||||
);
|
||||
currentDeficit -= saved;
|
||||
}
|
||||
}
|
||||
|
||||
return newEpisodes;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { createMockEnvironment } from '../testing/contextTestUtils.js';
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { SemanticCompressionProcessor } from './semanticCompressionProcessor.js';
|
||||
import type {
|
||||
Episode,
|
||||
UserPrompt,
|
||||
ToolExecution,
|
||||
AgentThought,
|
||||
} from '../ir/types.js';
|
||||
import type { ContextAccountingState } from '../pipeline.js';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
|
||||
describe('SemanticCompressionProcessor', () => {
|
||||
let processor: SemanticCompressionProcessor;
|
||||
let generateContentMock: ReturnType<typeof vi.fn>;
|
||||
|
||||
beforeEach(() => {
|
||||
generateContentMock = vi.fn().mockResolvedValue({
|
||||
candidates: [{ content: { parts: [{ text: 'Mocked Summary!' }] } }],
|
||||
});
|
||||
|
||||
const env = createMockEnvironment();
|
||||
env.getLlmClient = vi
|
||||
.fn()
|
||||
.mockReturnValue({ generateContent: generateContentMock }) as any;
|
||||
processor = new SemanticCompressionProcessor(env, {
|
||||
nodeThresholdTokens: 2000,
|
||||
});
|
||||
});
|
||||
|
||||
const getDummyState = (
|
||||
isSatisfied = false,
|
||||
deficit = 0,
|
||||
protectedIds = new Set<string>(),
|
||||
): ContextAccountingState => ({
|
||||
currentTokens: 5000,
|
||||
maxTokens: 10000,
|
||||
retainedTokens: 4000,
|
||||
deficitTokens: deficit,
|
||||
protectedEpisodeIds: protectedIds,
|
||||
isBudgetSatisfied: isSatisfied,
|
||||
});
|
||||
|
||||
const createDummyEpisode = (
|
||||
id: string,
|
||||
userText: string,
|
||||
thoughtText: string,
|
||||
toolObs: string,
|
||||
): Episode => ({
|
||||
id,
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: randomUUID(),
|
||||
type: 'USER_PROMPT',
|
||||
semanticParts: [{ type: 'text', text: userText }],
|
||||
metadata: {
|
||||
originalTokens: 3800,
|
||||
currentTokens: 3800,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
steps: [
|
||||
{
|
||||
id: randomUUID(),
|
||||
type: 'AGENT_THOUGHT',
|
||||
text: thoughtText,
|
||||
metadata: {
|
||||
originalTokens: 100,
|
||||
currentTokens: 100,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
{
|
||||
id: randomUUID(),
|
||||
type: 'TOOL_EXECUTION',
|
||||
toolName: 'test',
|
||||
intent: {},
|
||||
observation: toolObs,
|
||||
tokens: { intent: 10, observation: 3800 },
|
||||
metadata: {
|
||||
originalTokens: 3810,
|
||||
currentTokens: 3810,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
it('bypasses processing if budget is satisfied', async () => {
|
||||
const episodes = [createDummyEpisode('1', 'short', 'short', 'short')];
|
||||
const state = getDummyState(true);
|
||||
|
||||
await processor.process(episodes, state);
|
||||
expect(generateContentMock).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('skips protected episodes even if over budget', async () => {
|
||||
const massiveStr = 'M'.repeat(15000); // Exceeds threshold (10 * 4 = 40)
|
||||
const episodes = [
|
||||
createDummyEpisode('ep-1', massiveStr, massiveStr, massiveStr),
|
||||
];
|
||||
const state = getDummyState(false, 1000, new Set(['ep-1']));
|
||||
|
||||
await processor.process(episodes, state);
|
||||
expect(generateContentMock).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('summarizes unprotected UserPrompts, Thoughts, and Tool observations until deficit is met', async () => {
|
||||
const massiveStr = 'M'.repeat(15000);
|
||||
const episodes = [
|
||||
createDummyEpisode('ep-1', massiveStr, massiveStr, massiveStr),
|
||||
];
|
||||
const state = getDummyState(false, 50000, new Set()); // Massive deficit, forces all 3 to summarize
|
||||
|
||||
const result = await processor.process(episodes, state);
|
||||
expect(generateContentMock).toHaveBeenCalledTimes(3);
|
||||
|
||||
// Verify presentation layers were injected
|
||||
const userPart = (result[0].trigger as UserPrompt).semanticParts[0];
|
||||
const thoughtPart = result[0].steps[0] as AgentThought;
|
||||
const toolPart = result[0].steps[1] as ToolExecution;
|
||||
|
||||
expect(userPart.presentation).toBeDefined();
|
||||
expect(userPart.presentation!.text).toContain('Mocked Summary!');
|
||||
|
||||
expect(thoughtPart.presentation).toBeDefined();
|
||||
expect(thoughtPart.presentation!.text).toContain('Mocked Summary!');
|
||||
|
||||
expect(toolPart.presentation).toBeDefined();
|
||||
expect(
|
||||
(toolPart.presentation!.observation as Record<string, string>)['summary'],
|
||||
).toContain('Mocked Summary!');
|
||||
});
|
||||
|
||||
it('stops calling LLM when deficit hits zero', async () => {
|
||||
const massiveStr = 'M'.repeat(15000);
|
||||
const episodes = [
|
||||
createDummyEpisode('ep-1', massiveStr, massiveStr, massiveStr),
|
||||
];
|
||||
|
||||
// Set deficit low enough that ONE summary solves the problem
|
||||
const state = getDummyState(false, 5, new Set());
|
||||
|
||||
await processor.process(episodes, state);
|
||||
// It should only compress the UserPrompt and then stop
|
||||
expect(generateContentMock).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,202 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Episode } from '../ir/types.js';
|
||||
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { LlmRole } from '../../telemetry/types.js';
|
||||
import { getResponseText } from '../../utils/partUtils.js';
|
||||
import { estimateTokenCountSync } from '../../utils/tokenCalculation.js';
|
||||
|
||||
export class SemanticCompressionProcessor implements ContextProcessor {
|
||||
readonly name = 'SemanticCompression';
|
||||
private env: ContextEnvironment;
|
||||
private options: { nodeThresholdTokens: number };
|
||||
private modelToUse: string = 'chat-compression-2.5-flash-lite';
|
||||
|
||||
constructor(
|
||||
env: ContextEnvironment,
|
||||
options: { nodeThresholdTokens: number },
|
||||
) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
async process(
|
||||
episodes: Episode[],
|
||||
state: ContextAccountingState,
|
||||
): Promise<Episode[]> {
|
||||
require('fs').appendFileSync(
|
||||
'/tmp/debug2.json',
|
||||
'SEMANTIC PROCESS: First episode ID: ' +
|
||||
episodes[0]?.id +
|
||||
'\nProtected IDs: ' +
|
||||
Array.from(state.protectedEpisodeIds).join(', ') +
|
||||
'\n',
|
||||
);
|
||||
// If the budget is satisfied, or semantic compression isn't enabled
|
||||
if (state.isBudgetSatisfied) {
|
||||
return episodes;
|
||||
}
|
||||
|
||||
const semanticConfig = this.options;
|
||||
// We estimate 4 chars per token for truncation logic
|
||||
const thresholdChars = semanticConfig.nodeThresholdTokens * 4;
|
||||
this.modelToUse = 'gemini-2.5-flash';
|
||||
|
||||
let currentDeficit = state.deficitTokens;
|
||||
const newEpisodes = [...episodes];
|
||||
|
||||
// We scan backwards (oldest to newest would also work, but older is safer to degrade first)
|
||||
for (let i = 0; i < newEpisodes.length; i++) {
|
||||
if (currentDeficit <= 0) break;
|
||||
const ep = newEpisodes[i];
|
||||
if (state.protectedEpisodeIds.has(ep.id)) continue;
|
||||
|
||||
// 1. Compress User Prompts
|
||||
if (ep.trigger.type === 'USER_PROMPT') {
|
||||
for (const part of ep.trigger.semanticParts) {
|
||||
if (currentDeficit <= 0) break;
|
||||
if (part.type !== 'text') continue;
|
||||
// If it's already got a presentation, we don't want to re-summarize a summary
|
||||
if (part.presentation) continue;
|
||||
|
||||
if (part.text.length > thresholdChars) {
|
||||
const summary = await this.generateSummary(
|
||||
part.text,
|
||||
'User Prompt',
|
||||
);
|
||||
const newTokens = estimateTokenCountSync([{ text: summary }]);
|
||||
const oldTokens = estimateTokenCountSync([{ text: part.text }]);
|
||||
|
||||
if (newTokens < oldTokens) {
|
||||
part.presentation = { text: summary, tokens: newTokens };
|
||||
ep.trigger.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'SUMMARIZED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
currentDeficit -= oldTokens - newTokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Compress Model Thoughts
|
||||
for (const step of ep.steps) {
|
||||
if (currentDeficit <= 0) break;
|
||||
if (step.type === 'AGENT_THOUGHT') {
|
||||
if (step.presentation) continue;
|
||||
if (step.text.length > thresholdChars) {
|
||||
const summary = await this.generateSummary(
|
||||
step.text,
|
||||
'Agent Thought',
|
||||
);
|
||||
const newTokens = estimateTokenCountSync([{ text: summary }]);
|
||||
const oldTokens = estimateTokenCountSync([{ text: step.text }]);
|
||||
|
||||
if (newTokens < oldTokens) {
|
||||
step.presentation = { text: summary, tokens: newTokens };
|
||||
step.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'SUMMARIZED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
currentDeficit -= oldTokens - newTokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Compress Tool Observations
|
||||
if (step.type === 'TOOL_EXECUTION') {
|
||||
const rawObs = step.presentation?.observation ?? step.observation;
|
||||
|
||||
let stringifiedObs = '';
|
||||
if (typeof rawObs === 'string') {
|
||||
stringifiedObs = rawObs;
|
||||
} else {
|
||||
try {
|
||||
stringifiedObs = JSON.stringify(rawObs);
|
||||
} catch (_e) {
|
||||
stringifiedObs = String(rawObs);
|
||||
}
|
||||
}
|
||||
|
||||
if (
|
||||
stringifiedObs.length > thresholdChars &&
|
||||
!stringifiedObs.includes('<tool_output_masked>')
|
||||
) {
|
||||
const summary = await this.generateSummary(
|
||||
stringifiedObs,
|
||||
`Tool Output (${step.toolName})`,
|
||||
);
|
||||
|
||||
// Wrap the summary in an object so the Gemini API accepts it as a valid functionResponse.response
|
||||
const newObsObject = { summary };
|
||||
|
||||
const newObsTokens = estimateTokenCountSync([
|
||||
{
|
||||
functionResponse: {
|
||||
name: step.toolName,
|
||||
response: newObsObject as unknown as Record<string, unknown>, // eslint-disable-line @typescript-eslint/no-unsafe-type-assertion
|
||||
id: step.id,
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
const oldObsTokens =
|
||||
step.presentation?.tokens.observation ?? step.tokens.observation;
|
||||
const intentTokens =
|
||||
step.presentation?.tokens.intent ?? step.tokens.intent;
|
||||
|
||||
if (newObsTokens < oldObsTokens) {
|
||||
step.presentation = {
|
||||
intent: step.presentation?.intent ?? step.intent,
|
||||
observation: newObsObject,
|
||||
tokens: { intent: intentTokens, observation: newObsTokens },
|
||||
};
|
||||
step.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'SUMMARIZED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
currentDeficit -= oldObsTokens - newObsTokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return newEpisodes;
|
||||
}
|
||||
|
||||
private async generateSummary(
|
||||
content: string,
|
||||
contentType: string,
|
||||
abortSignal?: AbortSignal,
|
||||
): Promise<string> {
|
||||
const promptMessage = `You are compressing an old episodic context buffer for an AI assistant.\nSummarize this ${contentType} block in 2-3 highly technical sentences. Keep all critical facts, file names, dependencies, and architectural decisions. Discard conversational filler and boilerplate.\n\nContent:\n${content.slice(0, 30000)}`;
|
||||
|
||||
const client = this.env.getLlmClient();
|
||||
try {
|
||||
const response = await client.generateContent({
|
||||
modelConfigKey: { model: this.modelToUse },
|
||||
contents: [{ role: 'user', parts: [{ text: promptMessage }] }],
|
||||
promptId: 'local-context-compression-summary',
|
||||
role: LlmRole.UTILITY_COMPRESSOR,
|
||||
abortSignal: abortSignal ?? new AbortController().signal,
|
||||
});
|
||||
const text = getResponseText(response) ?? '';
|
||||
return `[Semantic Summary of old ${contentType}]\n${text.trim()}`;
|
||||
} catch (_e) {
|
||||
debugLogger.warn('Semantic compression LLM call failed: ' + String(_e));
|
||||
// If we fail to summarize, we just return the original truncated by 50% as a fail-safe, or the original.
|
||||
// Returning original is safer to prevent data loss on API failure.
|
||||
return content;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,130 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { createMockEnvironment } from '../testing/contextTestUtils.js';
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { ToolMaskingProcessor } from './toolMaskingProcessor.js';
|
||||
import type { Episode, ToolExecution } from '../ir/types.js';
|
||||
import type { ContextAccountingState } from '../pipeline.js';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import * as fsPromises from 'node:fs/promises';
|
||||
|
||||
vi.mock('node:fs/promises');
|
||||
|
||||
describe('ToolMaskingProcessor', () => {
|
||||
let processor: ToolMaskingProcessor;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
|
||||
processor = new ToolMaskingProcessor(createMockEnvironment(), {
|
||||
stringLengthThresholdTokens: 100,
|
||||
});
|
||||
});
|
||||
|
||||
const getDummyState = (
|
||||
isSatisfied = false,
|
||||
deficit = 0,
|
||||
protectedIds = new Set<string>(),
|
||||
): ContextAccountingState => ({
|
||||
currentTokens: 5000,
|
||||
maxTokens: 10000,
|
||||
retainedTokens: 4000,
|
||||
deficitTokens: deficit,
|
||||
protectedEpisodeIds: protectedIds,
|
||||
isBudgetSatisfied: isSatisfied,
|
||||
});
|
||||
|
||||
const createDummyEpisode = (
|
||||
id: string,
|
||||
intent: Record<string, unknown>,
|
||||
observation: Record<string, unknown>,
|
||||
): Episode => ({
|
||||
id,
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: randomUUID(),
|
||||
type: 'SYSTEM_EVENT',
|
||||
name: 'test',
|
||||
payload: {},
|
||||
metadata: { originalTokens: 10, currentTokens: 10, transformations: [] },
|
||||
},
|
||||
steps: [
|
||||
{
|
||||
id: randomUUID(),
|
||||
type: 'TOOL_EXECUTION',
|
||||
toolName: 'test_tool',
|
||||
intent,
|
||||
observation,
|
||||
tokens: { intent: 500, observation: 500 }, // Claim they are big enough to be masked
|
||||
metadata: {
|
||||
originalTokens: 1000,
|
||||
currentTokens: 1000,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
it('bypasses processing if budget is satisfied', async () => {
|
||||
const episodes = [
|
||||
createDummyEpisode('1', { arg: 'short' }, { out: 'short' }),
|
||||
];
|
||||
const state = getDummyState(true);
|
||||
|
||||
const result = await processor.process(episodes, state);
|
||||
require('fs').appendFileSync(
|
||||
'/tmp/debug.json',
|
||||
'\n\n' + JSON.stringify({ res: result[0].steps[0] }, null, 2),
|
||||
);
|
||||
|
||||
expect(result).toStrictEqual(episodes);
|
||||
expect((result[0].steps[0] as ToolExecution).presentation).toBeUndefined();
|
||||
});
|
||||
|
||||
it('deep masks massive string intents and observations', async () => {
|
||||
// We need strings > limitChars (100 tokens * 4 chars = 400 chars)
|
||||
const massiveIntentString = 'I'.repeat(500);
|
||||
const massiveObsString = 'O'.repeat(500);
|
||||
|
||||
const intentPayload = { args: { nested: [massiveIntentString, 'short'] } };
|
||||
const obsPayload = { result: massiveObsString, error: null };
|
||||
|
||||
const episodes = [createDummyEpisode('ep-1', intentPayload, obsPayload)];
|
||||
const state = getDummyState(false, 1000, new Set()); // Huge deficit
|
||||
|
||||
const result = await processor.process(episodes, state);
|
||||
|
||||
const toolStep = result[0].steps[0] as ToolExecution;
|
||||
|
||||
expect(toolStep.presentation).toBeDefined();
|
||||
|
||||
// Check intent was deep masked
|
||||
const maskedIntent = toolStep.presentation!.intent as Record<
|
||||
string,
|
||||
unknown
|
||||
>;
|
||||
expect((maskedIntent['args'] as { nested: string }).nested[0]).toContain(
|
||||
'<tool_output_masked>',
|
||||
);
|
||||
expect((maskedIntent['args'] as { nested: string }).nested[1]).toBe(
|
||||
'short',
|
||||
); // Unchanged
|
||||
|
||||
// Check observation was deep masked
|
||||
const maskedObs = toolStep.presentation!.observation as Record<
|
||||
string,
|
||||
unknown
|
||||
>;
|
||||
expect((maskedObs as { result: string }).result).toContain(
|
||||
'<tool_output_masked>',
|
||||
);
|
||||
expect((maskedObs as { error: string }).error).toBeNull();
|
||||
|
||||
// Check disk writes occurred
|
||||
expect(fsPromises.writeFile).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,243 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { estimateTokenCountSync } from '../../utils/tokenCalculation.js';
|
||||
import { sanitizeFilenamePart } from '../../utils/fileUtils.js';
|
||||
import * as fsPromises from 'node:fs/promises';
|
||||
import path from 'node:path';
|
||||
import {
|
||||
ACTIVATE_SKILL_TOOL_NAME,
|
||||
MEMORY_TOOL_NAME,
|
||||
ASK_USER_TOOL_NAME,
|
||||
ENTER_PLAN_MODE_TOOL_NAME,
|
||||
EXIT_PLAN_MODE_TOOL_NAME,
|
||||
} from '../../tools/tool-names.js';
|
||||
import type { Episode } from '../ir/types.js';
|
||||
|
||||
const UNMASKABLE_TOOLS = new Set([
|
||||
ACTIVATE_SKILL_TOOL_NAME,
|
||||
MEMORY_TOOL_NAME,
|
||||
ASK_USER_TOOL_NAME,
|
||||
ENTER_PLAN_MODE_TOOL_NAME,
|
||||
EXIT_PLAN_MODE_TOOL_NAME,
|
||||
]);
|
||||
|
||||
export class ToolMaskingProcessor implements ContextProcessor {
|
||||
readonly name = 'ToolMasking';
|
||||
private env: ContextEnvironment;
|
||||
private options: { stringLengthThresholdTokens: number };
|
||||
|
||||
constructor(
|
||||
env: ContextEnvironment,
|
||||
options: { stringLengthThresholdTokens: number },
|
||||
) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
async process(
|
||||
episodes: Episode[],
|
||||
state: ContextAccountingState,
|
||||
): Promise<Episode[]> {
|
||||
const maskingConfig = this.options;
|
||||
if (!maskingConfig) return episodes;
|
||||
if (state.isBudgetSatisfied) return episodes;
|
||||
|
||||
const newEpisodes = [...episodes];
|
||||
let currentDeficit = state.deficitTokens;
|
||||
const limitChars = maskingConfig.stringLengthThresholdTokens * 4;
|
||||
|
||||
let toolOutputsDir = path.join(
|
||||
this.env.getProjectTempDir(),
|
||||
'tool-outputs',
|
||||
);
|
||||
const sessionId = this.env.getSessionId();
|
||||
if (sessionId) {
|
||||
toolOutputsDir = path.join(
|
||||
toolOutputsDir,
|
||||
`session-${sanitizeFilenamePart(sessionId)}`,
|
||||
);
|
||||
}
|
||||
|
||||
// We only create the directory if we actually mask something
|
||||
let directoryCreated = false;
|
||||
|
||||
// Helper to extract string and write to disk
|
||||
const handleMasking = async (
|
||||
content: string,
|
||||
toolName: string,
|
||||
callId: string,
|
||||
nodeType: string,
|
||||
): Promise<string> => {
|
||||
if (!directoryCreated) {
|
||||
await fsPromises.mkdir(toolOutputsDir, { recursive: true });
|
||||
directoryCreated = true;
|
||||
}
|
||||
|
||||
const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${Math.random().toString(36).substring(7)}.txt`;
|
||||
const filePath = path.join(toolOutputsDir, fileName);
|
||||
|
||||
await fsPromises.writeFile(filePath, content, 'utf-8');
|
||||
|
||||
const fileSizeMB = (
|
||||
Buffer.byteLength(content, 'utf8') /
|
||||
1024 /
|
||||
1024
|
||||
).toFixed(2);
|
||||
const totalLines = content.split('\n').length;
|
||||
return `<tool_output_masked>\n[Tool ${nodeType} string (${fileSizeMB}MB, ${totalLines} lines) masked to preserve context window. Full string saved to: ${filePath}]\n</tool_output_masked>`;
|
||||
};
|
||||
|
||||
// Forward scan, looking for massive intents or observations to mask
|
||||
for (let i = 0; i < newEpisodes.length; i++) {
|
||||
if (currentDeficit <= 0) break;
|
||||
const ep = newEpisodes[i];
|
||||
if (!ep || !ep.steps || state.protectedEpisodeIds.has(ep.id)) continue;
|
||||
|
||||
for (let j = 0; j < ep.steps.length; j++) {
|
||||
if (currentDeficit <= 0) break;
|
||||
const step = ep.steps[j];
|
||||
if (step.type !== 'TOOL_EXECUTION') continue;
|
||||
|
||||
const toolName = step.toolName;
|
||||
if (toolName && UNMASKABLE_TOOLS.has(toolName)) continue;
|
||||
|
||||
// Ensure presentation object exists
|
||||
if (!step.presentation) {
|
||||
step.presentation = {
|
||||
intent: step.intent,
|
||||
observation: step.observation,
|
||||
tokens: step.tokens, // Fallback to raw tokens initially
|
||||
};
|
||||
}
|
||||
|
||||
const callId = step.id || Date.now().toString();
|
||||
|
||||
/* eslint-disable @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-assignment */
|
||||
|
||||
const maskAsync = async (
|
||||
obj: any,
|
||||
nodeType: string,
|
||||
): Promise<{ masked: any; changed: boolean }> => {
|
||||
if (typeof obj === 'string') {
|
||||
require('fs').appendFileSync(
|
||||
'/tmp/debug.json',
|
||||
'STRING FOUND. length: ' +
|
||||
obj.length +
|
||||
' limitChars: ' +
|
||||
limitChars +
|
||||
'\n',
|
||||
);
|
||||
if (obj.length > 1000)
|
||||
console.log(
|
||||
'Found string of length:',
|
||||
obj.length,
|
||||
'limitChars is:',
|
||||
limitChars,
|
||||
'isAlreadyMasked:',
|
||||
this.isAlreadyMasked(obj),
|
||||
);
|
||||
if (obj.length > limitChars && !this.isAlreadyMasked(obj)) {
|
||||
const newString = await handleMasking(
|
||||
obj,
|
||||
toolName,
|
||||
callId,
|
||||
nodeType,
|
||||
);
|
||||
return { masked: newString, changed: true };
|
||||
}
|
||||
return { masked: obj, changed: false };
|
||||
}
|
||||
if (Array.isArray(obj)) {
|
||||
let changed = false;
|
||||
const masked = [];
|
||||
for (const item of obj) {
|
||||
const res = await maskAsync(item, nodeType);
|
||||
if (res.changed) changed = true;
|
||||
masked.push(res.masked);
|
||||
}
|
||||
return { masked, changed };
|
||||
}
|
||||
if (typeof obj === 'object' && obj !== null) {
|
||||
let changed = false;
|
||||
const masked: Record<string, any> = {};
|
||||
for (const [key, value] of Object.entries(obj)) {
|
||||
const res = await maskAsync(value, nodeType);
|
||||
if (res.changed) changed = true;
|
||||
masked[key] = res.masked;
|
||||
}
|
||||
return { masked, changed };
|
||||
}
|
||||
return { masked: obj, changed: false };
|
||||
};
|
||||
|
||||
const intentRes = await maskAsync(
|
||||
step.presentation.intent ?? step.intent,
|
||||
'intent',
|
||||
);
|
||||
const obsRes = await maskAsync(
|
||||
step.presentation.observation ?? step.observation,
|
||||
'observation',
|
||||
);
|
||||
|
||||
if (intentRes.changed || obsRes.changed) {
|
||||
step.presentation.intent = intentRes.masked;
|
||||
step.presentation.observation = obsRes.masked;
|
||||
|
||||
// Recalculate tokens perfectly
|
||||
const newIntentTokens = estimateTokenCountSync([
|
||||
{
|
||||
functionCall: {
|
||||
name: toolName,
|
||||
args: intentRes.masked,
|
||||
id: callId,
|
||||
},
|
||||
},
|
||||
]);
|
||||
const newObsTokens = estimateTokenCountSync([
|
||||
{
|
||||
functionResponse: {
|
||||
name: toolName,
|
||||
response: obsRes.masked,
|
||||
id: callId,
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
const oldTotal =
|
||||
step.presentation.tokens?.intent !== undefined
|
||||
? step.presentation.tokens.intent +
|
||||
step.presentation.tokens.observation
|
||||
: step.tokens.intent + step.tokens.observation;
|
||||
|
||||
const newTotal = newIntentTokens + newObsTokens;
|
||||
const savings = oldTotal - newTotal;
|
||||
|
||||
if (savings > 0) {
|
||||
step.presentation.tokens = {
|
||||
intent: newIntentTokens,
|
||||
observation: newObsTokens,
|
||||
};
|
||||
step.metadata.transformations.push({
|
||||
processorName: 'ToolMasking',
|
||||
action: 'MASKED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
currentDeficit -= savings;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return newEpisodes;
|
||||
}
|
||||
|
||||
private isAlreadyMasked(content: string): boolean {
|
||||
return content.includes('<tool_output_masked>');
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import * as fs from 'node:fs';
|
||||
import type { Config } from '../../config/config.js';
|
||||
import type { SidecarConfig } from './types.js';
|
||||
import { defaultSidecarProfile } from './profiles.js';
|
||||
|
||||
export class SidecarLoader {
|
||||
/**
|
||||
* Generates a Sidecar JSON graph from the experimental config file path or defaults.
|
||||
*/
|
||||
static fromConfig(config: Config): SidecarConfig {
|
||||
const sidecarPath =
|
||||
typeof (config as any).getExperimentalContextSidecarConfig === 'function'
|
||||
? (config as any).getExperimentalContextSidecarConfig()
|
||||
: undefined;
|
||||
|
||||
if (sidecarPath && fs.existsSync(sidecarPath)) {
|
||||
try {
|
||||
const fileContent = fs.readFileSync(sidecarPath, 'utf8');
|
||||
return JSON.parse(fileContent) as SidecarConfig;
|
||||
} catch (error) {
|
||||
console.error(
|
||||
`Failed to parse Sidecar configuration file at ${sidecarPath}:`,
|
||||
error,
|
||||
);
|
||||
// Fallback to default
|
||||
}
|
||||
}
|
||||
|
||||
return defaultSidecarProfile;
|
||||
}
|
||||
|
||||
static fromLegacyConfig(config: Config): SidecarConfig {
|
||||
return SidecarLoader.fromConfig(config);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
import type { ContextTracer } from '../tracer.js';
|
||||
|
||||
export interface ContextEnvironment {
|
||||
getLlmClient(): BaseLlmClient;
|
||||
getSessionId(): string;
|
||||
getTraceDir(): string;
|
||||
getProjectTempDir(): string;
|
||||
getTracer(): ContextTracer;
|
||||
getCharsPerToken(): number;
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
import type { ContextTracer } from '../tracer.js';
|
||||
import type { ContextEnvironment } from './environment.js';
|
||||
|
||||
export class ContextEnvironmentImpl implements ContextEnvironment {
|
||||
constructor(
|
||||
private llmClient: BaseLlmClient,
|
||||
private sessionId: string,
|
||||
private traceDir: string,
|
||||
private tempDir: string,
|
||||
private tracer: ContextTracer,
|
||||
private charsPerToken: number,
|
||||
) {}
|
||||
|
||||
getLlmClient(): BaseLlmClient {
|
||||
return this.llmClient;
|
||||
}
|
||||
|
||||
getSessionId(): string {
|
||||
return this.sessionId;
|
||||
}
|
||||
|
||||
getTraceDir(): string {
|
||||
return this.traceDir;
|
||||
}
|
||||
|
||||
getProjectTempDir(): string {
|
||||
return this.tempDir;
|
||||
}
|
||||
|
||||
getTracer(): ContextTracer {
|
||||
return this.tracer;
|
||||
}
|
||||
|
||||
getCharsPerToken(): number {
|
||||
return this.charsPerToken;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { SidecarConfig } from './types.js';
|
||||
|
||||
/**
|
||||
* The standard default context management profile.
|
||||
* Optimized for safety, precision, and reliable summarization.
|
||||
*/
|
||||
export const defaultSidecarProfile: SidecarConfig = {
|
||||
budget: {
|
||||
retainedTokens: 65000,
|
||||
maxTokens: 150000,
|
||||
},
|
||||
gcBackstop: {
|
||||
strategy: 'truncate',
|
||||
target: 'incremental',
|
||||
freeTokensTarget: 10000,
|
||||
},
|
||||
pipelines: {
|
||||
eagerBackground: [
|
||||
{
|
||||
processorId: 'StateSnapshotWorker',
|
||||
options: { pollingIntervalMs: 5000 },
|
||||
},
|
||||
],
|
||||
retainedProcessingGraph: [
|
||||
{
|
||||
processorId: 'HistorySquashingProcessor',
|
||||
options: { maxTokensPerNode: 3000 },
|
||||
},
|
||||
],
|
||||
normalProcessingGraph: [
|
||||
{
|
||||
processorId: 'ToolMaskingProcessor',
|
||||
options: { stringLengthThresholdTokens: 8000 },
|
||||
},
|
||||
{
|
||||
processorId: 'BlobDegradationProcessor',
|
||||
options: {},
|
||||
},
|
||||
{
|
||||
processorId: 'SemanticCompressionProcessor',
|
||||
options: { nodeThresholdTokens: 5000, contextWindowPercentage: 0.2 },
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,42 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextProcessor } from '../pipeline.js';
|
||||
import type { AsyncContextWorker } from '../workers/asyncContextWorker.js';
|
||||
import type { ContextEnvironment } from './environment.js';
|
||||
|
||||
export interface ContextProcessorDef<
|
||||
TOptions extends Record<string, unknown> = any,
|
||||
> {
|
||||
readonly id: string;
|
||||
create(
|
||||
env: ContextEnvironment,
|
||||
options: TOptions,
|
||||
): ContextProcessor | AsyncContextWorker;
|
||||
}
|
||||
|
||||
/**
|
||||
* Registry for mapping declarative sidecar configs to running Processor instances.
|
||||
*/
|
||||
export class ProcessorRegistry {
|
||||
private static processors = new Map<string, ContextProcessorDef>();
|
||||
|
||||
static register(def: ContextProcessorDef) {
|
||||
this.processors.set(def.id, def);
|
||||
}
|
||||
|
||||
static get(id: string): ContextProcessorDef {
|
||||
const def = this.processors.get(id);
|
||||
if (!def) {
|
||||
throw new Error(`Context Processor [${id}] is not registered.`);
|
||||
}
|
||||
return def;
|
||||
}
|
||||
|
||||
static clear() {
|
||||
this.processors.clear();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
/**
|
||||
* Definition of a processor or worker to be instantiated in the graph.
|
||||
*/
|
||||
export interface ProcessorConfig {
|
||||
/** The registered ID of the processor (e.g. 'SemanticCompressionProcessor') */
|
||||
processorId: string;
|
||||
|
||||
/** Dynamic, processor-specific hyperparameters */
|
||||
options: Record<string, unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
* The Data-Driven Schema for the Context Manager.
|
||||
*/
|
||||
export interface SidecarConfig {
|
||||
/** Defines the token ceilings and limits for the pipeline. */
|
||||
budget: {
|
||||
retainedTokens: number;
|
||||
maxTokens: number;
|
||||
};
|
||||
|
||||
/** Defines what happens when the pipeline fails to compress under 'maxTokens' */
|
||||
gcBackstop: {
|
||||
strategy: 'truncate' | 'compress' | 'rollingSummarizer';
|
||||
target: 'incremental' | 'freeNTokens' | 'max';
|
||||
freeTokensTarget?: number;
|
||||
};
|
||||
|
||||
/** The execution graphs for context manipulation */
|
||||
pipelines: {
|
||||
/**
|
||||
* Eagerly executes in the background when the 'retainedTokens' boundary is crossed.
|
||||
* Contains AsyncContextWorkers (e.g. StateSnapshotWorker).
|
||||
*/
|
||||
eagerBackground: ProcessorConfig[];
|
||||
|
||||
/**
|
||||
* Executes sequentially to protect the pristine outliers within the retained window.
|
||||
* Contains ContextProcessors (e.g. HistorySquashingProcessor).
|
||||
*/
|
||||
retainedProcessingGraph: ProcessorConfig[];
|
||||
|
||||
/**
|
||||
* Executes sequentially to opportunistically degrade messages older than the retained window.
|
||||
* Contains ContextProcessors (e.g. ToolMaskingProcessor, SemanticCompressionProcessor).
|
||||
*/
|
||||
normalProcessingGraph: ProcessorConfig[];
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { vi } from 'vitest';
|
||||
import type { Config } from '../../config/config.js';
|
||||
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
|
||||
export function createMockEnvironment(): ContextEnvironment {
|
||||
return {
|
||||
getLlmClient: vi.fn().mockReturnValue({
|
||||
generateContent: vi.fn().mockResolvedValue({
|
||||
text: 'Mock LLM summary response',
|
||||
}),
|
||||
}) as any,
|
||||
getSessionId: vi.fn().mockReturnValue('mock-session'),
|
||||
getTraceDir: vi.fn().mockReturnValue('/tmp/.gemini/trace'),
|
||||
getProjectTempDir: vi.fn().mockReturnValue('/tmp'),
|
||||
getTracer: vi.fn().mockReturnValue({
|
||||
logEvent: vi.fn(),
|
||||
saveAsset: vi.fn().mockReturnValue('mock-asset-id'),
|
||||
}) as any,
|
||||
getCharsPerToken: vi.fn().mockReturnValue(1),
|
||||
};
|
||||
}
|
||||
|
||||
import type { Content } from '@google/genai';
|
||||
import { AgentChatHistory } from '../../core/agentChatHistory.js';
|
||||
import { ContextManager } from '../contextManager.js';
|
||||
|
||||
/**
|
||||
* Creates a block of synthetic conversation history designed to consume a specific number of tokens.
|
||||
* Assumes roughly 4 characters per token for standard English text.
|
||||
*/
|
||||
export function createSyntheticHistory(
|
||||
numTurns: number,
|
||||
tokensPerTurn: number,
|
||||
): Content[] {
|
||||
const history: Content[] = [];
|
||||
const charsPerTurn = tokensPerTurn * 1;
|
||||
|
||||
for (let i = 0; i < numTurns; i++) {
|
||||
history.push({
|
||||
role: 'user',
|
||||
parts: [{ text: `User turn ${i}. ` + 'A'.repeat(charsPerTurn) }],
|
||||
});
|
||||
history.push({
|
||||
role: 'model',
|
||||
parts: [{ text: `Model response ${i}. ` + 'B'.repeat(charsPerTurn) }],
|
||||
});
|
||||
}
|
||||
|
||||
return history;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a fully mocked Config object tailored for Context Component testing.
|
||||
*/
|
||||
export function createMockContextConfig(
|
||||
overrides?: Record<string, unknown>,
|
||||
llmClientOverride?: unknown,
|
||||
): Config {
|
||||
const defaultConfig = {
|
||||
isContextManagementEnabled: vi.fn().mockReturnValue(true),
|
||||
storage: {
|
||||
getProjectTempDir: vi.fn().mockReturnValue('/tmp/gemini-test'),
|
||||
},
|
||||
getBaseLlmClient: vi.fn().mockReturnValue(
|
||||
llmClientOverride || {
|
||||
generateContent: vi.fn().mockResolvedValue({
|
||||
text: '<mocked_snapshot>Synthesized state</mocked_snapshot>',
|
||||
}),
|
||||
},
|
||||
),
|
||||
getUsageStatisticsEnabled: vi.fn().mockReturnValue(false),
|
||||
getTargetDir: vi.fn().mockReturnValue('/tmp'),
|
||||
getSessionId: vi.fn().mockReturnValue('test-session'),
|
||||
};
|
||||
|
||||
return { ...defaultConfig, ...overrides } as unknown as Config;
|
||||
}
|
||||
|
||||
/**
|
||||
* Wires up a full ContextManager component with an AgentChatHistory and active background workers.
|
||||
*/
|
||||
import { ContextTracer } from '../tracer.js';
|
||||
import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js';
|
||||
import { SidecarLoader } from '../sidecar/SidecarLoader.js';
|
||||
|
||||
export function setupContextComponentTest(config: Config) {
|
||||
const chatHistory = new AgentChatHistory();
|
||||
const sidecar = SidecarLoader.fromLegacyConfig(config);
|
||||
const tracer = new ContextTracer('/tmp', 'test-session');
|
||||
const env = new ContextEnvironmentImpl(
|
||||
config.getBaseLlmClient() as any,
|
||||
'test-session',
|
||||
'/tmp',
|
||||
'/tmp/gemini-test',
|
||||
tracer,
|
||||
1,
|
||||
);
|
||||
const contextManager = new ContextManager(sidecar, env, tracer);
|
||||
|
||||
// The async worker is now internally managed by ContextManager
|
||||
|
||||
// Subscribe to history to enable the Eager/Opportunistic triggers
|
||||
contextManager.subscribeToHistory(chatHistory);
|
||||
|
||||
return { chatHistory, contextManager };
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import * as fs from 'node:fs';
|
||||
import * as path from 'node:path';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
|
||||
export class ContextTracer {
|
||||
private traceDir: string;
|
||||
private assetsDir: string;
|
||||
private enabled: boolean;
|
||||
|
||||
constructor(targetDir: string, sessionId: string) {
|
||||
this.enabled = process.env['GEMINI_CONTEXT_TRACE'] === 'true';
|
||||
this.traceDir = path.join(targetDir, '.gemini', 'context_trace', sessionId);
|
||||
this.assetsDir = path.join(this.traceDir, 'assets');
|
||||
|
||||
if (this.enabled) {
|
||||
try {
|
||||
fs.mkdirSync(this.assetsDir, { recursive: true });
|
||||
this.logEvent('SYSTEM', 'Context Tracer Initialized', { sessionId });
|
||||
} catch (e) {
|
||||
console.error('Failed to initialize ContextTracer', e);
|
||||
this.enabled = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logEvent(
|
||||
component: string,
|
||||
action: string,
|
||||
details?: Record<string, unknown>,
|
||||
) {
|
||||
if (!this.enabled) return;
|
||||
try {
|
||||
const timestamp = new Date().toISOString();
|
||||
const detailsStr = details
|
||||
? ` | Details: ${JSON.stringify(details)}`
|
||||
: '';
|
||||
const logLine = `[${timestamp}] [${component}] ${action}${detailsStr}\n`;
|
||||
fs.appendFileSync(
|
||||
path.join(this.traceDir, 'trace.log'),
|
||||
logLine,
|
||||
'utf-8',
|
||||
);
|
||||
} catch (e) {
|
||||
// fail silently in trace
|
||||
}
|
||||
}
|
||||
|
||||
saveAsset(component: string, assetName: string, data: unknown): string {
|
||||
if (!this.enabled) return 'asset-recording-disabled';
|
||||
try {
|
||||
const assetId = `${Date.now()}-${randomUUID().slice(0, 6)}-${assetName}.json`;
|
||||
const assetPath = path.join(this.assetsDir, assetId);
|
||||
|
||||
fs.writeFileSync(assetPath, JSON.stringify(data, null, 2), 'utf-8');
|
||||
this.logEvent(component, `Saved asset: ${assetName}`, { assetId });
|
||||
return assetId;
|
||||
} catch (e) {
|
||||
this.logEvent(component, `Failed to save asset: ${assetName}`, {
|
||||
error: String(e),
|
||||
});
|
||||
return 'asset-save-failed';
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Part } from '@google/genai';
|
||||
import { estimateTokenCountSync as baseEstimate } from '../../utils/tokenCalculation.js';
|
||||
|
||||
export function estimateContextTokenCountSync(
|
||||
parts: Part[],
|
||||
depth: number = 0,
|
||||
config?: { charsPerToken?: number },
|
||||
): number {
|
||||
if (config?.charsPerToken !== undefined && config.charsPerToken !== 4) {
|
||||
let totalTokens = 0;
|
||||
for (const part of parts) {
|
||||
if (typeof part.text === 'string') {
|
||||
totalTokens += Math.ceil(part.text.length / config.charsPerToken);
|
||||
} else {
|
||||
totalTokens += Math.ceil(
|
||||
JSON.stringify(part).length / config.charsPerToken,
|
||||
);
|
||||
}
|
||||
}
|
||||
return totalTokens;
|
||||
}
|
||||
|
||||
// The baseEstimate no longer accepts config because we forked it!
|
||||
return baseEstimate(parts, depth);
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextEventBus } from '../eventBus.js';
|
||||
|
||||
export interface AsyncContextWorker {
|
||||
/** The unique name of the worker (e.g., 'StateSnapshotWorker') */
|
||||
readonly name: string;
|
||||
|
||||
/** Starts listening to the ContextEventBus for background tasks */
|
||||
start(bus: ContextEventBus): void;
|
||||
|
||||
/** Stops listening and aborts any pending background tasks */
|
||||
stop(): void;
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import type { Episode, SnapshotVariant } from '../ir/types.js';
|
||||
import type { AsyncContextWorker } from './asyncContextWorker.js';
|
||||
import type {
|
||||
ContextEventBus,
|
||||
ContextConsolidationEvent,
|
||||
} from '../eventBus.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { estimateContextTokenCountSync as estimateTokenCountSync } from '../utils/contextTokenCalculator.js';
|
||||
import { IrMapper } from '../ir/mapper.js';
|
||||
import { LlmRole } from '../../telemetry/llmRole.js';
|
||||
import type { ContextTracer } from '../tracer.js';
|
||||
|
||||
export class StateSnapshotWorker implements AsyncContextWorker {
|
||||
name = 'StateSnapshotWorker';
|
||||
private bus?: ContextEventBus;
|
||||
private tracer?: ContextTracer;
|
||||
private isSynthesizing = false;
|
||||
|
||||
constructor(private readonly env: ContextEnvironment) {}
|
||||
|
||||
start(bus: ContextEventBus, tracer?: ContextTracer): void {
|
||||
console.log('Worker start() called with bus:', !!bus);
|
||||
this.bus = bus;
|
||||
this.tracer = tracer;
|
||||
this.bus.onConsolidationNeeded(this.handleConsolidation.bind(this));
|
||||
}
|
||||
|
||||
stop(): void {
|
||||
if (this.bus) {
|
||||
// In a real implementation we would `removeListener` here
|
||||
this.bus = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
private async handleConsolidation(
|
||||
event: ContextConsolidationEvent,
|
||||
): Promise<void> {
|
||||
console.log(
|
||||
`Worker handling consolidation. targetDeficit: ${event.targetDeficit}, isSynthesizing: ${this.isSynthesizing}`,
|
||||
);
|
||||
if (this.isSynthesizing || event.targetDeficit <= 0) return;
|
||||
|
||||
// Identify the "dying" block of episodes that need to be collected.
|
||||
// For now, we assume older episodes are at the front of the array.
|
||||
// We only want episodes that don't already have a snapshot variant computing/ready.
|
||||
const unprotectedOldest = event.episodes.filter(
|
||||
(ep) => !ep.variants?.['snapshot'],
|
||||
);
|
||||
|
||||
if (unprotectedOldest.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
let targetDeficit = event.targetDeficit;
|
||||
const episodesToSynthesize: Episode[] = [];
|
||||
let tokensToSynthesize = 0;
|
||||
|
||||
for (const ep of unprotectedOldest) {
|
||||
console.log('Worker considering episode:', ep.id);
|
||||
if (tokensToSynthesize >= targetDeficit) break;
|
||||
episodesToSynthesize.push(ep);
|
||||
// Rough estimate of tokens in this episode
|
||||
const epTokens = ep.steps.reduce(
|
||||
(sum, step) => sum + step.metadata.currentTokens,
|
||||
ep.trigger.metadata.currentTokens +
|
||||
(ep.yield?.metadata.currentTokens || 0),
|
||||
);
|
||||
tokensToSynthesize += epTokens;
|
||||
}
|
||||
|
||||
if (episodesToSynthesize.length === 0) return;
|
||||
|
||||
console.log(
|
||||
`Worker synthesized logic loop complete. Selected ${episodesToSynthesize.length} episodes for ~${tokensToSynthesize} tokens.`,
|
||||
);
|
||||
this.isSynthesizing = true;
|
||||
|
||||
try {
|
||||
debugLogger.log(
|
||||
`StateSnapshotWorker: Asynchronously synthesizing ${episodesToSynthesize.length} episodes to recover ~${tokensToSynthesize} tokens.`,
|
||||
);
|
||||
this.tracer?.logEvent(
|
||||
'StateSnapshotWorker',
|
||||
`Consolidation requested. Synthesizing ${episodesToSynthesize.length} episodes for ~${tokensToSynthesize} tokens.`,
|
||||
);
|
||||
|
||||
const client = this.env.getLlmClient();
|
||||
const rawContents = IrMapper.fromIr(episodesToSynthesize);
|
||||
const rawAssetId = this.tracer?.saveAsset(
|
||||
'StateSnapshotWorker',
|
||||
'episodes_to_synthesize',
|
||||
rawContents,
|
||||
);
|
||||
this.tracer?.logEvent(
|
||||
'StateSnapshotWorker',
|
||||
'Dispatching LLM request for snapshot generation',
|
||||
{ rawAssetId },
|
||||
);
|
||||
|
||||
const promptText = `
|
||||
You are a background memory consolidation worker for an AI assistant.
|
||||
Your task is to review the following block of the oldest conversation history and synthesize it into a highly dense, accurate "World State Snapshot".
|
||||
This snapshot will completely replace these old memories.
|
||||
Preserve all critical facts, technical decisions, file paths, and outstanding tasks. Discard all conversational filler.
|
||||
|
||||
Conversation History to Synthesize:
|
||||
${JSON.stringify(rawContents, null, 2).slice(0, 50000)}
|
||||
|
||||
Output the snapshot as a dense, structured summary.`;
|
||||
|
||||
const response = await client.generateContent({
|
||||
modelConfigKey: { model: 'gemini-2.5-flash' }, // Fast and cheap for background tasks
|
||||
contents: [{ role: 'user', parts: [{ text: promptText }] }],
|
||||
promptId: 'async-world-state-snapshot',
|
||||
role: LlmRole.UTILITY_COMPRESSOR,
|
||||
abortSignal: new AbortController().signal, // Run in background, could add cancellation logic later
|
||||
});
|
||||
|
||||
// Extract text safely from the GenAI response
|
||||
const snapshotText = response.text;
|
||||
const responseAssetId = this.tracer?.saveAsset(
|
||||
'StateSnapshotWorker',
|
||||
'snapshot_response',
|
||||
snapshotText || '',
|
||||
);
|
||||
this.tracer?.logEvent('StateSnapshotWorker', 'Received LLM response', {
|
||||
responseAssetId,
|
||||
});
|
||||
if (!snapshotText) {
|
||||
debugLogger.warn(
|
||||
'StateSnapshotWorker: LLM returned empty response for snapshot generation.',
|
||||
);
|
||||
}
|
||||
|
||||
const mockSnapshotText = `
|
||||
<world_state_snapshot>
|
||||
${snapshotText || '[Failed to generate snapshot]'}
|
||||
</world_state_snapshot>`;
|
||||
|
||||
const snapshotTokens = estimateTokenCountSync(
|
||||
[{ text: mockSnapshotText }],
|
||||
0,
|
||||
{ charsPerToken: this.env.getCharsPerToken() },
|
||||
);
|
||||
|
||||
const replacedEpisodeIds = episodesToSynthesize.map((e) => e.id);
|
||||
|
||||
const snapshotEpisode: Episode = {
|
||||
id: randomUUID(),
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: randomUUID(),
|
||||
type: 'SYSTEM_EVENT',
|
||||
name: 'world_state_snapshot',
|
||||
payload: {
|
||||
originalEpisodeCount: episodesToSynthesize.length,
|
||||
recoveredTokens: tokensToSynthesize,
|
||||
},
|
||||
metadata: {
|
||||
originalTokens: snapshotTokens,
|
||||
currentTokens: snapshotTokens,
|
||||
transformations: [
|
||||
{
|
||||
processorName: 'StateSnapshotWorker',
|
||||
action: 'SYNTHESIZED',
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
steps: [
|
||||
{
|
||||
id: randomUUID(),
|
||||
type: 'AGENT_THOUGHT',
|
||||
text: mockSnapshotText,
|
||||
metadata: {
|
||||
originalTokens: snapshotTokens,
|
||||
currentTokens: snapshotTokens,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const variant: SnapshotVariant = {
|
||||
type: 'snapshot',
|
||||
status: 'ready',
|
||||
recoveredTokens: tokensToSynthesize,
|
||||
episode: snapshotEpisode,
|
||||
replacedEpisodeIds,
|
||||
};
|
||||
|
||||
// Emit the variant for the MOST RECENT episode in the batch,
|
||||
// since the Opportunistic Swapper sweeps from newest to oldest.
|
||||
const targetId = replacedEpisodeIds[replacedEpisodeIds.length - 1];
|
||||
|
||||
if (this.bus) {
|
||||
this.tracer?.logEvent(
|
||||
'StateSnapshotWorker',
|
||||
`Emitting VARIANT_READY for targetId [${targetId}]`,
|
||||
);
|
||||
|
||||
this.bus.emitVariantReady({
|
||||
targetId,
|
||||
variantId: 'snapshot',
|
||||
variant,
|
||||
});
|
||||
} else {
|
||||
debugLogger.warn(
|
||||
'StateSnapshotWorker: Event bus disconnected before variant could be emitted.',
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
debugLogger.error(
|
||||
`StateSnapshotWorker: Critical failure during snapshot synthesis: ${error instanceof Error ? error.message : String(error)}`,
|
||||
);
|
||||
} finally {
|
||||
this.isSynthesizing = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Content } from '@google/genai';
|
||||
|
||||
export type HistoryEventType = 'PUSH' | 'SYNC_FULL' | 'CLEAR';
|
||||
|
||||
export interface HistoryEvent {
|
||||
type: HistoryEventType;
|
||||
payload: readonly Content[];
|
||||
}
|
||||
|
||||
export type HistoryListener = (event: HistoryEvent) => void;
|
||||
|
||||
export class AgentChatHistory {
|
||||
private history: Content[];
|
||||
private listeners: Set<HistoryListener> = new Set();
|
||||
|
||||
constructor(initialHistory: Content[] = []) {
|
||||
this.history = [...initialHistory];
|
||||
}
|
||||
|
||||
subscribe(listener: HistoryListener): () => void {
|
||||
this.listeners.add(listener);
|
||||
// Emit initial state to new subscriber
|
||||
listener({ type: 'SYNC_FULL', payload: this.history });
|
||||
return () => this.listeners.delete(listener);
|
||||
}
|
||||
|
||||
private notify(type: HistoryEventType, payload: readonly Content[]) {
|
||||
const event: HistoryEvent = { type, payload };
|
||||
for (const listener of this.listeners) {
|
||||
listener(event);
|
||||
}
|
||||
}
|
||||
|
||||
push(content: Content) {
|
||||
this.history.push(content);
|
||||
this.notify('PUSH', [content]);
|
||||
}
|
||||
|
||||
set(history: readonly Content[]) {
|
||||
this.history = [...history];
|
||||
this.notify('SYNC_FULL', this.history);
|
||||
}
|
||||
|
||||
clear() {
|
||||
this.history = [];
|
||||
this.notify('CLEAR', []);
|
||||
}
|
||||
|
||||
get(): readonly Content[] {
|
||||
return this.history;
|
||||
}
|
||||
|
||||
map(callback: (value: Content, index: number, array: Content[]) => Content) {
|
||||
this.history = this.history.map(callback);
|
||||
this.notify('SYNC_FULL', this.history);
|
||||
}
|
||||
|
||||
flatMap<U>(
|
||||
callback: (
|
||||
value: Content,
|
||||
index: number,
|
||||
array: Content[],
|
||||
) => U | readonly U[],
|
||||
): U[] {
|
||||
return this.history.flatMap(callback);
|
||||
}
|
||||
|
||||
get length(): number {
|
||||
return this.history.length;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user