Merge branch 'jl/cm-part-1' into gundermanc/compression-eval

This commit is contained in:
Christian Gunderman
2026-04-08 16:35:17 -07:00
56 changed files with 6528 additions and 0 deletions
+8
View File
@@ -700,6 +700,7 @@ export interface ConfigParameters {
experimentalJitContext?: boolean;
autoDistillation?: boolean;
experimentalMemoryManager?: boolean;
experimentalContextSidecarConfig?: string;
experimentalAgentHistoryTruncation?: boolean;
experimentalAgentHistoryTruncationThreshold?: number;
experimentalAgentHistoryRetainedMessages?: number;
@@ -942,6 +943,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;
@@ -1153,6 +1155,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,
@@ -2427,6 +2431,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,110 @@
/**
* @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.
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
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);
// 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.
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!!' }] },
];
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,73 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
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;
// 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,189 @@
/**
* @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 { ContextEventBus } from './eventBus.js';
import { ContextTokenCalculator } from './utils/contextTokenCalculator.js';
import type { Content } from '@google/genai';
import type { BaseLlmClient } from '../core/baseLlmClient.js';
import type { Episode } from './ir/types.js';
import type { SidecarConfig } from './sidecar/types.js';
import { ProcessorRegistry } from './sidecar/registry.js';
import { registerBuiltInProcessors } from './sidecar/builtins.js';
import { IrMapper } from './ir/mapper.js';
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),
getExperimentalContextSidecarConfig: vi.fn().mockReturnValue(undefined),
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 registry = new ProcessorRegistry();
registerBuiltInProcessors(registry);
const sidecar = SidecarLoader.fromConfig(mockConfig, registry);
const tracer = new ContextTracer({
targetDir: '/tmp',
sessionId: 'test-session',
});
const eventBus = new ContextEventBus();
const env = new ContextEnvironmentImpl(
{
generateContent: async () => ({}),
generateJson: async () => ({}),
} as unknown as BaseLlmClient,
'test-prompt-id',
'test',
'/tmp',
'/tmp',
tracer,
4,
eventBus,
);
contextManager = ContextManager.create(
sidecar,
env,
tracer,
undefined,
registry,
);
});
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 unknown as { pristineEpisodes: Episode[] }
).pristineEpisodes = IrMapper.toIr(history, new ContextTokenCalculator(4));
const result = await contextManager.projectCompressedHistory();
expect(result).toMatchSnapshot();
});
it('should not modify history when under budget', async () => {
const history = createLargeHistory();
(
contextManager as unknown as { pristineEpisodes: Episode[] }
).pristineEpisodes = IrMapper.toIr(history, new ContextTokenCalculator(4));
// 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({
targetDir: '/tmp',
sessionId: 'test2',
});
const eventBus2 = new ContextEventBus();
const env2 = new ContextEnvironmentImpl(
{
generateContent: async () => ({}),
generateJson: async () => ({}),
} as unknown as BaseLlmClient,
'test-prompt-id',
'test',
'/tmp',
'/tmp',
tracer2,
4,
eventBus2,
);
contextManager = ContextManager.create(
{
budget: { retainedTokens: 100000, maxTokens: 150000 },
pipelines: [],
} as unknown as SidecarConfig,
env2,
tracer2,
);
(
contextManager as unknown as { pristineEpisodes: Episode[] }
).pristineEpisodes = IrMapper.toIr(history, new ContextTokenCalculator(4));
const result = await contextManager.projectCompressedHistory();
expect(result.length).toEqual(history.length);
});
});
+184
View File
@@ -0,0 +1,184 @@
/**
* @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 type { Episode } from './ir/types.js';
import type { ContextEventBus } from './eventBus.js';
import type { ContextTracer } from './tracer.js';
import type { ContextEnvironment } from './sidecar/environment.js';
import type { SidecarConfig } from './sidecar/types.js';
import { PipelineOrchestrator } from './sidecar/orchestrator.js';
import { HistoryObserver } from './historyObserver.js';
import { generateWorkingBufferView } from './ir/graphUtils.js';
import { IrProjector } from './ir/projector.js';
import { registerBuiltInProcessors } from './sidecar/builtins.js';
import { ProcessorRegistry } from './sidecar/registry.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 readonly eventBus: ContextEventBus;
// Internal sub-components
// Synchronous processors are instantiated but effectively used as singletons within this class
private orchestrator: PipelineOrchestrator;
private historyObserver?: HistoryObserver;
static create(
sidecar: SidecarConfig,
env: ContextEnvironment,
tracer: ContextTracer,
orchestrator?: PipelineOrchestrator,
registry?: ProcessorRegistry,
): ContextManager {
if (!registry) {
registry = new ProcessorRegistry();
registerBuiltInProcessors(registry);
}
const orch =
orchestrator ||
new PipelineOrchestrator(sidecar, env, env.eventBus, tracer, registry);
return new ContextManager(sidecar, env, tracer, orch);
}
// Use ContextManager.create() instead
private constructor(
private sidecar: SidecarConfig,
private env: ContextEnvironment,
private readonly tracer: ContextTracer,
orchestrator: PipelineOrchestrator,
) {
this.eventBus = env.eventBus;
this.orchestrator = orchestrator;
this.eventBus.onPristineHistoryUpdated((event) => {
this.pristineEpisodes = event.episodes;
this.evaluateTriggers();
});
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}.`,
);
}
});
}
/**
* Safely stops background workers and clears event listeners.
*/
shutdown() {
this.orchestrator.shutdown();
if (this.historyObserver) {
this.historyObserver.stop();
}
}
/**
* Evaluates if the current working buffer exceeds configured budget thresholds,
* firing consolidation events if necessary.
*/
private evaluateTriggers() {
if (!this.sidecar.budget) return;
const workingBuffer = this.getWorkingBufferView();
const currentTokens =
this.env.tokenCalculator.calculateEpisodeListTokens(workingBuffer);
this.tracer.logEvent('ContextManager', 'Evaluated triggers', {
currentTokens,
retainedTokens: this.sidecar.budget.retainedTokens,
});
// 1. Eager Compute Trigger
this.eventBus.emitChunkReceived({ episodes: this.pristineEpisodes });
// 2. Budget Crossed Trigger
if (currentTokens > this.sidecar.budget.retainedTokens) {
const deficit = currentTokens - this.sidecar.budget.retainedTokens;
this.tracer.logEvent(
'ContextManager',
'Budget crossed. Emitting ConsolidationNeeded',
{ deficit },
);
this.eventBus.emitConsolidationNeeded({
episodes: workingBuffer,
targetDeficit: deficit,
});
}
}
/**
* Subscribes to the core AgentChatHistory to natively track all message events,
* converting them seamlessly into pristine Episodes.
*/
subscribeToHistory(chatHistory: AgentChatHistory) {
if (this.historyObserver) {
this.historyObserver.stop();
}
this.historyObserver = new HistoryObserver(
chatHistory,
this.eventBus,
this.tracer,
this.env.tokenCalculator,
);
this.historyObserver.start();
}
/**
* 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.
*/
getWorkingBufferView(): Episode[] {
return generateWorkingBufferView(
this.pristineEpisodes,
this.sidecar.budget.retainedTokens,
this.tracer,
this.env,
);
}
/**
* 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[]> {
this.tracer.logEvent('ContextManager', 'Projection requested.');
const protectedIds = new Set<string>();
if (this.pristineEpisodes.length > 0) {
protectedIds.add(this.pristineEpisodes[0].id); // Structural invariant
}
return IrProjector.project(
this.getWorkingBufferView(),
this.orchestrator,
this.sidecar,
this.tracer,
this.env,
protectedIds,
);
}
}
+63
View File
@@ -0,0 +1,63 @@
/**
* @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 PristineHistoryUpdatedEvent {
episodes: Episode[];
}
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 {
emitPristineHistoryUpdated(event: PristineHistoryUpdatedEvent) {
this.emit('PRISTINE_HISTORY_UPDATED', event);
}
onPristineHistoryUpdated(
listener: (event: PristineHistoryUpdatedEvent) => void,
) {
this.on('PRISTINE_HISTORY_UPDATED', listener);
}
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,62 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type {
AgentChatHistory,
HistoryEvent,
} from '../core/agentChatHistory.js';
import { IrMapper } from './ir/mapper.js';
import type { ContextTokenCalculator } from './utils/contextTokenCalculator.js';
import type { ContextEventBus } from './eventBus.js';
import type { ContextTracer } from './tracer.js';
/**
* Connects the raw AgentChatHistory to the ContextManager.
* It maps raw messages into Episodic Intermediate Representation (IR)
* and evaluates background triggers whenever history changes.
*/
export class HistoryObserver {
private unsubscribeHistory?: () => void;
constructor(
private readonly chatHistory: AgentChatHistory,
private readonly eventBus: ContextEventBus,
private readonly tracer: ContextTracer,
private readonly tokenCalculator: ContextTokenCalculator,
) {}
start() {
if (this.unsubscribeHistory) {
this.unsubscribeHistory();
}
this.unsubscribeHistory = this.chatHistory.subscribe(
(_event: HistoryEvent) => {
// Rebuild the pristine IR graph from the full source history on every change.
const pristineEpisodes = IrMapper.toIr(
this.chatHistory.get(),
this.tokenCalculator,
);
this.tracer.logEvent(
'HistoryObserver',
'Rebuilt pristine graph from chat history update',
{ episodeCount: pristineEpisodes.length },
);
this.eventBus.emitPristineHistoryUpdated({
episodes: pristineEpisodes,
});
},
);
}
stop() {
if (this.unsubscribeHistory) {
this.unsubscribeHistory();
this.unsubscribeHistory = undefined;
}
}
}
@@ -0,0 +1,137 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Episode } from './types.js';
export interface MutationRecord {
episodeId: string;
type: 'modified' | 'inserted' | 'replaced' | 'deleted';
action: string;
originalIds?: string[]; // If replaced
episode?: Episode; // For new or modified
}
export class EpisodeEditor {
private originalMap: Map<string, Episode>;
private workingOrder: string[];
private workingMap: Map<string, Episode>;
private mutations: MutationRecord[] = [];
constructor(episodes: Episode[]) {
this.originalMap = new Map(episodes.map((e) => [e.id, e]));
this.workingOrder = episodes.map((e) => e.id);
this.workingMap = new Map(episodes.map((e) => [e.id, e]));
}
/**
* Provides a readonly view of the current working state of the episodes.
* Processors should iterate over this to decide what to mutate.
*/
get episodes(): readonly Episode[] {
return this.workingOrder.map((id) => this.workingMap.get(id)!);
}
/**
* Safely edits an existing episode.
* The framework will handle deeply cloning the episode before passing it to the mutator,
* guaranteeing that original references are never modified.
*/
editEpisode(id: string, action: string, mutator: (draft: Episode) => void) {
const ep = this.workingMap.get(id);
if (!ep) return;
// Lazy deep clone only if it's the original reference
if (ep === this.originalMap.get(id)) {
const clone = structuredClone(ep);
this.workingMap.set(id, clone);
}
const draft = this.workingMap.get(id)!;
mutator(draft);
// Log mutation if not already tracked as modified/inserted/replaced
if (!this.mutations.find((m) => m.episodeId === id)) {
this.mutations.push({
episodeId: id,
type: 'modified',
action,
episode: draft,
});
}
}
/**
* Inserts a brand new episode into the graph at the specified index.
*/
insertEpisode(index: number, newEpisode: Episode, action: string) {
this.workingMap.set(newEpisode.id, newEpisode);
this.workingOrder.splice(index, 0, newEpisode.id);
this.mutations.push({
episodeId: newEpisode.id,
type: 'inserted',
action,
episode: newEpisode,
});
}
/**
* Replaces a set of older episodes with a single new episode (e.g., a Summary or Snapshot).
* It inserts the new episode at the lowest index of the removed episodes.
*/
replaceEpisodes(oldIds: string[], newEpisode: Episode, action: string) {
const indices = oldIds
.map((id) => this.workingOrder.indexOf(id))
.filter((i) => i !== -1);
if (indices.length === 0) return;
const insertIndex = Math.min(...indices);
// Remove old
this.workingOrder = this.workingOrder.filter((id) => !oldIds.includes(id));
for (const id of oldIds) {
this.workingMap.delete(id);
}
// Insert new
this.workingOrder.splice(insertIndex, 0, newEpisode.id);
this.workingMap.set(newEpisode.id, newEpisode);
this.mutations.push({
episodeId: newEpisode.id,
type: 'replaced',
action,
originalIds: oldIds,
episode: newEpisode,
});
}
/**
* Removes episodes from the graph completely (e.g., emergency truncation).
*/
removeEpisodes(oldIds: string[], action: string) {
this.workingOrder = this.workingOrder.filter((id) => !oldIds.includes(id));
for (const id of oldIds) {
this.workingMap.delete(id);
this.mutations.push({ episodeId: id, type: 'deleted', action });
}
}
/**
* Retrieves the final, finalized array of episodes.
* Called by the Orchestrator.
*/
getFinalEpisodes(): Episode[] {
return this.workingOrder.map((id) => this.workingMap.get(id)!);
}
/**
* Retrieves a log of all structural and property mutations performed by this editor.
* Called by the Orchestrator to emit VariantReady events.
*/
getMutations(): MutationRecord[] {
return this.mutations;
}
}
+107
View File
@@ -0,0 +1,107 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content, Part } from '@google/genai';
import type { Episode, EpisodeStep, UserPrompt, AgentYield } from './types.js';
export function fromIr(episodes: Episode[]): Content[] {
const history: Content[] = [];
for (const ep of episodes) {
if (ep.trigger.type === 'USER_PROMPT') {
const triggerContent = serializeTrigger(ep.trigger);
if (triggerContent) history.push(triggerContent);
}
const stepContents = serializeSteps(ep.steps);
history.push(...stepContents);
if (ep.yield) {
history.push(serializeYield(ep.yield));
}
}
return history;
}
function serializeTrigger(trigger: UserPrompt): Content | null {
const parts: Part[] = [];
for (const sp of 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') {
parts.push(sp.part);
}
}
return parts.length > 0 ? { role: 'user', parts } : null;
}
function serializeSteps(steps: EpisodeStep[]): Content[] {
const history: Content[] = [];
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 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,
id: step.id,
},
});
const observation = step.presentation
? step.presentation.observation
: step.observation;
pendingUserParts.push({
functionResponse: {
name: step.toolName,
response:
typeof observation === 'string'
? { message: observation }
: observation,
id: step.id,
},
});
}
}
flushPending();
return history;
}
function serializeYield(yieldNode: AgentYield): Content {
return {
role: 'model',
parts: [{ text: yieldNode.presentation?.text ?? yieldNode.text }],
};
}
@@ -0,0 +1,175 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { generateWorkingBufferView } from './graphUtils.js';
import {
createMockEnvironment,
createDummyEpisode,
} from '../testing/contextTestUtils.js';
import type { ContextEnvironment } from '../sidecar/environment.js';
import type { AgentThought, UserPrompt } from './types.js';
describe('graphUtils (View Generator)', () => {
let env: ContextEnvironment;
beforeEach(() => {
vi.resetAllMocks();
env = createMockEnvironment();
// Our token mock is 1 char = 1 token for simplicity
vi.spyOn(
env.tokenCalculator,
'calculateEpisodeListTokens',
).mockImplementation((eps) =>
eps.reduce(
(acc, ep) => acc + (ep.trigger.metadata.originalTokens || 100),
0,
),
);
});
it('returns pristine episodes untouched if under budget', () => {
const episodes = [
createDummyEpisode('ep-1', 'USER_PROMPT', [{ type: 'text', text: '1' }]),
createDummyEpisode('ep-2', 'USER_PROMPT', [{ type: 'text', text: '2' }]),
];
// We retain 5000 tokens. Total mock tokens = 200.
const view = generateWorkingBufferView(episodes, 5000, env.tracer, env);
expect(view).toHaveLength(2);
// Must be a deep copy! The view generator clones episodes.
expect(view).not.toBe(episodes);
expect(view[0].id).toBe('ep-1');
expect(view[1].id).toBe('ep-2');
});
it('swaps to Masked variant when over budget (rolling backwards)', () => {
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
{ text: '1', type: 'text' },
]);
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
{ text: '2', type: 'text' },
]);
ep1.variants = {
masked: {
type: 'masked',
status: 'ready',
text: '<MASKED>',
recoveredTokens: 10,
},
};
// We only retain 100 tokens.
// ep-2 (newest) takes 100 tokens.
// Now rolling = 100. Over budget!
// ep-1 is evaluated, and swapped for Masked.
const view = generateWorkingBufferView([ep1, ep2], 10, env.tracer, env);
expect(view).toHaveLength(2);
expect(view[1].id).toBe('ep-2'); // Unchanged (newest)
expect(view[0].id).toBe('ep-1');
expect(
(view[0].trigger as UserPrompt).semanticParts[0].presentation?.text,
).toBe('<MASKED>');
});
it('swaps to Summary variant when over budget', () => {
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
{ type: 'text', text: '1' },
]);
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
{ type: 'text', text: '2' },
]);
ep1.variants = {
summary: {
type: 'summary',
status: 'ready',
text: '<SUMMARY>',
recoveredTokens: 50,
},
};
const view = generateWorkingBufferView([ep1, ep2], 10, env.tracer, env);
expect(view).toHaveLength(2);
// The summary completely replaces the internal steps and clears the yield.
expect(view[0].steps).toHaveLength(1);
expect(view[0].steps[0].type).toBe('AGENT_THOUGHT');
expect((view[0].steps[0] as AgentThought).text).toBe('<SUMMARY>');
expect(view[0].yield).toBeUndefined();
});
it('handles complex N-to-1 Snapshot skipping gracefully', () => {
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
{ type: 'text', text: '1' },
]);
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
{ type: 'text', text: '2' },
]);
const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', [
{ type: 'text', text: '3' },
]);
const ep4 = createDummyEpisode('ep-4', 'USER_PROMPT', [
{ type: 'text', text: '4' },
]);
// ep-3 has a snapshot that replaces [ep-1, ep-2, ep-3]
const snapshotEp = createDummyEpisode('snap-1', 'SYSTEM_EVENT', []);
ep3.variants = {
snapshot: {
type: 'snapshot',
status: 'ready',
episode: snapshotEp,
replacedEpisodeIds: ['ep-1', 'ep-2', 'ep-3'],
},
};
// We only retain 5 tokens, forcing the sweep to use variants for EVERYTHING except ep4.
const view = generateWorkingBufferView(
[ep1, ep2, ep3, ep4],
5,
env.tracer,
env,
);
// Result should be exactly: [snapshot, ep-4]
expect(view).toHaveLength(2);
expect(view[0].id).toBe('snap-1');
expect(view[1].id).toBe('ep-4');
});
it('ignores variants that are not yet "ready"', () => {
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
{ type: 'text', text: '1' },
]);
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
{ type: 'text', text: '2' },
]);
ep1.variants = {
masked: {
type: 'masked',
status: 'computing',
text: '<MASKED>',
recoveredTokens: 10,
},
};
const view = generateWorkingBufferView([ep1, ep2], 10, env.tracer, env);
// Because the variant was computing, it must fall back to the raw pristine text.
expect(view).toHaveLength(2);
expect(
(view[0].trigger as UserPrompt).semanticParts[0].presentation,
).toBeUndefined();
});
});
+174
View File
@@ -0,0 +1,174 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Episode } from './types.js';
import type { ContextTracer } from '../tracer.js';
import { debugLogger } from '../../utils/debugLogger.js';
import type { ContextEnvironment } from '../sidecar/environment.js';
/**
* 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.
*/
export function generateWorkingBufferView(
pristineEpisodes: Episode[],
retainedTokens: number,
tracer: ContextTracer,
env: ContextEnvironment,
): Episode[] {
const currentEpisodes: Episode[] = [];
let rollingTokens = 0;
const skippedIds = new Set<string>();
tracer.logEvent('ViewGenerator', 'Generating Working Buffer View');
for (let i = pristineEpisodes.length - 1; i >= 0; i--) {
const ep = pristineEpisodes[i];
// If this episode was already replaced by an N-to-1 Snapshot injected earlier in the sweep, skip it entirely!
if (skippedIds.has(ep.id)) {
tracer.logEvent(
'ViewGenerator',
`Skipping episode [${ep.id}] due to N-to-1 replacement.`,
);
continue;
}
let projectedTrigger: typeof ep.trigger;
if (ep.trigger.type === 'USER_PROMPT') {
projectedTrigger = {
...ep.trigger,
metadata: {
...ep.trigger.metadata,
transformations: [...(ep.trigger.metadata?.transformations || [])],
},
semanticParts: ep.trigger.semanticParts.map((sp) => ({ ...sp })),
};
} else {
projectedTrigger = {
...ep.trigger,
metadata: {
...ep.trigger.metadata,
transformations: [...(ep.trigger.metadata?.transformations || [])],
},
};
}
let projectedEp: Episode = {
...ep,
trigger: projectedTrigger,
steps: ep.steps.map((step) => ({
...step,
metadata: {
...step.metadata,
transformations: [...(step.metadata?.transformations || [])],
},
})),
yield: ep.yield
? {
...ep.yield,
metadata: {
...ep.yield.metadata,
transformations: [...(ep.yield.metadata?.transformations || [])],
},
}
: undefined,
};
const epTokens = env.tokenCalculator.calculateEpisodeListTokens([
projectedEp,
]);
if (rollingTokens > retainedTokens && ep.variants) {
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;
// Mark all the episodes this snapshot covers to be skipped by the backwards sweep.
for (const id of snapshot.replacedEpisodeIds) {
skippedIds.add(id);
}
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 typeof projectedEp.steps;
projectedEp.yield = undefined;
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 &&
projectedEp.trigger.semanticParts.length > 0
) {
projectedEp.trigger.semanticParts[0].presentation = {
text: masked.text,
tokens: masked.recoveredTokens || 10,
};
}
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 += env.tokenCalculator.calculateEpisodeListTokens([
projectedEp,
]);
}
return currentEpisodes;
}
+271
View File
@@ -0,0 +1,271 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { IrMapper } from './mapper.js';
import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
import type { Content } from '@google/genai';
import type { UserPrompt, ToolExecution, AgentThought } 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 tokenCalculator = new ContextTokenCalculator(4);
const episodes = IrMapper.toIr(rawHistory, tokenCalculator);
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 correctly handle multi-tool-calls grouped within a single turn without dropping observations', () => {
const rawHistory: Content[] = [
{
role: 'user',
parts: [{ text: 'Examine both of these tools please.' }],
},
{
role: 'model',
parts: [
{ text: 'I will call them concurrently.' },
{
functionCall: {
id: 'c1',
name: 'tool_one',
args: { p: 1 },
},
},
{
functionCall: {
id: 'c2',
name: 'tool_two',
args: { p: 2 },
},
},
],
},
// Gemini forces the user turn to contain ALL function responses for that model turn
{
role: 'user',
parts: [
{
functionResponse: {
id: 'c1',
name: 'tool_one',
response: { r: 1 },
},
},
{
functionResponse: {
id: 'c2',
name: 'tool_two',
response: { r: 2 },
},
},
],
},
{
role: 'model',
parts: [{ text: 'Both complete.' }],
},
];
const tokenCalculator = new ContextTokenCalculator(4);
const episodes = IrMapper.toIr(rawHistory, tokenCalculator);
// It should collapse into a single episode
expect(episodes).toHaveLength(1);
const ep = episodes[0];
expect(ep.trigger.type).toBe('USER_PROMPT');
// The steps array should contain:
// 0: AgentThought ("I will call them concurrently")
// 1: ToolExecution(tool_one)
// 2: ToolExecution(tool_two)
expect(ep.steps).toHaveLength(3);
expect(ep.steps[0].type).toBe('AGENT_THOUGHT');
expect((ep.steps[0] as AgentThought).text).toBe(
'I will call them concurrently.',
);
expect(ep.steps[1].type).toBe('TOOL_EXECUTION');
expect((ep.steps[1] as ToolExecution).toolName).toBe('tool_one');
expect((ep.steps[1] as ToolExecution).intent).toEqual({ p: 1 });
expect((ep.steps[1] as ToolExecution).observation).toEqual({ r: 1 });
expect(ep.steps[2].type).toBe('TOOL_EXECUTION');
expect((ep.steps[2] as ToolExecution).toolName).toBe('tool_two');
expect((ep.steps[2] as ToolExecution).intent).toEqual({ p: 2 });
expect((ep.steps[2] as ToolExecution).observation).toEqual({ r: 2 });
// The final model turn should become the yield
expect(ep.yield).toBeDefined();
expect(ep.yield?.type).toBe('AGENT_YIELD');
expect(ep.yield?.text).toBe('Both complete.');
// Now verify we can reconstitute it without dropping the multiple calls
const reconstituted = IrMapper.fromIr(episodes);
// The reconstituted history should have exactly 4 turns, same as original
expect(reconstituted).toHaveLength(4);
// Check that the Model turn has both function calls
expect(reconstituted[1].role).toBe('model');
expect(reconstituted[1].parts).toHaveLength(3); // text + call1 + call2
expect(reconstituted[1].parts![1].functionCall?.name).toBe('tool_one');
expect(reconstituted[1].parts![2].functionCall?.name).toBe('tool_two');
// Check that the User turn has both function responses
expect(reconstituted[2].role).toBe('user');
expect(reconstituted[2].parts).toHaveLength(2); // response1 + response2
expect(reconstituted[2].parts![0].functionResponse?.name).toBe('tool_one');
expect(reconstituted[2].parts![1].functionResponse?.name).toBe('tool_two');
});
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 tokenCalculator = new ContextTokenCalculator(4);
const initialIr = IrMapper.toIr(history, tokenCalculator);
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, tokenCalculator);
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);
});
});
+31
View File
@@ -0,0 +1,31 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content } from '@google/genai';
import type { Episode } from './types.js';
import { toIr } from './toIr.js';
import { fromIr } from './fromIr.js';
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
export class IrMapper {
/**
* 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[],
tokenCalculator: ContextTokenCalculator,
): Episode[] {
return toIr(history, tokenCalculator);
}
/**
* Re-serializes the Episodic IR back into a flat Gemini Content[] array.
*/
static fromIr(episodes: Episode[]): Content[] {
return fromIr(episodes);
}
}
+92
View File
@@ -0,0 +1,92 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content } from '@google/genai';
import { IrMapper } from './mapper.js';
import type { Episode } from './types.js';
import { debugLogger } from '../../utils/debugLogger.js';
import type {
ContextEnvironment,
ContextTracer,
} from '../sidecar/environment.js';
import type { PipelineOrchestrator } from '../sidecar/orchestrator.js';
import type { SidecarConfig } from '../sidecar/types.js';
export class IrProjector {
/**
* Orchestrates the final projection: takes a working buffer view,
* applies the Immediate Sanitization pipeline, and enforces token boundaries.
*/
static async project(
workingBuffer: Episode[],
orchestrator: PipelineOrchestrator,
sidecar: SidecarConfig,
tracer: ContextTracer,
env: ContextEnvironment,
protectedIds: Set<string>,
): Promise<Content[]> {
if (!sidecar.budget) {
const contents = IrMapper.fromIr(workingBuffer);
tracer.logEvent('IrProjector', 'Projected Context to LLM (No Budget)', {
projectedContext: contents,
});
return contents;
}
const maxTokens = sidecar.budget.maxTokens;
const currentTokens =
env.tokenCalculator.calculateEpisodeListTokens(workingBuffer);
if (currentTokens <= maxTokens) {
tracer.logEvent(
'IrProjector',
`View is within maxTokens (${currentTokens} <= ${maxTokens}). Returning view.`,
);
const contents = IrMapper.fromIr(workingBuffer);
tracer.logEvent('IrProjector', 'Projected Context to LLM', {
projectedContext: contents,
});
return contents;
}
tracer.logEvent(
'IrProjector',
`View exceeds maxTokens (${currentTokens} > ${maxTokens}). Hitting Synchronous Pressure Barrier.`,
);
debugLogger.log(
`Context Manager Synchronous Barrier triggered: View at ${currentTokens} tokens (limit: ${maxTokens}).`,
);
const processedEpisodes = await orchestrator.executePipeline(
'Immediate Sanitization',
workingBuffer,
{
currentTokens,
maxTokens: sidecar.budget.maxTokens,
retainedTokens: sidecar.budget.retainedTokens,
deficitTokens: Math.max(0, currentTokens - sidecar.budget.maxTokens),
protectedEpisodeIds: protectedIds,
isBudgetSatisfied: currentTokens <= sidecar.budget.maxTokens,
},
);
const finalTokens =
env.tokenCalculator.calculateEpisodeListTokens(processedEpisodes);
tracer.logEvent(
'IrProjector',
`Finished projection. Final token count: ${finalTokens}.`,
);
debugLogger.log(
`Context Manager finished. Final actual token count: ${finalTokens}.`,
);
const contents = IrMapper.fromIr(processedEpisodes);
tracer.logEvent('IrProjector', 'Projected Sanitized Context to LLM', {
projectedContextSanitized: contents,
});
return contents;
}
}
+261
View File
@@ -0,0 +1,261 @@
/**
* @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,
SystemEvent,
} from './types.js';
import type { ContextTokenCalculator } 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>();
export function getStableId(obj: object): string {
let id = nodeIdentityMap.get(obj);
if (!id) {
id = randomUUID();
nodeIdentityMap.set(obj, id);
}
return id;
}
function isRecord(v: unknown): v is Record<string, unknown> {
return typeof v === 'object' && v !== null && !Array.isArray(v);
}
function isCompleteEpisode(ep: Partial<Episode>): ep is Episode {
return (
typeof ep.id === 'string' &&
typeof ep.timestamp === 'number' &&
!!ep.trigger &&
Array.isArray(ep.steps)
);
}
export function toIr(
history: readonly Content[],
tokenCalculator: ContextTokenCalculator,
): Episode[] {
const episodes: Episode[] = [];
let currentEpisode: Partial<Episode> | null = null;
const pendingCallParts: Map<string, Part> = new Map();
const createMetadata = (parts: Part[]): IrMetadata => {
const tokens = tokenCalculator.estimateTokensForParts(parts, 0);
return {
originalTokens: tokens,
currentTokens: tokens,
transformations: [],
};
};
const finalizeEpisode = () => {
if (currentEpisode && isCompleteEpisode(currentEpisode)) {
episodes.push(currentEpisode);
}
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) {
currentEpisode = parseToolResponses(
msg,
currentEpisode,
pendingCallParts,
tokenCalculator,
createMetadata,
);
}
if (hasUserParts) {
finalizeEpisode();
currentEpisode = parseUserParts(msg, createMetadata);
}
} else if (msg.role === 'model') {
currentEpisode = parseModelParts(
msg,
currentEpisode,
pendingCallParts,
createMetadata,
);
}
}
if (currentEpisode) {
finalizeYield(currentEpisode);
finalizeEpisode();
}
return episodes;
}
function parseToolResponses(
msg: Content,
currentEpisode: Partial<Episode> | null,
pendingCallParts: Map<string, Part>,
tokenCalculator: ContextTokenCalculator,
createMetadata: (parts: Part[]) => IrMetadata,
): Partial<Episode> {
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([]),
} as SystemEvent,
steps: [],
};
}
for (const part of msg.parts!) {
if (part.functionResponse) {
const callId = part.functionResponse.id || '';
const matchingCall = pendingCallParts.get(callId);
const intentTokens = matchingCall
? tokenCalculator.estimateTokensForParts([matchingCall])
: 0;
const obsTokens = tokenCalculator.estimateTokensForParts([part]);
const step: ToolExecution = {
id: getStableId(part),
type: 'TOOL_EXECUTION',
toolName: part.functionResponse.name || 'unknown',
intent: isRecord(matchingCall?.functionCall?.args)
? matchingCall.functionCall.args
: {},
observation: isRecord(part.functionResponse.response)
? part.functionResponse.response
: {},
tokens: {
intent: intentTokens,
observation: obsTokens,
},
metadata: {
originalTokens: intentTokens + obsTokens,
currentTokens: intentTokens + obsTokens,
transformations: [],
},
};
currentEpisode.steps!.push(step);
if (callId) pendingCallParts.delete(callId);
}
}
return currentEpisode;
}
function parseUserParts(
msg: Content,
createMetadata: (parts: Part[]) => IrMetadata,
): Partial<Episode> {
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)),
};
return {
id: getStableId(msg),
timestamp: Date.now(),
trigger,
steps: [],
};
}
function parseModelParts(
msg: Content,
currentEpisode: Partial<Episode> | null,
pendingCallParts: Map<string, Part>,
createMetadata: (parts: Part[]) => IrMetadata,
): Partial<Episode> {
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([]),
} as SystemEvent,
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);
}
}
return currentEpisode;
}
function finalizeYield(currentEpisode: Partial<Episode>) {
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;
}
}
}
+204
View File
@@ -0,0 +1,204 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Part } from '@google/genai';
/**
* 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: Part;
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;
}
+45
View File
@@ -0,0 +1,45 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { EpisodeEditor } from './ir/episodeEditor.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 via the provided EpisodeEditor, based on the current accounting state.
* Processors should safely mutate or replace episodes using the editor's API.
*/
process(editor: EpisodeEditor, state: ContextAccountingState): Promise<void>;
}
@@ -0,0 +1,97 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import {
createMockEnvironment,
createDummyState,
createDummyEpisode,
} from '../testing/contextTestUtils.js';
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { BlobDegradationProcessor } from './blobDegradationProcessor.js';
import { EpisodeEditor } from '../ir/episodeEditor.js';
import type { UserPrompt } from '../ir/types.js';
import type { ContextEnvironment } from '../sidecar/environment.js';
import type { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
describe('BlobDegradationProcessor', () => {
let processor: BlobDegradationProcessor;
let env: ContextEnvironment;
let fileSystem: InMemoryFileSystem;
beforeEach(() => {
vi.resetAllMocks();
env = createMockEnvironment();
fileSystem = env.fileSystem as InMemoryFileSystem;
processor = new BlobDegradationProcessor(env);
});
it('degrades inline_data into a text reference and saves to disk', async () => {
const dummyImageBase64 = Buffer.from('fake-image-data').toString('base64');
const ep = createDummyEpisode('ep-1', 'USER_PROMPT', [
{ type: 'text', text: 'Look at this image:' },
{
type: 'inline_data',
mimeType: 'image/png',
data: dummyImageBase64,
},
]);
const state = createDummyState(false, 500);
const editor = new EpisodeEditor([ep]);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
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',
);
// Verify it was written to fake FS
expect(fileSystem.getFiles().size).toBeGreaterThan(0);
const files = Array.from(fileSystem.getFiles().keys());
expect(files[0]).toContain(
'.gemini/tool-outputs/degraded-blobs/session-mock-session/blob_',
);
expect(result[0].trigger.metadata.transformations.length).toBe(1);
});
it('degrades file_data into a text reference without disk write', async () => {
const ep = createDummyEpisode('ep-2', 'USER_PROMPT', [
{
type: 'file_data',
mimeType: 'application/pdf',
fileUri: 'gs://fake-bucket/doc.pdf',
},
]);
const state = createDummyState(false, 500);
const editor = new EpisodeEditor([ep]);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
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(fileSystem.getFiles().size).toBe(0);
});
});
@@ -0,0 +1,142 @@
/**
* @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 { sanitizeFilenamePart } from '../../utils/fileUtils.js';
import type { EpisodeEditor } from '../ir/episodeEditor.js';
export type BlobDegradationProcessorOptions = Record<string, never>;
export class BlobDegradationProcessor implements ContextProcessor {
static create(
env: ContextEnvironment,
_options: BlobDegradationProcessorOptions,
): BlobDegradationProcessor {
return new BlobDegradationProcessor(env);
}
readonly id = 'BlobDegradationProcessor';
readonly name = 'BlobDegradationProcessor';
readonly options = {};
private env: ContextEnvironment;
constructor(env: ContextEnvironment) {
this.env = env;
}
async process(
editor: EpisodeEditor,
state: ContextAccountingState,
): Promise<void> {
if (state.isBudgetSatisfied) {
return;
}
let currentDeficit = state.deficitTokens;
let directoryCreated = false;
let blobOutputsDir = this.env.fileSystem.join(
this.env.projectTempDir,
'degraded-blobs',
);
const sessionId = this.env.sessionId;
if (sessionId) {
blobOutputsDir = this.env.fileSystem.join(
blobOutputsDir,
`session-${sanitizeFilenamePart(sessionId)}`,
);
}
const ensureDir = async () => {
if (!directoryCreated) {
await this.env.fileSystem.mkdir(blobOutputsDir, { recursive: true });
directoryCreated = true;
}
};
// Forward scan, looking for bloated non-text parts to degrade
for (const ep of editor.episodes) {
if (currentDeficit <= 0) break;
if (state.protectedEpisodeIds.has(ep.id)) continue;
if (ep.trigger.type === 'USER_PROMPT') {
for (let j = 0; j < ep.trigger.semanticParts.length; j++) {
const part = ep.trigger.semanticParts[j];
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()}_${this.env.idGenerator.generateId()}.${ext}`;
const filePath = this.env.fileSystem.join(blobOutputsDir, fileName);
// Base64 to buffer
const buffer = Buffer.from(part.data, 'base64');
await this.env.fileSystem.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 = this.env.tokenCalculator.estimateTokensForParts([
{ inlineData: { mimeType: part.mimeType, data: part.data } },
]);
const newTokens = this.env.tokenCalculator.estimateTokensForParts([
{ 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 = this.env.tokenCalculator.estimateTokensForParts([
{ fileData: { mimeType: part.mimeType, fileUri: part.fileUri } },
]);
const newTokens = this.env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
tokensSaved = oldTokens - newTokens;
} else if (part.type === 'raw_part') {
newText = `[Unknown Part degraded to text to preserve context window.]`;
const oldTokens = this.env.tokenCalculator.estimateTokensForParts([
part.part,
]);
const newTokens = this.env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
tokensSaved = oldTokens - newTokens;
}
if (newText && tokensSaved > 0) {
const newTokens = this.env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
editor.editEpisode(ep.id, 'DEGRADE_BLOB', (draft) => {
if (draft.trigger.type === 'USER_PROMPT') {
draft.trigger.semanticParts[j].presentation = {
text: newText,
tokens: newTokens,
};
draft.trigger.metadata.transformations.push({
processorName: this.name,
action: 'DEGRADED',
timestamp: Date.now(),
});
}
});
currentDeficit -= tokensSaved;
}
}
}
}
}
}
@@ -0,0 +1,129 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import {
createMockEnvironment,
createDummyState,
createDummyEpisode,
} from '../testing/contextTestUtils.js';
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { EmergencyTruncationProcessor } from './emergencyTruncationProcessor.js';
import { EpisodeEditor } from '../ir/episodeEditor.js';
import type { ContextEnvironment } from '../sidecar/environment.js';
describe('EmergencyTruncationProcessor', () => {
let processor: EmergencyTruncationProcessor;
let env: ContextEnvironment;
beforeEach(() => {
vi.resetAllMocks();
env = createMockEnvironment();
// Force token calculator to return exactly what we tell it for deterministic testing
vi.spyOn(
env.tokenCalculator,
'calculateEpisodeListTokens',
).mockImplementation((episodes) =>
// Just sum up the metadata originalTokens for our dummy episodes
episodes.reduce(
(acc, ep) => acc + (ep.trigger.metadata.originalTokens || 100),
0,
),
);
processor = new EmergencyTruncationProcessor(env, {});
});
it('bypasses processing if currentTokens <= maxTokens', async () => {
const episodes = [
createDummyEpisode('ep-1', 'USER_PROMPT', [
{ type: 'text', text: 'short' },
]),
];
// State says we are under budget (5000 < 10000)
const state = createDummyState(true, 0, new Set(), 5000, 10000);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
expect(result).toStrictEqual(episodes);
expect(result.length).toBe(1);
});
it('truncates episodes from the front (oldest) until targetTokens is met', async () => {
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
{ type: 'text', text: 'oldest' },
]);
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
{ type: 'text', text: 'middle' },
]);
const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', [
{ type: 'text', text: 'newest' },
]);
// Each is worth 100 tokens according to our mock
const episodes = [ep1, ep2, ep3];
// We have 300 tokens, but max is 200. We need to drop 100 tokens.
const state = createDummyState(false, 100, new Set(), 300, 200);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
// It should drop the FIRST episode (ep-1) and keep the rest.
expect(result.length).toBe(2);
expect(result[0].id).toBe('ep-2');
expect(result[1].id).toBe('ep-3');
});
it('never drops protected episodes (e.g. system instructions)', async () => {
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
{ type: 'text', text: 'protected system prompt' },
]);
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
{ type: 'text', text: 'middle' },
]);
const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', [
{ type: 'text', text: 'newest' },
]);
const episodes = [ep1, ep2, ep3];
// We have 300 tokens, max is 200. We need to drop 100 tokens.
// However, ep-1 is protected!
const state = createDummyState(false, 100, new Set(['ep-1']), 300, 200);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
// It should SKIP dropping ep-1 (protected) and drop ep-2 instead.
expect(result.length).toBe(2);
expect(result[0].id).toBe('ep-1'); // Protected, survived
expect(result[1].id).toBe('ep-3'); // Survivor
});
it('can drop multiple episodes if deficit is huge', async () => {
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', []);
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', []);
const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', []);
const episodes = [ep1, ep2, ep3];
// We have 300 tokens, max is 50. We need to drop 250 tokens!
const state = createDummyState(false, 250, new Set(), 300, 50);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
// It must drop ep1 (100t) and ep2 (100t).
// Remaining is ep3 (100t).
// Wait, if it drops ep1 (remaining=200) and ep2 (remaining=100),
// when it looks at ep3, remaining (100) > max (50), so it will drop ep3 too!
expect(result.length).toBe(0);
});
});
@@ -0,0 +1,60 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ContextProcessor, ContextAccountingState } from '../pipeline.js';
import type { ContextEnvironment } from '../sidecar/environment.js';
import type { EpisodeEditor } from '../ir/episodeEditor.js';
export type EmergencyTruncationProcessorOptions = Record<string, never>;
export class EmergencyTruncationProcessor implements ContextProcessor {
static create(
env: ContextEnvironment,
options: EmergencyTruncationProcessorOptions,
): EmergencyTruncationProcessor {
return new EmergencyTruncationProcessor(env, options);
}
readonly id = 'EmergencyTruncationProcessor';
readonly name = 'EmergencyTruncationProcessor';
readonly options: EmergencyTruncationProcessorOptions;
constructor(
private readonly _env: ContextEnvironment,
options: EmergencyTruncationProcessorOptions,
) {
this.options = options;
}
async process(
editor: EpisodeEditor,
state: ContextAccountingState,
): Promise<void> {
if (state.currentTokens <= state.maxTokens) return;
let remainingTokens = state.currentTokens;
const targetTokens = state.maxTokens;
const toRemove: string[] = [];
// We respect the global protected Episode IDs (like the system prompt at index 0)
for (const ep of editor.episodes) {
const epTokens = this._env.tokenCalculator.calculateEpisodeListTokens([
ep,
]);
if (
remainingTokens > targetTokens &&
!state.protectedEpisodeIds.has(ep.id)
) {
remainingTokens -= epTokens;
toRemove.push(ep.id);
}
}
if (toRemove.length > 0) {
editor.removeEpisodes(toRemove, 'TRUNCATED');
}
}
}
@@ -0,0 +1,157 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import {
createMockEnvironment,
createDummyState,
createDummyEpisode,
} from '../testing/contextTestUtils.js';
import { describe, it, expect, beforeEach } from 'vitest';
import { HistorySquashingProcessor } from './historySquashingProcessor.js';
import { EpisodeEditor } from '../ir/episodeEditor.js';
import type { UserPrompt, AgentThought, AgentYield } from '../ir/types.js';
import { randomUUID } from 'node:crypto';
describe('HistorySquashingProcessor', () => {
let processor: HistorySquashingProcessor;
beforeEach(() => {
processor = new HistorySquashingProcessor(createMockEnvironment(), {
maxTokensPerNode: 100,
});
});
const createThoughtEpisode = (
id: string,
userText: string,
modelThought: string,
) => {
const ep = createDummyEpisode(id, 'USER_PROMPT', [
{ type: 'text', text: userText },
]);
// Replace the tool steps with a thought step for this test
ep.steps = [
{
id: randomUUID(),
type: 'AGENT_THOUGHT',
text: modelThought,
metadata: {
originalTokens: 1000,
currentTokens: 1000,
transformations: [],
},
},
];
return ep;
};
it('bypasses processing if budget is satisfied', async () => {
const episodes = [createThoughtEpisode('1', 'short text', 'short thought')];
const state = createDummyState(true);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
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 = [createThoughtEpisode('ep-1', longText, 'short thought')];
const state = createDummyState(false, 100, new Set(['ep-1']));
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
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 = [createThoughtEpisode('ep-2', longUser, longModel)];
const state = createDummyState(false, 500); // High deficit, force truncation
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
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 = [
createThoughtEpisode('ep-3', longUser1, 'short'),
createThoughtEpisode('ep-4', longUser2, 'short'),
];
// Set deficit to exactly what ONE truncation will save
// Original = ~250 tokens. Limit = 100. Truncation saves ~150 tokens.
const state = createDummyState(false, 150);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
// 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 = createThoughtEpisode('ep-5', 'short', 'short');
ep.yield = {
id: randomUUID(),
type: 'AGENT_YIELD',
text: longYield,
metadata: {
originalTokens: 250,
currentTokens: 250,
transformations: [],
},
};
const state = createDummyState(false, 500);
const editor = new EpisodeEditor([ep]);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
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,189 @@
/**
* @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 { truncateProportionally } from '../truncation.js';
import type { EpisodeEditor } from '../ir/episodeEditor.js';
export interface HistorySquashingProcessorOptions {
maxTokensPerNode: number;
}
export class HistorySquashingProcessor implements ContextProcessor {
static create(
env: ContextEnvironment,
options: HistorySquashingProcessorOptions,
): HistorySquashingProcessor {
return new HistorySquashingProcessor(env, options);
}
static readonly schema = {
type: 'object',
properties: {
maxTokensPerNode: {
type: 'number',
description:
'The maximum tokens a node can have before being truncated.',
},
},
required: ['maxTokensPerNode'],
};
readonly id = 'HistorySquashingProcessor';
readonly name = 'HistorySquashingProcessor';
readonly options: HistorySquashingProcessorOptions;
constructor(
env: ContextEnvironment,
options: HistorySquashingProcessorOptions,
) {
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(
editor: EpisodeEditor,
state: ContextAccountingState,
): Promise<void> {
if (state.isBudgetSatisfied) {
return;
}
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;
for (const ep of editor.episodes) {
if (currentDeficit <= 0) break;
if (state.protectedEpisodeIds.has(ep.id)) continue;
// 1. Squash User Prompts
if (ep.trigger.type === 'USER_PROMPT') {
for (let j = 0; j < ep.trigger.semanticParts.length; j++) {
const part = ep.trigger.semanticParts[j];
if (part.type === 'text') {
const saved = this.tryApplySquash(
part.text,
limitChars,
currentDeficit,
(p) => {
editor.editEpisode(ep.id, 'SQUASH_PROMPT', (draft) => {
if (draft.trigger.type === 'USER_PROMPT') {
draft.trigger.semanticParts[j].presentation = p;
}
});
},
() => {
editor.editEpisode(ep.id, 'SQUASH_PROMPT', (draft) => {
draft.trigger.metadata.transformations.push({
processorName: this.name,
action: 'TRUNCATED',
timestamp: Date.now(),
});
});
},
);
currentDeficit -= saved;
}
}
}
// 2. Squash Model Thoughts
if (ep.steps) {
for (let j = 0; j < ep.steps.length; j++) {
const step = ep.steps[j];
if (currentDeficit <= 0) break;
if (step.type === 'AGENT_THOUGHT') {
const saved = this.tryApplySquash(
step.text,
limitChars,
currentDeficit,
(p) => {
editor.editEpisode(ep.id, 'SQUASH_THOUGHT', (draft) => {
const draftStep = draft.steps[j];
if (draftStep.type === 'AGENT_THOUGHT') {
draftStep.presentation = p;
}
});
},
() => {
editor.editEpisode(ep.id, 'SQUASH_THOUGHT', (draft) => {
const draftStep = draft.steps[j];
if (draftStep.type === 'AGENT_THOUGHT') {
draftStep.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) => {
editor.editEpisode(ep.id, 'SQUASH_YIELD', (draft) => {
if (draft.yield) draft.yield.presentation = p;
});
},
() => {
editor.editEpisode(ep.id, 'SQUASH_YIELD', (draft) => {
if (draft.yield) {
draft.yield.metadata.transformations.push({
processorName: this.name,
action: 'TRUNCATED',
timestamp: Date.now(),
});
}
});
},
);
currentDeficit -= saved;
}
}
}
}
@@ -0,0 +1,164 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import {
createMockEnvironment,
createDummyState,
createDummyEpisode,
} from '../testing/contextTestUtils.js';
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { SemanticCompressionProcessor } from './semanticCompressionProcessor.js';
import { EpisodeEditor } from '../ir/episodeEditor.js';
import type { UserPrompt, ToolExecution, AgentThought } from '../ir/types.js';
import { randomUUID } from 'node:crypto';
import type { BaseLlmClient } from 'src/core/baseLlmClient.js';
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();
// Re-mock llmClient properly
vi.spyOn(env, 'llmClient', 'get').mockReturnValue({
generateContent: generateContentMock,
} as unknown as BaseLlmClient);
processor = new SemanticCompressionProcessor(env, {
nodeThresholdTokens: 2000,
});
});
const createEpisodeWithThoughtsAndTools = (
id: string,
userText: string,
thoughtText: string,
toolObs: string,
) => {
const ep = createDummyEpisode(id, 'USER_PROMPT', [
{ type: 'text', text: userText },
]);
// We override metadata for threshold triggering
ep.trigger.metadata.currentTokens = 3800;
ep.steps = [
{
id: randomUUID(),
type: 'AGENT_THOUGHT',
text: thoughtText,
metadata: {
originalTokens: 3800,
currentTokens: 3800,
transformations: [],
},
},
{
id: randomUUID(),
type: 'TOOL_EXECUTION',
toolName: 'test',
intent: {},
observation: toolObs,
tokens: { intent: 10, observation: 3800 },
metadata: {
originalTokens: 3810,
currentTokens: 3810,
transformations: [],
},
},
];
return ep;
};
it('bypasses processing if budget is satisfied', async () => {
const episodes = [
createEpisodeWithThoughtsAndTools('1', 'short', 'short', 'short'),
];
const state = createDummyState(true);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
expect(generateContentMock).not.toHaveBeenCalled();
});
it('skips protected episodes even if over budget', async () => {
const massiveStr = 'M'.repeat(15000);
const episodes = [
createEpisodeWithThoughtsAndTools(
'ep-1',
massiveStr,
massiveStr,
massiveStr,
),
];
const state = createDummyState(false, 1000, new Set(['ep-1']));
const editor = new EpisodeEditor(episodes);
await processor.process(editor, 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 = [
createEpisodeWithThoughtsAndTools(
'ep-1',
massiveStr,
massiveStr,
massiveStr,
),
];
const state = createDummyState(false, 50000); // Massive deficit, forces all 3 to summarize
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
expect(generateContentMock).toHaveBeenCalledTimes(3);
// Verify presentation layers were injected
const result = editor.getFinalEpisodes();
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 = [
createEpisodeWithThoughtsAndTools(
'ep-1',
massiveStr,
massiveStr,
massiveStr,
),
];
// Set deficit low enough that ONE summary solves the problem
const state = createDummyState(false, 5);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
// It should only compress the UserPrompt and then stop
expect(generateContentMock).toHaveBeenCalledTimes(1);
});
});
@@ -0,0 +1,261 @@
/**
* @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 { debugLogger } from '../../utils/debugLogger.js';
import { LlmRole } from '../../telemetry/types.js';
import { getResponseText } from '../../utils/partUtils.js';
import type { EpisodeEditor } from '../ir/episodeEditor.js';
export interface SemanticCompressionProcessorOptions {
nodeThresholdTokens: number;
}
export class SemanticCompressionProcessor implements ContextProcessor {
static create(
env: ContextEnvironment,
options: SemanticCompressionProcessorOptions,
): SemanticCompressionProcessor {
return new SemanticCompressionProcessor(env, options);
}
static readonly schema = {
type: 'object',
properties: {
nodeThresholdTokens: {
type: 'number',
description: 'The token threshold above which nodes are summarized.',
},
},
required: ['nodeThresholdTokens'],
};
readonly id = 'SemanticCompressionProcessor';
readonly name = 'SemanticCompressionProcessor';
readonly options: SemanticCompressionProcessorOptions;
private env: ContextEnvironment;
private modelToUse: string = 'chat-compression-2.5-flash-lite';
constructor(
env: ContextEnvironment,
options: SemanticCompressionProcessorOptions,
) {
this.env = env;
this.options = options;
}
async process(
editor: EpisodeEditor,
state: ContextAccountingState,
): Promise<void> {
// If the budget is satisfied, or semantic compression isn't enabled
if (state.isBudgetSatisfied) {
return;
}
const semanticConfig = this.options;
const limitTokens = semanticConfig.nodeThresholdTokens;
const thresholdChars = this.env.tokenCalculator.tokensToChars(limitTokens);
this.modelToUse = 'gemini-2.5-flash';
let currentDeficit = state.deficitTokens;
// We scan backwards (oldest to newest would also work, but older is safer to degrade first)
for (const ep of editor.episodes) {
if (currentDeficit <= 0) break;
if (state.protectedEpisodeIds.has(ep.id)) continue;
// 1. Compress User Prompts
if (ep.trigger.type === 'USER_PROMPT') {
for (let j = 0; j < ep.trigger.semanticParts.length; j++) {
const part = ep.trigger.semanticParts[j];
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 = this.env.tokenCalculator.estimateTokensForParts([
{ text: summary },
]);
const oldTokens = this.env.tokenCalculator.estimateTokensForParts([
{ text: part.text },
]);
if (newTokens < oldTokens) {
editor.editEpisode(ep.id, 'SUMMARIZE_PROMPT', (draft) => {
if (draft.trigger.type === 'USER_PROMPT') {
draft.trigger.semanticParts[j].presentation = {
text: summary,
tokens: newTokens,
};
draft.trigger.metadata.transformations.push({
processorName: this.name,
action: 'SUMMARIZED',
timestamp: Date.now(),
});
}
});
currentDeficit -= oldTokens - newTokens;
}
}
}
}
// 2. Compress Model Thoughts
if (ep.steps) {
for (let j = 0; j < ep.steps.length; j++) {
const step = ep.steps[j];
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 = this.env.tokenCalculator.estimateTokensForParts(
[{ text: summary }],
);
const oldTokens = this.env.tokenCalculator.estimateTokensForParts(
[{ text: step.text }],
);
if (newTokens < oldTokens) {
editor.editEpisode(ep.id, 'SUMMARIZE_THOUGHT', (draft) => {
const draftStep = draft.steps[j];
if (draftStep.type === 'AGENT_THOUGHT') {
draftStep.presentation = {
text: summary,
tokens: newTokens,
};
draftStep.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 {
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 =
this.env.tokenCalculator.estimateTokensForParts([
{
functionResponse: {
name: step.toolName,
response: newObsObject,
id: step.id,
},
},
]);
const oldObsTokens =
step.presentation?.tokens?.observation ??
step.tokens?.observation ??
step.tokens;
const intentTokens =
step.presentation?.tokens?.intent ?? step.tokens?.intent ?? 0;
if (newObsTokens < oldObsTokens) {
editor.editEpisode(ep.id, 'SUMMARIZE_TOOL', (draft) => {
const draftStep = draft.steps[j];
if (draftStep.type === 'TOOL_EXECUTION') {
draftStep.presentation = {
intent:
draftStep.presentation?.intent ?? draftStep.intent,
observation: newObsObject,
tokens: {
intent: intentTokens,
observation: newObsTokens,
},
};
if (!draftStep.metadata) {
draftStep.metadata = {
transformations: [],
currentTokens: 0,
originalTokens: 0,
};
}
if (!draftStep.metadata.transformations) {
draftStep.metadata.transformations = [];
}
draftStep.metadata.transformations.push({
processorName: this.name,
action: 'SUMMARIZED',
timestamp: Date.now(),
});
}
});
currentDeficit -= oldObsTokens - newObsTokens;
}
}
}
}
}
}
}
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.llmClient;
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: ${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,118 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import {
createMockEnvironment,
createDummyState,
createDummyEpisode,
} from '../testing/contextTestUtils.js';
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { StateSnapshotProcessor } from './stateSnapshotProcessor.js';
import { EpisodeEditor } from '../ir/episodeEditor.js';
import type { ContextEnvironment } from '../sidecar/environment.js';
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
describe('StateSnapshotProcessor', () => {
let processor: StateSnapshotProcessor;
let env: ContextEnvironment;
let generateContentMock: ReturnType<typeof vi.fn>;
beforeEach(() => {
vi.resetAllMocks();
env = createMockEnvironment();
generateContentMock = vi.fn().mockResolvedValue({
text: 'Mocked Compressed State Snapshot!',
});
vi.spyOn(env, 'llmClient', 'get').mockReturnValue({
generateContent: generateContentMock,
} as unknown as BaseLlmClient);
// Override token calc for testing
vi.spyOn(env.tokenCalculator, 'estimateTokensForParts').mockReturnValue(
100,
);
processor = new StateSnapshotProcessor(env, {}, env.eventBus);
});
it('bypasses processing if deficit is <= 0', async () => {
const episodes = [
createDummyEpisode('ep-1', 'USER_PROMPT', [
{ type: 'text', text: 'hello' },
]),
];
// current: 100, max: 1000, retained: 200 (deficit 0)
const state = createDummyState(false, 0, new Set(), 100, 1000, 200);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
expect(result).toStrictEqual(episodes);
expect(generateContentMock).not.toHaveBeenCalled();
});
it('bypasses processing if not enough episodes to summarize (needs at least 2 inner episodes)', async () => {
const episodes = [
createDummyEpisode('ep-sys', 'SYSTEM_EVENT', []),
createDummyEpisode('ep-active', 'USER_PROMPT', [
{ type: 'text', text: 'help' },
]),
];
// current: 1000, max: 10000, retained: 500. Target deficit = 500
const state = createDummyState(false, 500, new Set(), 1000, 10000, 500);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
expect(result).toStrictEqual(episodes);
expect(generateContentMock).not.toHaveBeenCalled();
});
it('summarizes intermediate episodes into a single snapshot episode', async () => {
const episodes = [
createDummyEpisode('ep-0', 'SYSTEM_EVENT', []),
createDummyEpisode('ep-1', 'USER_PROMPT', [
{ type: 'text', text: 'old 1' },
]),
createDummyEpisode('ep-2', 'USER_PROMPT', [
{ type: 'text', text: 'old 2' },
]),
createDummyEpisode('ep-3', 'USER_PROMPT', [
{ type: 'text', text: 'current' },
]),
];
// Target deficit = 200
const state = createDummyState(false, 200, new Set(), 1000, 10000, 800);
const editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
// We started with 4 episodes.
// Episodes [1, 2] were synthesized into a single new Snapshot episode.
// Final array should be: [0, SNAPSHOT, 3] = length 3.
expect(result.length).toBe(3);
expect(result[0].id).toBe('ep-0');
const snapshotEp = result[1];
expect(snapshotEp.yield).toBeDefined();
expect(snapshotEp.yield!.text).toContain('<CONTEXT_SNAPSHOT>');
expect(snapshotEp.yield!.text).toContain(
'Mocked Compressed State Snapshot!',
);
expect(result[2].id).toBe('ep-3');
expect(generateContentMock).toHaveBeenCalledTimes(1);
const llmArgs = generateContentMock.mock.calls[0][0];
expect(llmArgs.contents[0].parts[0].text).toContain('old 1');
expect(llmArgs.contents[0].parts[0].text).toContain('old 2');
});
});
@@ -0,0 +1,182 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ContextProcessor, ContextAccountingState } from '../pipeline.js';
import type { Episode } from '../ir/types.js';
import type {
ContextEnvironment,
ContextEventBus,
} from '../sidecar/environment.js';
import { v4 as uuidv4 } from 'uuid';
import { LlmRole } from '../../telemetry/llmRole.js';
import { debugLogger } from 'src/utils/debugLogger.js';
import type { EpisodeEditor } from '../ir/episodeEditor.js';
export interface StateSnapshotProcessorOptions {
model?: string;
systemInstruction?: string;
triggerDeficitTokens?: number;
}
export class StateSnapshotProcessor implements ContextProcessor {
static create(
env: ContextEnvironment,
options: StateSnapshotProcessorOptions,
): StateSnapshotProcessor {
return new StateSnapshotProcessor(env, options, env.eventBus);
}
readonly id = 'StateSnapshotProcessor';
readonly name = 'StateSnapshotProcessor';
readonly options: StateSnapshotProcessorOptions;
private readonly env: ContextEnvironment;
private isSynthesizing = false;
constructor(
env: ContextEnvironment,
options: StateSnapshotProcessorOptions,
_eventBus: ContextEventBus,
) {
this.env = env;
this.options = options;
}
async process(
editor: EpisodeEditor,
state: ContextAccountingState,
): Promise<void> {
const targetDeficit = Math.max(
0,
state.currentTokens - state.retainedTokens,
);
if (this.isSynthesizing || targetDeficit <= 0) return;
this.isSynthesizing = true;
try {
let deficitAccumulator = 0;
const selectedEpisodes: Episode[] = [];
for (let i = 1; i < editor.episodes.length - 1; i++) {
const ep = editor.episodes[i];
selectedEpisodes.push(ep);
let triggerText = '';
if (ep.trigger?.type === 'USER_PROMPT') {
const firstPart = ep.trigger.semanticParts?.[0];
if (firstPart) {
triggerText =
firstPart.type === 'text'
? firstPart.text
: (firstPart.presentation?.text ?? '');
}
}
deficitAccumulator += this.env.tokenCalculator.estimateTokensForParts([
{ text: triggerText },
{ text: ep.yield?.text ?? '' },
]);
if (deficitAccumulator >= targetDeficit) break;
}
if (selectedEpisodes.length < 2) return; // Not enough context to summarize
// Optimization: Do NOT emit VariantComputing, let the Orchestrator handle caching the final result.
const snapshotEp: Episode =
await this.synthesizeSnapshot(selectedEpisodes);
const oldIds = selectedEpisodes.map((ep) => ep.id);
editor.replaceEpisodes(oldIds, snapshotEp, 'STATE_SNAPSHOT');
} finally {
this.isSynthesizing = false;
}
}
private async synthesizeSnapshot(episodes: Episode[]): Promise<Episode> {
const client = this.env.llmClient;
const systemPrompt =
this.options.systemInstruction ??
`You are an expert Context Memory Manager. You will be provided with a raw transcript of older conversation turns between a user and an AI assistant.
Your task is to synthesize these turns into a single, dense, factual snapshot that preserves all critical context, preferences, active tasks, and factual knowledge, but discards conversational filler, pleasantries, and redundant back-and-forth iterations.
Output ONLY the raw factual snapshot, formatted compactly. Do not include markdown wrappers, prefixes like "Here is the snapshot", or conversational elements.`;
let userPromptText = 'TRANSCRIPT TO SNAPSHOT:\n\n';
for (const ep of episodes) {
if (ep.trigger?.type === 'USER_PROMPT') {
const partsText = ep.trigger.semanticParts
.map((p) => {
if (p.type === 'text') return p.text;
if (p.presentation) return p.presentation.text;
return '';
})
.join('');
userPromptText += `USER: ${partsText}\n`;
} else if (ep.trigger?.type === 'SYSTEM_EVENT') {
userPromptText += `[SYSTEM EVENT: ${ep.trigger.name}]\n`;
}
for (const step of ep.steps) {
if (step.type === 'TOOL_EXECUTION') {
userPromptText += `[Tool Called: ${step.toolName}]\n`;
}
}
if (ep.yield) {
userPromptText += `ASSISTANT: ${ep.yield.text}\n`;
}
userPromptText += '\n';
}
try {
const response = await client.generateContent({
modelConfigKey: { model: 'state-snapshot-processor' },
contents: [{ role: 'user', parts: [{ text: userPromptText }] }],
systemInstruction: { role: 'system', parts: [{ text: systemPrompt }] },
promptId: this.env.promptId,
role: LlmRole.UTILITY_STATE_SNAPSHOT_PROCESSOR,
abortSignal: new AbortController().signal,
});
const snapshotText = response.text;
// Synthesize a new "Episode" representing this compressed block
const newId = uuidv4();
const contentTokens = this.env.tokenCalculator.estimateTokensForParts([
{ text: snapshotText },
]);
return {
id: newId,
timestamp: Date.now(),
trigger: {
id: `${newId}-t`,
type: 'USER_PROMPT',
semanticParts: [],
metadata: {
originalTokens: 0,
currentTokens: 0,
transformations: [],
},
},
steps: [],
yield: {
id: `${newId}-y`,
type: 'AGENT_YIELD',
text: `<CONTEXT_SNAPSHOT>\n${snapshotText}\n</CONTEXT_SNAPSHOT>`,
metadata: {
originalTokens: contentTokens,
currentTokens: contentTokens,
transformations: [
{
processorName: 'StateSnapshotProcessor',
action: 'SYNTHESIZED',
timestamp: Date.now(),
},
],
},
},
};
} catch (error) {
debugLogger.error('Failed to synthesize snapshot:', error);
throw error;
}
}
}
@@ -0,0 +1,134 @@
/**
* @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 { EpisodeEditor } from '../ir/episodeEditor.js';
import type { Episode, ToolExecution } from '../ir/types.js';
import type { ContextAccountingState } from '../pipeline.js';
import { randomUUID } from 'node:crypto';
import type { ContextEnvironment } from '../sidecar/environment.js';
import type { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
describe('ToolMaskingProcessor', () => {
let processor: ToolMaskingProcessor;
let env: ContextEnvironment;
let fileSystem: InMemoryFileSystem;
beforeEach(() => {
vi.resetAllMocks();
env = createMockEnvironment();
fileSystem = env.fileSystem as InMemoryFileSystem;
processor = new ToolMaskingProcessor(env, {
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 editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
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 editor = new EpisodeEditor(episodes);
await processor.process(editor, state);
const result = editor.getFinalEpisodes();
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 to fake FS
expect(fileSystem.getFiles().size).toBe(2);
});
});
@@ -0,0 +1,319 @@
/**
* @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 { sanitizeFilenamePart } from '../../utils/fileUtils.js';
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 { EpisodeEditor } from '../ir/episodeEditor.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 interface ToolMaskingProcessorOptions {
stringLengthThresholdTokens: number;
}
type MaskableValue =
| string
| number
| boolean
| null
| MaskableValue[]
| { [key: string]: MaskableValue };
function isMaskableValue(val: unknown): val is MaskableValue {
if (
val === null ||
typeof val === 'string' ||
typeof val === 'number' ||
typeof val === 'boolean'
) {
return true;
}
if (Array.isArray(val)) {
return val.every(isMaskableValue);
}
if (typeof val === 'object') {
return Object.values(val).every(isMaskableValue);
}
return false;
}
function isMaskableRecord(val: unknown): val is Record<string, MaskableValue> {
return (
typeof val === 'object' &&
val !== null &&
!Array.isArray(val) &&
isMaskableValue(val)
);
}
export class ToolMaskingProcessor implements ContextProcessor {
static create(
env: ContextEnvironment,
options: ToolMaskingProcessorOptions,
): ToolMaskingProcessor {
return new ToolMaskingProcessor(env, options);
}
static readonly schema = {
type: 'object',
properties: {
stringLengthThresholdTokens: {
type: 'number',
description:
'The token threshold above which tool intents/observations are masked.',
},
},
required: ['stringLengthThresholdTokens'],
};
readonly id = 'ToolMaskingProcessor';
readonly name = 'ToolMaskingProcessor';
readonly options: ToolMaskingProcessorOptions;
private env: ContextEnvironment;
constructor(env: ContextEnvironment, options: ToolMaskingProcessorOptions) {
this.env = env;
this.options = options;
}
async process(
editor: EpisodeEditor,
state: ContextAccountingState,
): Promise<void> {
const maskingConfig = this.options;
if (!maskingConfig) return;
if (state.isBudgetSatisfied) return;
let currentDeficit = state.deficitTokens;
const limitChars = this.env.tokenCalculator.tokensToChars(
maskingConfig.stringLengthThresholdTokens,
);
let toolOutputsDir = this.env.fileSystem.join(
this.env.projectTempDir,
'tool-outputs',
);
const sessionId = this.env.sessionId;
if (sessionId) {
toolOutputsDir = this.env.fileSystem.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 this.env.fileSystem.mkdir(toolOutputsDir, { recursive: true });
directoryCreated = true;
}
const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${this.env.idGenerator.generateId()}.txt`;
const filePath = this.env.fileSystem.join(toolOutputsDir, fileName);
await this.env.fileSystem.writeFile(filePath, content);
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 (const ep of editor.episodes) {
if (currentDeficit <= 0) break;
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();
const maskAsync = async (
obj: MaskableValue,
nodeType: string,
): Promise<{ masked: MaskableValue; changed: boolean }> => {
if (typeof obj === 'string') {
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: MaskableValue[] = [];
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, MaskableValue> = {};
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 rawIntent = step.presentation?.intent ?? step.intent;
const rawObs = step.presentation?.observation ?? step.observation;
if (!isMaskableRecord(rawIntent)) {
throw new Error(
`ToolMaskingProcessor: step intent is not a valid JSON record. CallID: ${callId}`,
);
}
if (!isMaskableValue(rawObs)) {
throw new Error(
`ToolMaskingProcessor: step observation is not a valid JSON value. CallID: ${callId}`,
);
}
const intentRes = await maskAsync(rawIntent, 'intent');
const obsRes = await maskAsync(rawObs, 'observation');
if (intentRes.changed || obsRes.changed) {
const maskedIntent = isMaskableRecord(intentRes.masked)
? (intentRes.masked as Record<string, unknown>)
: undefined;
const maskedObs = isMaskableRecord(obsRes.masked)
? (obsRes.masked as Record<string, unknown>)
: undefined;
// Recalculate tokens perfectly
const newIntentTokens =
this.env.tokenCalculator.estimateTokensForParts([
{
functionCall: {
name: toolName,
args: maskedIntent,
id: callId,
},
},
]);
const newObsTokens = this.env.tokenCalculator.estimateTokensForParts([
{
functionResponse: {
name: toolName,
response:
typeof obsRes.masked === 'string'
? { message: obsRes.masked }
: maskedObs,
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) {
currentDeficit -= savings;
this.env.tracer.logEvent(
'ToolMaskingProcessor',
`Masked tool ${toolName}`,
{ recoveredTokens: savings },
);
editor.editEpisode(ep.id, 'MASK_TOOL', (draft) => {
const draftStep = draft.steps[j];
if (draftStep.type !== 'TOOL_EXECUTION') return;
if (!draftStep.presentation) {
draftStep.presentation = {
intent: draftStep.intent,
observation: draftStep.observation,
tokens: draftStep.tokens,
};
}
draftStep.presentation.intent = maskedIntent ?? {};
draftStep.presentation.observation =
typeof obsRes.masked === 'string'
? { message: obsRes.masked }
: (maskedObs ?? {});
draftStep.presentation.tokens = {
intent: newIntentTokens,
observation: newObsTokens,
};
draftStep.metadata = {
...draftStep.metadata,
transformations: [
...(draftStep.metadata?.transformations || []),
{
processorName: 'ToolMasking',
action: 'MASKED',
timestamp: Date.now(),
},
],
};
});
}
}
}
}
}
private isAlreadyMasked(content: string): boolean {
return content.includes('<tool_output_masked>');
}
}
@@ -0,0 +1,67 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { ProcessorRegistry } from './registry.js';
import { registerBuiltInProcessors } from './builtins.js';
import { describe, it, expect, beforeEach } from 'vitest';
import { SidecarLoader } from './SidecarLoader.js';
import { defaultSidecarProfile } from './profiles.js';
import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
import type { Config } from 'src/config/config.js';
describe('SidecarLoader (Fake FS)', () => {
let fileSystem: InMemoryFileSystem;
let registry: ProcessorRegistry;
beforeEach(() => {
fileSystem = new InMemoryFileSystem();
registry = new ProcessorRegistry();
registerBuiltInProcessors(registry);
});
const mockConfig = {
getExperimentalContextSidecarConfig: () => '/path/to/sidecar.json',
} as unknown as Config;
it('returns default profile if file does not exist', () => {
const result = SidecarLoader.fromConfig(mockConfig, registry, fileSystem);
expect(result).toBe(defaultSidecarProfile);
});
it('returns default profile if file exists but is 0 bytes', () => {
fileSystem.setFile('/path/to/sidecar.json', '');
const result = SidecarLoader.fromConfig(mockConfig, registry, fileSystem);
expect(result).toBe(defaultSidecarProfile);
});
it('throws an error if file is empty whitespace', () => {
fileSystem.setFile('/path/to/sidecar.json', ' \n ');
expect(() =>
SidecarLoader.fromConfig(mockConfig, registry, fileSystem),
).toThrow('is empty');
});
it('returns parsed config if file is valid', () => {
const validConfig = {
budget: { retainedTokens: 1000, maxTokens: 2000 },
gcBackstop: { strategy: 'truncate', target: 'max' },
pipelines: [],
};
fileSystem.setFile('/path/to/sidecar.json', JSON.stringify(validConfig));
const result = SidecarLoader.fromConfig(mockConfig, registry, fileSystem);
expect(result.budget.maxTokens).toBe(2000);
});
it('throws validation error if file is invalid', () => {
const invalidConfig = {
budget: { retainedTokens: 1000 }, // missing maxTokens
};
fileSystem.setFile('/path/to/sidecar.json', JSON.stringify(invalidConfig));
expect(() =>
SidecarLoader.fromConfig(mockConfig, registry, fileSystem),
).toThrow('Validation error:');
});
});
@@ -0,0 +1,87 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Config } from '../../config/config.js';
import type { SidecarConfig } from './types.js';
import { defaultSidecarProfile } from './profiles.js';
import { SchemaValidator } from '../../utils/schemaValidator.js';
import { getSidecarConfigSchema } from './schema.js';
import type { IFileSystem } from '../system/IFileSystem.js';
import { NodeFileSystem } from '../system/NodeFileSystem.js';
import type { ProcessorRegistry } from './registry.js';
export class SidecarLoader {
/**
* Loads and validates a sidecar config from a specific file path.
* Throws an error if the file cannot be read, parsed, or fails schema validation.
*/
static loadFromFile(
sidecarPath: string,
registry: ProcessorRegistry,
fileSystem: IFileSystem = new NodeFileSystem(),
): SidecarConfig {
const fileContent = fileSystem.readFileSync(sidecarPath, 'utf8');
if (!fileContent.trim()) {
throw new Error(`Sidecar configuration file at ${sidecarPath} is empty.`);
}
let parsed: unknown;
try {
parsed = JSON.parse(fileContent);
} catch (error) {
throw new Error(
`Failed to parse Sidecar configuration file at ${sidecarPath}: ${
error instanceof Error ? error.message : String(error)
}`,
);
}
const validationError = SchemaValidator.validate(
getSidecarConfigSchema(registry),
parsed,
);
if (validationError) {
throw new Error(
`Invalid sidecar configuration in ${sidecarPath}. Validation error: ${validationError}`,
);
}
// Schema has been validated.
const isSidecarConfig = (val: unknown): val is SidecarConfig => true;
if (isSidecarConfig(parsed)) {
return parsed;
}
throw new Error(
'Unreachable: schema validation passed but type predicate failed.',
);
}
/**
* Generates a Sidecar JSON graph from the experimental config file path or defaults.
* If a config file is present but invalid, this will THROW to prevent silent misconfiguration.
*/
static fromConfig(
config: Config,
registry: ProcessorRegistry,
fileSystem: IFileSystem = new NodeFileSystem(),
): SidecarConfig {
const sidecarPath = config.getExperimentalContextSidecarConfig();
if (sidecarPath && fileSystem.existsSync(sidecarPath)) {
const size = fileSystem.statSyncSize(sidecarPath);
// If the file exists but is completely empty (0 bytes), it's safe to fallback.
if (size === 0) {
return defaultSidecarProfile;
}
// If the file has content, enforce strict validation and throw on failure.
return this.loadFromFile(sidecarPath, registry, fileSystem);
}
return defaultSidecarProfile;
}
}
@@ -0,0 +1,127 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ProcessorRegistry } from './registry.js';
import {
ToolMaskingProcessor,
type ToolMaskingProcessorOptions,
} from '../processors/toolMaskingProcessor.js';
import { BlobDegradationProcessor } from '../processors/blobDegradationProcessor.js';
import {
SemanticCompressionProcessor,
type SemanticCompressionProcessorOptions,
} from '../processors/semanticCompressionProcessor.js';
import {
HistorySquashingProcessor,
type HistorySquashingProcessorOptions,
} from '../processors/historySquashingProcessor.js';
import {
StateSnapshotProcessor,
type StateSnapshotProcessorOptions,
} from '../processors/stateSnapshotProcessor.js';
import {
EmergencyTruncationProcessor,
type EmergencyTruncationProcessorOptions,
} from '../processors/emergencyTruncationProcessor.js';
export function registerBuiltInProcessors(registry: ProcessorRegistry) {
registry.register<ToolMaskingProcessorOptions>({
id: 'ToolMaskingProcessor',
schema: {
type: 'object',
properties: {
processorId: { const: 'ToolMaskingProcessor' },
options: {
type: 'object',
properties: { stringLengthThresholdTokens: { type: 'number' } },
required: ['stringLengthThresholdTokens'],
},
},
required: ['processorId', 'options'],
},
create: (env, opts) => new ToolMaskingProcessor(env, opts),
});
registry.register<Record<string, never>>({
id: 'BlobDegradationProcessor',
schema: {
type: 'object',
properties: {
processorId: { const: 'BlobDegradationProcessor' },
options: { type: 'object' },
},
required: ['processorId'],
},
create: (env) => new BlobDegradationProcessor(env),
});
registry.register<SemanticCompressionProcessorOptions>({
id: 'SemanticCompressionProcessor',
schema: {
type: 'object',
properties: {
processorId: { const: 'SemanticCompressionProcessor' },
options: {
type: 'object',
properties: { nodeThresholdTokens: { type: 'number' } },
required: ['nodeThresholdTokens'],
},
},
required: ['processorId', 'options'],
},
create: (env, opts) => new SemanticCompressionProcessor(env, opts),
});
registry.register<HistorySquashingProcessorOptions>({
id: 'HistorySquashingProcessor',
schema: {
type: 'object',
properties: {
processorId: { const: 'HistorySquashingProcessor' },
options: {
type: 'object',
properties: { maxTokensPerNode: { type: 'number' } },
required: ['maxTokensPerNode'],
},
},
required: ['processorId', 'options'],
},
create: (env, opts) => new HistorySquashingProcessor(env, opts),
});
registry.register<StateSnapshotProcessorOptions>({
id: 'StateSnapshotProcessor',
schema: {
type: 'object',
properties: {
processorId: { const: 'StateSnapshotProcessor' },
options: {
type: 'object',
properties: {
model: { type: 'string' },
systemInstruction: { type: 'string' },
triggerDeficitTokens: { type: 'number' },
},
},
},
required: ['processorId'],
},
create: (env, opts) => StateSnapshotProcessor.create(env, opts),
});
registry.register<EmergencyTruncationProcessorOptions>({
id: 'EmergencyTruncationProcessor',
schema: {
type: 'object',
properties: {
processorId: { const: 'EmergencyTruncationProcessor' },
options: { type: 'object' },
},
required: ['processorId'],
},
create: (env, opts) => EmergencyTruncationProcessor.create(env, opts),
});
}
@@ -0,0 +1,27 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
import type { ContextEventBus } from '../eventBus.js';
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
import type { ContextTracer } from '../tracer.js';
import type { IFileSystem } from '../system/IFileSystem.js';
import type { IIdGenerator } from '../system/IIdGenerator.js';
export type { ContextTracer, ContextEventBus };
export interface ContextEnvironment {
readonly llmClient: BaseLlmClient;
readonly promptId: string;
readonly sessionId: string;
readonly traceDir: string;
readonly projectTempDir: string;
readonly tracer: ContextTracer;
readonly charsPerToken: number;
readonly tokenCalculator: ContextTokenCalculator;
readonly fileSystem: IFileSystem;
readonly idGenerator: IIdGenerator;
readonly eventBus: ContextEventBus;
}
@@ -0,0 +1,38 @@
/**
* @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';
import type { ContextEventBus } from '../eventBus.js';
import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
import type { IFileSystem } from '../system/IFileSystem.js';
import { NodeFileSystem } from '../system/NodeFileSystem.js';
import type { IIdGenerator } from '../system/IIdGenerator.js';
import { NodeIdGenerator } from '../system/NodeIdGenerator.js';
export class ContextEnvironmentImpl implements ContextEnvironment {
readonly tokenCalculator: ContextTokenCalculator;
readonly fileSystem: IFileSystem;
readonly idGenerator: IIdGenerator;
constructor(
readonly llmClient: BaseLlmClient,
readonly sessionId: string,
readonly promptId: string,
readonly traceDir: string,
readonly projectTempDir: string,
readonly tracer: ContextTracer,
readonly charsPerToken: number,
readonly eventBus: ContextEventBus,
fileSystem?: IFileSystem,
idGenerator?: IIdGenerator,
) {
this.tokenCalculator = new ContextTokenCalculator(this.charsPerToken);
this.fileSystem = fileSystem || new NodeFileSystem();
this.idGenerator = idGenerator || new NodeIdGenerator();
}
}
@@ -0,0 +1,292 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
import { PipelineOrchestrator } from './orchestrator.js';
import { ProcessorRegistry } from './registry.js';
import {
createMockEnvironment,
createDummyState,
createDummyEpisode,
} from '../testing/contextTestUtils.js';
import type { ContextEnvironment } from './environment.js';
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
import type { PipelineDef, ProcessorConfig, SidecarConfig } from './types.js';
import type { ContextEventBus } from '../eventBus.js';
import type { EpisodeEditor } from '../ir/episodeEditor.js';
// Create a Dummy Processor for testing Orchestration routing
class DummySyncProcessor implements ContextProcessor {
static create() {
return new DummySyncProcessor();
}
constructor() {}
readonly name = 'DummySync';
readonly id = 'DummySync';
readonly options = {};
async process(editor: EpisodeEditor, _state: ContextAccountingState) {
editor.editEpisode(
editor.episodes[0].id,
'DUMMY_EDIT',
(draft: unknown) => {
(draft as Record<string, unknown>)['dummyModified'] = true;
},
);
}
}
class DummyAsyncProcessor implements ContextProcessor {
static create() {
return new DummyAsyncProcessor();
}
constructor() {}
readonly name = 'DummyAsync';
readonly id = 'DummyAsync';
readonly options = {};
async process(editor: EpisodeEditor, _state: ContextAccountingState) {
editor.editEpisode(
editor.episodes[0].id,
'DUMMY_EDIT',
(draft: unknown) => {
(draft as Record<string, unknown>)['dummyAsyncModified'] = true;
},
);
}
}
class ThrowingProcessor implements ContextProcessor {
static create() {
return new ThrowingProcessor();
}
constructor() {}
readonly name = 'Throwing';
readonly id = 'Throwing';
readonly options = {};
async process(
_editor: EpisodeEditor,
_state: ContextAccountingState,
): Promise<void> {
throw new Error('Processor failed intentionally');
}
}
describe('PipelineOrchestrator (Component)', () => {
let env: ContextEnvironment;
let eventBus: ContextEventBus;
let registry: ProcessorRegistry;
beforeEach(() => {
vi.resetAllMocks();
env = createMockEnvironment();
eventBus = env.eventBus;
registry = new ProcessorRegistry();
// Register our test processors
registry.register({
id: 'DummySyncProcessor',
schema: {},
create: () => new DummySyncProcessor(),
});
registry.register({
id: 'DummyAsyncProcessor',
schema: {},
create: () => new DummyAsyncProcessor(),
});
registry.register({
id: 'ThrowingProcessor',
schema: {},
create: () => new ThrowingProcessor(),
});
});
afterEach(() => {
// Cleanup registry to not pollute other tests
registry.clear();
});
const createConfig = (pipelines: PipelineDef[]): SidecarConfig => ({
budget: { maxTokens: 100, retainedTokens: 50 },
gcBackstop: { strategy: 'truncate', target: 'max' },
pipelines,
});
it('instantiates processors from the registry on initialization', () => {
const config = createConfig([
{
name: 'Sync',
execution: 'blocking',
triggers: [],
processors: [
{ processorId: 'DummySyncProcessor' } as unknown as ProcessorConfig,
],
},
]);
const orchestrator = new PipelineOrchestrator(
config,
env,
eventBus,
env.tracer,
registry,
);
expect(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(orchestrator as any).instantiatedProcessors.has('DummySyncProcessor'),
).toBe(true);
});
it('throws an error if a config requests an unknown processor', () => {
const config = createConfig([
{
name: 'Bad',
execution: 'blocking',
triggers: [],
processors: [
{ processorId: 'DoesNotExist' } as unknown as ProcessorConfig,
],
},
]);
expect(
() =>
new PipelineOrchestrator(config, env, eventBus, env.tracer, registry),
).toThrow('Context Processor [DoesNotExist] is not registered.');
});
it('executes blocking pipelines synchronously and returns the modified array', async () => {
const config = createConfig([
{
name: 'SyncPipe',
execution: 'blocking',
triggers: [],
processors: [
{ processorId: 'DummySyncProcessor' } as unknown as ProcessorConfig,
],
},
]);
const orchestrator = new PipelineOrchestrator(
config,
env,
eventBus,
env.tracer,
registry,
);
const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])];
const state = createDummyState(false);
const result = await orchestrator.executePipeline(
'SyncPipe',
episodes,
state,
);
expect(result).toHaveLength(1);
expect(
(result[0] as unknown as { dummyModified: boolean }).dummyModified,
).toBe(true);
});
it('executes background pipelines asynchronously without blocking the return', async () => {
const config = createConfig([
{
name: 'AsyncPipe',
execution: 'background',
triggers: [],
processors: [
{ processorId: 'DummyAsyncProcessor' } as unknown as ProcessorConfig,
],
},
]);
const orchestrator = new PipelineOrchestrator(
config,
env,
eventBus,
env.tracer,
registry,
);
const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])];
const state = createDummyState(false);
// This should resolve immediately with the UNMODIFIED array because execution is background
const result = await orchestrator.executePipeline(
'AsyncPipe',
episodes,
state,
);
expect(result).toHaveLength(1);
expect(
(result[0] as unknown as { asyncModified: unknown }).asyncModified,
).toBeUndefined(); // Not modified yet!
// Wait for the background task to complete (50ms delay in DummyAsyncProcessor)
await new Promise((resolve) => setTimeout(resolve, 60));
});
it('gracefully handles and swallows processor errors in synchronous pipelines', async () => {
const config = createConfig([
{
name: 'ThrowingPipe',
execution: 'blocking',
triggers: [],
processors: [
{ processorId: 'ThrowingProcessor' } as unknown as ProcessorConfig,
],
},
]);
const orchestrator = new PipelineOrchestrator(
config,
env,
eventBus,
env.tracer,
registry,
);
const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])];
const state = createDummyState(false);
// It should not throw! It should swallow the error and return the unmodified array.
const result = await orchestrator.executePipeline(
'ThrowingPipe',
episodes,
state,
);
expect(result).toHaveLength(1);
expect(result).toStrictEqual(episodes);
});
it('automatically binds to budget_exceeded trigger via EventBus', () => {
const config = createConfig([
{
name: 'PressureRelief',
execution: 'background',
triggers: ['budget_exceeded'],
processors: [
{ processorId: 'DummyAsyncProcessor' } as unknown as ProcessorConfig,
],
},
]);
// Spy on the private method to see if the trigger fires it
const executeSpy = vi.spyOn(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
PipelineOrchestrator.prototype as any,
'executePipelineAsync',
);
new PipelineOrchestrator(config, env, eventBus, env.tracer, registry);
const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])];
// Emit the trigger
eventBus.emitConsolidationNeeded({ episodes, targetDeficit: 100 });
expect(executeSpy).toHaveBeenCalled();
});
});
@@ -0,0 +1,234 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Episode } from '../ir/types.js';
import type { ContextProcessor, ContextAccountingState } from '../pipeline.js';
import type { SidecarConfig, PipelineDef } from './types.js';
import type {
ContextEnvironment,
ContextEventBus,
ContextTracer,
} from './environment.js';
import type { ProcessorRegistry } from './registry.js';
import { debugLogger } from '../../utils/debugLogger.js';
import { EpisodeEditor } from '../ir/episodeEditor.js';
export class PipelineOrchestrator {
private activeTimers: NodeJS.Timeout[] = [];
private readonly instantiatedProcessors = new Map<string, ContextProcessor>();
constructor(
private readonly config: SidecarConfig,
private readonly env: ContextEnvironment,
private readonly eventBus: ContextEventBus,
private readonly tracer: ContextTracer,
private readonly registry: ProcessorRegistry,
) {
this.instantiateProcessors();
this.registerTriggers();
}
/**
* Pre-loads and configures all processors defined in the sidecar config.
*/
private instantiateProcessors() {
for (const pipeline of this.config.pipelines) {
for (const procDef of pipeline.processors) {
if (!this.instantiatedProcessors.has(procDef.processorId)) {
const processorClass = this.registry.get(procDef.processorId);
if (!processorClass) {
throw new Error(
`Context Processor [${procDef.processorId}] is not registered.`,
);
}
// The Orchestrator injects standard dependencies required by processors
// If a processor needs the eventBus (like Snapshot), it expects it via constructor.
const instance = processorClass.create(
this.env,
procDef.options ?? {},
);
this.instantiatedProcessors.set(procDef.processorId, instance);
}
}
}
}
/**
* Sets up listeners for the triggers defined in the SidecarConfig.
*/
private registerTriggers() {
for (const pipeline of this.config.pipelines) {
for (const trigger of pipeline.triggers) {
if (typeof trigger === 'object' && trigger.type === 'timer') {
const timer = setInterval(() => {
// For background timers, we need a way to get the latest state
// But timers are generally disabled right now via the triggers config.
// If needed, we will pass it via event bus.
}, trigger.intervalMs);
this.activeTimers.push(timer);
} else if (trigger === 'budget_exceeded') {
this.eventBus.onConsolidationNeeded((event) => {
const state: ContextAccountingState = {
currentTokens: 0,
retainedTokens: this.config.budget.retainedTokens,
maxTokens: this.config.budget.maxTokens,
isBudgetSatisfied: false,
deficitTokens: event.targetDeficit,
protectedEpisodeIds: new Set(),
};
void this.executePipelineAsync(pipeline, event.episodes, state);
});
}
}
}
}
shutdown() {
for (const timer of this.activeTimers) {
clearInterval(timer);
}
}
/**
* Executes a pipeline asynchronously in the background. This is the "Eventual Consistency" path.
* When the pipeline resolves, it emits a VariantReady event to cache the new graph.
*/
/**
* Executes a pipeline based on its configured execution strategy ('blocking' or 'background').
*/
async executePipeline(
pipelineName: string,
episodes: Episode[],
state: ContextAccountingState,
): Promise<Episode[]> {
const pipeline = this.config.pipelines.find((p) => p.name === pipelineName);
if (!pipeline) return episodes;
if (pipeline.execution === 'background') {
this.executePipelineAsync(pipeline, episodes, state).catch((e) => {
debugLogger.error(`Background pipeline ${pipeline.name} failed:`, e);
});
return episodes; // Return immediately
}
// Blocking execution
this.tracer.logEvent(
'Orchestrator',
`Triggering synchronous pipeline: ${pipeline.name}`,
);
let currentEpisodes = [...episodes];
for (let i = 0; i < pipeline.processors.length; i++) {
const procDef = pipeline.processors[i];
const processor = this.instantiatedProcessors.get(procDef.processorId);
if (!processor) continue;
try {
this.tracer.logEvent(
'Orchestrator',
`Executing processor: ${procDef.processorId}`,
);
const editor = new EpisodeEditor(currentEpisodes);
await processor.process(editor, state);
currentEpisodes = editor.getFinalEpisodes();
} catch (error) {
debugLogger.error(
`Pipeline ${pipeline.name} failed synchronously at ${procDef.processorId}:`,
error,
);
return currentEpisodes; // Return what we have so far
}
}
return currentEpisodes;
}
/**
* Internal method for running a pipeline entirely in the background.
*/
private async executePipelineAsync(
pipeline: PipelineDef,
currentState: Episode[],
state: ContextAccountingState,
) {
this.tracer.logEvent(
'Orchestrator',
`Triggering async pipeline: ${pipeline.name}`,
);
if (!currentState || currentState.length === 0) return;
let currentEpisodes = [...currentState];
for (const procDef of pipeline.processors) {
const processor = this.instantiatedProcessors.get(procDef.processorId);
if (!processor) continue;
try {
this.tracer.logEvent(
'Orchestrator',
`Executing processor: ${procDef.processorId} (async)`,
);
const editor = new EpisodeEditor(currentEpisodes);
await processor.process(editor, state);
currentEpisodes = editor.getFinalEpisodes();
// Synthesize VariantReady events for anything that changed or was newly created
for (const mutation of editor.getMutations()) {
// We only broadcast modifications or replacements
// (Insertions without replacement and deletions are not tracked as variants on an existing node)
if (mutation.type === 'modified' || mutation.type === 'replaced') {
const variantId = `v-${procDef.processorId.toLowerCase()}`;
let vType: 'snapshot' | 'summary' | 'masked' = 'masked';
if (procDef.processorId.includes('Snapshot')) vType = 'snapshot';
else if (procDef.processorId.includes('Semantic'))
vType = 'summary';
const ep = mutation.episode!;
let fallbackText = '';
if (ep.yield?.text) fallbackText = ep.yield.text;
else if (ep.trigger?.type === 'USER_PROMPT') {
const firstPart = ep.trigger.semanticParts?.[0];
if (firstPart) {
fallbackText =
firstPart.type === 'text'
? firstPart.presentation?.text || firstPart.text
: '';
}
}
this.eventBus.emitVariantReady({
targetId:
mutation.type === 'replaced' ? mutation.originalIds![0] : ep.id,
variantId,
variant:
vType === 'snapshot'
? {
status: 'ready',
type: 'snapshot',
episode: ep,
recoveredTokens: ep.yield?.metadata?.currentTokens || 10,
replacedEpisodeIds: mutation.originalIds || [],
}
: {
status: 'ready',
type: vType,
text: fallbackText,
recoveredTokens: ep.yield?.metadata?.currentTokens || 10,
},
});
}
}
} catch (error) {
debugLogger.error(
`Pipeline ${pipeline.name} failed at ${procDef.processorId}:`,
error,
);
return; // Halt pipeline
}
}
}
}
@@ -0,0 +1,54 @@
/**
* @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: [
{
name: 'Immediate Sanitization',
triggers: ['on_turn'],
execution: 'blocking',
processors: [
{
processorId: 'ToolMaskingProcessor',
options: { stringLengthThresholdTokens: 8000 },
},
{ processorId: 'BlobDegradationProcessor', options: {} },
{
processorId: 'SemanticCompressionProcessor',
options: { nodeThresholdTokens: 5000 },
},
{ processorId: 'EmergencyTruncationProcessor', options: {} },
],
},
{
name: 'Deep Background Compression',
triggers: [{ type: 'timer', intervalMs: 5000 }, 'budget_exceeded'],
execution: 'background',
processors: [
{
processorId: 'HistorySquashingProcessor',
options: { maxTokensPerNode: 3000 },
},
{ processorId: 'StateSnapshotProcessor', options: {} },
],
},
],
};
@@ -0,0 +1,47 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ContextProcessor } from '../pipeline.js';
import type { ContextEnvironment } from './environment.js';
export interface ContextProcessorDef<TOptions = object> {
readonly id: string;
readonly schema: object;
create(env: ContextEnvironment, options: TOptions): ContextProcessor;
}
/**
* Registry for mapping declarative sidecar configs to running Processor instances.
*/
export class ProcessorRegistry {
private processors = new Map<string, ContextProcessorDef<unknown>>();
register<TOptions>(def: ContextProcessorDef<TOptions>) {
this.processors.set(def.id, def);
}
get(id: string): ContextProcessorDef {
const def = this.processors.get(id);
if (!def) {
throw new Error(`Context Processor [${id}] is not registered.`);
}
return def;
}
getSchemas(): object[] {
const schemas: object[] = [];
for (const def of this.processors.values()) {
if (def.schema) {
schemas.push(def.schema);
}
}
return schemas;
}
clear() {
this.processors.clear();
}
}
+103
View File
@@ -0,0 +1,103 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ProcessorRegistry } from './registry.js';
import './builtins.js';
export function getSidecarConfigSchema(registry: ProcessorRegistry) {
return {
$schema: 'http://json-schema.org/draft-07/schema#',
title: 'SidecarConfig',
description: 'The Data-Driven Schema for the Context Manager.',
type: 'object',
required: ['budget', 'gcBackstop', 'pipelines'],
properties: {
budget: {
type: 'object',
description: 'Defines the token ceilings and limits for the pipeline.',
required: ['retainedTokens', 'maxTokens'],
properties: {
retainedTokens: {
type: 'number',
description:
'The ideal token count the pipeline tries to shrink down to.',
},
maxTokens: {
type: 'number',
description:
'The absolute maximum token count allowed before synchronous truncation kicks in.',
},
},
},
gcBackstop: {
type: 'object',
description:
"Defines what happens when the pipeline fails to compress under 'maxTokens'",
required: ['strategy', 'target'],
properties: {
strategy: {
type: 'string',
enum: ['truncate', 'compress', 'rollingSummarizer'],
},
target: {
type: 'string',
enum: ['incremental', 'freeNTokens', 'max'],
},
freeTokensTarget: {
type: 'number',
},
},
},
pipelines: {
type: 'array',
description: 'The execution graphs for context manipulation.',
items: {
type: 'object',
required: ['name', 'triggers', 'execution', 'processors'],
properties: {
name: {
type: 'string',
},
triggers: {
type: 'array',
items: {
anyOf: [
{
type: 'string',
enum: ['on_turn', 'post_turn', 'budget_exceeded'],
},
{
type: 'object',
required: ['type', 'intervalMs'],
properties: {
type: {
type: 'string',
const: 'timer',
},
intervalMs: {
type: 'number',
},
},
},
],
},
},
execution: {
type: 'string',
enum: ['blocking', 'background'],
},
processors: {
type: 'array',
items: {
oneOf: registry.getSchemas(),
},
},
},
},
},
},
};
}
@@ -0,0 +1,67 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { StateSnapshotProcessorOptions } from '../processors/stateSnapshotProcessor.js';
/**
* Definition of a processor or worker to be instantiated in the graph.
*/
export type ProcessorConfig =
| {
processorId: 'ToolMaskingProcessor';
options: { stringLengthThresholdTokens: number };
}
| { processorId: 'BlobDegradationProcessor'; options?: object }
| {
processorId: 'SemanticCompressionProcessor';
options: { nodeThresholdTokens: number };
}
| {
processorId: 'HistorySquashingProcessor';
options: { maxTokensPerNode: number };
}
| {
processorId: 'StateSnapshotProcessor';
options: StateSnapshotProcessorOptions;
}
| {
processorId: 'EmergencyTruncationProcessor';
options?: Record<string, unknown>;
};
export type PipelineTrigger =
| 'on_turn'
| 'post_turn'
| 'budget_exceeded'
| { type: 'timer'; intervalMs: number };
export interface PipelineDef {
name: string;
triggers: PipelineTrigger[];
execution: 'blocking' | 'background';
processors: ProcessorConfig[];
}
/**
* 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: PipelineDef[];
}
@@ -0,0 +1,189 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { ContextManager } from '../contextManager.js';
import { AgentChatHistory } from '../../core/agentChatHistory.js';
import type { Content } from '@google/genai';
import type { SidecarConfig } from '../sidecar/types.js';
import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js';
import { ContextTracer } from '../tracer.js';
import { ContextEventBus } from '../eventBus.js';
import { PipelineOrchestrator } from '../sidecar/orchestrator.js';
import { registerBuiltInProcessors } from '../sidecar/builtins.js';
import { debugLogger } from '../../utils/debugLogger.js';
import { ProcessorRegistry } from '../sidecar/registry.js';
import { DeterministicIdGenerator } from '../system/DeterministicIdGenerator.js';
import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
export interface TurnSummary {
turnIndex: number;
tokensBeforeBackground: number;
tokensAfterBackground: number;
}
export class SimulationHarness {
readonly chatHistory: AgentChatHistory;
contextManager!: ContextManager;
env!: ContextEnvironmentImpl;
orchestrator!: PipelineOrchestrator;
readonly eventBus: ContextEventBus;
config!: SidecarConfig;
private tracer!: ContextTracer;
private currentTurnIndex = 0;
private tokenTrajectory: TurnSummary[] = [];
static async create(
config: SidecarConfig,
mockLlmClient: BaseLlmClient,
mockTempDir = '/tmp/sim',
): Promise<SimulationHarness> {
const harness = new SimulationHarness();
await harness.init(config, mockLlmClient, mockTempDir);
return harness;
}
private constructor() {
this.chatHistory = new AgentChatHistory();
this.eventBus = new ContextEventBus();
}
private async init(
config: SidecarConfig,
mockLlmClient: BaseLlmClient,
mockTempDir: string,
) {
this.config = config;
const registry = new ProcessorRegistry();
// Register all standard processors
registerBuiltInProcessors(registry);
this.tracer = new ContextTracer({
targetDir: mockTempDir,
sessionId: 'sim-session',
});
this.env = new ContextEnvironmentImpl(
mockLlmClient,
'sim-prompt',
'sim-session',
mockTempDir,
mockTempDir,
this.tracer,
4, // 4 chars per token average
this.eventBus,
new InMemoryFileSystem(),
new DeterministicIdGenerator(),
);
this.orchestrator = new PipelineOrchestrator(
config,
this.env,
this.eventBus,
this.tracer,
registry,
);
this.contextManager = ContextManager.create(
config,
this.env,
this.tracer,
this.orchestrator,
registry,
);
this.contextManager.subscribeToHistory(this.chatHistory);
}
/**
* Simulates a single "Turn" (User input + Model/Tool outputs)
* A turn might consist of multiple Content messages (e.g. user prompt -> model call -> user response -> model answer)
*/
async simulateTurn(messages: Content[]) {
// 1. Append the new messages
const currentHistory = this.chatHistory.get();
this.chatHistory.set([...currentHistory, ...messages]);
// 2. Measure tokens immediately after append (Before background processing)
const tokensBefore = this.env.tokenCalculator.calculateEpisodeListTokens(
this.contextManager.getWorkingBufferView(),
);
debugLogger.log(
`[Turn ${this.currentTurnIndex}] Tokens BEFORE: ${tokensBefore}`,
);
// 3. Yield to event loop to allow internal async subscribers and orchestrator to finish
await new Promise((resolve) => setTimeout(resolve, 50));
// 3.1 Simulate what projectCompressedHistory does with the sync handlers
let currentView = this.contextManager.getWorkingBufferView();
const currentTokens =
this.env.tokenCalculator.calculateEpisodeListTokens(currentView);
if (this.config.budget && currentTokens > this.config.budget.maxTokens) {
debugLogger.log(
`[Turn ${this.currentTurnIndex}] Sync panic triggered! ${currentTokens} > ${this.config.budget.maxTokens}`,
);
const syncPipelines = this.config.pipelines.filter(
(p) => p.execution === 'blocking',
);
const orchestrator = this.orchestrator;
for (const pipe of syncPipelines) {
await orchestrator.executePipeline(pipe.name, currentView, {
currentTokens,
maxTokens: this.config.budget.maxTokens,
retainedTokens: this.config.budget.retainedTokens,
isBudgetSatisfied: false,
deficitTokens: currentTokens - this.config.budget.maxTokens,
protectedEpisodeIds: new Set(),
});
currentView = this.contextManager.getWorkingBufferView();
}
// Inject the truncated view back into the graph
for (let i = 0; i < currentView.length; i++) {
const ep = currentView[i];
if (
!this.contextManager
.getWorkingBufferView()
.find((c) => c.id === ep.id)
) {
this.eventBus.emitVariantReady({
targetId: ep.id,
variantId: 'v-emergency',
variant: {
status: 'ready',
type: 'masked', // Truncation is technically a mask
text: ep.yield?.text || '',
recoveredTokens: 0,
},
});
}
}
// Wait for variant propagation
await new Promise((resolve) => setTimeout(resolve, 50));
}
// 4. Measure tokens after background processors have (hopefully) emitted variants
const tokensAfter = this.env.tokenCalculator.calculateEpisodeListTokens(
this.contextManager.getWorkingBufferView(),
);
debugLogger.log(
`[Turn ${this.currentTurnIndex}] Tokens AFTER: ${tokensAfter}`,
);
this.tokenTrajectory.push({
turnIndex: this.currentTurnIndex++,
tokensBeforeBackground: tokensBefore,
tokensAfterBackground: tokensAfter,
});
}
async getGoldenState() {
const finalProjection =
await this.contextManager.projectCompressedHistory();
return {
tokenTrajectory: this.tokenTrajectory,
finalProjection,
};
}
}
File diff suppressed because one or more lines are too long
@@ -0,0 +1,146 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeAll, afterAll } from 'vitest';
import { SimulationHarness } from './SimulationHarness.js';
import type { SidecarConfig } from '../sidecar/types.js';
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
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,
) ||
/^\/tmp\/sim/.test(val)), // Mask temp directories and UUIDs
print: (val) =>
typeof val === 'string' && /^\/tmp\/sim/.test(val)
? '"<MOCKED_DIR>"'
: '"<UUID>"',
});
describe('System Lifecycle Golden Tests', () => {
beforeAll(() => {
vi.spyOn(Math, 'random').mockReturnValue(0.5);
});
afterAll(() => {
vi.restoreAllMocks();
});
const getAggressiveConfig = (): SidecarConfig => ({
budget: { maxTokens: 4000, retainedTokens: 2000 }, // Extremely tight limits
gcBackstop: { strategy: 'truncate', target: 'max' },
pipelines: [
{
name: 'Pressure Relief', // Emits from eventBus 'budget_exceeded'
execution: 'background',
triggers: ['budget_exceeded'],
processors: [
{ processorId: 'BlobDegradationProcessor' },
{
processorId: 'ToolMaskingProcessor',
options: { stringLengthThresholdTokens: 50 },
}, // Mask any tool string > 200 chars
{ processorId: 'StateSnapshotProcessor', options: {} }, // Squash old history
],
},
{
name: 'Immediate Sanitization', // The magic string the projector is hardcoded to use
execution: 'blocking',
triggers: ['budget_exceeded'],
processors: [
{ processorId: 'EmergencyTruncationProcessor', options: {} },
],
},
],
});
const mockLlmClient = {
generateContent: vi.fn().mockResolvedValue({
text: '<MOCKED_STATE_SNAPSHOT_SUMMARY>',
}),
} as unknown as BaseLlmClient;
it('Scenario 1: Organic Growth with Huge Tool Output & Images', async () => {
const harness = await SimulationHarness.create(
getAggressiveConfig(),
mockLlmClient,
);
// Turn 0: System Prompt
await harness.simulateTurn([
{ role: 'user', parts: [{ text: 'System Instructions' }] },
{ role: 'model', parts: [{ text: 'Ack.' }] },
]);
// Turn 1: Normal conversation
await harness.simulateTurn([
{ role: 'user', parts: [{ text: 'Hello!' }] },
{ role: 'model', parts: [{ text: 'Hi, how can I help?' }] },
]);
// Turn 2: Massive Tool Output (Should trigger ToolMaskingProcessor in background)
await harness.simulateTurn([
{ role: 'user', parts: [{ text: 'Read the logs.' }] },
{
role: 'model',
parts: [
{
functionCall: {
name: 'run_shell_command',
args: { cmd: 'cat server.log' },
},
},
],
},
{
role: 'user',
parts: [
{
functionResponse: {
name: 'run_shell_command',
response: { output: 'LOG '.repeat(5000) },
},
},
],
},
{ role: 'model', parts: [{ text: 'The logs are very long.' }] },
]);
// Turn 3: Multi-modal blob (Should trigger BlobDegradationProcessor)
await harness.simulateTurn([
{
role: 'user',
parts: [
{ text: 'Look at this architecture diagram:' },
{
inlineData: {
mimeType: 'image/png',
data: 'fake_base64_data_'.repeat(1000),
},
},
],
},
{ role: 'model', parts: [{ text: 'Nice diagram.' }] },
]);
// Turn 4: More conversation to trigger StateSnapshot
await harness.simulateTurn([
{ role: 'user', parts: [{ text: 'Can we refactor?' }] },
{ role: 'model', parts: [{ text: 'Yes we can.' }] },
]);
// Get final state
const goldenState = await harness.getGoldenState();
// In a perfectly functioning opportunistic system, the token trajectory should show
// the massive spikes in Turn 2 and 3 being immediately resolved by the background tasks.
// The final projection should fit neatly under the Max Tokens limit.
expect(goldenState).toMatchSnapshot();
});
});
@@ -0,0 +1,18 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { IIdGenerator } from './IIdGenerator.js';
export class DeterministicIdGenerator implements IIdGenerator {
private counter = 0;
constructor(private prefix: string = 'id-') {}
generateId(): string {
this.counter++;
return `${this.prefix}${this.counter}`;
}
}
@@ -0,0 +1,20 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
export interface IFileSystem {
existsSync(path: string): boolean;
statSyncSize(path: string): number;
readFileSync(path: string, encoding: 'utf8'): string;
writeFileSync(path: string, data: string | Buffer, encoding?: 'utf-8'): void;
appendFileSync(path: string, data: string, encoding: 'utf-8'): void;
mkdirSync(path: string, options?: { recursive?: boolean }): void;
writeFile(path: string, data: string | Buffer): Promise<void>;
mkdir(path: string, options?: { recursive?: boolean }): Promise<void>;
join(...paths: string[]): string;
dirname(path: string): string;
}
@@ -0,0 +1,9 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
export interface IIdGenerator {
generateId(): string;
}
@@ -0,0 +1,79 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { IFileSystem } from './IFileSystem.js';
export class InMemoryFileSystem implements IFileSystem {
private files = new Map<string, string | Buffer>();
getFiles(): ReadonlyMap<string, string | Buffer> {
return this.files;
}
setFile(path: string, content: string | Buffer) {
this.files.set(this.normalize(path), content);
}
private normalize(p: string): string {
return p.replace(/\/+/g, '/');
}
existsSync(p: string): boolean {
return this.files.has(this.normalize(p));
}
statSyncSize(p: string): number {
const content = this.files.get(this.normalize(p));
if (content === undefined) {
throw new Error(`ENOENT: no such file or directory, stat '${p}'`);
}
return Buffer.isBuffer(content)
? content.byteLength
: Buffer.byteLength(content, 'utf8');
}
readFileSync(p: string, encoding: 'utf8'): string {
const content = this.files.get(this.normalize(p));
if (content === undefined) {
throw new Error(`ENOENT: no such file or directory, open '${p}'`);
}
if (Buffer.isBuffer(content)) {
return content.toString(encoding);
}
return content;
}
writeFileSync(p: string, data: string | Buffer, _encoding?: 'utf-8'): void {
this.files.set(this.normalize(p), data);
}
appendFileSync(p: string, data: string, _encoding: 'utf-8'): void {
const norm = this.normalize(p);
const existing = this.files.get(norm) || '';
const existingStr = Buffer.isBuffer(existing)
? existing.toString('utf8')
: existing;
this.files.set(norm, existingStr + data);
}
mkdirSync(_p: string, _options?: { recursive?: boolean }): void {}
async writeFile(p: string, data: string | Buffer): Promise<void> {
this.writeFileSync(p, data);
}
async mkdir(_p: string, _options?: { recursive?: boolean }): Promise<void> {}
join(...paths: string[]): string {
return this.normalize(paths.join('/'));
}
dirname(p: string): string {
const parts = this.normalize(p).split('/');
parts.pop();
return parts.length === 0 ? '.' : parts.join('/') || '/';
}
}
@@ -0,0 +1,56 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import * as fs from 'node:fs';
import * as fsPromises from 'node:fs/promises';
import * as path from 'node:path';
import type { IFileSystem } from './IFileSystem.js';
export class NodeFileSystem implements IFileSystem {
existsSync(p: string): boolean {
return fs.existsSync(p);
}
statSyncSize(p: string): number {
return fs.statSync(p).size;
}
readFileSync(p: string, encoding: 'utf8'): string {
return fs.readFileSync(p, encoding);
}
writeFileSync(p: string, data: string | Buffer, encoding?: 'utf-8'): void {
if (Buffer.isBuffer(data)) {
fs.writeFileSync(p, data);
} else {
fs.writeFileSync(p, data, encoding);
}
}
appendFileSync(p: string, data: string, encoding: 'utf-8'): void {
fs.appendFileSync(p, data, encoding);
}
mkdirSync(p: string, options?: { recursive?: boolean }): void {
fs.mkdirSync(p, options);
}
async writeFile(p: string, data: string | Buffer): Promise<void> {
await fsPromises.writeFile(p, data);
}
async mkdir(p: string, options?: { recursive?: boolean }): Promise<void> {
await fsPromises.mkdir(p, options);
}
join(...paths: string[]): string {
return path.join(...paths);
}
dirname(p: string): string {
return path.dirname(p);
}
}
@@ -0,0 +1,14 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import type { IIdGenerator } from './IIdGenerator.js';
export class NodeIdGenerator implements IIdGenerator {
generateId(): string {
return randomUUID();
}
}
@@ -0,0 +1,221 @@
/**
* @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';
import type { Content } from '@google/genai';
import { AgentChatHistory } from '../../core/agentChatHistory.js';
import { ContextManager } from '../contextManager.js';
import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
import { DeterministicIdGenerator } from '../system/DeterministicIdGenerator.js';
import type {
Episode,
UserPrompt,
SystemEvent,
SemanticPart,
} from '../ir/types.js';
import type { ContextAccountingState } from '../pipeline.js';
import { randomUUID } from 'node:crypto';
export function createDummyState(
isSatisfied = false,
deficit = 0,
protectedIds = new Set<string>(),
currentTokens = 5000,
maxTokens = 10000,
retainedTokens = 4000,
): ContextAccountingState {
return {
currentTokens,
maxTokens,
retainedTokens,
deficitTokens: deficit,
protectedEpisodeIds: protectedIds,
isBudgetSatisfied: isSatisfied,
};
}
export function createDummyEpisode(
id: string,
type: 'USER_PROMPT' | 'SYSTEM_EVENT',
parts: SemanticPart[] = [],
toolSteps: Array<{
intent: Record<string, unknown>;
observation: Record<string, unknown>;
toolName?: string;
tokens?: { intent: number; observation: number };
}> = [],
): Episode {
let trigger: UserPrompt | SystemEvent;
if (type === 'USER_PROMPT') {
trigger = {
id: randomUUID(),
type: 'USER_PROMPT',
semanticParts: parts,
metadata: {
originalTokens: 100,
currentTokens: 100,
transformations: [],
},
};
} else {
trigger = {
id: randomUUID(),
type: 'SYSTEM_EVENT',
name: 'dummy_event',
payload: {},
metadata: {
originalTokens: 100,
currentTokens: 100,
transformations: [],
},
};
}
return {
id,
timestamp: Date.now(),
trigger,
steps: toolSteps.map((step) => ({
id: randomUUID(),
type: 'TOOL_EXECUTION',
toolName: step.toolName || 'test_tool',
intent: step.intent,
observation: step.observation,
tokens: step.tokens || { intent: 50, observation: 50 },
metadata: {
originalTokens: 100,
currentTokens: 100,
transformations: [],
},
})),
};
}
export function createMockEnvironment(): ContextEnvironment {
return {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
llmClient: vi.fn().mockReturnValue({
generateContent: vi.fn().mockResolvedValue({
text: 'Mock LLM summary response',
}),
})() as unknown as BaseLlmClient,
promptId: 'mock-prompt-id',
sessionId: 'mock-session',
traceDir: '/tmp/.gemini/trace',
projectTempDir: '/tmp/.gemini/tool-outputs',
eventBus: new ContextEventBus(),
tracer: new ContextTracer({ targetDir: '/tmp', sessionId: 'mock-session' }),
charsPerToken: 1,
tokenCalculator: new ContextTokenCalculator(1),
fileSystem: new InMemoryFileSystem(),
idGenerator: new DeterministicIdGenerator('mock-uuid-'),
};
}
/**
* 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'),
getExperimentalContextSidecarConfig: vi.fn().mockReturnValue(undefined),
};
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
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';
import { ContextEventBus } from '../eventBus.js';
import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
import type { BaseLlmClient } from 'src/core/baseLlmClient.js';
import { ProcessorRegistry } from '../sidecar/registry.js';
import { registerBuiltInProcessors } from '../sidecar/builtins.js';
export function setupContextComponentTest(config: Config) {
const chatHistory = new AgentChatHistory();
const registry = new ProcessorRegistry();
registerBuiltInProcessors(registry);
const sidecar = SidecarLoader.fromConfig(config, registry);
const tracer = new ContextTracer({
targetDir: '/tmp',
sessionId: 'test-session',
});
const eventBus = new ContextEventBus();
const env = new ContextEnvironmentImpl(
config.getBaseLlmClient(),
'test prompt-id',
'test-session',
'/tmp',
'/tmp/gemini-test',
tracer,
1,
eventBus,
);
const contextManager = ContextManager.create(
sidecar,
env,
tracer,
undefined,
registry,
);
// The async worker is now internally managed by ContextManager
// Subscribe to history to enable the Eager/Opportunistic triggers
contextManager.subscribeToHistory(chatHistory);
return { chatHistory, contextManager };
}
+85
View File
@@ -0,0 +1,85 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, beforeEach, vi } from 'vitest';
import { ContextTracer } from './tracer.js';
import { InMemoryFileSystem } from './system/InMemoryFileSystem.js';
import { DeterministicIdGenerator } from './system/DeterministicIdGenerator.js';
describe('ContextTracer (Fake FS & ID Gen)', () => {
let fileSystem: InMemoryFileSystem;
let idGenerator: DeterministicIdGenerator;
beforeEach(() => {
fileSystem = new InMemoryFileSystem();
idGenerator = new DeterministicIdGenerator('mock-uuid-');
// We must mock Date.now() to ensure asset file names are perfectly deterministic
vi.useFakeTimers();
vi.setSystemTime(new Date('2026-01-01T12:00:00Z'));
});
it('initializes, logs events, and auto-saves large assets deterministically', () => {
const tracer = new ContextTracer(
{ enabled: true, targetDir: '/fake/target', sessionId: 'test-session' },
fileSystem,
idGenerator,
);
// Verify Initialization
const initTraceLog = fileSystem.readFileSync(
'/fake/target/.gemini/context_trace/test-session/trace.log',
'utf8',
);
expect(initTraceLog).toContain('[SYSTEM] Context Tracer Initialized');
// Small logging: shouldn't trigger saveAsset
tracer.logEvent('TestComponent', 'TestAction', { key: 'value' });
const smallTraceLog = fileSystem.readFileSync(
'/fake/target/.gemini/context_trace/test-session/trace.log',
'utf8',
);
expect(smallTraceLog).toContain('[TestComponent] TestAction');
expect(smallTraceLog).toContain('{"key":"value"}');
// Large logging: should trigger auto-asset save
const hugeString = 'a'.repeat(2000);
tracer.logEvent('TestComponent', 'LargeAction', { largeKey: hugeString });
// 1767268800000 is 2026-01-01T12:00:00Z
const expectedAssetPath =
'/fake/target/.gemini/context_trace/test-session/assets/1767268800000-mock-uuid-1-largeKey.json';
// Assert asset was written to FS
expect(fileSystem.existsSync(expectedAssetPath)).toBe(true);
const largeTraceLog = fileSystem.readFileSync(
'/fake/target/.gemini/context_trace/test-session/trace.log',
'utf8',
);
expect(largeTraceLog).toContain('[TestComponent] LargeAction');
expect(largeTraceLog).toContain(
`{"largeKey":{"$asset":"1767268800000-mock-uuid-1-largeKey.json"}}`,
);
});
it('silently ignores logging when disabled', () => {
const tracer = new ContextTracer(
{ enabled: false, targetDir: '/fake/target', sessionId: 'test-session' },
fileSystem,
idGenerator,
);
tracer.logEvent('TestComponent', 'TestAction');
const hugeString = 'a'.repeat(2000);
tracer.logEvent('TestComponent', 'LargeAction', { largeKey: hugeString });
// FS should be completely empty
expect(fileSystem.getFiles().size).toBe(0);
});
});
+120
View File
@@ -0,0 +1,120 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { debugLogger } from '../utils/debugLogger.js';
import type { IFileSystem } from './system/IFileSystem.js';
import { NodeFileSystem } from './system/NodeFileSystem.js';
import type { IIdGenerator } from './system/IIdGenerator.js';
import { NodeIdGenerator } from './system/NodeIdGenerator.js';
export interface ContextTracerOptions {
enabled?: boolean;
targetDir: string;
sessionId: string;
}
export class ContextTracer {
private traceDir: string;
private assetsDir: string;
private enabled: boolean;
private fileSystem: IFileSystem;
private idGenerator: IIdGenerator;
private readonly MAX_INLINE_SIZE = 1000;
constructor(
options: ContextTracerOptions,
fileSystem: IFileSystem = new NodeFileSystem(),
idGenerator: IIdGenerator = new NodeIdGenerator(),
) {
this.enabled = options.enabled ?? false;
this.fileSystem = fileSystem;
this.idGenerator = idGenerator;
this.traceDir = this.fileSystem.join(
options.targetDir,
'.gemini',
'context_trace',
options.sessionId,
);
this.assetsDir = this.fileSystem.join(this.traceDir, 'assets');
if (this.enabled) {
try {
this.fileSystem.mkdirSync(this.assetsDir, { recursive: true });
this.logEvent('SYSTEM', 'Context Tracer Initialized', {
sessionId: options.sessionId,
});
} catch (e) {
debugLogger.error('Failed to initialize ContextTracer', e);
this.enabled = false;
}
}
}
logEvent(
component: string,
action: string,
details?: Record<string, unknown>,
) {
if (!this.enabled) return;
try {
let processedDetails: Record<string, unknown> | undefined;
if (details) {
processedDetails = {};
for (const [key, value] of Object.entries(details)) {
const strValue =
typeof value === 'string' ? value : JSON.stringify(value);
if (strValue && strValue.length > this.MAX_INLINE_SIZE) {
const assetId = this.saveAsset(component, key, value);
processedDetails[key] = { $asset: assetId };
} else {
processedDetails[key] = value;
}
}
}
const timestamp = new Date().toISOString();
const detailsStr = processedDetails
? ` | Details: ${JSON.stringify(processedDetails)}`
: '';
const logLine = `[${timestamp}] [${component}] ${action}${detailsStr}\n`;
this.fileSystem.appendFileSync(
this.fileSystem.join(this.traceDir, 'trace.log'),
logLine,
'utf-8',
);
} catch (e) {
debugLogger.warn(`Tracing failed: ${e}`);
}
}
private saveAsset(
component: string,
assetName: string,
data: unknown,
): string {
if (!this.enabled) return 'asset-recording-disabled';
try {
const assetId = `${Date.now()}-${this.idGenerator.generateId()}-${assetName}.json`;
const assetPath = this.fileSystem.join(this.assetsDir, assetId);
this.fileSystem.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,71 @@
/**
* @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';
import type { Episode } from '../ir/types.js';
/**
* The flat token cost assigned to a single multi-modal asset (like an image tile)
* by the Gemini API. We use this as a baseline heuristic for inlineData/fileData.
*/
const BASE_MULTIMODAL_TOKEN_COST = 258;
export class ContextTokenCalculator {
constructor(private readonly charsPerToken: number) {}
/**
* Fast, simple heuristic estimation for a raw string.
*/
estimateTokensForString(text: string): number {
return Math.ceil(text.length / this.charsPerToken);
}
/**
* Fast, simple heuristic conversion from tokens to expected character length.
* Useful for calculating truncation thresholds.
*/
tokensToChars(tokens: number): number {
return tokens * this.charsPerToken;
}
/**
* Calculates the total token count for a complete Episodic IR graph.
* This is fast because it relies on pre-computed metadata where available.
*/
calculateEpisodeListTokens(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;
}
/**
* Slower, precise estimation for a Gemini Content/Part graph.
* Deeply inspects the nested structure and uses the base tokenization math.
*/
estimateTokensForParts(parts: Part[], depth: number = 0): number {
let totalTokens = 0;
for (const part of parts) {
if (typeof part.text === 'string') {
totalTokens += Math.ceil(part.text.length / this.charsPerToken);
} else if (part.inlineData !== undefined || part.fileData !== undefined) {
totalTokens += BASE_MULTIMODAL_TOKEN_COST;
} else {
totalTokens += Math.ceil(
JSON.stringify(part).length / this.charsPerToken,
);
}
}
// Also include structural overhead
return totalTokens + baseEstimate(parts, depth);
}
}
@@ -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;
}
}
+1
View File
@@ -16,4 +16,5 @@ export enum LlmRole {
UTILITY_EDIT_CORRECTOR = 'utility_edit_corrector',
UTILITY_AUTOCOMPLETE = 'utility_autocomplete',
UTILITY_FAST_ACK_HELPER = 'utility_fast_ack_helper',
UTILITY_STATE_SNAPSHOT_PROCESSOR = 'utility_state_snapshot_processr',
}