feat(core): introduce decoupled ContextManager and Sidecar architecture (#24752)

This commit is contained in:
joshualitt 2026-04-13 15:02:22 -07:00 committed by GitHub
parent 706d4d4707
commit daf5006237
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
54 changed files with 6454 additions and 0 deletions

View file

@ -699,6 +699,7 @@ export interface ConfigParameters {
experimentalJitContext?: boolean;
autoDistillation?: boolean;
experimentalMemoryManager?: boolean;
experimentalContextManagementConfig?: string;
experimentalAgentHistoryTruncation?: boolean;
experimentalAgentHistoryTruncationThreshold?: number;
experimentalAgentHistoryRetainedMessages?: number;
@ -939,6 +940,7 @@ export class Config implements McpContext, AgentLoopContext {
private readonly adminSkillsEnabled: boolean;
private readonly experimentalJitContext: boolean;
private readonly experimentalMemoryManager: boolean;
private readonly experimentalContextManagementConfig?: string;
private readonly memoryBoundaryMarkers: readonly string[];
private readonly topicUpdateNarration: boolean;
private readonly disableLLMCorrection: boolean;
@ -1150,6 +1152,8 @@ export class Config implements McpContext, AgentLoopContext {
this.experimentalJitContext = params.experimentalJitContext ?? false;
this.experimentalMemoryManager = params.experimentalMemoryManager ?? false;
this.experimentalContextManagementConfig =
params.experimentalContextManagementConfig;
this.memoryBoundaryMarkers = params.memoryBoundaryMarkers ?? ['.git'];
this.contextManagement = {
enabled: params.contextManagement?.enabled ?? false,
@ -2434,6 +2438,10 @@ export class Config implements McpContext, AgentLoopContext {
return this.experimentalMemoryManager;
}
getExperimentalContextManagementConfig(): string | undefined {
return this.experimentalContextManagementConfig;
}
getContextManagementConfig(): ContextManagementConfig {
return this.contextManagement;
}

View file

@ -0,0 +1,94 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import { loadContextManagementConfig } from './configLoader.js';
import { defaultContextProfile } from './profiles.js';
import { ContextProcessorRegistry } from './registry.js';
import * as fs from 'node:fs/promises';
import * as path from 'node:path';
import * as os from 'node:os';
import type { Config } from '../../config/config.js';
import type { JSONSchemaType } from 'ajv';
describe('SidecarLoader (Real FS)', () => {
let tmpDir: string;
let registry: ContextProcessorRegistry;
let sidecarPath: string;
let mockConfig: Config;
beforeEach(async () => {
tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'gemini-sidecar-test-'));
sidecarPath = path.join(tmpDir, 'sidecar.json');
registry = new ContextProcessorRegistry();
registry.registerProcessor({
id: 'NodeTruncation',
schema: {
type: 'object',
properties: { maxTokens: { type: 'number' } },
required: ['maxTokens'],
} as unknown as JSONSchemaType<{ maxTokens: number }>,
});
mockConfig = {
getExperimentalContextManagementConfig: () => sidecarPath,
} as unknown as Config;
});
afterEach(async () => {
await fs.rm(tmpDir, { recursive: true, force: true });
});
it('returns default profile if file does not exist', async () => {
const result = await loadContextManagementConfig(mockConfig, registry);
expect(result).toBe(defaultContextProfile);
});
it('returns default profile if file exists but is 0 bytes', async () => {
await fs.writeFile(sidecarPath, '');
const result = await loadContextManagementConfig(mockConfig, registry);
expect(result).toBe(defaultContextProfile);
});
it('returns parsed config if file is valid', async () => {
const validConfig = {
budget: { retainedTokens: 1000, maxTokens: 2000 },
processorOptions: {
myTruncation: {
type: 'NodeTruncation',
options: { maxTokens: 500 },
},
},
};
await fs.writeFile(sidecarPath, JSON.stringify(validConfig));
const result = await loadContextManagementConfig(mockConfig, registry);
expect(result.config.budget?.maxTokens).toBe(2000);
expect(result.config.processorOptions?.['myTruncation']).toBeDefined();
});
it('throws validation error if processorOptions contains invalid data for the schema', async () => {
const invalidConfig = {
budget: { retainedTokens: 1000, maxTokens: 2000 },
processorOptions: {
myTruncation: {
type: 'NodeTruncation',
options: { maxTokens: 'this should be a number' },
},
},
};
await fs.writeFile(sidecarPath, JSON.stringify(invalidConfig));
await expect(
loadContextManagementConfig(mockConfig, registry),
).rejects.toThrow('Validation error');
});
it('throws validation error if file is empty whitespace', async () => {
await fs.writeFile(sidecarPath, ' \n ');
await expect(
loadContextManagementConfig(mockConfig, registry),
).rejects.toThrow('Unexpected end of JSON input');
});
});

View file

@ -0,0 +1,90 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Config } from '../../config/config.js';
import * as fsSync from 'node:fs';
import * as fs from 'node:fs/promises';
import type { ContextManagementConfig } from './types.js';
import { defaultContextProfile, type ContextProfile } from './profiles.js';
import { SchemaValidator } from '../../utils/schemaValidator.js';
import { getContextManagementConfigSchema } from './schema.js';
import type { ContextProcessorRegistry } from './registry.js';
import { getErrorMessage } from '../../utils/errors.js';
/**
* 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.
*/
async function loadConfigFromFile(
sidecarPath: string,
registry: ContextProcessorRegistry,
): Promise<ContextProfile> {
const fileContent = await fs.readFile(sidecarPath, 'utf8');
let parsed: unknown;
try {
parsed = JSON.parse(fileContent);
} catch (error) {
throw new Error(
`Failed to parse Sidecar configuration file at ${sidecarPath}: ${getErrorMessage(
error,
)}`,
);
}
// Validate the complete structure, including deep options
const validationError = SchemaValidator.validate(
getContextManagementConfigSchema(registry),
parsed,
);
if (validationError) {
throw new Error(
`Invalid sidecar configuration in ${sidecarPath}. Validation error: ${validationError}`,
);
}
// Extract strictly what we need.
// Why this unsafe cast is acceptable:
// SchemaValidator just ran \`getSidecarConfigSchema(registry)\` against \`parsed\`.
// That function dynamically maps the \`processorOptions\` to strict JSON schema definitions,
// so we know with absolute certainty at runtime that \`parsed\` conforms to this shape.
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const validConfig = parsed as ContextManagementConfig;
return {
...defaultContextProfile,
config: {
...defaultContextProfile.config,
...(validConfig.budget ? { budget: validConfig.budget } : {}),
...(validConfig.processorOptions
? { processorOptions: validConfig.processorOptions }
: {}),
},
};
}
/**
* 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.
*/
export async function loadContextManagementConfig(
config: Config,
registry: ContextProcessorRegistry,
): Promise<ContextProfile> {
const sidecarPath = config.getExperimentalContextManagementConfig();
if (sidecarPath && fsSync.existsSync(sidecarPath)) {
const size = fsSync.statSync(sidecarPath).size;
// If the file exists but is completely empty (0 bytes), it's safe to fallback.
if (size === 0) {
return defaultContextProfile;
}
// If the file has content, enforce strict validation and throw on failure.
return loadConfigFromFile(sidecarPath, registry);
}
return defaultContextProfile;
}

View file

@ -0,0 +1,145 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type {
AsyncPipelineDef,
ContextManagementConfig,
PipelineDef,
} from './types.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
// Import factories
import { createToolMaskingProcessor } from '../processors/toolMaskingProcessor.js';
import { createBlobDegradationProcessor } from '../processors/blobDegradationProcessor.js';
import { createNodeTruncationProcessor } from '../processors/nodeTruncationProcessor.js';
import { createNodeDistillationProcessor } from '../processors/nodeDistillationProcessor.js';
import { createStateSnapshotProcessor } from '../processors/stateSnapshotProcessor.js';
import { createStateSnapshotAsyncProcessor } from '../processors/stateSnapshotAsyncProcessor.js';
/**
* Helper to safely merge static default options with dynamically loaded
* JSON overrides from the SidecarConfig.
*
* Why the unsafe cast is acceptable here:
* Before the \`config\` object ever reaches this function, \`SidecarLoader.ts\`
* passes the raw JSON through \`SchemaValidator\`. The schema dynamically generates
* a \`oneOf\` map linking every \`type\` discriminator to its corresponding processor
* schema definition. By the time we access \`options\` here, its shape has been
* strictly validated against the corresponding Zod/JSONSchema definition at runtime,
* making the generic cast to \`<T>\` structurally safe.
*/
function resolveProcessorOptions<T>(
config: ContextManagementConfig | undefined,
id: string,
defaultOptions: T,
): T {
if (config?.processorOptions && config.processorOptions[id]) {
return {
...defaultOptions,
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
...(config.processorOptions[id].options as T),
};
}
return defaultOptions;
}
export interface ContextProfile {
config: ContextManagementConfig;
buildPipelines: (
env: ContextEnvironment,
config?: ContextManagementConfig,
) => PipelineDef[];
buildAsyncPipelines: (
env: ContextEnvironment,
config?: ContextManagementConfig,
) => AsyncPipelineDef[];
}
/**
* The standard default context management profile.
* Optimized for safety, precision, and reliable summarization.
*/
export const defaultContextProfile: ContextProfile = {
config: {
budget: {
retainedTokens: 65000,
maxTokens: 150000,
},
},
buildPipelines: (
env: ContextEnvironment,
config?: ContextManagementConfig,
): PipelineDef[] =>
// Helper to merge default options with dynamically loaded processorOptions by ID
[
{
name: 'Immediate Sanitization',
triggers: ['new_message'],
processors: [
createToolMaskingProcessor(
'ToolMasking',
env,
resolveProcessorOptions(config, 'ToolMasking', {
stringLengthThresholdTokens: 8000,
}),
),
createBlobDegradationProcessor('BlobDegradation', env), // No options
],
},
{
name: 'Normalization',
triggers: ['retained_exceeded'],
processors: [
createNodeTruncationProcessor(
'NodeTruncation',
env,
resolveProcessorOptions(config, 'NodeTruncation', {
maxTokensPerNode: 3000,
}),
),
createNodeDistillationProcessor(
'NodeDistillation',
env,
resolveProcessorOptions(config, 'NodeDistillation', {
nodeThresholdTokens: 5000,
}),
),
],
},
{
name: 'Emergency Backstop',
triggers: ['gc_backstop'],
processors: [
createStateSnapshotProcessor(
'StateSnapshotSync',
env,
resolveProcessorOptions(config, 'StateSnapshotSync', {
target: 'max',
}),
),
],
},
],
buildAsyncPipelines: (
env: ContextEnvironment,
config?: ContextManagementConfig,
): AsyncPipelineDef[] => [
{
name: 'Async Background GC',
triggers: ['nodes_aged_out'],
processors: [
createStateSnapshotAsyncProcessor(
'StateSnapshotAsync',
env,
resolveProcessorOptions(config, 'StateSnapshotAsync', {
type: 'accumulate',
}),
),
],
},
],
};

View file

@ -0,0 +1,42 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { JSONSchemaType } from 'ajv';
export interface ContextProcessorDef<T = unknown> {
readonly id: string;
readonly schema: JSONSchemaType<T>;
}
/**
* Registry for validating declarative sidecar configuration schemas.
* (Dynamic instantiation has been replaced by static ContextProfiles)
*/
export class ContextProcessorRegistry {
private readonly processors = new Map<string, ContextProcessorDef>();
registerProcessor<T>(def: ContextProcessorDef<T>) {
// Erasing the type.
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
this.processors.set(def.id, def as unknown as ContextProcessorDef<unknown>);
}
getSchema(id: string): object | undefined {
return this.processors.get(id)?.schema;
}
getSchemaDefs(): ContextProcessorDef[] {
const defs = [];
for (const def of this.processors.values()) {
if (def.schema) defs.push({ id: def.id, schema: def.schema });
}
return defs;
}
clear() {
this.processors.clear();
}
}

View file

@ -0,0 +1,55 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ContextProcessorRegistry } from './registry.js';
export function getContextManagementConfigSchema(
registry: ContextProcessorRegistry,
) {
// We use a registry to deeply validate processor overrides.
// We do this by generating a `oneOf` list that matches the `type` discriminator
// to the specific processor `options` schema.
const processorOptionSchemas = registry.getSchemaDefs().map((def) => ({
type: 'object',
required: ['type', 'options'],
properties: {
type: { const: def.id },
options: def.schema,
},
}));
return {
$schema: 'http://json-schema.org/draft-07/schema#',
title: 'ContextManagementConfig',
description: 'The Hyperparameter schema for a Context Profile.',
type: 'object',
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.',
},
},
},
processorOptions: {
type: 'object',
description:
'Named hyperparameter configurations for ContextProcessors and AsyncProcessors.',
additionalProperties: { oneOf: processorOptionSchemas },
},
},
};
}

View file

@ -0,0 +1,46 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ContextProcessor, AsyncContextProcessor } from '../pipeline.js';
export type PipelineTrigger =
| 'new_message'
| 'retained_exceeded'
| 'gc_backstop'
| 'nodes_added'
| 'nodes_aged_out'
| { type: 'timer'; intervalMs: number };
export interface PipelineDef {
name: string;
triggers: PipelineTrigger[];
processors: ContextProcessor[];
}
export interface AsyncPipelineDef {
name: string;
triggers: PipelineTrigger[];
processors: AsyncContextProcessor[];
}
export interface ContextBudget {
retainedTokens: number;
maxTokens: number;
}
/**
* The Data-Driven Schema for the Context Manager.
*/
export interface ContextManagementConfig {
/** Defines the token ceilings and limits for the pipeline. */
budget: ContextBudget;
/**
* Dynamic hyperparameter overrides for individual ContextProcessors and AsyncProcessors.
* Keys are named identifiers (e.g. "gentleTruncation").
*/
processorOptions?: Record<string, { type: string; options: unknown }>;
}

View file

@ -0,0 +1,82 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { testTruncateProfile } from './testing/testProfile.js';
import {
createSyntheticHistory,
createMockContextConfig,
setupContextComponentTest,
} from './testing/contextTestUtils.js';
describe('ContextManager Sync Pressure Barrier Tests', () => {
beforeEach(() => {
vi.useFakeTimers();
});
afterEach(() => {
vi.useRealTimers();
vi.restoreAllMocks();
});
it('should instantly truncate history when maxTokens is exceeded using truncate strategy', async () => {
// 1. Setup
const config = createMockContextConfig();
const { chatHistory, contextManager } = setupContextComponentTest(
config,
testTruncateProfile,
);
// 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.renderHistory();
// 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');
// Filter out synthetic Yield nodes (they are model responses without actual tool/text bodies)
const contentNodes = projection.filter(
(p) =>
p.parts && p.parts.some((part) => part.text && part.text !== 'Yield'),
);
// Verify the latest turn is perfectly preserved at the back
const lastUser = contentNodes[contentNodes.length - 2];
const lastModel = contentNodes[contentNodes.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.');
});
});

View file

@ -0,0 +1,170 @@
/**
* @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 type { ConcreteNode } from './graph/types.js';
import type { ContextEventBus } from './eventBus.js';
import type { ContextTracer } from './tracer.js';
import type { ContextEnvironment } from './pipeline/environment.js';
import type { ContextProfile } from './config/profiles.js';
import type { PipelineOrchestrator } from './pipeline/orchestrator.js';
import { HistoryObserver } from './historyObserver.js';
import { render } from './graph/render.js';
import { ContextWorkingBufferImpl } from './pipeline/contextWorkingBuffer.js';
export class ContextManager {
// The master state containing the pristine graph and current active graph.
private buffer: ContextWorkingBufferImpl =
ContextWorkingBufferImpl.initialize([]);
private readonly eventBus: ContextEventBus;
// Internal sub-components
private readonly orchestrator: PipelineOrchestrator;
private readonly historyObserver: HistoryObserver;
constructor(
private readonly sidecar: ContextProfile,
private readonly env: ContextEnvironment,
private readonly tracer: ContextTracer,
orchestrator: PipelineOrchestrator,
chatHistory: AgentChatHistory,
) {
this.eventBus = env.eventBus;
this.orchestrator = orchestrator;
this.historyObserver = new HistoryObserver(
chatHistory,
this.env.eventBus,
this.tracer,
this.env.tokenCalculator,
this.env.graphMapper,
);
this.historyObserver.start();
this.eventBus.onPristineHistoryUpdated((event) => {
const existingIds = new Set(this.buffer.nodes.map((n) => n.id));
const newIds = new Set(event.nodes.map((n) => n.id));
const addedNodes = event.nodes.filter((n) => !existingIds.has(n.id));
// Prune any pristine nodes that were dropped from the upstream history
this.buffer = this.buffer.prunePristineNodes(newIds);
if (addedNodes.length > 0) {
this.buffer = this.buffer.appendPristineNodes(addedNodes);
}
this.evaluateTriggers(event.newNodes);
});
}
/**
* Safely stops background async pipelines and clears event listeners.
*/
shutdown() {
this.orchestrator.shutdown();
this.historyObserver.stop();
}
/**
* Evaluates if the current working buffer exceeds configured budget thresholds,
* firing consolidation events if necessary.
*/
private evaluateTriggers(newNodes: Set<string>) {
if (!this.sidecar.config.budget) return;
if (newNodes.size > 0) {
this.eventBus.emitChunkReceived({
nodes: this.buffer.nodes,
targetNodeIds: newNodes,
});
}
const currentTokens = this.env.tokenCalculator.calculateConcreteListTokens(
this.buffer.nodes,
);
if (currentTokens > this.sidecar.config.budget.retainedTokens) {
const agedOutNodes = new Set<string>();
let rollingTokens = 0;
// Walk backwards finding nodes that fall out of the retained budget
for (let i = this.buffer.nodes.length - 1; i >= 0; i--) {
const node = this.buffer.nodes[i];
rollingTokens += this.env.tokenCalculator.calculateConcreteListTokens([
node,
]);
if (rollingTokens > this.sidecar.config.budget.retainedTokens) {
agedOutNodes.add(node.id);
}
}
if (agedOutNodes.size > 0) {
this.env.tokenCalculator.garbageCollectCache(
new Set(this.buffer.nodes.map((n) => n.id)),
);
this.eventBus.emitConsolidationNeeded({
nodes: this.buffer.nodes,
targetDeficit:
currentTokens - this.sidecar.config.budget.retainedTokens,
targetNodeIds: agedOutNodes,
});
}
}
}
/**
* Retrieves the raw, uncompressed Episodic Context Graph graph.
* Useful for internal tool rendering (like the trace viewer).
* Note: This is an expensive, deep clone operation.
*/
getPristineGraph(): readonly ConcreteNode[] {
const pristineSet = new Map<string, ConcreteNode>();
for (const node of this.buffer.nodes) {
const roots = this.buffer.getPristineNodes(node.id);
for (const root of roots) {
pristineSet.set(root.id, root);
}
}
// We sort them by timestamp to ensure they are returned in chronological order
return Array.from(pristineSet.values()).sort(
(a, b) => a.timestamp - b.timestamp,
);
}
/**
* Generates a virtual view of the pristine graph, substituting in variants
* up to the configured token budget.
* This is the view that will eventually be projected back to the LLM.
*/
getNodes(): readonly ConcreteNode[] {
return [...this.buffer.nodes];
}
/**
* Executes the final 'gc_backstop' pipeline if necessary, enforcing the token budget,
* and maps the Episodic Context Graph back into a raw Gemini Content[] array for transmission.
* This is the primary method called by the agent framework before sending a request.
*/
async renderHistory(
activeTaskIds: Set<string> = new Set(),
): Promise<Content[]> {
this.tracer.logEvent('ContextManager', 'Starting rendering of LLM context');
// Apply final GC Backstop pressure barrier synchronously before mapping
const finalHistory = await render(
this.buffer.nodes,
this.orchestrator,
this.sidecar,
this.tracer,
this.env,
activeTaskIds,
);
this.tracer.logEvent('ContextManager', 'Finished rendering');
return finalHistory;
}
}

View file

@ -0,0 +1,52 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { EventEmitter } from 'node:events';
import type { ConcreteNode } from './graph/types.js';
export interface PristineHistoryUpdatedEvent {
nodes: readonly ConcreteNode[];
newNodes: Set<string>;
}
export interface ContextConsolidationEvent {
nodes: readonly ConcreteNode[];
targetDeficit: number;
targetNodeIds: Set<string>;
}
export interface ChunkReceivedEvent {
nodes: readonly ConcreteNode[];
targetNodeIds: Set<string>;
}
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: ChunkReceivedEvent) {
this.emit('IR_CHUNK_RECEIVED', event);
}
onChunkReceived(listener: (event: ChunkReceivedEvent) => 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);
}
}

View file

@ -0,0 +1,43 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content, Part } from '@google/genai';
import type { ConcreteNode } from './types.js';
export interface NodeSerializationWriter {
appendContent(content: Content): void;
appendModelPart(part: Part): void;
appendUserPart(part: Part): void;
flushModelParts(): void;
}
export interface NodeBehavior<T extends ConcreteNode = ConcreteNode> {
readonly type: T['type'];
/** Serializes the node into the Gemini Content structure. */
serialize(node: T, writer: NodeSerializationWriter): void;
/**
* Generates a structural representation of the node for the purpose
* of estimating its token cost.
*/
getEstimatableParts(node: T): Part[];
}
export class NodeBehaviorRegistry {
private readonly behaviors = new Map<string, NodeBehavior<ConcreteNode>>();
register<T extends ConcreteNode>(behavior: NodeBehavior<T>) {
this.behaviors.set(behavior.type, behavior);
}
get(type: string): NodeBehavior<ConcreteNode> {
const behavior = this.behaviors.get(type);
if (!behavior) {
throw new Error(`Unregistered Node type: ${type}`);
}
return behavior;
}
}

View file

@ -0,0 +1,172 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Part } from '@google/genai';
import type { NodeBehavior, NodeBehaviorRegistry } from './behaviorRegistry.js';
import type {
UserPrompt,
AgentThought,
ToolExecution,
MaskedTool,
AgentYield,
Snapshot,
RollingSummary,
SystemEvent,
} from './types.js';
export const UserPromptBehavior: NodeBehavior<UserPrompt> = {
type: 'USER_PROMPT',
getEstimatableParts(prompt) {
const parts: Part[] = [];
for (const sp of prompt.semanticParts) {
switch (sp.type) {
case 'text':
parts.push({ text: sp.text });
break;
case 'inline_data':
parts.push({ inlineData: { mimeType: sp.mimeType, data: sp.data } });
break;
case 'file_data':
parts.push({
fileData: { mimeType: sp.mimeType, fileUri: sp.fileUri },
});
break;
case 'raw_part':
parts.push(sp.part);
break;
default:
break;
}
}
return parts;
},
serialize(prompt, writer) {
const parts = this.getEstimatableParts(prompt);
if (parts.length > 0) {
writer.flushModelParts();
writer.appendContent({ role: 'user', parts });
}
},
};
export const AgentThoughtBehavior: NodeBehavior<AgentThought> = {
type: 'AGENT_THOUGHT',
getEstimatableParts(thought) {
return [{ text: thought.text }];
},
serialize(thought, writer) {
writer.appendModelPart({ text: thought.text });
},
};
export const ToolExecutionBehavior: NodeBehavior<ToolExecution> = {
type: 'TOOL_EXECUTION',
getEstimatableParts(tool) {
return [
{ functionCall: { id: tool.id, name: tool.toolName, args: tool.intent } },
{
functionResponse: {
id: tool.id,
name: tool.toolName,
response:
typeof tool.observation === 'string'
? { message: tool.observation }
: tool.observation,
},
},
];
},
serialize(tool, writer) {
const parts = this.getEstimatableParts(tool);
writer.appendModelPart(parts[0]);
writer.flushModelParts();
writer.appendUserPart(parts[1]);
},
};
export const MaskedToolBehavior: NodeBehavior<MaskedTool> = {
type: 'MASKED_TOOL',
getEstimatableParts(tool) {
return [
{
functionCall: {
id: tool.id,
name: tool.toolName,
args: tool.intent ?? {},
},
},
{
functionResponse: {
id: tool.id,
name: tool.toolName,
response:
typeof tool.observation === 'string'
? { message: tool.observation }
: (tool.observation ?? {}),
},
},
];
},
serialize(tool, writer) {
const parts = this.getEstimatableParts(tool);
writer.appendModelPart(parts[0]);
writer.flushModelParts();
writer.appendUserPart(parts[1]);
},
};
export const AgentYieldBehavior: NodeBehavior<AgentYield> = {
type: 'AGENT_YIELD',
getEstimatableParts(yieldNode) {
return [{ text: yieldNode.text }];
},
serialize(yieldNode, writer) {
writer.appendModelPart({ text: yieldNode.text });
writer.flushModelParts();
},
};
export const SystemEventBehavior: NodeBehavior<SystemEvent> = {
type: 'SYSTEM_EVENT',
getEstimatableParts() {
return [];
},
serialize(node, writer) {
writer.flushModelParts();
},
};
export const SnapshotBehavior: NodeBehavior<Snapshot> = {
type: 'SNAPSHOT',
getEstimatableParts(node) {
return [{ text: node.text }];
},
serialize(node, writer) {
writer.flushModelParts();
writer.appendUserPart({ text: node.text });
},
};
export const RollingSummaryBehavior: NodeBehavior<RollingSummary> = {
type: 'ROLLING_SUMMARY',
getEstimatableParts(node) {
return [{ text: node.text }];
},
serialize(node, writer) {
writer.flushModelParts();
writer.appendUserPart({ text: node.text });
},
};
export function registerBuiltInBehaviors(registry: NodeBehaviorRegistry) {
registry.register(UserPromptBehavior);
registry.register(AgentThoughtBehavior);
registry.register(ToolExecutionBehavior);
registry.register(MaskedToolBehavior);
registry.register(AgentYieldBehavior);
registry.register(SystemEventBehavior);
registry.register(SnapshotBehavior);
registry.register(RollingSummaryBehavior);
}

View file

@ -0,0 +1,54 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content, Part } from '@google/genai';
import type { ConcreteNode } from './types.js';
import type {
NodeSerializationWriter,
NodeBehaviorRegistry,
} from './behaviorRegistry.js';
class NodeSerializer implements NodeSerializationWriter {
private history: Content[] = [];
private currentModelParts: Part[] = [];
appendContent(content: Content) {
this.flushModelParts();
this.history.push(content);
}
appendModelPart(part: Part) {
this.currentModelParts.push(part);
}
appendUserPart(part: Part) {
this.flushModelParts();
this.history.push({ role: 'user', parts: [part] });
}
flushModelParts() {
if (this.currentModelParts.length > 0) {
this.history.push({ role: 'model', parts: [...this.currentModelParts] });
this.currentModelParts = [];
}
}
getContents(): Content[] {
this.flushModelParts();
return this.history;
}
}
export function fromGraph(
nodes: readonly ConcreteNode[],
registry: NodeBehaviorRegistry,
): Content[] {
const writer = new NodeSerializer();
for (const node of nodes) {
const behavior = registry.get(node.type);
behavior.serialize(node, writer);
}
return writer.getContents();
}

View file

@ -0,0 +1,28 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content } from '@google/genai';
import type { Episode, ConcreteNode } from './types.js';
import { toGraph } from './toGraph.js';
import { fromGraph } from './fromGraph.js';
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
import type { NodeBehaviorRegistry } from './behaviorRegistry.js';
export class ContextGraphMapper {
private readonly nodeIdentityMap = new WeakMap<object, string>();
constructor(private readonly registry: NodeBehaviorRegistry) {}
toGraph(
history: readonly Content[],
tokenCalculator: ContextTokenCalculator,
): Episode[] {
return toGraph(history, tokenCalculator, this.nodeIdentityMap);
}
fromGraph(nodes: readonly ConcreteNode[]): Content[] {
return fromGraph(nodes, this.registry);
}
}

View file

@ -0,0 +1,115 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content } from '@google/genai';
import type { ConcreteNode } from './types.js';
import { debugLogger } from '../../utils/debugLogger.js';
import type {
ContextEnvironment,
ContextTracer,
} from '../pipeline/environment.js';
import type { PipelineOrchestrator } from '../pipeline/orchestrator.js';
import type { ContextProfile } from '../config/profiles.js';
/**
* Orchestrates the final render: takes a working buffer view (The Nodes),
* applies the Immediate Sanitization pipeline, and enforces token boundaries.
*/
export async function render(
nodes: readonly ConcreteNode[],
orchestrator: PipelineOrchestrator,
sidecar: ContextProfile,
tracer: ContextTracer,
env: ContextEnvironment,
protectedIds: Set<string>,
): Promise<Content[]> {
if (!sidecar.config.budget) {
const contents = env.graphMapper.fromGraph(nodes);
tracer.logEvent('Render', 'Render Context to LLM (No Budget)', {
renderedContext: contents,
});
return contents;
}
const maxTokens = sidecar.config.budget.maxTokens;
const currentTokens = env.tokenCalculator.calculateConcreteListTokens(nodes);
// V0: Always protect the first node (System Prompt) and the last turn
if (nodes.length > 0) {
protectedIds.add(nodes[0].id);
if (nodes[0].logicalParentId) protectedIds.add(nodes[0].logicalParentId);
const lastNode = nodes[nodes.length - 1];
protectedIds.add(lastNode.id);
if (lastNode.logicalParentId) protectedIds.add(lastNode.logicalParentId);
}
if (currentTokens <= maxTokens) {
tracer.logEvent(
'Render',
`View is within maxTokens (${currentTokens} <= ${maxTokens}). Returning view.`,
);
const contents = env.graphMapper.fromGraph(nodes);
tracer.logEvent('Render', 'Render Context for LLM', {
renderedContext: contents,
});
return contents;
}
tracer.logEvent(
'Render',
`View exceeds maxTokens (${currentTokens} > ${maxTokens}). Hitting Synchronous Pressure Barrier.`,
);
debugLogger.log(
`Context Manager Synchronous Barrier triggered: View at ${currentTokens} tokens (limit: ${maxTokens}).`,
);
// Calculate exactly which nodes aged out of the retainedTokens budget to form our target delta
const agedOutNodes = new Set<string>();
let rollingTokens = 0;
// Start from newest and count backwards
for (let i = nodes.length - 1; i >= 0; i--) {
const node = nodes[i];
const nodeTokens = env.tokenCalculator.calculateConcreteListTokens([node]);
rollingTokens += nodeTokens;
if (rollingTokens > sidecar.config.budget.retainedTokens) {
agedOutNodes.add(node.id);
}
}
const processedNodes = await orchestrator.executeTriggerSync(
'gc_backstop',
nodes,
agedOutNodes,
protectedIds,
);
const finalTokens =
env.tokenCalculator.calculateConcreteListTokens(processedNodes);
tracer.logEvent(
'Render',
`Finished rendering. Final token count: ${finalTokens}.`,
);
debugLogger.log(
`Context Manager finished. Final actual token count: ${finalTokens}.`,
);
// Apply skipList logic to abstract over summarized nodes
const skipList = new Set<string>();
for (const node of processedNodes) {
if (node.abstractsIds) {
for (const id of node.abstractsIds) skipList.add(id);
}
}
const visibleNodes = processedNodes.filter((n) => !skipList.has(n.id));
const contents = env.graphMapper.fromGraph(visibleNodes);
tracer.logEvent('Render', 'Render Sanitized Context for LLM', {
renderedContextSanitized: contents,
});
return contents;
}

View file

@ -0,0 +1,236 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content, Part } from '@google/genai';
import type {
Episode,
SemanticPart,
ToolExecution,
AgentThought,
AgentYield,
UserPrompt,
} from './types.js';
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
import { randomUUID } from 'node:crypto';
import { isRecord } from '../../utils/markdownUtils.js';
// We remove the global nodeIdentityMap and instead rely on one passed from ContextGraphMapper
export function getStableId(
obj: object,
nodeIdentityMap: WeakMap<object, string>,
): string {
let id = nodeIdentityMap.get(obj);
if (!id) {
id = randomUUID();
nodeIdentityMap.set(obj, id);
}
return id;
}
function isCompleteEpisode(ep: Partial<Episode>): ep is Episode {
return (
typeof ep.id === 'string' &&
Array.isArray(ep.concreteNodes) &&
ep.concreteNodes.length > 0
);
}
export function toGraph(
history: readonly Content[],
tokenCalculator: ContextTokenCalculator,
nodeIdentityMap: WeakMap<object, string>,
): Episode[] {
const episodes: Episode[] = [];
let currentEpisode: Partial<Episode> | null = null;
const pendingCallParts: Map<string, Part> = new Map();
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,
nodeIdentityMap,
);
}
if (hasUserParts) {
finalizeEpisode();
currentEpisode = parseUserParts(msg, nodeIdentityMap);
}
} else if (msg.role === 'model') {
currentEpisode = parseModelParts(
msg,
currentEpisode,
pendingCallParts,
nodeIdentityMap,
);
}
}
if (currentEpisode) {
finalizeYield(currentEpisode);
finalizeEpisode();
}
return episodes;
}
function parseToolResponses(
msg: Content,
currentEpisode: Partial<Episode> | null,
pendingCallParts: Map<string, Part>,
tokenCalculator: ContextTokenCalculator,
nodeIdentityMap: WeakMap<object, string>,
): Partial<Episode> {
if (!currentEpisode) {
currentEpisode = {
id: getStableId(msg, nodeIdentityMap),
concreteNodes: [],
};
}
const parts = msg.parts || [];
for (const part of 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, nodeIdentityMap),
timestamp: Date.now(),
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,
},
};
currentEpisode.concreteNodes = [
...(currentEpisode.concreteNodes || []),
step,
];
if (callId) pendingCallParts.delete(callId);
}
}
return currentEpisode;
}
function parseUserParts(
msg: Content,
nodeIdentityMap: WeakMap<object, string>,
): Partial<Episode> {
const semanticParts: SemanticPart[] = [];
const parts = msg.parts || [];
for (const p of 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 baseObj = parts.length > 0 ? parts[0] : msg;
const trigger: UserPrompt = {
id: getStableId(baseObj, nodeIdentityMap),
timestamp: Date.now(),
type: 'USER_PROMPT',
semanticParts,
};
return {
id: getStableId(msg, nodeIdentityMap),
concreteNodes: [trigger],
};
}
function parseModelParts(
msg: Content,
currentEpisode: Partial<Episode> | null,
pendingCallParts: Map<string, Part>,
nodeIdentityMap: WeakMap<object, string>,
): Partial<Episode> {
if (!currentEpisode) {
currentEpisode = {
id: getStableId(msg, nodeIdentityMap),
concreteNodes: [],
};
}
const parts = msg.parts || [];
for (const part of 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, nodeIdentityMap),
timestamp: Date.now(),
type: 'AGENT_THOUGHT',
text: part.text,
};
currentEpisode.concreteNodes = [
...(currentEpisode.concreteNodes || []),
thought,
];
}
}
return currentEpisode;
}
function finalizeYield(currentEpisode: Partial<Episode>) {
if (currentEpisode.concreteNodes && currentEpisode.concreteNodes.length > 0) {
const yieldNode: AgentYield = {
id: randomUUID(),
timestamp: Date.now(),
type: 'AGENT_YIELD',
text: 'Yield', // Synthesized yield since we don't have the original concrete node
};
const existingNodes = currentEpisode.concreteNodes || [];
currentEpisode.concreteNodes = [...existingNodes, yieldNode];
}
}

View file

@ -0,0 +1,222 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Part } from '@google/genai';
export type NodeType =
// Organic Concrete Nodes
| 'USER_PROMPT'
| 'SYSTEM_EVENT'
| 'AGENT_THOUGHT'
| 'TOOL_EXECUTION'
| 'AGENT_YIELD'
// Synthetic Concrete Nodes
| 'SNAPSHOT'
| 'ROLLING_SUMMARY'
| 'MASKED_TOOL'
// Logical Nodes
| 'TASK'
| 'EPISODE';
/** Base interface for all nodes in the Episodic Context Graph */
export interface Node {
readonly id: string;
readonly type: NodeType;
}
/**
* Concrete Nodes: The atomic, renderable pieces of data.
* These are the actual "planks" of the Nodes of Theseus.
*/
export interface BaseConcreteNode extends Node {
readonly timestamp: number;
/** The ID of the Logical Node (e.g., Episode) that structurally owns this node */
readonly logicalParentId?: string;
/** If this node replaced a single node 1:1 (e.g., masking), this points to the original */
readonly replacesId?: string;
/** If this node is a synthetic summary of N nodes, this points to the original IDs */
readonly abstractsIds?: readonly string[];
}
/**
* Semantic Parts for User Prompts
*/
export interface SemanticTextPart {
readonly type: 'text';
readonly text: string;
}
export interface SemanticInlineDataPart {
readonly type: 'inline_data';
readonly mimeType: string;
readonly data: string;
}
export interface SemanticFileDataPart {
readonly type: 'file_data';
readonly mimeType: string;
readonly fileUri: string;
}
export interface SemanticRawPart {
readonly type: 'raw_part';
readonly part: Part;
}
export type SemanticPart =
| SemanticTextPart
| SemanticInlineDataPart
| SemanticFileDataPart
| SemanticRawPart;
/**
* Trigger Nodes
* Events that wake the agent up and initiate an Episode.
*/
export interface UserPrompt extends BaseConcreteNode {
readonly type: 'USER_PROMPT';
readonly semanticParts: readonly SemanticPart[];
}
export interface SystemEvent extends BaseConcreteNode {
readonly type: 'SYSTEM_EVENT';
readonly name: string;
readonly 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 BaseConcreteNode {
readonly type: 'AGENT_THOUGHT';
readonly text: string;
}
export interface ToolExecution extends BaseConcreteNode {
readonly type: 'TOOL_EXECUTION';
readonly toolName: string;
readonly intent: Record<string, unknown>;
readonly observation: string | Record<string, unknown>;
readonly tokens: {
readonly intent: number;
readonly observation: number;
};
}
export interface MaskedTool extends BaseConcreteNode {
readonly type: 'MASKED_TOOL';
readonly toolName: string;
readonly intent?: Record<string, unknown>;
readonly observation?: string | Record<string, unknown>;
readonly tokens: {
readonly intent: number;
readonly observation: number;
};
}
export type EpisodeStep = AgentThought | ToolExecution | MaskedTool;
/**
* Resolution Node
* The final message where the agent yields control back to the user.
*/
export interface AgentYield extends BaseConcreteNode {
readonly type: 'AGENT_YIELD';
readonly text: string;
}
/**
* Synthetic Leaf Interfaces
* Processors that generate summaries emit explicit synthetic nodes.
*/
export interface Snapshot extends BaseConcreteNode {
readonly type: 'SNAPSHOT';
readonly text: string;
}
export interface RollingSummary extends BaseConcreteNode {
readonly type: 'ROLLING_SUMMARY';
readonly text: string;
}
export type SyntheticLeaf = Snapshot | RollingSummary;
export type ConcreteNode =
| UserPrompt
| SystemEvent
| AgentThought
| ToolExecution
| MaskedTool
| AgentYield
| Snapshot
| RollingSummary;
/**
* Logical Nodes
* These define hierarchy and grouping. They do not directly render to Gemini.
*/
export interface Episode extends Node {
readonly type: 'EPISODE';
/** References to the Concrete Node IDs that conceptually belong to this Episode. */
concreteNodes: readonly ConcreteNode[];
}
export interface Task extends Node {
readonly type: 'TASK';
readonly goal: string;
readonly status: 'active' | 'completed' | 'failed';
/** References to the Episode IDs that belong to this task */
readonly episodeIds: readonly string[];
}
export type LogicalNode = Task | Episode;
export function isEpisode(node: Node): node is Episode {
return node.type === 'EPISODE';
}
export function isTask(node: Node): node is Task {
return node.type === 'TASK';
}
export function isAgentThought(node: Node): node is AgentThought {
return node.type === 'AGENT_THOUGHT';
}
export function isAgentYield(node: Node): node is AgentYield {
return node.type === 'AGENT_YIELD';
}
export function isToolExecution(node: Node): node is ToolExecution {
return node.type === 'TOOL_EXECUTION';
}
export function isMaskedTool(node: Node): node is MaskedTool {
return node.type === 'MASKED_TOOL';
}
export function isUserPrompt(node: Node): node is UserPrompt {
return node.type === 'USER_PROMPT';
}
export function isSystemEvent(node: Node): node is SystemEvent {
return node.type === 'SYSTEM_EVENT';
}
export function isSnapshot(node: Node): node is Snapshot {
return node.type === 'SNAPSHOT';
}
export function isRollingSummary(node: Node): node is RollingSummary {
return node.type === 'ROLLING_SUMMARY';
}

View file

@ -0,0 +1,88 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type {
AgentChatHistory,
HistoryEvent,
} from '../core/agentChatHistory.js';
import type { ContextGraphMapper } from './graph/mapper.js';
import type { ContextTokenCalculator } from './utils/contextTokenCalculator.js';
import type { ContextEventBus } from './eventBus.js';
import type { ContextTracer } from './tracer.js';
import type { ConcreteNode } from './graph/types.js';
/**
* Connects the raw AgentChatHistory to the ContextManager.
* It maps raw messages into Episodic Intermediate Representation (Context Graph)
* and evaluates background triggers whenever history changes.
*/
export class HistoryObserver {
private unsubscribeHistory?: () => void;
private readonly seenNodeIds = new Set<string>();
constructor(
private readonly chatHistory: AgentChatHistory,
private readonly eventBus: ContextEventBus,
private readonly tracer: ContextTracer,
private readonly tokenCalculator: ContextTokenCalculator,
private readonly graphMapper: ContextGraphMapper,
) {}
start() {
if (this.unsubscribeHistory) {
this.unsubscribeHistory();
}
this.unsubscribeHistory = this.chatHistory.subscribe(
(_event: HistoryEvent) => {
// Rebuild the pristine Context Graph graph from the full source history on every change.
// Wait, toGraph still returns an Episode[].
// We actually need to map the Episode[] to a flat ConcreteNode[] here to form the 'nodes'.
const pristineEpisodes = this.graphMapper.toGraph(
this.chatHistory.get(),
this.tokenCalculator,
);
const nodes: ConcreteNode[] = [];
for (const ep of pristineEpisodes) {
if (ep.concreteNodes) {
for (const child of ep.concreteNodes) {
nodes.push(child);
}
}
}
const newNodes = new Set<string>();
for (const node of nodes) {
if (!this.seenNodeIds.has(node.id)) {
newNodes.add(node.id);
this.seenNodeIds.add(node.id);
}
}
this.tracer.logEvent(
'HistoryObserver',
'Rebuilt pristine graph from chat history update',
{ nodesSize: nodes.length, newNodesCount: newNodes.size },
);
this.eventBus.emitPristineHistoryUpdated({
nodes,
newNodes,
});
},
);
}
stop() {
if (this.unsubscribeHistory) {
this.unsubscribeHistory();
this.unsubscribeHistory = undefined;
}
}
}

View file

@ -0,0 +1,61 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ConcreteNode } from './graph/types.js';
export interface InboxMessage<T = unknown> {
id: string;
topic: string;
payload: T;
timestamp: number;
}
export interface InboxSnapshot {
getMessages<T = unknown>(topic: string): ReadonlyArray<InboxMessage<T>>;
consume(messageId: string): void;
}
export interface GraphMutation {
readonly processorId: string;
readonly timestamp: number;
readonly removedIds: readonly string[];
readonly addedNodes: readonly ConcreteNode[];
}
export interface ContextWorkingBuffer {
readonly nodes: readonly ConcreteNode[];
getPristineNodes(id: string): readonly ConcreteNode[];
getLineage(id: string): readonly ConcreteNode[];
getAuditLog(): readonly GraphMutation[];
}
export interface ProcessArgs {
readonly buffer: ContextWorkingBuffer;
readonly targets: readonly ConcreteNode[];
readonly inbox: InboxSnapshot;
}
/**
* A ContextProcessor is a pure, closure-based object that returns a modified subset of nodes
* (or the original targets if no changes are needed).
* The Orchestrator will use this to generate a new graph delta.
*/
export interface ContextProcessor {
readonly id: string;
readonly name: string;
process(args: ProcessArgs): Promise<readonly ConcreteNode[]>;
}
export interface AsyncContextProcessor {
readonly id: string;
readonly name: string;
process(args: ProcessArgs): Promise<void>;
}
export interface BackstopTargetOptions {
target?: 'incremental' | 'freeNTokens' | 'max';
freeTokensTarget?: number;
}

View file

@ -0,0 +1,162 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { ContextWorkingBufferImpl } from './contextWorkingBuffer.js';
import { createDummyNode } from '../testing/contextTestUtils.js';
describe('ContextWorkingBufferImpl', () => {
it('should initialize with a pristine graph correctly', () => {
const pristine1 = createDummyNode(
'ep1',
'USER_PROMPT',
10,
undefined,
'p1',
);
const pristine2 = createDummyNode(
'ep1',
'AGENT_THOUGHT',
10,
undefined,
'p2',
);
const buffer = ContextWorkingBufferImpl.initialize([pristine1, pristine2]);
expect(buffer.nodes).toHaveLength(2);
expect(buffer.getAuditLog()).toHaveLength(0);
// Pristine nodes should point to themselves
expect(buffer.getPristineNodes('p1')).toEqual([pristine1]);
expect(buffer.getPristineNodes('p2')).toEqual([pristine2]);
});
it('should track 1:1 replacements (e.g., masking) and append to audit log', () => {
const pristine1 = createDummyNode(
'ep1',
'USER_PROMPT',
10,
undefined,
'p1',
);
let buffer = ContextWorkingBufferImpl.initialize([pristine1]);
const maskedNode = createDummyNode(
'ep1',
'USER_PROMPT',
5,
undefined,
'm1',
);
// Simulate what a processor does
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(maskedNode as any).replacesId = 'p1';
buffer = buffer.applyProcessorResult(
'ToolMasking',
[pristine1],
[maskedNode],
);
expect(buffer.nodes).toHaveLength(1);
expect(buffer.nodes[0].id).toBe('m1');
const log = buffer.getAuditLog();
expect(log).toHaveLength(1);
expect(log[0].processorId).toBe('ToolMasking');
expect(log[0].removedIds).toEqual(['p1']);
expect(log[0].addedNodes[0].id).toBe('m1');
// Provenance lookup: the masked node should resolve back to the pristine root
expect(buffer.getPristineNodes('m1')).toEqual([pristine1]);
});
it('should track N:1 abstractions (e.g., rolling summaries)', () => {
const p1 = createDummyNode('ep1', 'USER_PROMPT', 10, undefined, 'p1');
const p2 = createDummyNode('ep1', 'AGENT_THOUGHT', 10, undefined, 'p2');
const p3 = createDummyNode('ep1', 'USER_PROMPT', 10, undefined, 'p3');
let buffer = ContextWorkingBufferImpl.initialize([p1, p2, p3]);
const summaryNode = createDummyNode(
'ep1',
'ROLLING_SUMMARY',
15,
undefined,
's1',
);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(summaryNode as any).abstractsIds = ['p1', 'p2'];
buffer = buffer.applyProcessorResult('Summarizer', [p1, p2], [summaryNode]);
// p1 and p2 are removed, p3 remains, s1 is added
expect(buffer.nodes.map((n) => n.id)).toEqual(['p3', 's1']);
// Provenance lookup: The summary node should resolve to both p1 and p2!
const roots = buffer.getPristineNodes('s1');
expect(roots).toHaveLength(2);
expect(roots).toContain(p1);
expect(roots).toContain(p2);
});
it('should track multi-generation provenance correctly', () => {
const p1 = createDummyNode('ep1', 'USER_PROMPT', 10, undefined, 'p1');
let buffer = ContextWorkingBufferImpl.initialize([p1]);
// Gen 1: Masked
const gen1 = createDummyNode('ep1', 'USER_PROMPT', 8, undefined, 'gen1');
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(gen1 as any).replacesId = 'p1';
buffer = buffer.applyProcessorResult('Masking', [p1], [gen1]);
// Gen 2: Summarized
const gen2 = createDummyNode(
'ep1',
'ROLLING_SUMMARY',
5,
undefined,
'gen2',
);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(gen2 as any).abstractsIds = ['gen1'];
buffer = buffer.applyProcessorResult('Summarizer', [gen1], [gen2]);
expect(buffer.nodes).toHaveLength(1);
expect(buffer.nodes[0].id).toBe('gen2');
// Audit log should show sequence
const log = buffer.getAuditLog();
expect(log).toHaveLength(2);
expect(log[0].processorId).toBe('Masking');
expect(log[1].processorId).toBe('Summarizer');
// Multi-gen Provenance lookup: gen2 -> gen1 -> p1
expect(buffer.getPristineNodes('gen2')).toEqual([p1]);
});
it('should handle net-new injected nodes without throwing', () => {
const p1 = createDummyNode('ep1', 'USER_PROMPT', 10, undefined, 'p1');
let buffer = ContextWorkingBufferImpl.initialize([p1]);
const injected = createDummyNode(
'ep1',
'SYSTEM_EVENT',
5,
undefined,
'injected1',
);
// No replacesId or abstractsIds
buffer = buffer.applyProcessorResult('Injector', [], [injected]);
expect(buffer.nodes.map((n) => n.id)).toEqual(['p1', 'injected1']);
// It should root to itself
expect(buffer.getPristineNodes('injected1')).toEqual([injected]);
});
});

View file

@ -0,0 +1,270 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ContextWorkingBuffer, GraphMutation } from '../pipeline.js';
import type { ConcreteNode } from '../graph/types.js';
export class ContextWorkingBufferImpl implements ContextWorkingBuffer {
// The current active graph
readonly nodes: readonly ConcreteNode[];
// The AOT pre-calculated provenance index (Current ID -> Pristine IDs)
private readonly provenanceMap: ReadonlyMap<string, ReadonlySet<string>>;
// The original immutable pristine nodes mapping
private readonly pristineNodesMap: ReadonlyMap<string, ConcreteNode>;
// The historical linked list of changes
private readonly history: readonly GraphMutation[];
private constructor(
nodes: readonly ConcreteNode[],
pristineNodesMap: ReadonlyMap<string, ConcreteNode>,
provenanceMap: ReadonlyMap<string, ReadonlySet<string>>,
history: readonly GraphMutation[],
) {
this.nodes = nodes;
this.pristineNodesMap = pristineNodesMap;
this.provenanceMap = provenanceMap;
this.history = history;
}
/**
* Initializes a brand new ContextWorkingBuffer from a pristine graph.
* Every node's provenance points to itself.
*/
static initialize(
pristineNodes: readonly ConcreteNode[],
): ContextWorkingBufferImpl {
const pristineMap = new Map<string, ConcreteNode>();
const initialProvenance = new Map<string, ReadonlySet<string>>();
for (const node of pristineNodes) {
pristineMap.set(node.id, node);
initialProvenance.set(node.id, new Set([node.id]));
}
return new ContextWorkingBufferImpl(
pristineNodes,
pristineMap,
initialProvenance,
[], // Empty history
);
}
/**
* Appends newly observed pristine nodes (e.g. from a user message) to the working buffer.
* Ensures they are tracked in the pristine map and point to themselves in provenance.
*/
appendPristineNodes(
newNodes: readonly ConcreteNode[],
): ContextWorkingBufferImpl {
if (newNodes.length === 0) return this;
const newPristineMap = new Map<string, ConcreteNode>(this.pristineNodesMap);
const newProvenanceMap = new Map(this.provenanceMap);
for (const node of newNodes) {
newPristineMap.set(node.id, node);
newProvenanceMap.set(node.id, new Set([node.id]));
}
return new ContextWorkingBufferImpl(
[...this.nodes, ...newNodes],
newPristineMap,
newProvenanceMap,
[...this.history],
);
}
/**
* Generates an entirely new buffer instance by calculating the delta between the processor's input and output.
*/
applyProcessorResult(
processorId: string,
inputTargets: readonly ConcreteNode[],
outputNodes: readonly ConcreteNode[],
): ContextWorkingBufferImpl {
const outputIds = new Set(outputNodes.map((n) => n.id));
const inputIds = new Set(inputTargets.map((n) => n.id));
// Calculate diffs
const removedIds = inputTargets
.filter((n) => !outputIds.has(n.id))
.map((n) => n.id);
const addedNodes = outputNodes.filter((n) => !inputIds.has(n.id));
// Create mutation record
const mutation: GraphMutation = {
processorId,
timestamp: Date.now(),
removedIds,
addedNodes,
};
// Calculate new node array
const removedSet = new Set(removedIds);
const retainedNodes = this.nodes.filter((n) => !removedSet.has(n.id));
const newGraph = [...retainedNodes];
// We append the output nodes in the same general position if possible,
// but in a complex graph we just ensure they exist. V2 graph uses timestamps for order.
// For simplicity, we just push added nodes to the end of the retained array
newGraph.push(...addedNodes);
// Calculate new provenance map
const newProvenanceMap = new Map(this.provenanceMap);
let finalPristineMap = this.pristineNodesMap;
// Map the new synthetic nodes back to their pristine roots
for (const added of addedNodes) {
const roots = new Set<string>();
// 1:1 Replacement (e.g. Masked Node)
if (added.replacesId) {
const inheritedRoots = this.provenanceMap.get(added.replacesId);
if (inheritedRoots) {
for (const rootId of inheritedRoots) roots.add(rootId);
}
}
// N:1 Abstraction (e.g. Rolling Summary)
if (added.abstractsIds) {
for (const abstractId of added.abstractsIds) {
const inheritedRoots = this.provenanceMap.get(abstractId);
if (inheritedRoots) {
for (const rootId of inheritedRoots) roots.add(rootId);
}
}
}
// If it has no links back to the original graph, it is its own root
// (e.g., a system-injected instruction)
if (roots.size === 0) {
roots.add(added.id);
// It acts as a net-new pristine root.
if (!finalPristineMap.has(added.id)) {
const mutableMap = new Map<string, ConcreteNode>(finalPristineMap);
mutableMap.set(added.id, added);
finalPristineMap = mutableMap;
}
}
newProvenanceMap.set(added.id, roots);
}
// GC the Caches
// We only want to keep provenance and pristine entries that are reachable
// from the nodes in 'newGraph'.
const reachablePristineIds = new Set<string>();
const reachableCurrentIds = new Set<string>();
for (const node of newGraph) {
reachableCurrentIds.add(node.id);
const roots = newProvenanceMap.get(node.id);
if (roots) {
for (const root of roots) {
reachablePristineIds.add(root);
}
}
}
// Prune Provenance Map
for (const [id] of newProvenanceMap) {
if (!reachableCurrentIds.has(id)) {
newProvenanceMap.delete(id);
}
}
// Prune Pristine Map
const prunedPristineMap = new Map<string, ConcreteNode>();
for (const id of reachablePristineIds) {
const node = finalPristineMap.get(id);
if (node) prunedPristineMap.set(id, node);
}
finalPristineMap = prunedPristineMap;
return new ContextWorkingBufferImpl(
newGraph,
finalPristineMap,
newProvenanceMap,
[...this.history, mutation],
);
}
/** Removes nodes from the working buffer that were completely dropped from the upstream pristine history */
prunePristineNodes(
retainedIds: ReadonlySet<string>,
): ContextWorkingBufferImpl {
const newGraph = this.nodes.filter(
(n) => retainedIds.has(n.id) || !this.pristineNodesMap.has(n.id),
);
const newProvenanceMap = new Map(this.provenanceMap);
const reachablePristineIds = new Set<string>();
const reachableCurrentIds = new Set<string>();
for (const node of newGraph) {
reachableCurrentIds.add(node.id);
const roots = newProvenanceMap.get(node.id);
if (roots) {
for (const root of roots) {
if (retainedIds.has(root) || !this.pristineNodesMap.has(root)) {
reachablePristineIds.add(root);
}
}
}
}
for (const [id] of newProvenanceMap) {
if (!reachableCurrentIds.has(id)) {
newProvenanceMap.delete(id);
}
}
const prunedPristineMap = new Map<string, ConcreteNode>();
for (const id of reachablePristineIds) {
const node = this.pristineNodesMap.get(id);
if (node) prunedPristineMap.set(id, node);
}
return new ContextWorkingBufferImpl(
newGraph,
prunedPristineMap,
newProvenanceMap,
[...this.history],
);
}
getPristineNodes(id: string): readonly ConcreteNode[] {
const pristineIds = this.provenanceMap.get(id);
if (!pristineIds) return [];
return Array.from(pristineIds).map(
(pid) => this.pristineNodesMap.get(pid)!,
);
}
getAuditLog(): readonly GraphMutation[] {
return this.history;
}
getLineage(id: string): readonly ConcreteNode[] {
const lineage: ConcreteNode[] = [];
const currentNodesMap = new Map(this.nodes.map((n) => [n.id, n]));
let current = currentNodesMap.get(id);
while (current) {
lineage.push(current);
if (current.logicalParentId && current.logicalParentId !== current.id) {
current = currentNodesMap.get(current.logicalParentId);
} else {
break;
}
}
return lineage;
}
}

View file

@ -0,0 +1,29 @@
/**
* @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 { LiveInbox } from './inbox.js';
import type { NodeBehaviorRegistry } from '../graph/behaviorRegistry.js';
import type { ContextGraphMapper } from '../graph/mapper.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 eventBus: ContextEventBus;
readonly inbox: LiveInbox;
readonly behaviorRegistry: NodeBehaviorRegistry;
readonly graphMapper: ContextGraphMapper;
}

View file

@ -0,0 +1,44 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { ContextEnvironmentImpl } from './environmentImpl.js';
import { ContextTracer } from '../tracer.js';
import { ContextEventBus } from '../eventBus.js';
import { createMockLlmClient } from '../testing/contextTestUtils.js';
describe('ContextEnvironmentImpl', () => {
it('should initialize with defaults correctly', () => {
const tracer = new ContextTracer({ targetDir: '/tmp', sessionId: 'mock' });
const eventBus = new ContextEventBus();
const mockLlmClient = createMockLlmClient();
const env = new ContextEnvironmentImpl(
mockLlmClient,
'mock-session',
'mock-prompt',
'/tmp/trace',
'/tmp/temp',
tracer,
4,
eventBus,
);
expect(env.llmClient).toBe(mockLlmClient);
expect(env.sessionId).toBe('mock-session');
expect(env.promptId).toBe('mock-prompt');
expect(env.traceDir).toBe('/tmp/trace');
expect(env.projectTempDir).toBe('/tmp/temp');
expect(env.tracer).toBe(tracer);
expect(env.charsPerToken).toBe(4);
expect(env.eventBus).toBe(eventBus);
// Default internals
expect(env.behaviorRegistry).toBeDefined();
expect(env.tokenCalculator).toBeDefined();
expect(env.inbox).toBeDefined();
expect(env.graphMapper).toBeDefined();
});
});

View file

@ -0,0 +1,42 @@
/**
* @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 { LiveInbox } from './inbox.js';
import { NodeBehaviorRegistry } from '../graph/behaviorRegistry.js';
import { registerBuiltInBehaviors } from '../graph/builtinBehaviors.js';
import { ContextGraphMapper } from '../graph/mapper.js';
export class ContextEnvironmentImpl implements ContextEnvironment {
readonly tokenCalculator: ContextTokenCalculator;
readonly inbox: LiveInbox;
readonly behaviorRegistry: NodeBehaviorRegistry;
readonly graphMapper: ContextGraphMapper;
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,
) {
this.behaviorRegistry = new NodeBehaviorRegistry();
registerBuiltInBehaviors(this.behaviorRegistry);
this.tokenCalculator = new ContextTokenCalculator(
this.charsPerToken,
this.behaviorRegistry,
);
this.inbox = new LiveInbox();
this.graphMapper = new ContextGraphMapper(this.behaviorRegistry);
}
}

View file

@ -0,0 +1,45 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { LiveInbox, InboxSnapshotImpl } from './inbox.js';
describe('Inbox', () => {
it('should publish messages and provide snapshots', () => {
const inbox = new LiveInbox();
inbox.publish('test-topic', { data: 'hello' });
inbox.publish('other-topic', { data: 'world' });
const messages = inbox.getMessages();
expect(messages.length).toBe(2);
expect(messages[0].topic).toBe('test-topic');
expect(messages[0].payload).toEqual({ data: 'hello' });
});
it('should drain consumed messages from the snapshot', () => {
const inbox = new LiveInbox();
inbox.publish('test-topic', { data: 'hello' });
inbox.publish('other-topic', { data: 'world' });
const messages = inbox.getMessages();
const snapshot = new InboxSnapshotImpl(messages);
const filtered = snapshot.getMessages<{ data: string }>('test-topic');
expect(filtered.length).toBe(1);
expect(filtered[0].payload.data).toBe('hello');
// Consume the message
snapshot.consume(filtered[0].id);
// Provide the consumed IDs to the real inbox to drain them
inbox.drainConsumed(snapshot.getConsumedIds());
const finalMessages = inbox.getMessages();
expect(finalMessages.length).toBe(1);
expect(finalMessages[0].topic).toBe('other-topic');
});
});

View file

@ -0,0 +1,61 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import type { InboxMessage, InboxSnapshot } from '../pipeline.js';
export class LiveInbox {
private messages: InboxMessage[] = [];
publish<T>(topic: string, payload: T): void {
this.messages.push({
id: randomUUID(),
topic,
payload,
timestamp: Date.now(),
});
}
getMessages(): readonly InboxMessage[] {
return [...this.messages];
}
drainConsumed(consumedIds: Set<string>): void {
this.messages = this.messages.filter((m) => !consumedIds.has(m.id));
}
}
export class InboxSnapshotImpl implements InboxSnapshot {
private messages: readonly InboxMessage[];
private consumedIds = new Set<string>();
constructor(messages: readonly InboxMessage[]) {
this.messages = messages;
}
getMessages<T = unknown>(topic: string): ReadonlyArray<InboxMessage<T>> {
const raw = this.messages.filter((m) => m.topic === topic);
/*
* Architectural Justification for Unchecked Cast:
* The Inbox is a heterogeneous event bus designed to support arbitrary, declarative
* routing via configuration files (where topics are just strings). Because TypeScript
* completely erases generic type information (<T>) at runtime, the central array
* can only hold `unknown` payloads. To enforce strict type safety without a central
* registry (which would break decoupling) or heavy runtime validation (Zod schemas),
* we must assert the type boundary here. The contract relies on the async pipeline and Processor
* agreeing on the payload structure associated with the configured topic string.
*/
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
return raw as ReadonlyArray<InboxMessage<T>>;
}
consume(messageId: string): void {
this.consumedIds.add(messageId);
}
getConsumedIds(): Set<string> {
return this.consumedIds;
}
}

View file

@ -0,0 +1,224 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import assert from 'node:assert';
import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
import { PipelineOrchestrator } from './orchestrator.js';
import {
createMockEnvironment,
createDummyNode,
} from '../testing/contextTestUtils.js';
import type { ContextEnvironment } from './environment.js';
import type {
ContextProcessor,
AsyncContextProcessor,
ProcessArgs,
} from '../pipeline.js';
import type { PipelineDef, AsyncPipelineDef } from '../config/types.js';
import type { ContextEventBus } from '../eventBus.js';
import type { ConcreteNode, UserPrompt } from '../graph/types.js';
// A realistic mock processor that modifies the text of the first target node
function createModifyingProcessor(id: string): ContextProcessor {
return {
id,
name: 'ModifyingProcessor',
process: async (args: ProcessArgs) => {
const newTargets = [...args.targets];
if (newTargets.length > 0 && newTargets[0].type === 'USER_PROMPT') {
const prompt = newTargets[0];
const newParts = [...prompt.semanticParts];
if (newParts.length > 0 && newParts[0].type === 'text') {
newParts[0] = {
...newParts[0],
text: newParts[0].text + ' [modified]',
};
}
newTargets[0] = {
...prompt,
id: prompt.id + '-modified',
replacesId: prompt.id,
semanticParts: newParts,
};
}
return newTargets;
},
};
}
// A processor that just throws an error
function createThrowingProcessor(id: string): ContextProcessor {
return {
id,
name: 'Throwing',
process: async (): Promise<readonly ConcreteNode[]> => {
throw new Error('Processor failed intentionally');
},
};
}
// A mock async processor that signals it ran
function createMockAsyncProcessor(
id: string,
executeSpy: ReturnType<typeof vi.fn>,
): AsyncContextProcessor {
return {
id,
name: 'MockAsyncProcessor',
process: async (args: ProcessArgs) => {
executeSpy(args);
},
};
}
describe('PipelineOrchestrator (Component)', () => {
let env: ContextEnvironment;
let eventBus: ContextEventBus;
beforeEach(() => {
env = createMockEnvironment();
eventBus = env.eventBus;
});
afterEach(() => {
vi.restoreAllMocks();
});
const setupOrchestrator = (
pipelines: PipelineDef[],
asyncPipelines: AsyncPipelineDef[] = [],
) => {
const orchestrator = new PipelineOrchestrator(
pipelines,
asyncPipelines,
env,
eventBus,
env.tracer,
);
return orchestrator;
};
describe('Synchronous Pipeline Execution', () => {
it('applies processors in sequence on matching trigger', async () => {
const pipelines: PipelineDef[] = [
{
name: 'TestPipeline',
triggers: ['new_message'],
processors: [createModifyingProcessor('Mod')],
},
];
const orchestrator = setupOrchestrator(pipelines);
const originalNode = createDummyNode('ep1', 'USER_PROMPT', 50, {
semanticParts: [{ type: 'text', text: 'Original' }],
});
const processed = await orchestrator.executeTriggerSync(
'new_message',
[originalNode],
new Set([originalNode.id]),
new Set(),
);
expect(processed.length).toBe(1);
const resultingNode = processed[0] as UserPrompt;
assert(resultingNode.semanticParts[0].type === 'text');
expect(resultingNode.semanticParts[0].text).toBe('Original [modified]');
expect(resultingNode.replacesId).toBe(originalNode.id);
});
it('bypasses pipelines that do not match the trigger', async () => {
const pipelines: PipelineDef[] = [
{
name: 'TestPipeline',
triggers: ['gc_backstop'], // Different trigger
processors: [createModifyingProcessor('Mod')],
},
];
const orchestrator = setupOrchestrator(pipelines);
const originalNode = createDummyNode('ep1', 'USER_PROMPT', 50, {
semanticParts: [{ type: 'text', text: 'Original' }],
});
const processed = await orchestrator.executeTriggerSync(
'new_message',
[originalNode],
new Set([originalNode.id]),
new Set(),
);
expect(processed).toEqual([originalNode]); // Untouched
});
it('gracefully handles a failing processor without crashing the pipeline', async () => {
const pipelines: PipelineDef[] = [
{
name: 'FailingPipeline',
triggers: ['new_message'],
processors: [
createThrowingProcessor('Thrower'),
createModifyingProcessor('Mod'),
],
},
];
const orchestrator = setupOrchestrator(pipelines);
const originalNode = createDummyNode('ep1', 'USER_PROMPT', 50, {
semanticParts: [{ type: 'text', text: 'Original' }],
});
// The throwing processor should be caught and logged, allowing Mod to still run.
const processed = await orchestrator.executeTriggerSync(
'new_message',
[originalNode],
new Set([originalNode.id]),
new Set(),
);
expect(processed.length).toBe(1);
const resultingNode = processed[0] as UserPrompt;
assert(resultingNode.semanticParts[0].type === 'text');
expect(resultingNode.semanticParts[0].text).toBe('Original [modified]');
});
});
describe('Asynchronous async pipeline Events', () => {
it('routes emitChunkReceived to async pipelines with nodes_added trigger', async () => {
const executeSpy = vi.fn();
const asyncProcessor = createMockAsyncProcessor(
'MyAsyncProcessor',
executeSpy,
);
setupOrchestrator(
[],
[
{
name: 'TestAsync',
triggers: ['nodes_added'],
processors: [asyncProcessor],
},
],
);
const node1 = createDummyNode('ep1', 'USER_PROMPT', 10);
const node2 = createDummyNode('ep1', 'AGENT_THOUGHT', 20);
eventBus.emitChunkReceived({
nodes: [node1, node2],
targetNodeIds: new Set([node2.id]),
});
// Yield event loop
await new Promise((resolve) => setTimeout(resolve, 0));
expect(executeSpy).toHaveBeenCalledTimes(1);
const callArgs = executeSpy.mock.calls[0][0];
expect(callArgs.targets).toEqual([node2]); // AsyncProcessors only get the target nodes
});
});
});

View file

@ -0,0 +1,218 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ConcreteNode } from '../graph/types.js';
import type {
AsyncPipelineDef,
PipelineDef,
PipelineTrigger,
} from '../config/types.js';
import type {
ContextEnvironment,
ContextEventBus,
ContextTracer,
} from './environment.js';
import { debugLogger } from '../../utils/debugLogger.js';
import { InboxSnapshotImpl } from './inbox.js';
import { ContextWorkingBufferImpl } from './contextWorkingBuffer.js';
export class PipelineOrchestrator {
private activeTimers: NodeJS.Timeout[] = [];
constructor(
private readonly pipelines: PipelineDef[],
private readonly asyncPipelines: AsyncPipelineDef[],
private readonly env: ContextEnvironment,
private readonly eventBus: ContextEventBus,
private readonly tracer: ContextTracer,
) {
this.setupTriggers();
}
private isNodeAllowed(
node: ConcreteNode,
triggerTargets: ReadonlySet<string>,
protectedLogicalIds: ReadonlySet<string> = new Set(),
): boolean {
return (
triggerTargets.has(node.id) &&
!protectedLogicalIds.has(node.id) &&
(!node.logicalParentId || !protectedLogicalIds.has(node.logicalParentId))
);
}
private setupTriggers() {
const bindTriggers = <P extends PipelineDef | AsyncPipelineDef>(
pipelines: P[],
executeFn: (
pipeline: P,
nodes: readonly ConcreteNode[],
targets: ReadonlySet<string>,
protectedIds: ReadonlySet<string>,
) => void,
) => {
for (const pipeline of pipelines) {
for (const trigger of pipeline.triggers) {
if (typeof trigger === 'object' && trigger.type === 'timer') {
const timer = setInterval(() => {
// Background timers not fully implemented in V1 yet
}, trigger.intervalMs);
this.activeTimers.push(timer);
} else if (
trigger === 'retained_exceeded' ||
trigger === 'nodes_aged_out'
) {
this.eventBus.onConsolidationNeeded((event) => {
executeFn(pipeline, event.nodes, event.targetNodeIds, new Set());
});
} else if (trigger === 'new_message' || trigger === 'nodes_added') {
this.eventBus.onChunkReceived((event) => {
executeFn(pipeline, event.nodes, event.targetNodeIds, new Set());
});
}
}
}
};
bindTriggers(this.pipelines, (pipeline, nodes, targets, protectedIds) => {
void this.executePipelineAsync(
pipeline,
nodes,
new Set(targets),
new Set(protectedIds),
);
});
bindTriggers(this.asyncPipelines, (pipeline, nodes, targetIds) => {
const inboxSnapshot = new InboxSnapshotImpl(
this.env.inbox.getMessages() || [],
);
const targets = nodes.filter((n) => targetIds.has(n.id));
for (const processor of pipeline.processors) {
processor
.process({
targets,
inbox: inboxSnapshot,
buffer: ContextWorkingBufferImpl.initialize(nodes),
})
.catch((e: unknown) =>
debugLogger.error(`AsyncProcessor ${processor.name} failed:`, e),
);
}
});
}
shutdown() {
for (const timer of this.activeTimers) {
clearInterval(timer);
}
}
async executeTriggerSync(
trigger: PipelineTrigger,
nodes: readonly ConcreteNode[],
triggerTargets: ReadonlySet<string>,
protectedLogicalIds: ReadonlySet<string> = new Set(),
): Promise<readonly ConcreteNode[]> {
let currentBuffer = ContextWorkingBufferImpl.initialize(nodes);
const triggerPipelines = this.pipelines.filter((p) =>
p.triggers.includes(trigger),
);
// Freeze the inbox for this pipeline run
const inboxSnapshot = new InboxSnapshotImpl(
this.env.inbox.getMessages() || [],
);
for (const pipeline of triggerPipelines) {
for (const processor of pipeline.processors) {
try {
this.tracer.logEvent(
'Orchestrator',
`Executing processor synchronously: ${processor.id}`,
);
const allowedTargets = currentBuffer.nodes.filter((n) =>
this.isNodeAllowed(n, triggerTargets, protectedLogicalIds),
);
const returnedNodes = await processor.process({
buffer: currentBuffer,
targets: allowedTargets,
inbox: inboxSnapshot,
});
currentBuffer = currentBuffer.applyProcessorResult(
processor.id,
allowedTargets,
returnedNodes,
);
} catch (error) {
debugLogger.error(
`Synchronous processor ${processor.id} failed:`,
error,
);
}
}
}
// Success! Drain consumed messages
this.env.inbox.drainConsumed(inboxSnapshot.getConsumedIds());
return currentBuffer.nodes;
}
private async executePipelineAsync(
pipeline: PipelineDef,
nodes: readonly ConcreteNode[],
triggerTargets: Set<string>,
protectedLogicalIds: ReadonlySet<string> = new Set(),
) {
this.tracer.logEvent(
'Orchestrator',
`Triggering async pipeline: ${pipeline.name}`,
);
if (!nodes || nodes.length === 0) return;
let currentBuffer = ContextWorkingBufferImpl.initialize(nodes);
const inboxSnapshot = new InboxSnapshotImpl(
this.env.inbox.getMessages() || [],
);
for (const processor of pipeline.processors) {
try {
this.tracer.logEvent(
'Orchestrator',
`Executing processor: ${processor.id} (async)`,
);
const allowedTargets = currentBuffer.nodes.filter((n) =>
this.isNodeAllowed(n, triggerTargets, protectedLogicalIds),
);
const returnedNodes = await processor.process({
buffer: currentBuffer,
targets: allowedTargets,
inbox: inboxSnapshot,
});
currentBuffer = currentBuffer.applyProcessorResult(
processor.id,
allowedTargets,
returnedNodes,
);
} catch (error) {
debugLogger.error(
`Pipeline ${pipeline.name} failed async at ${processor.id}:`,
error,
);
return;
}
}
this.env.inbox.drainConsumed(inboxSnapshot.getConsumedIds());
}
}

View file

@ -0,0 +1,110 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import assert from 'node:assert';
import { describe, it, expect } from 'vitest';
import { createBlobDegradationProcessor } from './blobDegradationProcessor.js';
import {
createMockProcessArgs,
createMockEnvironment,
createDummyNode,
} from '../testing/contextTestUtils.js';
import type { UserPrompt, SemanticPart, ConcreteNode } from '../graph/types.js';
describe('BlobDegradationProcessor', () => {
it('should ignore text parts and only target inline_data and file_data', async () => {
const env = createMockEnvironment();
// charsPerToken = 1
// We want the degraded text to be cheaper than the original blob.
// Degraded text is ~100 chars ("...degraded to text...").
// So we make the blob data 200 chars.
const fakeData = 'A'.repeat(200);
const processor = createBlobDegradationProcessor(
'BlobDegradationProcessor',
env,
);
const parts: SemanticPart[] = [
{ type: 'text', text: 'Hello' },
{ type: 'inline_data', mimeType: 'image/png', data: fakeData },
{ type: 'text', text: 'World' },
];
const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, {
semanticParts: parts,
}) as UserPrompt;
const targets = [prompt];
const result = await processor.process(createMockProcessArgs(targets));
expect(result.length).toBe(1);
const modifiedPrompt = result[0] as UserPrompt;
expect(modifiedPrompt.id).not.toBe(prompt.id);
expect(modifiedPrompt.semanticParts.length).toBe(3);
// Text parts should be untouched
expect(modifiedPrompt.semanticParts[0]).toEqual(parts[0]);
expect(modifiedPrompt.semanticParts[2]).toEqual(parts[2]);
// The inline_data part should be replaced with text
const degradedPart = modifiedPrompt.semanticParts[1];
expect(degradedPart.type).toBe('text');
assert(degradedPart.type === 'text');
expect(degradedPart.text).toContain(
'[Multi-Modal Blob (image/png, 0.00MB) degraded to text',
);
});
it('should degrade all blobs unconditionally', async () => {
const env = createMockEnvironment();
const processor = createBlobDegradationProcessor(
'BlobDegradationProcessor',
env,
);
// Tokens for fileData = 258.
// Degraded text = "[File Reference (video/mp4) degraded to text to preserve context window. Original URI: gs://test1]"
// Degraded text length ~100 characters.
// Since charsPerToken=1, degraded text = 100 tokens.
// Tokens saved = 258 - 100 = 158. This is > 0, so it WILL degrade it!
const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, {
semanticParts: [
{ type: 'file_data', mimeType: 'video/mp4', fileUri: 'gs://test1' },
{ type: 'file_data', mimeType: 'video/mp4', fileUri: 'gs://test2' },
],
}) as UserPrompt;
const targets = [prompt];
const result = await processor.process(createMockProcessArgs(targets));
const modifiedPrompt = result[0] as UserPrompt;
expect(modifiedPrompt.semanticParts.length).toBe(2);
// Both parts should be degraded
expect(modifiedPrompt.semanticParts[0].type).toBe('text');
expect(modifiedPrompt.semanticParts[1].type).toBe('text');
});
it('should return exactly the targets array if targets are empty', async () => {
const env = createMockEnvironment();
const processor = createBlobDegradationProcessor(
'BlobDegradationProcessor',
env,
);
const targets: ConcreteNode[] = [];
const result = await processor.process(createMockProcessArgs(targets));
expect(result).toBe(targets);
});
});

View file

@ -0,0 +1,153 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import type { JSONSchemaType } from 'ajv';
import type { ProcessArgs, ContextProcessor } from '../pipeline.js';
import * as fs from 'node:fs/promises';
import * as path from 'node:path';
import type { ConcreteNode, UserPrompt } from '../graph/types.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { sanitizeFilenamePart } from '../../utils/fileUtils.js';
export type BlobDegradationProcessorOptions = Record<string, never>;
export const BlobDegradationProcessorOptionsSchema: JSONSchemaType<BlobDegradationProcessorOptions> =
{
type: 'object',
properties: {},
required: [],
};
export function createBlobDegradationProcessor(
id: string,
env: ContextEnvironment,
): ContextProcessor {
return {
id,
name: 'BlobDegradationProcessor',
process: async ({ targets }: ProcessArgs) => {
if (targets.length === 0) {
return targets;
}
let directoryCreated = false;
let blobOutputsDir = path.join(env.projectTempDir, 'degraded-blobs');
const sessionId = env.sessionId;
if (sessionId) {
blobOutputsDir = path.join(
blobOutputsDir,
`session-${sanitizeFilenamePart(sessionId)}`,
);
}
const ensureDir = async () => {
if (!directoryCreated) {
await fs.mkdir(blobOutputsDir, { recursive: true });
directoryCreated = true;
}
};
const returnedNodes: ConcreteNode[] = [];
// Forward scan, looking for bloated non-text parts to degrade
for (const node of targets) {
switch (node.type) {
case 'USER_PROMPT': {
let modified = false;
const newParts = [...node.semanticParts];
for (let j = 0; j < node.semanticParts.length; j++) {
const part = node.semanticParts[j];
if (part.type === 'text') continue;
let newText = '';
let tokensSaved = 0;
switch (part.type) {
case 'inline_data': {
await ensureDir();
const ext = part.mimeType.split('/')[1] || 'bin';
const fileName = `blob_${Date.now()}_${randomUUID()}.${ext}`;
const filePath = path.join(blobOutputsDir, fileName);
const buffer = Buffer.from(part.data, 'base64');
await fs.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}]`;
const oldTokens = env.tokenCalculator.estimateTokensForParts([
{
inlineData: { mimeType: part.mimeType, data: part.data },
},
]);
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
tokensSaved = oldTokens - newTokens;
break;
}
case 'file_data': {
newText = `[File Reference (${part.mimeType}) degraded to text to preserve context window. Original URI: ${part.fileUri}]`;
const oldTokens = env.tokenCalculator.estimateTokensForParts([
{
fileData: {
mimeType: part.mimeType,
fileUri: part.fileUri,
},
},
]);
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
tokensSaved = oldTokens - newTokens;
break;
}
case 'raw_part': {
newText = `[Unknown Part degraded to text to preserve context window.]`;
const oldTokens = env.tokenCalculator.estimateTokensForParts([
part.part,
]);
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: newText },
]);
tokensSaved = oldTokens - newTokens;
break;
}
default:
break;
}
if (newText && tokensSaved > 0) {
newParts[j] = { type: 'text', text: newText };
modified = true;
}
}
if (modified) {
const degradedNode: UserPrompt = {
...node,
id: randomUUID(),
semanticParts: newParts,
replacesId: node.id,
};
returnedNodes.push(degradedNode);
} else {
returnedNodes.push(node);
}
break;
}
default:
returnedNodes.push(node);
break;
}
}
return returnedNodes;
},
};
}

View file

@ -0,0 +1,82 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type {
ContextProcessor,
BackstopTargetOptions,
ProcessArgs,
} from '../pipeline.js';
import type { ConcreteNode } from '../graph/types.js';
import type { JSONSchemaType } from 'ajv';
import type { ContextEnvironment } from '../pipeline/environment.js';
export type HistoryTruncationProcessorOptions = BackstopTargetOptions;
export const HistoryTruncationProcessorOptionsSchema: JSONSchemaType<HistoryTruncationProcessorOptions> =
{
type: 'object',
properties: {
target: {
type: 'string',
enum: ['incremental', 'freeNTokens', 'max'],
nullable: true,
},
freeTokensTarget: { type: 'number', nullable: true },
},
required: [],
};
export function createHistoryTruncationProcessor(
id: string,
env: ContextEnvironment,
options: HistoryTruncationProcessorOptions,
): ContextProcessor {
return {
id,
name: 'HistoryTruncationProcessor',
process: async ({ targets }: ProcessArgs) => {
const strategy = options.target ?? 'max';
const keptNodes: ConcreteNode[] = [];
if (strategy === 'incremental') {
// 'incremental' simply drops the single oldest node in the targets, ignoring tokens.
let removedNodes = 0;
for (const node of targets) {
if (removedNodes < 1) {
removedNodes++;
continue;
}
keptNodes.push(node);
}
return keptNodes;
}
let targetTokensToRemove = 0;
if (strategy === 'freeNTokens') {
targetTokensToRemove = options.freeTokensTarget ?? 0;
if (targetTokensToRemove <= 0) return targets;
} else if (strategy === 'max') {
// 'max' means we remove all targets without stopping early
targetTokensToRemove = Infinity;
}
let removedTokens = 0;
// The targets are sequentially ordered from oldest to newest.
// We want to delete the oldest targets first.
for (const node of targets) {
if (removedTokens >= targetTokensToRemove) {
keptNodes.push(node);
continue;
}
removedTokens += env.tokenCalculator.getTokenCost(node);
}
return keptNodes;
},
};
}

View file

@ -0,0 +1,152 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import assert from 'node:assert';
import { describe, it, expect } from 'vitest';
import { createNodeDistillationProcessor } from './nodeDistillationProcessor.js';
import {
createMockProcessArgs,
createMockEnvironment,
createDummyNode,
createDummyToolNode,
createMockLlmClient,
} from '../testing/contextTestUtils.js';
import type {
UserPrompt,
AgentThought,
ToolExecution,
} from '../graph/types.js';
describe('NodeDistillationProcessor', () => {
it('should trigger summarization via LLM for long text parts', async () => {
const mockLlmClient = createMockLlmClient(['Mocked Summary!']);
// Use charsPerToken=1 naturally.
const env = createMockEnvironment({
llmClient: mockLlmClient,
});
const processor = createNodeDistillationProcessor(
'NodeDistillationProcessor',
env,
{
nodeThresholdTokens: 10,
},
);
const longText = 'A'.repeat(50); // 50 chars
const prompt = createDummyNode(
'ep1',
'USER_PROMPT',
50,
{
semanticParts: [{ type: 'text', text: longText }],
},
'prompt-id',
) as UserPrompt;
const thought = createDummyNode(
'ep1',
'AGENT_THOUGHT',
50,
{
text: longText,
},
'thought-id',
) as AgentThought;
const tool = createDummyToolNode(
'ep1',
5,
500,
{
observation: { result: 'A'.repeat(500) },
},
'tool-id',
);
const targets = [prompt, thought, tool];
const result = await processor.process(createMockProcessArgs(targets));
expect(result.length).toBe(3);
// 1. User Prompt
const compressedPrompt = result[0] as UserPrompt;
expect(compressedPrompt.id).not.toBe(prompt.id);
expect(compressedPrompt.semanticParts[0].type).toBe('text');
assert(compressedPrompt.semanticParts[0].type === 'text');
expect(compressedPrompt.semanticParts[0].text).toBe('Mocked Summary!');
// 2. Agent Thought
const compressedThought = result[1] as AgentThought;
expect(compressedThought.id).not.toBe(thought.id);
expect(compressedThought.text).toBe('Mocked Summary!');
// 3. Tool Execution
const compressedTool = result[2] as ToolExecution;
expect(compressedTool.id).not.toBe(tool.id);
expect(compressedTool.observation).toEqual({ summary: 'Mocked Summary!' });
expect(mockLlmClient.generateContent).toHaveBeenCalledTimes(3);
});
it('should ignore nodes that are below the threshold', async () => {
const mockLlmClient = createMockLlmClient(['S']); // length = 1
const env = createMockEnvironment({
llmClient: mockLlmClient,
});
const processor = createNodeDistillationProcessor(
'NodeDistillationProcessor',
env,
{
nodeThresholdTokens: 100, // Very high threshold
},
);
const shortText = 'Short text'; // 10 chars
const prompt = createDummyNode(
'ep1',
'USER_PROMPT',
10,
{
semanticParts: [{ type: 'text', text: shortText }],
},
'prompt-id',
) as UserPrompt;
const thought = createDummyNode(
'ep1',
'AGENT_THOUGHT',
13,
{
text: 'Short thought',
},
'thought-id',
) as AgentThought;
const targets = [prompt, thought];
const result = await processor.process(createMockProcessArgs(targets));
expect(result.length).toBe(2);
// 1. User Prompt (NOT compressed)
const untouchedPrompt = result[0] as UserPrompt;
expect(untouchedPrompt.id).toBe(prompt.id);
// 2. Agent Thought (NOT compressed)
const untouchedThought = result[1] as AgentThought;
expect(untouchedThought.id).toBe(thought.id);
// LLM should not have been called
expect(mockLlmClient.generateContent).toHaveBeenCalledTimes(0);
});
});

View file

@ -0,0 +1,203 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import type { ConcreteNode } from '../graph/types.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { debugLogger } from '../../utils/debugLogger.js';
import { getResponseText } from '../../utils/partUtils.js';
import { LlmRole } from '../../telemetry/llmRole.js';
export interface NodeDistillationProcessorOptions {
nodeThresholdTokens: number;
}
export const NodeDistillationProcessorOptionsSchema: JSONSchemaType<NodeDistillationProcessorOptions> =
{
type: 'object',
properties: {
nodeThresholdTokens: { type: 'number' },
},
required: ['nodeThresholdTokens'],
};
export function createNodeDistillationProcessor(
id: string,
env: ContextEnvironment,
options: NodeDistillationProcessorOptions,
): ContextProcessor {
const generateSummary = async (
text: string,
contextInfo: string,
): Promise<string> => {
try {
const response = await env.llmClient.generateContent({
role: LlmRole.UTILITY_COMPRESSOR,
modelConfigKey: { model: 'gemini-3-flash-base' },
promptId: env.promptId,
abortSignal: new AbortController().signal,
contents: [
{
role: 'user',
parts: [{ text }],
},
],
systemInstruction: {
role: 'system',
parts: [
{
text: `You are an expert context compressor. Your job is to drastically shorten the following ${contextInfo} while preserving the absolute core semantic meaning, facts, and intent. Omit all conversational filler, pleasantries, or redundant information. Return ONLY the compressed summary.`,
},
],
},
});
return getResponseText(response) || text;
} catch (e) {
debugLogger.warn(
`NodeDistillationProcessor failed to summarize ${contextInfo}`,
e,
);
return text; // Fallback to original text on API failure
}
};
return {
id,
name: 'NodeDistillationProcessor',
process: async ({ targets }: ProcessArgs) => {
const semanticConfig = options;
const limitTokens = semanticConfig.nodeThresholdTokens;
const thresholdChars = env.tokenCalculator.tokensToChars(limitTokens);
const returnedNodes: ConcreteNode[] = [];
// Scan the target working buffer and unconditionally apply the configured hyperparameter threshold
for (const node of targets) {
switch (node.type) {
case 'USER_PROMPT': {
let modified = false;
const newParts = [...node.semanticParts];
for (let j = 0; j < node.semanticParts.length; j++) {
const part = node.semanticParts[j];
if (part.type !== 'text') continue;
if (part.text.length > thresholdChars) {
const summary = await generateSummary(part.text, 'User Prompt');
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: summary },
]);
const oldTokens = env.tokenCalculator.estimateTokensForParts([
{ text: part.text },
]);
if (newTokens < oldTokens) {
newParts[j] = { type: 'text', text: summary };
modified = true;
}
}
}
if (modified) {
returnedNodes.push({
...node,
id: randomUUID(),
semanticParts: newParts,
replacesId: node.id,
});
} else {
returnedNodes.push(node);
}
break;
}
case 'AGENT_THOUGHT': {
if (node.text.length > thresholdChars) {
const summary = await generateSummary(node.text, 'Agent Thought');
const newTokens = env.tokenCalculator.estimateTokensForParts([
{ text: summary },
]);
const oldTokens = env.tokenCalculator.getTokenCost(node);
if (newTokens < oldTokens) {
returnedNodes.push({
...node,
id: randomUUID(),
text: summary,
replacesId: node.id,
});
break;
}
}
returnedNodes.push(node);
break;
}
case 'TOOL_EXECUTION': {
const rawObs = node.observation;
let stringifiedObs = '';
if (typeof rawObs === 'string') {
stringifiedObs = rawObs;
} else {
try {
stringifiedObs = JSON.stringify(rawObs);
} catch {
stringifiedObs = String(rawObs);
}
}
if (stringifiedObs.length > thresholdChars) {
const summary = await generateSummary(
stringifiedObs,
node.toolName || 'unknown',
);
const newObsObject = { summary };
const newObsTokens = env.tokenCalculator.estimateTokensForParts([
{
functionResponse: {
name: node.toolName || 'unknown',
response: newObsObject,
id: node.id,
},
},
]);
const oldObsTokens =
node.tokens?.observation ??
env.tokenCalculator.getTokenCost(node);
const intentTokens = node.tokens?.intent ?? 0;
if (newObsTokens < oldObsTokens) {
returnedNodes.push({
...node,
id: randomUUID(),
observation: newObsObject,
tokens: {
intent: intentTokens,
observation: newObsTokens,
},
replacesId: node.id,
});
break;
}
}
returnedNodes.push(node);
break;
}
default:
returnedNodes.push(node);
break;
}
}
return returnedNodes;
},
};
}

View file

@ -0,0 +1,136 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import assert from 'node:assert';
import { describe, it, expect } from 'vitest';
import { createNodeTruncationProcessor } from './nodeTruncationProcessor.js';
import {
createMockProcessArgs,
createMockEnvironment,
createDummyNode,
} from '../testing/contextTestUtils.js';
import type { UserPrompt, AgentThought, AgentYield } from '../graph/types.js';
describe('NodeTruncationProcessor', () => {
it('should truncate nodes that exceed maxTokensPerNode', async () => {
// env.tokenCalculator uses charsPerToken=1 natively.
const env = createMockEnvironment();
const processor = createNodeTruncationProcessor(
'NodeTruncationProcessor',
env,
{
maxTokensPerNode: 10, // 10 chars limit
},
);
const longText = 'A'.repeat(50); // 50 tokens
const prompt = createDummyNode(
'ep1',
'USER_PROMPT',
50,
{
semanticParts: [{ type: 'text', text: longText }],
},
'prompt-id',
) as UserPrompt;
const thought = createDummyNode(
'ep1',
'AGENT_THOUGHT',
50,
{
text: longText,
},
'thought-id',
) as AgentThought;
const yieldNode = createDummyNode(
'ep1',
'AGENT_YIELD',
50,
{
text: longText,
},
'yield-id',
) as AgentYield;
const targets = [prompt, thought, yieldNode];
const result = await processor.process(createMockProcessArgs(targets));
expect(result.length).toBe(3);
// 1. User Prompt
const squashedPrompt = result[0] as UserPrompt;
expect(squashedPrompt.id).not.toBe(prompt.id);
expect(squashedPrompt.semanticParts[0].type).toBe('text');
assert(squashedPrompt.semanticParts[0].type === 'text');
expect(squashedPrompt.semanticParts[0].text).toContain('[... OMITTED');
// 2. Agent Thought
const squashedThought = result[1] as AgentThought;
expect(squashedThought.id).not.toBe(thought.id);
expect(squashedThought.text).toContain('[... OMITTED');
// 3. Agent Yield
const squashedYield = result[2] as AgentYield;
expect(squashedYield.id).not.toBe(yieldNode.id);
expect(squashedYield.text).toContain('[... OMITTED');
});
it('should ignore nodes that are below maxTokensPerNode', async () => {
const env = createMockEnvironment();
const processor = createNodeTruncationProcessor(
'NodeTruncationProcessor',
env,
{
maxTokensPerNode: 100, // 100 chars limit
},
);
const shortText = 'Short text'; // 10 chars
const prompt = createDummyNode(
'ep1',
'USER_PROMPT',
10,
{
semanticParts: [{ type: 'text', text: shortText }],
},
'prompt-id',
) as UserPrompt;
const thought = createDummyNode(
'ep1',
'AGENT_THOUGHT',
13,
{
text: 'Short thought', // 13 chars
},
'thought-id',
) as AgentThought;
const targets = [prompt, thought];
const result = await processor.process(createMockProcessArgs(targets));
expect(result.length).toBe(2);
// 1. User Prompt (untouched)
const squashedPrompt = result[0] as UserPrompt;
expect(squashedPrompt.id).toBe(prompt.id);
assert(squashedPrompt.semanticParts[0].type === 'text');
expect(squashedPrompt.semanticParts[0].text).not.toContain('[... OMITTED');
// 2. Agent Thought (untouched)
const untouchedThought = result[1] as AgentThought;
expect(untouchedThought.id).toBe(thought.id);
expect(untouchedThought.text).not.toContain('[... OMITTED');
});
});

View file

@ -0,0 +1,144 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { truncateProportionally } from '../truncation.js';
import type { ConcreteNode } from '../graph/types.js';
export interface NodeTruncationProcessorOptions {
maxTokensPerNode: number;
}
export const NodeTruncationProcessorOptionsSchema: JSONSchemaType<NodeTruncationProcessorOptions> =
{
type: 'object',
properties: {
maxTokensPerNode: { type: 'number' },
},
required: ['maxTokensPerNode'],
};
export function createNodeTruncationProcessor(
id: string,
env: ContextEnvironment,
options: NodeTruncationProcessorOptions,
): ContextProcessor {
const tryApplySquash = (
text: string,
limitChars: number,
): {
text: string;
newTokens: number;
oldTokens: number;
tokensSaved: number;
} | null => {
const originalLength = text.length;
if (originalLength <= limitChars) return null;
const newText = truncateProportionally(
text,
limitChars,
`\n\n[... OMITTED ${originalLength - limitChars} chars ...]\n\n`,
);
if (newText !== text) {
// Using accurate TokenCalculator instead of simple math
const newTokens = env.tokenCalculator.estimateTokensForString(newText);
const oldTokens = env.tokenCalculator.estimateTokensForString(text);
const tokensSaved = oldTokens - newTokens;
if (tokensSaved > 0) {
return { text: newText, newTokens, oldTokens, tokensSaved };
}
}
return null;
};
return {
id,
name: 'NodeTruncationProcessor',
process: async ({ targets }: ProcessArgs) => {
if (targets.length === 0) {
return targets;
}
const { maxTokensPerNode } = options;
const limitChars = env.tokenCalculator.tokensToChars(maxTokensPerNode);
const returnedNodes: ConcreteNode[] = [];
for (const node of targets) {
switch (node.type) {
case 'USER_PROMPT': {
let modified = false;
const newParts = [...node.semanticParts];
for (let j = 0; j < node.semanticParts.length; j++) {
const part = node.semanticParts[j];
if (part.type === 'text') {
const squashResult = tryApplySquash(part.text, limitChars);
if (squashResult) {
newParts[j] = { type: 'text', text: squashResult.text };
modified = true;
}
}
}
if (modified) {
returnedNodes.push({
...node,
id: randomUUID(),
semanticParts: newParts,
replacesId: node.id,
});
} else {
returnedNodes.push(node);
}
break;
}
case 'AGENT_THOUGHT': {
const squashResult = tryApplySquash(node.text, limitChars);
if (squashResult) {
returnedNodes.push({
...node,
id: randomUUID(),
text: squashResult.text,
replacesId: node.id,
});
} else {
returnedNodes.push(node);
}
break;
}
case 'AGENT_YIELD': {
const squashResult = tryApplySquash(node.text, limitChars);
if (squashResult) {
returnedNodes.push({
...node,
id: randomUUID(),
text: squashResult.text,
replacesId: node.id,
});
} else {
returnedNodes.push(node);
}
break;
}
default:
returnedNodes.push(node);
break;
}
}
return returnedNodes;
},
};
}

View file

@ -0,0 +1,98 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { createRollingSummaryProcessor } from './rollingSummaryProcessor.js';
import {
createMockProcessArgs,
createMockEnvironment,
createDummyNode,
} from '../testing/contextTestUtils.js';
describe('RollingSummaryProcessor', () => {
it('should initialize with correct default options', () => {
const env = createMockEnvironment();
const processor = createRollingSummaryProcessor(
'RollingSummaryProcessor',
env,
{
target: 'incremental',
},
);
expect(processor.id).toBe('RollingSummaryProcessor');
});
it('should summarize older nodes when the deficit exceeds the threshold', async () => {
// env.tokenCalculator uses charsPerToken=1 based on createMockEnvironment
const env = createMockEnvironment();
// We want to free exactly 100 tokens.
// We will supply nodes that cost 50 tokens each.
const processor = createRollingSummaryProcessor(
'RollingSummaryProcessor',
env,
{
target: 'freeNTokens',
freeTokensTarget: 100,
},
);
const text50 = 'A'.repeat(50);
const targets = [
createDummyNode(
'ep1',
'USER_PROMPT',
50,
{ semanticParts: [{ type: 'text', text: text50 }] },
'id1',
),
createDummyNode('ep1', 'AGENT_THOUGHT', 50, { text: text50 }, 'id2'),
createDummyNode('ep1', 'AGENT_YIELD', 50, { text: text50 }, 'id3'),
];
const result = await processor.process(createMockProcessArgs(targets));
// 3 nodes at 50 cost each.
// The first node (id1) is the initial USER_PROMPT and is always skipped by RollingSummaryProcessor.
// Node id2 adds 50 deficit. Node id3 adds 50 deficit. Total = 100 deficit, which hits the target break point.
// Thus, id2 and id3 are summarized into a new ROLLING_SUMMARY node.
expect(result.length).toBe(2);
expect(result[0].type).toBe('USER_PROMPT');
expect(result[1].type).toBe('ROLLING_SUMMARY');
});
it('should preserve targets if deficit does not trigger summary', async () => {
const env = createMockEnvironment();
// We want to free 100 tokens, but our nodes will only cost 10 tokens each.
const processor = createRollingSummaryProcessor(
'RollingSummaryProcessor',
env,
{
target: 'freeNTokens',
freeTokensTarget: 100,
},
);
const text10 = 'A'.repeat(10);
const targets = [
createDummyNode(
'ep1',
'USER_PROMPT',
10,
{ semanticParts: [{ type: 'text', text: text10 }] },
'id1',
),
createDummyNode('ep1', 'AGENT_THOUGHT', 10, { text: text10 }, 'id2'),
];
const result = await processor.process(createMockProcessArgs(targets));
// Deficit accumulator reaches 10. This is < 100 limit, and total summarizable nodes < 2 anyway.
expect(result.length).toBe(2);
expect(result[0].type).toBe('USER_PROMPT');
expect(result[1].type).toBe('AGENT_THOUGHT');
});
});

View file

@ -0,0 +1,157 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import type { JSONSchemaType } from 'ajv';
import type {
ContextProcessor,
ProcessArgs,
BackstopTargetOptions,
} from '../pipeline.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import type { ConcreteNode, RollingSummary } from '../graph/types.js';
import { debugLogger } from '../../utils/debugLogger.js';
import { LlmRole } from '../../telemetry/llmRole.js';
export interface RollingSummaryProcessorOptions extends BackstopTargetOptions {
systemInstruction?: string;
}
export const RollingSummaryProcessorOptionsSchema: JSONSchemaType<RollingSummaryProcessorOptions> =
{
type: 'object',
properties: {
target: {
type: 'string',
enum: ['incremental', 'freeNTokens', 'max'],
nullable: true,
},
freeTokensTarget: { type: 'number', nullable: true },
maxRollingSummaries: { type: 'number', nullable: true },
systemInstruction: { type: 'string', nullable: true },
},
required: [],
};
export function createRollingSummaryProcessor(
id: string,
env: ContextEnvironment,
options: RollingSummaryProcessorOptions,
): ContextProcessor {
const generateRollingSummary = async (
nodes: ConcreteNode[],
): Promise<string> => {
let transcript = '';
for (const node of nodes) {
let nodeContent = '';
if ('text' in node && typeof node.text === 'string') {
nodeContent = node.text;
} else if ('semanticParts' in node) {
nodeContent = JSON.stringify(node.semanticParts);
} else if ('observation' in node) {
nodeContent =
typeof node.observation === 'string'
? node.observation
: JSON.stringify(node.observation);
}
transcript += `[${node.type}]: ${nodeContent}\n`;
}
const systemPrompt =
options.systemInstruction ??
`You are an expert context compressor. Your job is to drastically shorten the provided conversational transcript while preserving the absolute core semantic meaning, facts, and intent. Omit all conversational filler, pleasantries, or redundant information. Return ONLY the compressed summary.`;
const response = await env.llmClient.generateContent({
role: LlmRole.UTILITY_COMPRESSOR,
modelConfigKey: { model: 'gemini-3-flash-base' },
promptId: env.promptId,
abortSignal: new AbortController().signal,
contents: [{ role: 'user', parts: [{ text: transcript }] }],
systemInstruction: { role: 'system', parts: [{ text: systemPrompt }] },
});
const candidate = response.candidates?.[0];
const textPart = candidate?.content?.parts?.[0];
return textPart?.text || '';
};
return {
id,
name: 'RollingSummaryProcessor',
process: async ({ targets }: ProcessArgs) => {
if (targets.length === 0) return targets;
const strategy = options.target ?? 'max';
const nodesToSummarize: ConcreteNode[] = [];
if (strategy === 'incremental') {
// 'incremental' simply summarizes the minimum viable chunk (the oldest 2 nodes), ignoring token math.
for (const node of targets) {
if (node.id === targets[0].id && node.type === 'USER_PROMPT') {
continue; // Keep system prompt
}
nodesToSummarize.push(node);
if (nodesToSummarize.length >= 2) break; // We have enough for a minimum rolling summary
}
} else {
let targetTokensToRemove = 0;
if (strategy === 'freeNTokens') {
targetTokensToRemove = options.freeTokensTarget ?? Infinity;
} else if (strategy === 'max') {
targetTokensToRemove = Infinity;
}
if (targetTokensToRemove > 0) {
let deficitAccumulator = 0;
for (const node of targets) {
if (node.id === targets[0].id && node.type === 'USER_PROMPT') {
continue; // Keep system prompt
}
nodesToSummarize.push(node);
deficitAccumulator += env.tokenCalculator.getTokenCost(node);
if (deficitAccumulator >= targetTokensToRemove) break;
}
}
}
if (nodesToSummarize.length < 2) return targets; // Not enough context to summarize
try {
// Synthesize the rolling summary synchronously
const snapshotText = await generateRollingSummary(nodesToSummarize);
const newId = randomUUID();
const summaryNode: RollingSummary = {
id: newId,
logicalParentId: newId,
type: 'ROLLING_SUMMARY',
timestamp: Date.now(),
text: snapshotText,
abstractsIds: nodesToSummarize.map((n) => n.id),
};
const consumedIds = nodesToSummarize.map((n) => n.id);
const returnedNodes = targets.filter(
(t) => !consumedIds.includes(t.id),
);
const firstRemovedIdx = targets.findIndex((t) =>
consumedIds.includes(t.id),
);
if (firstRemovedIdx !== -1) {
const idx = Math.max(0, firstRemovedIdx);
returnedNodes.splice(idx, 0, summaryNode);
} else {
returnedNodes.unshift(summaryNode);
}
return returnedNodes;
} catch (e) {
debugLogger.error('RollingSummaryProcessor failed sync backstop', e);
return targets;
}
},
};
}

View file

@ -0,0 +1,125 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi } from 'vitest';
import { createStateSnapshotAsyncProcessor } from './stateSnapshotAsyncProcessor.js';
import {
createMockEnvironment,
createDummyNode,
createMockProcessArgs,
} from '../testing/contextTestUtils.js';
import type { InboxMessage } from '../pipeline.js';
import type { InboxSnapshotImpl } from '../pipeline/inbox.js';
describe('StateSnapshotAsyncProcessor', () => {
it('should generate a snapshot and publish it to the inbox', async () => {
const env = createMockEnvironment();
// Spy on the publish method
const publishSpy = vi.spyOn(env.inbox, 'publish');
const worker = createStateSnapshotAsyncProcessor(
'StateSnapshotAsyncProcessor',
env,
{ type: 'point-in-time' },
);
const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A');
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
const targets = [nodeA, nodeB];
await worker.process(createMockProcessArgs(targets, targets, []));
// Ensure generateContent was called
expect(env.llmClient.generateContent).toHaveBeenCalled();
// Verify it published to the inbox
expect(publishSpy).toHaveBeenCalledWith(
'PROPOSED_SNAPSHOT',
expect.objectContaining({
newText: 'Mock LLM summary response',
consumedIds: ['node-A', 'node-B'],
type: 'point-in-time',
}),
);
});
it('should pull previous accumulate snapshot from inbox and append new targets', async () => {
const env = createMockEnvironment();
const publishSpy = vi.spyOn(env.inbox, 'publish');
const drainSpy = vi.spyOn(env.inbox, 'drainConsumed');
const worker = createStateSnapshotAsyncProcessor(
'StateSnapshotAsyncProcessor',
env,
{ type: 'accumulate' },
);
const nodeC = createDummyNode('ep2', 'USER_PROMPT', 50, {}, 'node-C');
const targets = [nodeC];
const inboxMessages: InboxMessage[] = [
{
id: 'draft-1',
topic: 'PROPOSED_SNAPSHOT',
timestamp: Date.now() - 1000,
payload: {
consumedIds: ['node-A', 'node-B'],
newText: '<old snapshot>',
type: 'accumulate',
},
},
];
const args = createMockProcessArgs(targets, targets, inboxMessages);
await worker.process(args);
// The old draft should be consumed
expect(
(args.inbox as InboxSnapshotImpl).getConsumedIds().has('draft-1'),
).toBe(true);
expect(drainSpy).toHaveBeenCalledWith(expect.any(Set));
// The new publish should contain ALL consumed IDs (old + new)
expect(publishSpy).toHaveBeenCalledWith(
'PROPOSED_SNAPSHOT',
expect.objectContaining({
newText: 'Mock LLM summary response',
consumedIds: ['node-A', 'node-B', 'node-C'], // Aggregated!
type: 'accumulate',
}),
);
// Verify the LLM was called with the old snapshot prepended
expect(env.llmClient.generateContent).toHaveBeenCalledWith(
expect.objectContaining({
contents: expect.arrayContaining([
expect.objectContaining({
parts: expect.arrayContaining([
expect.objectContaining({
text: expect.stringContaining('<old snapshot>'),
}),
]),
}),
]),
}),
);
});
it('should ignore empty targets', async () => {
const env = createMockEnvironment();
const publishSpy = vi.spyOn(env.inbox, 'publish');
const worker = createStateSnapshotAsyncProcessor(
'StateSnapshotAsyncProcessor',
env,
{ type: 'accumulate' },
);
await worker.process(createMockProcessArgs([], [], []));
expect(env.llmClient.generateContent).not.toHaveBeenCalled();
expect(publishSpy).not.toHaveBeenCalled();
});
});

View file

@ -0,0 +1,113 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import type { JSONSchemaType } from 'ajv';
import type { AsyncContextProcessor, ProcessArgs } from '../pipeline.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import type { ConcreteNode } from '../graph/types.js';
import { SnapshotGenerator } from '../utils/snapshotGenerator.js';
import { debugLogger } from '../../utils/debugLogger.js';
export interface StateSnapshotAsyncProcessorOptions {
type?: 'accumulate' | 'point-in-time';
systemInstruction?: string;
}
export const StateSnapshotAsyncProcessorOptionsSchema: JSONSchemaType<StateSnapshotAsyncProcessorOptions> =
{
type: 'object',
properties: {
type: {
type: 'string',
enum: ['accumulate', 'point-in-time'],
nullable: true,
},
systemInstruction: { type: 'string', nullable: true },
},
required: [],
};
export function createStateSnapshotAsyncProcessor(
id: string,
env: ContextEnvironment,
options: StateSnapshotAsyncProcessorOptions,
): AsyncContextProcessor {
const generator = new SnapshotGenerator(env);
return {
id,
name: 'StateSnapshotAsyncProcessor',
process: async ({ targets, inbox }: ProcessArgs): Promise<void> => {
if (targets.length === 0) return;
try {
let nodesToSummarize = [...targets];
let previousConsumedIds: string[] = [];
const processorType = options.type ?? 'point-in-time';
if (processorType === 'accumulate') {
// Look for the most recent unconsumed accumulate snapshot in the inbox
const proposedSnapshots = inbox.getMessages<{
newText: string;
consumedIds: string[];
type: string;
}>('PROPOSED_SNAPSHOT');
const accumulateSnapshots = proposedSnapshots.filter(
(s) => s.payload.type === 'accumulate',
);
if (accumulateSnapshots.length > 0) {
// Sort to find the most recent
const latest = [...accumulateSnapshots].sort(
(a, b) => b.timestamp - a.timestamp,
)[0];
// Consume the old draft so the inbox doesn't fill up with stale drafts
inbox.consume(latest.id);
// And we must persist its consumption back to the live inbox immediately,
// because we are effectively "taking" it from the shelf to modify.
env.inbox.drainConsumed(new Set([latest.id]));
previousConsumedIds = latest.payload.consumedIds;
// Prepend a synthetic node representing the previous rolling state
const previousStateNode: ConcreteNode = {
id: randomUUID(),
logicalParentId: '',
type: 'SNAPSHOT',
timestamp: latest.timestamp,
text: latest.payload.newText,
};
nodesToSummarize = [previousStateNode, ...targets];
}
}
const snapshotText = await generator.synthesizeSnapshot(
nodesToSummarize,
options.systemInstruction,
);
const newConsumedIds = [
...previousConsumedIds,
...targets.map((t) => t.id),
];
// In V2, async pipelines communicate their work to the inbox, and the processor picks it up.
env.inbox.publish('PROPOSED_SNAPSHOT', {
newText: snapshotText,
consumedIds: newConsumedIds,
type: processorType,
});
} catch (e) {
debugLogger.error(
'StateSnapshotAsyncProcessor failed to generate snapshot',
e,
);
}
},
};
}

View file

@ -0,0 +1,131 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { createStateSnapshotProcessor } from './stateSnapshotProcessor.js';
import {
createMockEnvironment,
createDummyNode,
createMockProcessArgs,
} from '../testing/contextTestUtils.js';
import type { InboxSnapshotImpl } from '../pipeline/inbox.js';
describe('StateSnapshotProcessor', () => {
it('should ignore if budget is satisfied', async () => {
const env = createMockEnvironment();
const processor = createStateSnapshotProcessor(
'StateSnapshotProcessor',
env,
{
target: 'incremental',
},
);
const targets = [createDummyNode('ep1', 'USER_PROMPT')];
const result = await processor.process(createMockProcessArgs(targets));
expect(result).toBe(targets); // Strict equality
});
it('should apply a valid snapshot from the Inbox (Fast Path)', async () => {
const env = createMockEnvironment();
const processor = createStateSnapshotProcessor(
'StateSnapshotProcessor',
env,
{
target: 'incremental',
},
);
const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A');
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
const nodeC = createDummyNode('ep2', 'USER_PROMPT', 50, {}, 'node-C');
const targets = [nodeA, nodeB, nodeC];
// The async background pipeline created a snapshot of A and B
const messages = [
{
id: 'msg-1',
topic: 'PROPOSED_SNAPSHOT',
timestamp: Date.now(),
payload: {
consumedIds: ['node-A', 'node-B'],
newText: '<compressed A and B>',
type: 'point-in-time',
},
},
];
const processArgs = createMockProcessArgs(targets, [], messages);
const result = await processor.process(processArgs);
// Should remove A and B, insert Snapshot, keep C
expect(result.length).toBe(2);
expect(result[0].type).toBe('SNAPSHOT');
expect(result[1].id).toBe('node-C');
// Should consume the message
expect(
(processArgs.inbox as InboxSnapshotImpl).getConsumedIds().has('msg-1'),
).toBe(true);
});
it('should reject a snapshot if the nodes were modified/deleted (Cache Invalidated)', async () => {
const env = createMockEnvironment();
const processor = createStateSnapshotProcessor(
'StateSnapshotProcessor',
env,
{
target: 'incremental',
},
);
// Make deficit 0 so we don't fall through to the sync backstop and fail the test that way
// node-A is MISSING (user deleted it)
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
const targets = [nodeB];
const messages = [
{
id: 'msg-1',
topic: 'PROPOSED_SNAPSHOT',
timestamp: Date.now(),
payload: {
consumedIds: ['node-A', 'node-B'],
newText: '<compressed A and B>',
},
},
];
const processArgs = createMockProcessArgs(targets, [], messages);
const result = await processor.process(processArgs);
// Because deficit is 0, and Inbox was rejected, nothing should change
expect(result.length).toBe(1);
expect(result[0].id).toBe('node-B');
expect(
(processArgs.inbox as InboxSnapshotImpl).getConsumedIds().has('msg-1'),
).toBe(false);
});
it('should fall back to sync backstop if inbox is empty', async () => {
const env = createMockEnvironment();
const processor = createStateSnapshotProcessor(
'StateSnapshotProcessor',
env,
{ target: 'max' },
); // Summarize all
const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A');
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
const nodeC = createDummyNode('ep2', 'USER_PROMPT', 50, {}, 'node-C');
const targets = [nodeA, nodeB, nodeC];
const result = await processor.process(createMockProcessArgs(targets));
// Should synthesize a new snapshot synchronously
expect(env.llmClient.generateContent).toHaveBeenCalled();
expect(result.length).toBe(2); // nodeA is skipped as "system prompt", snapshot + nodeA
expect(result[1].type).toBe('SNAPSHOT');
});
});

View file

@ -0,0 +1,185 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import type { JSONSchemaType } from 'ajv';
import type {
ContextProcessor,
ProcessArgs,
BackstopTargetOptions,
} from '../pipeline.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import type { ConcreteNode, Snapshot } from '../graph/types.js';
import { SnapshotGenerator } from '../utils/snapshotGenerator.js';
import { debugLogger } from '../../utils/debugLogger.js';
export interface StateSnapshotProcessorOptions extends BackstopTargetOptions {
model?: string;
systemInstruction?: string;
}
export const StateSnapshotProcessorOptionsSchema: JSONSchemaType<StateSnapshotProcessorOptions> =
{
type: 'object',
properties: {
target: {
type: 'string',
enum: ['incremental', 'freeNTokens', 'max'],
nullable: true,
},
freeTokensTarget: { type: 'number', nullable: true },
model: { type: 'string', nullable: true },
systemInstruction: { type: 'string', nullable: true },
},
required: [],
};
export function createStateSnapshotProcessor(
id: string,
env: ContextEnvironment,
options: StateSnapshotProcessorOptions,
): ContextProcessor {
const generator = new SnapshotGenerator(env);
return {
id,
name: 'StateSnapshotProcessor',
process: async ({ targets, inbox }: ProcessArgs) => {
if (targets.length === 0) {
return targets;
}
// Determine what mode we are looking for: 'incremental' -> 'point-in-time', 'max' -> 'accumulate'
const strategy = options.target ?? 'max';
const expectedType =
strategy === 'incremental' ? 'point-in-time' : 'accumulate';
// 1. Check Inbox for a completed Snapshot (The Fast Path)
const proposedSnapshots = inbox.getMessages<{
newText: string;
consumedIds: string[];
type: string;
}>('PROPOSED_SNAPSHOT');
if (proposedSnapshots.length > 0) {
// Filter for the snapshot type that matches our processor mode
const matchingSnapshots = proposedSnapshots.filter(
(s) => s.payload.type === expectedType,
);
// Sort by newest timestamp first (we want the most accumulated snapshot)
const sorted = [...matchingSnapshots].sort(
(a, b) => b.timestamp - a.timestamp,
);
for (const proposed of sorted) {
const { consumedIds, newText } = proposed.payload;
// Verify all consumed IDs still exist sequentially in targets
const targetIds = new Set(targets.map((t) => t.id));
const isValid = consumedIds.every((id) => targetIds.has(id));
if (isValid) {
// If valid, apply it!
const newId = randomUUID();
const snapshotNode: Snapshot = {
id: newId,
logicalParentId: newId,
type: 'SNAPSHOT',
timestamp: Date.now(),
text: newText,
abstractsIds: consumedIds,
};
// Remove the consumed nodes and insert the snapshot at the earliest index
const returnedNodes = targets.filter(
(t) => !consumedIds.includes(t.id),
);
const firstRemovedIdx = targets.findIndex((t) =>
consumedIds.includes(t.id),
);
if (firstRemovedIdx !== -1) {
const idx = Math.max(0, firstRemovedIdx);
returnedNodes.splice(idx, 0, snapshotNode);
} else {
returnedNodes.unshift(snapshotNode);
}
inbox.consume(proposed.id);
return returnedNodes;
}
}
}
// 2. The Synchronous Backstop (The Slow Path)
let targetTokensToRemove = 0;
if (strategy === 'incremental') {
targetTokensToRemove = Infinity; // incremental implies removing as much as possible if no state is passed
} else if (strategy === 'freeNTokens') {
targetTokensToRemove = options.freeTokensTarget ?? Infinity;
} else if (strategy === 'max') {
targetTokensToRemove = Infinity;
}
let deficitAccumulator = 0;
const nodesToSummarize: ConcreteNode[] = [];
// Scan oldest to newest
for (const node of targets) {
if (node.id === targets[0].id && node.type === 'USER_PROMPT') {
// Keep system prompt if it's the very first node
// In a real system, system prompt is protected, but we double check
continue;
}
nodesToSummarize.push(node);
deficitAccumulator += env.tokenCalculator.getTokenCost(node);
if (deficitAccumulator >= targetTokensToRemove) break;
}
if (nodesToSummarize.length < 2) return targets; // Not enough context
try {
const snapshotText = await generator.synthesizeSnapshot(
nodesToSummarize,
options.systemInstruction,
);
const newId = randomUUID();
const snapshotNode: Snapshot = {
id: newId,
logicalParentId: newId,
type: 'SNAPSHOT',
timestamp: Date.now(),
text: snapshotText,
abstractsIds: nodesToSummarize.map((n) => n.id),
};
const consumedIds = nodesToSummarize.map((n) => n.id);
const returnedNodes = targets.filter(
(t) => !consumedIds.includes(t.id),
);
const firstRemovedIdx = targets.findIndex((t) =>
consumedIds.includes(t.id),
);
if (firstRemovedIdx !== -1) {
const idx = Math.max(0, firstRemovedIdx);
returnedNodes.splice(idx, 0, snapshotNode);
} else {
returnedNodes.unshift(snapshotNode);
}
return returnedNodes;
} catch (e) {
debugLogger.error('StateSnapshotProcessor failed sync backstop', e);
return targets;
}
},
};
}

View file

@ -0,0 +1,68 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { createToolMaskingProcessor } from './toolMaskingProcessor.js';
import {
createMockProcessArgs,
createMockEnvironment,
createDummyToolNode,
} from '../testing/contextTestUtils.js';
import type { ToolExecution } from '../graph/types.js';
describe('ToolMaskingProcessor', () => {
it('should write large strings to disk and replace them with a masked pointer', async () => {
const env = createMockEnvironment();
// env uses charsPerToken=1 natively.
// original string lengths > stringLengthThresholdTokens (which is 10) will be masked
const processor = createToolMaskingProcessor('ToolMaskingProcessor', env, {
stringLengthThresholdTokens: 10,
});
const longString = 'A'.repeat(500); // 500 chars
const toolStep = createDummyToolNode('ep1', 50, 500, {
observation: {
result: longString,
metadata: 'short', // 5 chars, will not be masked
},
});
const result = await processor.process(createMockProcessArgs([toolStep]));
expect(result.length).toBe(1);
const masked = result[0] as ToolExecution;
// It should have generated a new ID because it modified it
expect(masked.id).not.toBe(toolStep.id);
// It should have masked the observation
const obs = masked.observation as { result: string; metadata: string };
expect(obs.result).toContain('<tool_output_masked>');
expect(obs.metadata).toBe('short'); // Untouched
});
it('should skip unmaskable tools', async () => {
const env = createMockEnvironment();
const processor = createToolMaskingProcessor('ToolMaskingProcessor', env, {
stringLengthThresholdTokens: 10,
});
const toolStep = createDummyToolNode('ep1', 10, 10, {
toolName: 'activate_skill',
observation: {
result:
'this is a really long string that normally would get masked but wont because of the tool name',
},
});
const result = await processor.process(createMockProcessArgs([toolStep]));
// Returned the exact same object reference
expect(result[0]).toBe(toolStep);
});
});

View file

@ -0,0 +1,271 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import * as fs from 'node:fs/promises';
import * as path from 'node:path';
import type { ConcreteNode, ToolExecution } from '../graph/types.js';
import type { ContextEnvironment } from '../pipeline/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 { Part } from '@google/genai';
export interface ToolMaskingProcessorOptions {
stringLengthThresholdTokens: number;
}
export const ToolMaskingProcessorOptionsSchema: JSONSchemaType<ToolMaskingProcessorOptions> =
{
type: 'object',
properties: {
stringLengthThresholdTokens: { type: 'number' },
},
required: ['stringLengthThresholdTokens'],
};
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,
]);
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 function createToolMaskingProcessor(
id: string,
env: ContextEnvironment,
options: ToolMaskingProcessorOptions,
): ContextProcessor {
const isAlreadyMasked = (text: string): boolean =>
text.includes('<tool_output_masked>');
return {
id,
name: 'ToolMaskingProcessor',
process: async ({ targets }: ProcessArgs) => {
const maskingConfig = options;
if (!maskingConfig) return targets;
if (targets.length === 0) return targets;
const limitChars = env.tokenCalculator.tokensToChars(
maskingConfig.stringLengthThresholdTokens,
);
let toolOutputsDir = path.join(env.projectTempDir, 'tool-outputs');
const sessionId = env.sessionId;
if (sessionId) {
toolOutputsDir = path.join(
toolOutputsDir,
`session-${sanitizeFilenamePart(sessionId)}`,
);
}
let directoryCreated = false;
const handleMasking = async (
content: string,
toolName: string,
callId: string,
nodeType: string,
): Promise<string> => {
if (!directoryCreated) {
await fs.mkdir(toolOutputsDir, { recursive: true });
directoryCreated = true;
}
const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${randomUUID()}.txt`;
const filePath = path.join(toolOutputsDir, fileName);
await fs.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>`;
};
const returnedNodes: ConcreteNode[] = [];
for (const node of targets) {
switch (node.type) {
case 'TOOL_EXECUTION': {
const toolName = node.toolName;
if (toolName && UNMASKABLE_TOOLS.has(toolName)) {
returnedNodes.push(node);
break;
}
const callId = node.id || Date.now().toString();
const maskAsync = async (
obj: MaskableValue,
nodeType: string,
): Promise<{ masked: MaskableValue; changed: boolean }> => {
if (typeof obj === 'string') {
if (obj.length > limitChars && !isAlreadyMasked(obj)) {
const newString = await handleMasking(
obj,
toolName || 'unknown',
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 = node.intent;
const rawObs = node.observation;
if (!isMaskableRecord(rawIntent) || !isMaskableValue(rawObs)) {
returnedNodes.push(node);
break;
}
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;
// Handle observation explicitly as string vs object
const maskedObs =
typeof obsRes.masked === 'string'
? ({ message: obsRes.masked } as Record<string, unknown>)
: isMaskableRecord(obsRes.masked)
? (obsRes.masked as Record<string, unknown>)
: undefined;
const newIntentTokens =
env.tokenCalculator.estimateTokensForParts([
{
functionCall: {
name: toolName || 'unknown',
args: maskedIntent,
id: callId,
},
},
]);
let obsPart: Record<string, unknown> = {};
if (maskedObs) {
obsPart = {
functionResponse: {
name: toolName || 'unknown',
response: maskedObs,
id: callId,
},
};
}
const newObsTokens = env.tokenCalculator.estimateTokensForParts([
obsPart as Part,
]);
const tokensSaved =
env.tokenCalculator.getTokenCost(node) -
(newIntentTokens + newObsTokens);
if (tokensSaved > 0) {
const maskedNode: ToolExecution = {
...node,
id: randomUUID(), // Modified, so generate new ID
intent: maskedIntent ?? node.intent,
observation: maskedObs ?? node.observation,
tokens: {
intent: newIntentTokens,
observation: newObsTokens,
},
replacesId: node.id,
};
returnedNodes.push(maskedNode);
} else {
returnedNodes.push(node);
}
} else {
returnedNodes.push(node);
}
break;
}
default:
returnedNodes.push(node);
break;
}
}
return returnedNodes;
},
};
}

File diff suppressed because one or more lines are too long

View file

@ -0,0 +1,246 @@
/**
* @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 { createMockLlmClient } from '../testing/contextTestUtils.js';
import type { ContextProfile } from '../config/profiles.js';
import { createToolMaskingProcessor } from '../processors/toolMaskingProcessor.js';
import { createBlobDegradationProcessor } from '../processors/blobDegradationProcessor.js';
import { createStateSnapshotProcessor } from '../processors/stateSnapshotProcessor.js';
import { createHistoryTruncationProcessor } from '../processors/historyTruncationProcessor.js';
import { createStateSnapshotAsyncProcessor } from '../processors/stateSnapshotAsyncProcessor.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 = (): ContextProfile => ({
config: {
budget: { maxTokens: 1000, retainedTokens: 500 }, // Extremely tight limits
},
buildPipelines: (env) => [
{
name: 'Pressure Relief', // Emits from eventBus 'retained_exceeded'
triggers: ['retained_exceeded'],
processors: [
createBlobDegradationProcessor('BlobDegradationProcessor', env),
createToolMaskingProcessor('ToolMaskingProcessor', env, {
stringLengthThresholdTokens: 50,
}),
createStateSnapshotProcessor('StateSnapshotProcessor', env, {}),
],
},
{
name: 'Immediate Sanitization', // The magic string the projector is hardcoded to use
triggers: ['retained_exceeded'],
processors: [
createHistoryTruncationProcessor(
'HistoryTruncationProcessor',
env,
{},
),
],
},
],
buildAsyncPipelines: (env) => [
{
name: 'Async',
triggers: ['nodes_aged_out'],
processors: [
createStateSnapshotAsyncProcessor(
'StateSnapshotAsyncProcessor',
env,
{},
),
],
},
],
});
const mockLlmClient = createMockLlmClient([
'<MOCKED_STATE_SNAPSHOT_SUMMARY>',
]);
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();
});
it('Scenario 2: Under Budget (No Modifications)', async () => {
const generousConfig: ContextProfile = {
config: {
budget: { maxTokens: 100000, retainedTokens: 50000 },
},
buildPipelines: () => [],
buildAsyncPipelines: () => [],
};
const harness = await SimulationHarness.create(
generousConfig,
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?' }] },
]);
const goldenState = await harness.getGoldenState();
// Total tokens should cleanly match character count with no synthetic nodes
expect(goldenState).toMatchSnapshot();
});
it('Scenario 3: Async-Driven Background GC', async () => {
const gcConfig: ContextProfile = {
config: {
budget: { maxTokens: 200, retainedTokens: 100 },
},
buildPipelines: () => [],
buildAsyncPipelines: (env) => [
{
name: 'Async',
triggers: ['nodes_aged_out'],
processors: [
createStateSnapshotAsyncProcessor(
'StateSnapshotAsyncProcessor',
env,
{},
),
],
},
],
};
const harness = await SimulationHarness.create(gcConfig, mockLlmClient);
// Turn 0
await harness.simulateTurn([
{ role: 'user', parts: [{ text: 'A'.repeat(50) }] },
{ role: 'model', parts: [{ text: 'B'.repeat(50) }] },
]);
// Turn 1 (Should trigger StateSnapshotasync pipeline because we exceed 100 retainedTokens)
await harness.simulateTurn([
{ role: 'user', parts: [{ text: 'C'.repeat(50) }] },
{ role: 'model', parts: [{ text: 'D'.repeat(50) }] },
]);
// Give the async background pipeline an extra beat to complete its async execution and emit variants
await new Promise((resolve) => setTimeout(resolve, 50));
// Turn 2
await harness.simulateTurn([
{ role: 'user', parts: [{ text: 'E'.repeat(50) }] },
{ role: 'model', parts: [{ text: 'F'.repeat(50) }] },
]);
const goldenState = await harness.getGoldenState();
// We should see ROLLING_SUMMARY nodes injected into the graph, proving the async pipeline ran in the background
expect(goldenState).toMatchSnapshot();
});
});

View file

@ -0,0 +1,157 @@
/**
* @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 { ContextProfile } from '../config/profiles.js';
import { ContextEnvironmentImpl } from '../pipeline/environmentImpl.js';
import { ContextTracer } from '../tracer.js';
import { ContextEventBus } from '../eventBus.js';
import { PipelineOrchestrator } from '../pipeline/orchestrator.js';
import { debugLogger } from '../../utils/debugLogger.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!: ContextProfile;
private tracer!: ContextTracer;
private currentTurnIndex = 0;
private tokenTrajectory: TurnSummary[] = [];
static async create(
config: ContextProfile,
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: ContextProfile,
mockLlmClient: BaseLlmClient,
mockTempDir: string,
) {
this.config = config;
this.tracer = new ContextTracer({
targetDir: mockTempDir,
sessionId: 'sim-session',
});
this.env = new ContextEnvironmentImpl(
mockLlmClient,
'sim-prompt',
'sim-session',
mockTempDir,
mockTempDir,
this.tracer,
1, // 1 char per token average
this.eventBus,
);
this.orchestrator = new PipelineOrchestrator(
config.buildPipelines(this.env),
config.buildAsyncPipelines(this.env),
this.env,
this.eventBus,
this.tracer,
);
this.contextManager = new ContextManager(
config,
this.env,
this.tracer,
this.orchestrator,
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.calculateConcreteListTokens(
this.contextManager.getNodes(),
);
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.getNodes();
const currentTokens =
this.env.tokenCalculator.calculateConcreteListTokens(currentView);
if (
this.config.config.budget &&
currentTokens > this.config.config.budget.maxTokens
) {
debugLogger.log(
`[Turn ${this.currentTurnIndex}] Sync panic triggered! ${currentTokens} > ${this.config.config.budget.maxTokens}`,
);
const orchestrator = this.orchestrator;
// In the V2 simulation, we trigger the 'gc_backstop' to simulate emergency pressure.
// Since contextManager owns its buffer natively, the simulation now properly matches reality
// where the manager runs the orchestrator and keeps the resulting modified view.
const modifiedView = await orchestrator.executeTriggerSync(
'gc_backstop',
currentView,
new Set(currentView.map((e) => e.id)),
new Set<string>(),
);
// In the real system, ContextManager triggers this and retains it.
// We will emulate that behavior internally in the test loop for token counting.
currentView = modifiedView;
}
// 4. Measure tokens after background processors have processed inboxes
const tokensAfter = this.env.tokenCalculator.calculateConcreteListTokens(
this.contextManager.getNodes(),
);
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.renderHistory();
return {
tokenTrajectory: this.tokenTrajectory,
finalProjection,
};
}
}

View file

@ -0,0 +1,278 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { vi } from 'vitest';
import { AgentChatHistory } from '../../core/agentChatHistory.js';
import { ContextManager } from '../contextManager.js';
import { randomUUID } from 'node:crypto';
import { ContextTracer } from '../tracer.js';
import { ContextEnvironmentImpl } from '../pipeline/environmentImpl.js';
import { ContextEventBus } from '../eventBus.js';
import { PipelineOrchestrator } from '../pipeline/orchestrator.js';
import type { ConcreteNode, ToolExecution } from '../graph/types.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import type { Config } from '../../config/config.js';
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
import type { Content, GenerateContentResponse } from '@google/genai';
import { InboxSnapshotImpl } from '../pipeline/inbox.js';
import type { InboxMessage, ProcessArgs } from '../pipeline.js';
import type { ContextProfile } from '../config/profiles.js';
import type { Mock } from 'vitest';
import { ContextWorkingBufferImpl } from '../pipeline/contextWorkingBuffer.js';
import { testTruncateProfile } from './testProfile.js';
/**
* Creates a valid mock GenerateContentResponse with the provided text.
* Used to avoid having to manually construct the deeply nested candidate/content/part structure.
*/
export const createMockGenerateContentResponse = (
text: string,
): GenerateContentResponse =>
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
({
candidates: [{ content: { role: 'model', parts: [{ text }] }, index: 0 }],
}) as GenerateContentResponse;
export function createDummyNode(
logicalParentId: string,
type: ConcreteNode['type'],
tokens = 100,
overrides?: Partial<ConcreteNode>,
id?: string,
): ConcreteNode {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
return {
id: id || randomUUID(),
episodeId: logicalParentId,
logicalParentId,
type,
timestamp: Date.now(),
text: `Dummy ${type}`,
name: type === 'SYSTEM_EVENT' ? 'dummy_event' : undefined,
payload: type === 'SYSTEM_EVENT' ? {} : undefined,
semanticParts: [],
metadata: {
originalTokens: tokens,
currentTokens: tokens,
transformations: [],
},
...overrides,
} as unknown as ConcreteNode;
}
export function createDummyToolNode(
logicalParentId: string,
intentTokens = 100,
obsTokens = 200,
overrides?: Partial<ToolExecution>,
id?: string,
): ToolExecution {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
return {
id: id || randomUUID(),
episodeId: logicalParentId,
logicalParentId,
type: 'TOOL_EXECUTION',
timestamp: Date.now(),
toolName: 'dummy_tool',
intent: { action: 'test' },
observation: { result: 'ok' },
tokens: {
intent: intentTokens,
observation: obsTokens,
},
metadata: {
originalTokens: intentTokens + obsTokens,
currentTokens: intentTokens + obsTokens,
transformations: [],
},
...overrides,
} as unknown as ToolExecution;
}
export interface MockLlmClient extends BaseLlmClient {
generateContent: Mock;
}
export function createMockLlmClient(
responses?: Array<string | GenerateContentResponse>,
): MockLlmClient {
const generateContentMock = vi.fn();
if (responses && responses.length > 0) {
for (const response of responses) {
if (typeof response === 'string') {
generateContentMock.mockResolvedValueOnce(
createMockGenerateContentResponse(response),
);
} else {
generateContentMock.mockResolvedValueOnce(response);
}
}
// Fallback to the last response for any subsequent calls
const lastResponse = responses[responses.length - 1];
if (typeof lastResponse === 'string') {
generateContentMock.mockResolvedValue(
createMockGenerateContentResponse(lastResponse),
);
} else {
generateContentMock.mockResolvedValue(lastResponse);
}
} else {
// Default fallback
generateContentMock.mockResolvedValue(
createMockGenerateContentResponse('Mock LLM response'),
);
}
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
return {
generateContent: generateContentMock,
} as unknown as MockLlmClient;
}
export function createMockEnvironment(
overrides?: Partial<ContextEnvironment>,
): ContextEnvironment {
const llmClient = createMockLlmClient(['Mock LLM summary response']);
const tracer = new ContextTracer({
targetDir: '/tmp',
sessionId: 'mock-session',
});
const eventBus = new ContextEventBus();
const env = new ContextEnvironmentImpl(
llmClient,
'mock-session',
'mock-prompt-id',
'/tmp/.gemini/trace',
'/tmp/.gemini/tool-outputs',
tracer,
1,
eventBus,
);
if (overrides) {
Object.assign(env, overrides);
}
return env;
}
/**
* 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 createMockProcessArgs(
targets: ConcreteNode[],
bufferNodes: ConcreteNode[] = [],
inboxMessages: InboxMessage[] = [],
): ProcessArgs {
return {
targets,
buffer: ContextWorkingBufferImpl.initialize(
bufferNodes.length ? bufferNodes : targets,
),
inbox: new InboxSnapshotImpl(inboxMessages),
};
}
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'),
getExperimentalContextManagementConfig: 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 async pipelines.
*/
export function setupContextComponentTest(
config: Config,
sidecarOverride?: ContextProfile,
): { chatHistory: AgentChatHistory; contextManager: ContextManager } {
const chatHistory = new AgentChatHistory();
const sidecar = sidecarOverride || testTruncateProfile;
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 orchestrator = new PipelineOrchestrator(
sidecar.buildPipelines(env),
sidecar.buildAsyncPipelines(env),
env,
eventBus,
tracer,
);
const contextManager = new ContextManager(
sidecar,
env,
tracer,
orchestrator,
chatHistory,
);
// The async async pipeline is now internally managed by ContextManager
return { chatHistory, contextManager };
}

View file

@ -0,0 +1,28 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ContextProfile } from '../config/profiles.js';
import type { PipelineDef } from '../config/types.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { createHistoryTruncationProcessor } from '../processors/historyTruncationProcessor.js';
export const testTruncateProfile: ContextProfile = {
config: {
budget: {
retainedTokens: 65000,
maxTokens: 150000,
},
},
buildPipelines: (env: ContextEnvironment): PipelineDef[] => [
{
name: 'Emergency Backstop (Truncate Only)',
triggers: ['gc_backstop', 'retained_exceeded'],
processors: [
createHistoryTruncationProcessor('HistoryTruncation', env, {}),
],
},
],
buildAsyncPipelines: () => [],
};

View file

@ -0,0 +1,94 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import { ContextTracer } from './tracer.js';
import * as fs from 'node:fs/promises';
import { existsSync, readFileSync } from 'node:fs';
import * as path from 'node:path';
import * as os from 'node:os';
vi.mock('node:crypto', () => {
let count = 0;
return {
randomUUID: vi.fn(() => `mock-uuid-${++count}`),
};
});
describe('ContextTracer (Real FS & Mock ID Gen)', () => {
let tmpDir: string;
beforeEach(async () => {
tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'gemini-tracer-test-'));
vi.useFakeTimers();
vi.setSystemTime(new Date('2026-01-01T12:00:00Z'));
});
afterEach(async () => {
vi.useRealTimers();
await fs.rm(tmpDir, { recursive: true, force: true });
});
it('initializes, logs events, and auto-saves large assets deterministically', async () => {
const tracer = new ContextTracer({
enabled: true,
targetDir: tmpDir,
sessionId: 'test-session',
});
vi.advanceTimersByTime(10);
await Promise.resolve(); // allow async mkdir to happen in constructor
// Verify Initialization
const traceLogPath = path.join(
tmpDir,
'.gemini/context_trace/test-session/trace.log',
);
const initTraceLog = readFileSync(traceLogPath, 'utf-8');
expect(initTraceLog).toContain('[SYSTEM] Context Tracer Initialized');
tracer.logEvent('TestComponent', 'TestAction', { key: 'value' });
vi.advanceTimersByTime(10);
await Promise.resolve();
const smallTraceLog = readFileSync(traceLogPath, 'utf-8');
expect(smallTraceLog).toContain('[TestComponent] TestAction');
expect(smallTraceLog).toContain('{"key":"value"}');
const hugeString = 'a'.repeat(2000);
tracer.logEvent('TestComponent', 'LargeAction', { largeKey: hugeString });
vi.advanceTimersByTime(10);
await Promise.resolve();
const expectedAssetPath = path.join(
tmpDir,
'.gemini/context_trace/test-session/assets/1767268800020-mock-uuid-1-largeKey.json',
);
expect(existsSync(expectedAssetPath)).toBe(true);
const largeTraceLog = readFileSync(traceLogPath, 'utf-8');
expect(largeTraceLog).toContain('[TestComponent] LargeAction');
expect(largeTraceLog).toContain(
`{"largeKey":{"$asset":"1767268800020-mock-uuid-1-largeKey.json"}}`,
);
});
it('silently ignores logging when disabled', async () => {
const tracer = new ContextTracer({
enabled: false,
targetDir: tmpDir,
sessionId: 'test-session',
});
tracer.logEvent('TestComponent', 'TestAction');
const hugeString = 'a'.repeat(2000);
tracer.logEvent('TestComponent', 'LargeAction', { largeKey: hugeString });
// Nothing should be written
const traceDir = path.join(tmpDir, '.gemini');
expect(existsSync(traceDir)).toBe(false);
});
});

View file

@ -0,0 +1,107 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { debugLogger } from '../utils/debugLogger.js';
import * as fsSync from 'node:fs';
import * as path from 'node:path';
import { randomUUID } from 'node:crypto';
export interface ContextTracerOptions {
enabled?: boolean;
targetDir: string;
sessionId: string;
}
export class ContextTracer {
private traceDir: string;
private assetsDir: string;
private enabled: boolean;
private readonly MAX_INLINE_SIZE = 1000;
constructor(options: ContextTracerOptions) {
this.enabled = options.enabled ?? false;
this.traceDir = path.join(
options.targetDir,
'.gemini',
'context_trace',
options.sessionId,
);
this.assetsDir = path.join(this.traceDir, 'assets');
if (this.enabled) {
try {
fsSync.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`;
fsSync.appendFileSync(
path.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()}-${randomUUID()}-${assetName}.json`;
const assetPath = path.join(this.assetsDir, assetId);
fsSync.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';
}
}
}

View file

@ -0,0 +1,107 @@
/**
* @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 { ConcreteNode } from '../graph/types.js';
import type { NodeBehaviorRegistry } from '../graph/behaviorRegistry.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.
*/
export class ContextTokenCalculator {
private readonly tokenCache = new Map<string, number>();
constructor(
private readonly charsPerToken: number,
private readonly registry: NodeBehaviorRegistry,
) {}
/**
* Estimates tokens for a simple string based on character count.
* Fast, but inherently inaccurate compared to real model tokenization.
*/
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;
}
/**
* Pre-calculates and caches the token cost of a newly minted node.
* Because nodes are immutable, this cost never changes for this node ID.
*/
/**
* Removes cached token counts for any nodes that are no longer in the given live set.
* This prevents unbounded memory growth during long sessions.
*/
garbageCollectCache(liveNodeIds: ReadonlySet<string>): void {
for (const [id] of this.tokenCache) {
if (!liveNodeIds.has(id)) {
this.tokenCache.delete(id);
}
}
}
cacheNodeTokens(node: ConcreteNode): number {
const behavior = this.registry.get(node.type);
const parts = behavior.getEstimatableParts(node);
const tokens = this.estimateTokensForParts(parts);
this.tokenCache.set(node.id, tokens);
return tokens;
}
/**
* Retrieves the token cost of a single node from the cache.
* If it misses the cache, it computes it and caches it.
*/
getTokenCost(node: ConcreteNode): number {
const cached = this.tokenCache.get(node.id);
if (cached !== undefined) return cached;
return this.cacheNodeTokens(node);
}
/**
* Fast calculation for a flat array of ConcreteNodes (The Nodes).
* It relies entirely on the O(1) sidecar token cache.
*/
calculateConcreteListTokens(nodes: readonly ConcreteNode[]): number {
let tokens = 0;
for (const node of nodes) {
tokens += this.getTokenCost(node);
}
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 += 258;
} else {
totalTokens += Math.ceil(
JSON.stringify(part).length / this.charsPerToken,
);
}
}
// Also include structural overhead
return totalTokens + baseEstimate(parts, depth);
}
}

View file

@ -0,0 +1,54 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { ConcreteNode } from '../graph/types.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { LlmRole } from '../../telemetry/llmRole.js';
export class SnapshotGenerator {
constructor(private readonly env: ContextEnvironment) {}
async synthesizeSnapshot(
nodes: readonly ConcreteNode[],
systemInstruction?: string,
): Promise<string> {
const systemPrompt =
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 node of nodes) {
let nodeContent = '';
if ('text' in node && typeof node.text === 'string') {
nodeContent = node.text;
} else if ('semanticParts' in node) {
nodeContent = JSON.stringify(node.semanticParts);
} else if ('observation' in node) {
nodeContent =
typeof node.observation === 'string'
? node.observation
: JSON.stringify(node.observation);
}
userPromptText += `[${node.type}]: ${nodeContent}\n`;
}
const response = await this.env.llmClient.generateContent({
role: LlmRole.UTILITY_STATE_SNAPSHOT_PROCESSOR,
modelConfigKey: { model: 'gemini-3-flash-base' },
contents: [{ role: 'user', parts: [{ text: userPromptText }] }],
systemInstruction: { role: 'system', parts: [{ text: systemPrompt }] },
promptId: this.env.promptId,
abortSignal: new AbortController().signal,
});
const candidate = response.candidates?.[0];
const textPart = candidate?.content?.parts?.[0];
return textPart?.text || '';
}
}

View file

@ -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;
}
}

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_processor',
}