mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-06-20 00:06:49 -07:00
Merge branch 'jl/cm-part-1' into gundermanc/compression-eval
This commit is contained in:
@@ -700,6 +700,7 @@ export interface ConfigParameters {
|
||||
experimentalJitContext?: boolean;
|
||||
autoDistillation?: boolean;
|
||||
experimentalMemoryManager?: boolean;
|
||||
experimentalContextSidecarConfig?: string;
|
||||
experimentalAgentHistoryTruncation?: boolean;
|
||||
experimentalAgentHistoryTruncationThreshold?: number;
|
||||
experimentalAgentHistoryRetainedMessages?: number;
|
||||
@@ -942,6 +943,7 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
private readonly adminSkillsEnabled: boolean;
|
||||
private readonly experimentalJitContext: boolean;
|
||||
private readonly experimentalMemoryManager: boolean;
|
||||
private readonly experimentalContextSidecarConfig?: string;
|
||||
private readonly memoryBoundaryMarkers: readonly string[];
|
||||
private readonly topicUpdateNarration: boolean;
|
||||
private readonly disableLLMCorrection: boolean;
|
||||
@@ -1153,6 +1155,8 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
|
||||
this.experimentalJitContext = params.experimentalJitContext ?? false;
|
||||
this.experimentalMemoryManager = params.experimentalMemoryManager ?? false;
|
||||
this.experimentalContextSidecarConfig =
|
||||
params.experimentalContextSidecarConfig;
|
||||
this.memoryBoundaryMarkers = params.memoryBoundaryMarkers ?? ['.git'];
|
||||
this.contextManagement = {
|
||||
enabled: params.contextManagement?.enabled ?? false,
|
||||
@@ -2427,6 +2431,10 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
return this.experimentalMemoryManager;
|
||||
}
|
||||
|
||||
getExperimentalContextSidecarConfig(): string | undefined {
|
||||
return this.experimentalContextSidecarConfig;
|
||||
}
|
||||
|
||||
getContextManagementConfig(): ContextManagementConfig {
|
||||
return this.contextManagement;
|
||||
}
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,110 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import {
|
||||
createMockContextConfig,
|
||||
setupContextComponentTest,
|
||||
} from './testing/contextTestUtils.js';
|
||||
|
||||
describe('ContextManager Barrier Tests', () => {
|
||||
it('Soft Barrier (retainedTokens): should inject ready variants and shrink projection', async () => {
|
||||
const config = createMockContextConfig();
|
||||
const { chatHistory, contextManager } = setupContextComponentTest(config);
|
||||
|
||||
// 1. Shrink limits: 1 char = 1 token. RetainedTokens = 10. MaxTokens = 100.
|
||||
|
||||
contextManager['sidecar'].budget.retainedTokens = 5;
|
||||
contextManager['sidecar'].budget.maxTokens = 100;
|
||||
|
||||
// 2. Build tiny history: 5 turns (10 messages). 2 tokens per turn.
|
||||
const tinyHistory = [];
|
||||
for (let i = 0; i < 5; i++) {
|
||||
tinyHistory.push({ role: 'user', parts: [{ text: `U${i}` }] });
|
||||
tinyHistory.push({ role: 'model', parts: [{ text: `M${i}` }] });
|
||||
}
|
||||
|
||||
// Set history directly to avoid event races
|
||||
chatHistory.set(tinyHistory);
|
||||
|
||||
// 3. Pre-verify baseline length.
|
||||
const baseline = await contextManager.projectCompressedHistory();
|
||||
expect(baseline.length).toBe(10);
|
||||
|
||||
// 4. Emit a fake snapshot covering the first 3 pairs (6 messages)
|
||||
const targetEp = contextManager['pristineEpisodes'][2];
|
||||
const replacedIds = contextManager['pristineEpisodes']
|
||||
.slice(0, 3)
|
||||
.map((ep) => ep.id);
|
||||
|
||||
contextManager['eventBus'].emitVariantReady({
|
||||
targetId: targetEp.id,
|
||||
variantId: 'snapshot',
|
||||
variant: {
|
||||
status: 'ready',
|
||||
type: 'snapshot',
|
||||
replacedEpisodeIds: replacedIds,
|
||||
episode: {
|
||||
id: 'snapshot-ep',
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: 't1',
|
||||
type: 'USER_PROMPT',
|
||||
semanticParts: [],
|
||||
metadata: {
|
||||
originalTokens: 0,
|
||||
currentTokens: 0,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
yield: {
|
||||
id: 'y1',
|
||||
type: 'AGENT_YIELD',
|
||||
text: '<SNAP>',
|
||||
metadata: {
|
||||
originalTokens: 5,
|
||||
currentTokens: 5,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
steps: [],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// 5. Verify Projection shrinks: 6 original messages replaced by 1 snapshot episode (1 text part) -> length 5.
|
||||
const projection = await contextManager.projectCompressedHistory();
|
||||
expect(projection.length).toBe(5);
|
||||
// projection[0] should be the snapshot yield
|
||||
expect(projection[0].parts![0].text).toBe('<SNAP>');
|
||||
});
|
||||
|
||||
it('Hard Barrier (maxTokens): should ruthlessly truncate unprotected episodes', async () => {
|
||||
const config = createMockContextConfig();
|
||||
const { chatHistory, contextManager } = setupContextComponentTest(config);
|
||||
|
||||
// 1. Shrink limits: maxTokens = 15.
|
||||
|
||||
contextManager['sidecar'].budget.maxTokens = 15;
|
||||
|
||||
// 2. Build history: 2 turns. Total = 24 tokens.
|
||||
const history = [
|
||||
{ role: 'user', parts: [{ text: 'U0' }] },
|
||||
{ role: 'model', parts: [{ text: 'M0_LARGE!!' }] },
|
||||
{ role: 'user', parts: [{ text: 'U1' }] },
|
||||
{ role: 'model', parts: [{ text: 'M1_LARGE!!' }] },
|
||||
];
|
||||
chatHistory.set(history);
|
||||
|
||||
const projection = await contextManager.projectCompressedHistory();
|
||||
|
||||
// Because Turn 0 is architecturally protected (system prompt/initialization), it SURVIVES!
|
||||
// Turn 1 is dropped to satisfy the maxTokens constraint.
|
||||
expect(projection.length).toBe(2);
|
||||
expect(projection[0].parts![0].text).toBe('U0');
|
||||
expect(projection[1].parts![0].text).toBe('M0_LARGE!!');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,73 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||
|
||||
import {
|
||||
createSyntheticHistory,
|
||||
createMockContextConfig,
|
||||
setupContextComponentTest,
|
||||
} from './testing/contextTestUtils.js';
|
||||
|
||||
describe('ContextManager Sync Pressure Barrier Tests', () => {
|
||||
beforeEach(() => {
|
||||
vi.useFakeTimers();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it('should instantly truncate history when maxTokens is exceeded using truncate strategy', async () => {
|
||||
// 1. Setup
|
||||
const config = createMockContextConfig();
|
||||
const { chatHistory, contextManager } = setupContextComponentTest(config);
|
||||
|
||||
// 2. Add System Prompt (Episode 0 - Protected)
|
||||
chatHistory.set([
|
||||
{ role: 'user', parts: [{ text: 'System prompt' }] },
|
||||
{ role: 'model', parts: [{ text: 'Understood.' }] },
|
||||
]);
|
||||
|
||||
// 3. Add massive history that blows past the 150k maxTokens limit
|
||||
// 20 turns * 10,000 tokens/turn = ~200,000 tokens
|
||||
const massiveHistory = createSyntheticHistory(20, 35000);
|
||||
chatHistory.set([...chatHistory.get(), ...massiveHistory]);
|
||||
|
||||
// 4. Add the Latest Turn (Protected)
|
||||
chatHistory.set([
|
||||
...chatHistory.get(),
|
||||
{ role: 'user', parts: [{ text: 'Final question.' }] },
|
||||
{ role: 'model', parts: [{ text: 'Final answer.' }] },
|
||||
]);
|
||||
|
||||
const rawHistoryLength = chatHistory.get().length;
|
||||
|
||||
// 5. Project History (Triggers Sync Barrier)
|
||||
const projection = await contextManager.projectCompressedHistory();
|
||||
|
||||
// 6. Assertions
|
||||
// The barrier should have dropped several older episodes to get under 150k.
|
||||
|
||||
expect(projection.length).toBeLessThan(rawHistoryLength);
|
||||
|
||||
// Verify Episode 0 (System) is perfectly preserved at the front
|
||||
|
||||
expect(projection[0].role).toBe('user');
|
||||
expect(projection[0].parts![0].text).toBe('System prompt');
|
||||
|
||||
// Verify the latest turn is perfectly preserved at the back
|
||||
const lastUser = projection[projection.length - 2];
|
||||
const lastModel = projection[projection.length - 1];
|
||||
|
||||
expect(lastUser.role).toBe('user');
|
||||
expect(lastUser.parts![0].text).toBe('Final question.');
|
||||
|
||||
expect(lastModel.role).toBe('model');
|
||||
expect(lastModel.parts![0].text).toBe('Final answer.');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,189 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
describe,
|
||||
it,
|
||||
expect,
|
||||
vi,
|
||||
beforeEach,
|
||||
beforeAll,
|
||||
afterAll,
|
||||
} from 'vitest';
|
||||
import { ContextManager } from './contextManager.js';
|
||||
import { ContextEnvironmentImpl } from './sidecar/environmentImpl.js';
|
||||
import { SidecarLoader } from './sidecar/SidecarLoader.js';
|
||||
import { ContextTracer } from './tracer.js';
|
||||
import { ContextEventBus } from './eventBus.js';
|
||||
import { ContextTokenCalculator } from './utils/contextTokenCalculator.js';
|
||||
import type { Content } from '@google/genai';
|
||||
import type { BaseLlmClient } from '../core/baseLlmClient.js';
|
||||
import type { Episode } from './ir/types.js';
|
||||
import type { SidecarConfig } from './sidecar/types.js';
|
||||
import { ProcessorRegistry } from './sidecar/registry.js';
|
||||
import { registerBuiltInProcessors } from './sidecar/builtins.js';
|
||||
import { IrMapper } from './ir/mapper.js';
|
||||
|
||||
expect.addSnapshotSerializer({
|
||||
test: (val) =>
|
||||
typeof val === 'string' &&
|
||||
/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i.test(val),
|
||||
print: () => '"<UUID>"',
|
||||
});
|
||||
|
||||
describe('ContextManager Golden Tests', () => {
|
||||
beforeAll(() => {
|
||||
vi.useFakeTimers();
|
||||
vi.setSystemTime(new Date(2026, 3, 2).getTime());
|
||||
vi.spyOn(Math, 'random').mockReturnValue(0.5);
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
vi.useRealTimers();
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
let mockConfig: any; // eslint-disable-line @typescript-eslint/no-explicit-any
|
||||
let contextManager: ContextManager;
|
||||
|
||||
beforeEach(() => {
|
||||
mockConfig = {
|
||||
isContextManagementEnabled: vi.fn().mockReturnValue(true),
|
||||
getExperimentalContextSidecarConfig: vi.fn().mockReturnValue(undefined),
|
||||
getTargetDir: vi.fn().mockReturnValue('/tmp'),
|
||||
getSessionId: vi.fn().mockReturnValue('test-session'),
|
||||
getToolOutputMaskingConfig: vi.fn().mockResolvedValue({
|
||||
enabled: true,
|
||||
minPrunableThresholdTokens: 50,
|
||||
protectLatestTurn: false,
|
||||
protectionThresholdTokens: 100,
|
||||
}),
|
||||
storage: { getProjectTempDir: vi.fn().mockReturnValue('/tmp') },
|
||||
getUsageStatisticsEnabled: vi.fn().mockReturnValue(false),
|
||||
getBaseLlmClient: vi.fn().mockReturnValue({
|
||||
generateJson: vi.fn().mockResolvedValue({
|
||||
'test_file.txt': { level: 'SUMMARY' },
|
||||
}),
|
||||
generateContent: vi.fn().mockResolvedValue({
|
||||
candidates: [
|
||||
{ content: { parts: [{ text: 'This is a summary.' }] } },
|
||||
],
|
||||
}),
|
||||
}),
|
||||
};
|
||||
|
||||
const registry = new ProcessorRegistry();
|
||||
registerBuiltInProcessors(registry);
|
||||
|
||||
const sidecar = SidecarLoader.fromConfig(mockConfig, registry);
|
||||
const tracer = new ContextTracer({
|
||||
targetDir: '/tmp',
|
||||
sessionId: 'test-session',
|
||||
});
|
||||
const eventBus = new ContextEventBus();
|
||||
const env = new ContextEnvironmentImpl(
|
||||
{
|
||||
generateContent: async () => ({}),
|
||||
generateJson: async () => ({}),
|
||||
} as unknown as BaseLlmClient,
|
||||
'test-prompt-id',
|
||||
'test',
|
||||
'/tmp',
|
||||
'/tmp',
|
||||
tracer,
|
||||
4,
|
||||
eventBus,
|
||||
);
|
||||
contextManager = ContextManager.create(
|
||||
sidecar,
|
||||
env,
|
||||
tracer,
|
||||
undefined,
|
||||
registry,
|
||||
);
|
||||
});
|
||||
|
||||
const createLargeHistory = (): Content[] => [
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{ text: 'A long long time ago, '.repeat(500) }, // Squashing target
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'model',
|
||||
parts: [{ text: 'in a galaxy far far away...' }],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'some_tool',
|
||||
response: { output: 'TOOL OUTPUT DATA '.repeat(500) }, // Masking target
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{ text: '--- test_file.txt ---\n' + 'FILE DATA '.repeat(1000) }, // Semantic target
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
it('should process history and match golden snapshot', async () => {
|
||||
const history = createLargeHistory();
|
||||
(
|
||||
contextManager as unknown as { pristineEpisodes: Episode[] }
|
||||
).pristineEpisodes = IrMapper.toIr(history, new ContextTokenCalculator(4));
|
||||
const result = await contextManager.projectCompressedHistory();
|
||||
expect(result).toMatchSnapshot();
|
||||
});
|
||||
|
||||
it('should not modify history when under budget', async () => {
|
||||
const history = createLargeHistory();
|
||||
(
|
||||
contextManager as unknown as { pristineEpisodes: Episode[] }
|
||||
).pristineEpisodes = IrMapper.toIr(history, new ContextTokenCalculator(4));
|
||||
// In Golden Tests, we just want to ensure the logic doesn't throw or alter unprotected history in weird ways.
|
||||
// Since we're skipping processors due to being under budget, it should equal history.
|
||||
const tracer2 = new ContextTracer({
|
||||
targetDir: '/tmp',
|
||||
sessionId: 'test2',
|
||||
});
|
||||
const eventBus2 = new ContextEventBus();
|
||||
const env2 = new ContextEnvironmentImpl(
|
||||
{
|
||||
generateContent: async () => ({}),
|
||||
generateJson: async () => ({}),
|
||||
} as unknown as BaseLlmClient,
|
||||
'test-prompt-id',
|
||||
'test',
|
||||
'/tmp',
|
||||
'/tmp',
|
||||
tracer2,
|
||||
4,
|
||||
eventBus2,
|
||||
);
|
||||
contextManager = ContextManager.create(
|
||||
{
|
||||
budget: { retainedTokens: 100000, maxTokens: 150000 },
|
||||
pipelines: [],
|
||||
} as unknown as SidecarConfig,
|
||||
env2,
|
||||
tracer2,
|
||||
);
|
||||
|
||||
(
|
||||
contextManager as unknown as { pristineEpisodes: Episode[] }
|
||||
).pristineEpisodes = IrMapper.toIr(history, new ContextTokenCalculator(4));
|
||||
const result = await contextManager.projectCompressedHistory();
|
||||
|
||||
expect(result.length).toEqual(history.length);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,184 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Content } from '@google/genai';
|
||||
import type { AgentChatHistory } from '../core/agentChatHistory.js';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
import type { Episode } from './ir/types.js';
|
||||
import type { ContextEventBus } from './eventBus.js';
|
||||
import type { ContextTracer } from './tracer.js';
|
||||
import type { ContextEnvironment } from './sidecar/environment.js';
|
||||
import type { SidecarConfig } from './sidecar/types.js';
|
||||
import { PipelineOrchestrator } from './sidecar/orchestrator.js';
|
||||
import { HistoryObserver } from './historyObserver.js';
|
||||
import { generateWorkingBufferView } from './ir/graphUtils.js';
|
||||
import { IrProjector } from './ir/projector.js';
|
||||
import { registerBuiltInProcessors } from './sidecar/builtins.js';
|
||||
import { ProcessorRegistry } from './sidecar/registry.js';
|
||||
|
||||
export class ContextManager {
|
||||
// The stateful, pristine Episodic Intermediate Representation graph.
|
||||
// This allows the agent to remember and summarize continuously without losing data across turns.
|
||||
private pristineEpisodes: Episode[] = [];
|
||||
private readonly eventBus: ContextEventBus;
|
||||
|
||||
// Internal sub-components
|
||||
// Synchronous processors are instantiated but effectively used as singletons within this class
|
||||
private orchestrator: PipelineOrchestrator;
|
||||
private historyObserver?: HistoryObserver;
|
||||
|
||||
static create(
|
||||
sidecar: SidecarConfig,
|
||||
env: ContextEnvironment,
|
||||
tracer: ContextTracer,
|
||||
orchestrator?: PipelineOrchestrator,
|
||||
registry?: ProcessorRegistry,
|
||||
): ContextManager {
|
||||
if (!registry) {
|
||||
registry = new ProcessorRegistry();
|
||||
registerBuiltInProcessors(registry);
|
||||
}
|
||||
const orch =
|
||||
orchestrator ||
|
||||
new PipelineOrchestrator(sidecar, env, env.eventBus, tracer, registry);
|
||||
return new ContextManager(sidecar, env, tracer, orch);
|
||||
}
|
||||
|
||||
// Use ContextManager.create() instead
|
||||
private constructor(
|
||||
private sidecar: SidecarConfig,
|
||||
private env: ContextEnvironment,
|
||||
private readonly tracer: ContextTracer,
|
||||
orchestrator: PipelineOrchestrator,
|
||||
) {
|
||||
this.eventBus = env.eventBus;
|
||||
this.orchestrator = orchestrator;
|
||||
|
||||
this.eventBus.onPristineHistoryUpdated((event) => {
|
||||
this.pristineEpisodes = event.episodes;
|
||||
this.evaluateTriggers();
|
||||
});
|
||||
|
||||
this.eventBus.onVariantReady((event) => {
|
||||
// Find the target episode in the pristine graph
|
||||
const targetEp = this.pristineEpisodes.find(
|
||||
(ep) => ep.id === event.targetId,
|
||||
);
|
||||
if (targetEp) {
|
||||
if (!targetEp.variants) {
|
||||
targetEp.variants = {};
|
||||
}
|
||||
targetEp.variants[event.variantId] = event.variant;
|
||||
this.tracer.logEvent(
|
||||
'ContextManager',
|
||||
`Received async variant [${event.variantId}] for Episode ${event.targetId}`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`ContextManager: Received async variant [${event.variantId}] for Episode ${event.targetId}.`,
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Safely stops background workers and clears event listeners.
|
||||
*/
|
||||
shutdown() {
|
||||
this.orchestrator.shutdown();
|
||||
if (this.historyObserver) {
|
||||
this.historyObserver.stop();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Evaluates if the current working buffer exceeds configured budget thresholds,
|
||||
* firing consolidation events if necessary.
|
||||
*/
|
||||
private evaluateTriggers() {
|
||||
if (!this.sidecar.budget) return;
|
||||
|
||||
const workingBuffer = this.getWorkingBufferView();
|
||||
const currentTokens =
|
||||
this.env.tokenCalculator.calculateEpisodeListTokens(workingBuffer);
|
||||
|
||||
this.tracer.logEvent('ContextManager', 'Evaluated triggers', {
|
||||
currentTokens,
|
||||
retainedTokens: this.sidecar.budget.retainedTokens,
|
||||
});
|
||||
|
||||
// 1. Eager Compute Trigger
|
||||
this.eventBus.emitChunkReceived({ episodes: this.pristineEpisodes });
|
||||
|
||||
// 2. Budget Crossed Trigger
|
||||
if (currentTokens > this.sidecar.budget.retainedTokens) {
|
||||
const deficit = currentTokens - this.sidecar.budget.retainedTokens;
|
||||
this.tracer.logEvent(
|
||||
'ContextManager',
|
||||
'Budget crossed. Emitting ConsolidationNeeded',
|
||||
{ deficit },
|
||||
);
|
||||
this.eventBus.emitConsolidationNeeded({
|
||||
episodes: workingBuffer,
|
||||
targetDeficit: deficit,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Subscribes to the core AgentChatHistory to natively track all message events,
|
||||
* converting them seamlessly into pristine Episodes.
|
||||
*/
|
||||
subscribeToHistory(chatHistory: AgentChatHistory) {
|
||||
if (this.historyObserver) {
|
||||
this.historyObserver.stop();
|
||||
}
|
||||
|
||||
this.historyObserver = new HistoryObserver(
|
||||
chatHistory,
|
||||
this.eventBus,
|
||||
this.tracer,
|
||||
this.env.tokenCalculator,
|
||||
);
|
||||
this.historyObserver.start();
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a computed view of the pristine log.
|
||||
* Sweeps backwards (newest to oldest), tracking rolling tokens.
|
||||
* When rollingTokens > retainedTokens, it injects the "best" available ready variant
|
||||
* (snapshot > summary > masked) instead of the raw text.
|
||||
* Handles N-to-1 variant skipping automatically.
|
||||
*/
|
||||
getWorkingBufferView(): Episode[] {
|
||||
return generateWorkingBufferView(
|
||||
this.pristineEpisodes,
|
||||
this.sidecar.budget.retainedTokens,
|
||||
this.tracer,
|
||||
this.env,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a temporary, compressed Content[] array to be used exclusively for the LLM request.
|
||||
* This does NOT mutate the pristine episodic graph.
|
||||
*/
|
||||
async projectCompressedHistory(): Promise<Content[]> {
|
||||
this.tracer.logEvent('ContextManager', 'Projection requested.');
|
||||
const protectedIds = new Set<string>();
|
||||
if (this.pristineEpisodes.length > 0) {
|
||||
protectedIds.add(this.pristineEpisodes[0].id); // Structural invariant
|
||||
}
|
||||
|
||||
return IrProjector.project(
|
||||
this.getWorkingBufferView(),
|
||||
this.orchestrator,
|
||||
this.sidecar,
|
||||
this.tracer,
|
||||
this.env,
|
||||
protectedIds,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { EventEmitter } from 'node:events';
|
||||
import type { Episode, Variant } from './ir/types.js';
|
||||
|
||||
export interface PristineHistoryUpdatedEvent {
|
||||
episodes: Episode[];
|
||||
}
|
||||
|
||||
export interface ContextConsolidationEvent {
|
||||
episodes: Episode[];
|
||||
targetDeficit: number;
|
||||
}
|
||||
|
||||
export interface IrChunkReceivedEvent {
|
||||
episodes: Episode[];
|
||||
}
|
||||
|
||||
export interface VariantReadyEvent {
|
||||
targetId: string; // The Episode or Step ID this variant attaches to
|
||||
variantId: string; // A unique ID for the variant itself
|
||||
variant: Variant;
|
||||
}
|
||||
|
||||
export class ContextEventBus extends EventEmitter {
|
||||
emitPristineHistoryUpdated(event: PristineHistoryUpdatedEvent) {
|
||||
this.emit('PRISTINE_HISTORY_UPDATED', event);
|
||||
}
|
||||
|
||||
onPristineHistoryUpdated(
|
||||
listener: (event: PristineHistoryUpdatedEvent) => void,
|
||||
) {
|
||||
this.on('PRISTINE_HISTORY_UPDATED', listener);
|
||||
}
|
||||
|
||||
emitChunkReceived(event: IrChunkReceivedEvent) {
|
||||
this.emit('IR_CHUNK_RECEIVED', event);
|
||||
}
|
||||
|
||||
onChunkReceived(listener: (event: IrChunkReceivedEvent) => void) {
|
||||
this.on('IR_CHUNK_RECEIVED', listener);
|
||||
}
|
||||
|
||||
emitConsolidationNeeded(event: ContextConsolidationEvent) {
|
||||
this.emit('BUDGET_RETAINED_CROSSED', event);
|
||||
}
|
||||
|
||||
onConsolidationNeeded(listener: (event: ContextConsolidationEvent) => void) {
|
||||
this.on('BUDGET_RETAINED_CROSSED', listener);
|
||||
}
|
||||
|
||||
emitVariantReady(event: VariantReadyEvent) {
|
||||
this.emit('VARIANT_READY', event);
|
||||
}
|
||||
|
||||
onVariantReady(listener: (event: VariantReadyEvent) => void) {
|
||||
this.on('VARIANT_READY', listener);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type {
|
||||
AgentChatHistory,
|
||||
HistoryEvent,
|
||||
} from '../core/agentChatHistory.js';
|
||||
import { IrMapper } from './ir/mapper.js';
|
||||
import type { ContextTokenCalculator } from './utils/contextTokenCalculator.js';
|
||||
import type { ContextEventBus } from './eventBus.js';
|
||||
import type { ContextTracer } from './tracer.js';
|
||||
|
||||
/**
|
||||
* Connects the raw AgentChatHistory to the ContextManager.
|
||||
* It maps raw messages into Episodic Intermediate Representation (IR)
|
||||
* and evaluates background triggers whenever history changes.
|
||||
*/
|
||||
export class HistoryObserver {
|
||||
private unsubscribeHistory?: () => void;
|
||||
|
||||
constructor(
|
||||
private readonly chatHistory: AgentChatHistory,
|
||||
private readonly eventBus: ContextEventBus,
|
||||
private readonly tracer: ContextTracer,
|
||||
private readonly tokenCalculator: ContextTokenCalculator,
|
||||
) {}
|
||||
|
||||
start() {
|
||||
if (this.unsubscribeHistory) {
|
||||
this.unsubscribeHistory();
|
||||
}
|
||||
|
||||
this.unsubscribeHistory = this.chatHistory.subscribe(
|
||||
(_event: HistoryEvent) => {
|
||||
// Rebuild the pristine IR graph from the full source history on every change.
|
||||
const pristineEpisodes = IrMapper.toIr(
|
||||
this.chatHistory.get(),
|
||||
this.tokenCalculator,
|
||||
);
|
||||
this.tracer.logEvent(
|
||||
'HistoryObserver',
|
||||
'Rebuilt pristine graph from chat history update',
|
||||
{ episodeCount: pristineEpisodes.length },
|
||||
);
|
||||
|
||||
this.eventBus.emitPristineHistoryUpdated({
|
||||
episodes: pristineEpisodes,
|
||||
});
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
stop() {
|
||||
if (this.unsubscribeHistory) {
|
||||
this.unsubscribeHistory();
|
||||
this.unsubscribeHistory = undefined;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Episode } from './types.js';
|
||||
|
||||
export interface MutationRecord {
|
||||
episodeId: string;
|
||||
type: 'modified' | 'inserted' | 'replaced' | 'deleted';
|
||||
action: string;
|
||||
originalIds?: string[]; // If replaced
|
||||
episode?: Episode; // For new or modified
|
||||
}
|
||||
|
||||
export class EpisodeEditor {
|
||||
private originalMap: Map<string, Episode>;
|
||||
private workingOrder: string[];
|
||||
private workingMap: Map<string, Episode>;
|
||||
private mutations: MutationRecord[] = [];
|
||||
|
||||
constructor(episodes: Episode[]) {
|
||||
this.originalMap = new Map(episodes.map((e) => [e.id, e]));
|
||||
this.workingOrder = episodes.map((e) => e.id);
|
||||
this.workingMap = new Map(episodes.map((e) => [e.id, e]));
|
||||
}
|
||||
|
||||
/**
|
||||
* Provides a readonly view of the current working state of the episodes.
|
||||
* Processors should iterate over this to decide what to mutate.
|
||||
*/
|
||||
get episodes(): readonly Episode[] {
|
||||
return this.workingOrder.map((id) => this.workingMap.get(id)!);
|
||||
}
|
||||
|
||||
/**
|
||||
* Safely edits an existing episode.
|
||||
* The framework will handle deeply cloning the episode before passing it to the mutator,
|
||||
* guaranteeing that original references are never modified.
|
||||
*/
|
||||
editEpisode(id: string, action: string, mutator: (draft: Episode) => void) {
|
||||
const ep = this.workingMap.get(id);
|
||||
if (!ep) return;
|
||||
|
||||
// Lazy deep clone only if it's the original reference
|
||||
if (ep === this.originalMap.get(id)) {
|
||||
const clone = structuredClone(ep);
|
||||
this.workingMap.set(id, clone);
|
||||
}
|
||||
|
||||
const draft = this.workingMap.get(id)!;
|
||||
mutator(draft);
|
||||
|
||||
// Log mutation if not already tracked as modified/inserted/replaced
|
||||
if (!this.mutations.find((m) => m.episodeId === id)) {
|
||||
this.mutations.push({
|
||||
episodeId: id,
|
||||
type: 'modified',
|
||||
action,
|
||||
episode: draft,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Inserts a brand new episode into the graph at the specified index.
|
||||
*/
|
||||
insertEpisode(index: number, newEpisode: Episode, action: string) {
|
||||
this.workingMap.set(newEpisode.id, newEpisode);
|
||||
this.workingOrder.splice(index, 0, newEpisode.id);
|
||||
this.mutations.push({
|
||||
episodeId: newEpisode.id,
|
||||
type: 'inserted',
|
||||
action,
|
||||
episode: newEpisode,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Replaces a set of older episodes with a single new episode (e.g., a Summary or Snapshot).
|
||||
* It inserts the new episode at the lowest index of the removed episodes.
|
||||
*/
|
||||
replaceEpisodes(oldIds: string[], newEpisode: Episode, action: string) {
|
||||
const indices = oldIds
|
||||
.map((id) => this.workingOrder.indexOf(id))
|
||||
.filter((i) => i !== -1);
|
||||
if (indices.length === 0) return;
|
||||
|
||||
const insertIndex = Math.min(...indices);
|
||||
|
||||
// Remove old
|
||||
this.workingOrder = this.workingOrder.filter((id) => !oldIds.includes(id));
|
||||
for (const id of oldIds) {
|
||||
this.workingMap.delete(id);
|
||||
}
|
||||
|
||||
// Insert new
|
||||
this.workingOrder.splice(insertIndex, 0, newEpisode.id);
|
||||
this.workingMap.set(newEpisode.id, newEpisode);
|
||||
|
||||
this.mutations.push({
|
||||
episodeId: newEpisode.id,
|
||||
type: 'replaced',
|
||||
action,
|
||||
originalIds: oldIds,
|
||||
episode: newEpisode,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Removes episodes from the graph completely (e.g., emergency truncation).
|
||||
*/
|
||||
removeEpisodes(oldIds: string[], action: string) {
|
||||
this.workingOrder = this.workingOrder.filter((id) => !oldIds.includes(id));
|
||||
for (const id of oldIds) {
|
||||
this.workingMap.delete(id);
|
||||
this.mutations.push({ episodeId: id, type: 'deleted', action });
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves the final, finalized array of episodes.
|
||||
* Called by the Orchestrator.
|
||||
*/
|
||||
getFinalEpisodes(): Episode[] {
|
||||
return this.workingOrder.map((id) => this.workingMap.get(id)!);
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves a log of all structural and property mutations performed by this editor.
|
||||
* Called by the Orchestrator to emit VariantReady events.
|
||||
*/
|
||||
getMutations(): MutationRecord[] {
|
||||
return this.mutations;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Content, Part } from '@google/genai';
|
||||
import type { Episode, EpisodeStep, UserPrompt, AgentYield } from './types.js';
|
||||
|
||||
export function fromIr(episodes: Episode[]): Content[] {
|
||||
const history: Content[] = [];
|
||||
|
||||
for (const ep of episodes) {
|
||||
if (ep.trigger.type === 'USER_PROMPT') {
|
||||
const triggerContent = serializeTrigger(ep.trigger);
|
||||
if (triggerContent) history.push(triggerContent);
|
||||
}
|
||||
|
||||
const stepContents = serializeSteps(ep.steps);
|
||||
history.push(...stepContents);
|
||||
|
||||
if (ep.yield) {
|
||||
history.push(serializeYield(ep.yield));
|
||||
}
|
||||
}
|
||||
|
||||
return history;
|
||||
}
|
||||
|
||||
function serializeTrigger(trigger: UserPrompt): Content | null {
|
||||
const parts: Part[] = [];
|
||||
for (const sp of trigger.semanticParts) {
|
||||
if (sp.presentation) {
|
||||
parts.push({ text: sp.presentation.text });
|
||||
} else if (sp.type === 'text') {
|
||||
parts.push({ text: sp.text });
|
||||
} else if (sp.type === 'inline_data') {
|
||||
parts.push({
|
||||
inlineData: { mimeType: sp.mimeType, data: sp.data },
|
||||
});
|
||||
} else if (sp.type === 'file_data') {
|
||||
parts.push({
|
||||
fileData: { mimeType: sp.mimeType, fileUri: sp.fileUri },
|
||||
});
|
||||
} else if (sp.type === 'raw_part') {
|
||||
parts.push(sp.part);
|
||||
}
|
||||
}
|
||||
return parts.length > 0 ? { role: 'user', parts } : null;
|
||||
}
|
||||
|
||||
function serializeSteps(steps: EpisodeStep[]): Content[] {
|
||||
const history: Content[] = [];
|
||||
let pendingModelParts: Part[] = [];
|
||||
let pendingUserParts: Part[] = [];
|
||||
|
||||
const flushPending = () => {
|
||||
if (pendingModelParts.length > 0) {
|
||||
history.push({ role: 'model', parts: [...pendingModelParts] });
|
||||
pendingModelParts = [];
|
||||
}
|
||||
if (pendingUserParts.length > 0) {
|
||||
history.push({ role: 'user', parts: [...pendingUserParts] });
|
||||
pendingUserParts = [];
|
||||
}
|
||||
};
|
||||
|
||||
for (const step of steps) {
|
||||
if (step.type === 'AGENT_THOUGHT') {
|
||||
if (pendingUserParts.length > 0) flushPending();
|
||||
pendingModelParts.push({
|
||||
text: step.presentation?.text ?? step.text,
|
||||
});
|
||||
} else if (step.type === 'TOOL_EXECUTION') {
|
||||
pendingModelParts.push({
|
||||
functionCall: {
|
||||
name: step.toolName,
|
||||
args: step.intent,
|
||||
id: step.id,
|
||||
},
|
||||
});
|
||||
const observation = step.presentation
|
||||
? step.presentation.observation
|
||||
: step.observation;
|
||||
pendingUserParts.push({
|
||||
functionResponse: {
|
||||
name: step.toolName,
|
||||
response:
|
||||
typeof observation === 'string'
|
||||
? { message: observation }
|
||||
: observation,
|
||||
id: step.id,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
flushPending();
|
||||
|
||||
return history;
|
||||
}
|
||||
|
||||
function serializeYield(yieldNode: AgentYield): Content {
|
||||
return {
|
||||
role: 'model',
|
||||
parts: [{ text: yieldNode.presentation?.text ?? yieldNode.text }],
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { generateWorkingBufferView } from './graphUtils.js';
|
||||
import {
|
||||
createMockEnvironment,
|
||||
createDummyEpisode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import type { AgentThought, UserPrompt } from './types.js';
|
||||
|
||||
describe('graphUtils (View Generator)', () => {
|
||||
let env: ContextEnvironment;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
env = createMockEnvironment();
|
||||
// Our token mock is 1 char = 1 token for simplicity
|
||||
vi.spyOn(
|
||||
env.tokenCalculator,
|
||||
'calculateEpisodeListTokens',
|
||||
).mockImplementation((eps) =>
|
||||
eps.reduce(
|
||||
(acc, ep) => acc + (ep.trigger.metadata.originalTokens || 100),
|
||||
0,
|
||||
),
|
||||
);
|
||||
});
|
||||
|
||||
it('returns pristine episodes untouched if under budget', () => {
|
||||
const episodes = [
|
||||
createDummyEpisode('ep-1', 'USER_PROMPT', [{ type: 'text', text: '1' }]),
|
||||
createDummyEpisode('ep-2', 'USER_PROMPT', [{ type: 'text', text: '2' }]),
|
||||
];
|
||||
|
||||
// We retain 5000 tokens. Total mock tokens = 200.
|
||||
const view = generateWorkingBufferView(episodes, 5000, env.tracer, env);
|
||||
|
||||
expect(view).toHaveLength(2);
|
||||
// Must be a deep copy! The view generator clones episodes.
|
||||
expect(view).not.toBe(episodes);
|
||||
expect(view[0].id).toBe('ep-1');
|
||||
expect(view[1].id).toBe('ep-2');
|
||||
});
|
||||
|
||||
it('swaps to Masked variant when over budget (rolling backwards)', () => {
|
||||
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
|
||||
{ text: '1', type: 'text' },
|
||||
]);
|
||||
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
|
||||
{ text: '2', type: 'text' },
|
||||
]);
|
||||
|
||||
ep1.variants = {
|
||||
masked: {
|
||||
type: 'masked',
|
||||
status: 'ready',
|
||||
text: '<MASKED>',
|
||||
recoveredTokens: 10,
|
||||
},
|
||||
};
|
||||
|
||||
// We only retain 100 tokens.
|
||||
// ep-2 (newest) takes 100 tokens.
|
||||
// Now rolling = 100. Over budget!
|
||||
// ep-1 is evaluated, and swapped for Masked.
|
||||
const view = generateWorkingBufferView([ep1, ep2], 10, env.tracer, env);
|
||||
|
||||
expect(view).toHaveLength(2);
|
||||
expect(view[1].id).toBe('ep-2'); // Unchanged (newest)
|
||||
|
||||
expect(view[0].id).toBe('ep-1');
|
||||
expect(
|
||||
(view[0].trigger as UserPrompt).semanticParts[0].presentation?.text,
|
||||
).toBe('<MASKED>');
|
||||
});
|
||||
|
||||
it('swaps to Summary variant when over budget', () => {
|
||||
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
|
||||
{ type: 'text', text: '1' },
|
||||
]);
|
||||
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
|
||||
{ type: 'text', text: '2' },
|
||||
]);
|
||||
|
||||
ep1.variants = {
|
||||
summary: {
|
||||
type: 'summary',
|
||||
status: 'ready',
|
||||
text: '<SUMMARY>',
|
||||
recoveredTokens: 50,
|
||||
},
|
||||
};
|
||||
|
||||
const view = generateWorkingBufferView([ep1, ep2], 10, env.tracer, env);
|
||||
|
||||
expect(view).toHaveLength(2);
|
||||
|
||||
// The summary completely replaces the internal steps and clears the yield.
|
||||
expect(view[0].steps).toHaveLength(1);
|
||||
expect(view[0].steps[0].type).toBe('AGENT_THOUGHT');
|
||||
expect((view[0].steps[0] as AgentThought).text).toBe('<SUMMARY>');
|
||||
expect(view[0].yield).toBeUndefined();
|
||||
});
|
||||
|
||||
it('handles complex N-to-1 Snapshot skipping gracefully', () => {
|
||||
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
|
||||
{ type: 'text', text: '1' },
|
||||
]);
|
||||
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
|
||||
{ type: 'text', text: '2' },
|
||||
]);
|
||||
const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', [
|
||||
{ type: 'text', text: '3' },
|
||||
]);
|
||||
const ep4 = createDummyEpisode('ep-4', 'USER_PROMPT', [
|
||||
{ type: 'text', text: '4' },
|
||||
]);
|
||||
|
||||
// ep-3 has a snapshot that replaces [ep-1, ep-2, ep-3]
|
||||
const snapshotEp = createDummyEpisode('snap-1', 'SYSTEM_EVENT', []);
|
||||
|
||||
ep3.variants = {
|
||||
snapshot: {
|
||||
type: 'snapshot',
|
||||
status: 'ready',
|
||||
episode: snapshotEp,
|
||||
replacedEpisodeIds: ['ep-1', 'ep-2', 'ep-3'],
|
||||
},
|
||||
};
|
||||
|
||||
// We only retain 5 tokens, forcing the sweep to use variants for EVERYTHING except ep4.
|
||||
const view = generateWorkingBufferView(
|
||||
[ep1, ep2, ep3, ep4],
|
||||
5,
|
||||
env.tracer,
|
||||
env,
|
||||
);
|
||||
|
||||
// Result should be exactly: [snapshot, ep-4]
|
||||
expect(view).toHaveLength(2);
|
||||
expect(view[0].id).toBe('snap-1');
|
||||
expect(view[1].id).toBe('ep-4');
|
||||
});
|
||||
|
||||
it('ignores variants that are not yet "ready"', () => {
|
||||
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
|
||||
{ type: 'text', text: '1' },
|
||||
]);
|
||||
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
|
||||
{ type: 'text', text: '2' },
|
||||
]);
|
||||
|
||||
ep1.variants = {
|
||||
masked: {
|
||||
type: 'masked',
|
||||
status: 'computing',
|
||||
text: '<MASKED>',
|
||||
recoveredTokens: 10,
|
||||
},
|
||||
};
|
||||
|
||||
const view = generateWorkingBufferView([ep1, ep2], 10, env.tracer, env);
|
||||
|
||||
// Because the variant was computing, it must fall back to the raw pristine text.
|
||||
expect(view).toHaveLength(2);
|
||||
expect(
|
||||
(view[0].trigger as UserPrompt).semanticParts[0].presentation,
|
||||
).toBeUndefined();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,174 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Episode } from './types.js';
|
||||
import type { ContextTracer } from '../tracer.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
|
||||
/**
|
||||
* Generates a computed view of the pristine log.
|
||||
* Sweeps backwards (newest to oldest), tracking rolling tokens.
|
||||
* When rollingTokens > retainedTokens, it injects the "best" available ready variant
|
||||
* (snapshot > summary > masked) instead of the raw text.
|
||||
* Handles N-to-1 variant skipping automatically.
|
||||
*/
|
||||
|
||||
export function generateWorkingBufferView(
|
||||
pristineEpisodes: Episode[],
|
||||
retainedTokens: number,
|
||||
tracer: ContextTracer,
|
||||
env: ContextEnvironment,
|
||||
): Episode[] {
|
||||
const currentEpisodes: Episode[] = [];
|
||||
let rollingTokens = 0;
|
||||
const skippedIds = new Set<string>();
|
||||
tracer.logEvent('ViewGenerator', 'Generating Working Buffer View');
|
||||
|
||||
for (let i = pristineEpisodes.length - 1; i >= 0; i--) {
|
||||
const ep = pristineEpisodes[i];
|
||||
|
||||
// If this episode was already replaced by an N-to-1 Snapshot injected earlier in the sweep, skip it entirely!
|
||||
if (skippedIds.has(ep.id)) {
|
||||
tracer.logEvent(
|
||||
'ViewGenerator',
|
||||
`Skipping episode [${ep.id}] due to N-to-1 replacement.`,
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
let projectedTrigger: typeof ep.trigger;
|
||||
|
||||
if (ep.trigger.type === 'USER_PROMPT') {
|
||||
projectedTrigger = {
|
||||
...ep.trigger,
|
||||
metadata: {
|
||||
...ep.trigger.metadata,
|
||||
transformations: [...(ep.trigger.metadata?.transformations || [])],
|
||||
},
|
||||
semanticParts: ep.trigger.semanticParts.map((sp) => ({ ...sp })),
|
||||
};
|
||||
} else {
|
||||
projectedTrigger = {
|
||||
...ep.trigger,
|
||||
metadata: {
|
||||
...ep.trigger.metadata,
|
||||
transformations: [...(ep.trigger.metadata?.transformations || [])],
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
let projectedEp: Episode = {
|
||||
...ep,
|
||||
trigger: projectedTrigger,
|
||||
steps: ep.steps.map((step) => ({
|
||||
...step,
|
||||
metadata: {
|
||||
...step.metadata,
|
||||
transformations: [...(step.metadata?.transformations || [])],
|
||||
},
|
||||
})),
|
||||
yield: ep.yield
|
||||
? {
|
||||
...ep.yield,
|
||||
metadata: {
|
||||
...ep.yield.metadata,
|
||||
transformations: [...(ep.yield.metadata?.transformations || [])],
|
||||
},
|
||||
}
|
||||
: undefined,
|
||||
};
|
||||
|
||||
const epTokens = env.tokenCalculator.calculateEpisodeListTokens([
|
||||
projectedEp,
|
||||
]);
|
||||
|
||||
if (rollingTokens > retainedTokens && ep.variants) {
|
||||
const snapshot = ep.variants['snapshot'];
|
||||
const summary = ep.variants['summary'];
|
||||
const masked = ep.variants['masked'];
|
||||
|
||||
if (
|
||||
snapshot &&
|
||||
snapshot.status === 'ready' &&
|
||||
snapshot.type === 'snapshot'
|
||||
) {
|
||||
projectedEp = snapshot.episode;
|
||||
// Mark all the episodes this snapshot covers to be skipped by the backwards sweep.
|
||||
for (const id of snapshot.replacedEpisodeIds) {
|
||||
skippedIds.add(id);
|
||||
}
|
||||
tracer.logEvent(
|
||||
'ViewGenerator',
|
||||
`Episode [${ep.id}] has SnapshotVariant. Selecting variant over raw text. Added [${snapshot.replacedEpisodeIds.join(',')}] to skippedIds.`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`Opportunistically swapped Episodes [${snapshot.replacedEpisodeIds.join(', ')}] for pre-computed Snapshot variant.`,
|
||||
);
|
||||
} else if (
|
||||
summary &&
|
||||
summary.status === 'ready' &&
|
||||
summary.type === 'summary'
|
||||
) {
|
||||
projectedEp.steps = [
|
||||
{
|
||||
id: ep.id + '-summary',
|
||||
type: 'AGENT_THOUGHT',
|
||||
text: summary.text,
|
||||
metadata: {
|
||||
originalTokens: epTokens,
|
||||
currentTokens: summary.recoveredTokens || 50,
|
||||
transformations: [
|
||||
{
|
||||
processorName: 'AsyncSemanticCompressor',
|
||||
action: 'SUMMARIZED',
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
] as typeof projectedEp.steps;
|
||||
projectedEp.yield = undefined;
|
||||
tracer.logEvent(
|
||||
'ViewGenerator',
|
||||
`Episode [${ep.id}] has SummaryVariant. Selecting variant over raw text.`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`Opportunistically swapped Episode ${ep.id} for pre-computed Summary variant.`,
|
||||
);
|
||||
} else if (
|
||||
masked &&
|
||||
masked.status === 'ready' &&
|
||||
masked.type === 'masked'
|
||||
) {
|
||||
if (
|
||||
projectedEp.trigger.type === 'USER_PROMPT' &&
|
||||
projectedEp.trigger.semanticParts &&
|
||||
projectedEp.trigger.semanticParts.length > 0
|
||||
) {
|
||||
projectedEp.trigger.semanticParts[0].presentation = {
|
||||
text: masked.text,
|
||||
tokens: masked.recoveredTokens || 10,
|
||||
};
|
||||
}
|
||||
tracer.logEvent(
|
||||
'ViewGenerator',
|
||||
`Episode [${ep.id}] has MaskedVariant. Selecting variant over raw text.`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`Opportunistically swapped Episode ${ep.id} for pre-computed Masked variant.`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
currentEpisodes.unshift(projectedEp);
|
||||
rollingTokens += env.tokenCalculator.calculateEpisodeListTokens([
|
||||
projectedEp,
|
||||
]);
|
||||
}
|
||||
|
||||
return currentEpisodes;
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { IrMapper } from './mapper.js';
|
||||
import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
|
||||
import type { Content } from '@google/genai';
|
||||
import type { UserPrompt, ToolExecution, AgentThought } from './types.js';
|
||||
|
||||
describe('IrMapper', () => {
|
||||
it('should correctly map a complex conversation into Episodes and back', () => {
|
||||
const rawHistory: Content[] = [
|
||||
{ role: 'user', parts: [{ text: 'Can you read file A and B?' }] },
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{ text: 'Let me check those files.' },
|
||||
{
|
||||
functionCall: {
|
||||
id: 'call_1',
|
||||
name: 'read_file',
|
||||
args: { filepath: 'A.txt' },
|
||||
},
|
||||
},
|
||||
{
|
||||
functionCall: {
|
||||
id: 'call_2',
|
||||
name: 'read_file',
|
||||
args: { filepath: 'B.txt' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
id: 'call_1',
|
||||
name: 'read_file',
|
||||
response: { output: 'Contents of A' },
|
||||
},
|
||||
},
|
||||
{
|
||||
functionResponse: {
|
||||
id: 'call_2',
|
||||
name: 'read_file',
|
||||
response: { output: 'Contents of B' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{ text: 'Thanks. Now I will compile.' },
|
||||
{
|
||||
functionCall: {
|
||||
id: 'call_3',
|
||||
name: 'shell',
|
||||
args: { cmd: 'make' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
id: 'call_3',
|
||||
name: 'shell',
|
||||
response: { output: 'success' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{ role: 'model', parts: [{ text: 'Everything is done!' }] },
|
||||
];
|
||||
|
||||
const tokenCalculator = new ContextTokenCalculator(4);
|
||||
const episodes = IrMapper.toIr(rawHistory, tokenCalculator);
|
||||
|
||||
expect(episodes).toHaveLength(1);
|
||||
const ep = episodes[0];
|
||||
|
||||
expect(ep.trigger.type).toBe('USER_PROMPT');
|
||||
expect(
|
||||
((ep.trigger as UserPrompt).semanticParts[0] as { text: string }).text,
|
||||
).toBe('Can you read file A and B?');
|
||||
|
||||
// Steps should be: Thought, ToolExecution(A), ToolExecution(B), Thought, ToolExecution(make)
|
||||
expect(ep.steps).toHaveLength(5);
|
||||
expect(ep.steps[0].type).toBe('AGENT_THOUGHT');
|
||||
expect(ep.steps[1].type).toBe('TOOL_EXECUTION');
|
||||
expect((ep.steps[1] as ToolExecution).toolName).toBe('read_file');
|
||||
expect((ep.steps[1] as ToolExecution).intent).toEqual({
|
||||
filepath: 'A.txt',
|
||||
});
|
||||
expect((ep.steps[1] as ToolExecution).observation).toEqual({
|
||||
output: 'Contents of A',
|
||||
});
|
||||
|
||||
expect(ep.steps[2].type).toBe('TOOL_EXECUTION');
|
||||
expect((ep.steps[2] as ToolExecution).intent).toEqual({
|
||||
filepath: 'B.txt',
|
||||
});
|
||||
|
||||
expect(ep.steps[3].type).toBe('AGENT_THOUGHT');
|
||||
|
||||
expect(ep.steps[4].type).toBe('TOOL_EXECUTION');
|
||||
expect((ep.steps[4] as ToolExecution).toolName).toBe('shell');
|
||||
|
||||
expect(ep.yield?.type).toBe('AGENT_YIELD');
|
||||
expect(ep.yield?.text).toBe('Everything is done!');
|
||||
|
||||
// Test Re-serialization
|
||||
const reconstituted = IrMapper.fromIr(episodes);
|
||||
|
||||
// Compare basic structure (the reconstituted version might have slightly different grouping of calls/responses
|
||||
// based on flush logic, but semantically equivalent)
|
||||
expect(reconstituted[0]).toEqual(rawHistory[0]);
|
||||
// Reconstituted history is identical except tool IDs will be reassigned because IrMapper discards string IDs in favor of deterministic object hash IDs
|
||||
expect(reconstituted[1].parts![0]).toEqual(rawHistory[1].parts![0]);
|
||||
|
||||
// The exact structural equivalence isn't mathematically perfect because Gemini allows mixing text and calls
|
||||
// in one Content block, but the flat representation is semantically identical.
|
||||
});
|
||||
|
||||
it('should correctly handle multi-tool-calls grouped within a single turn without dropping observations', () => {
|
||||
const rawHistory: Content[] = [
|
||||
{
|
||||
role: 'user',
|
||||
parts: [{ text: 'Examine both of these tools please.' }],
|
||||
},
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{ text: 'I will call them concurrently.' },
|
||||
{
|
||||
functionCall: {
|
||||
id: 'c1',
|
||||
name: 'tool_one',
|
||||
args: { p: 1 },
|
||||
},
|
||||
},
|
||||
{
|
||||
functionCall: {
|
||||
id: 'c2',
|
||||
name: 'tool_two',
|
||||
args: { p: 2 },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
// Gemini forces the user turn to contain ALL function responses for that model turn
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
id: 'c1',
|
||||
name: 'tool_one',
|
||||
response: { r: 1 },
|
||||
},
|
||||
},
|
||||
{
|
||||
functionResponse: {
|
||||
id: 'c2',
|
||||
name: 'tool_two',
|
||||
response: { r: 2 },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'model',
|
||||
parts: [{ text: 'Both complete.' }],
|
||||
},
|
||||
];
|
||||
|
||||
const tokenCalculator = new ContextTokenCalculator(4);
|
||||
const episodes = IrMapper.toIr(rawHistory, tokenCalculator);
|
||||
|
||||
// It should collapse into a single episode
|
||||
expect(episodes).toHaveLength(1);
|
||||
const ep = episodes[0];
|
||||
|
||||
expect(ep.trigger.type).toBe('USER_PROMPT');
|
||||
|
||||
// The steps array should contain:
|
||||
// 0: AgentThought ("I will call them concurrently")
|
||||
// 1: ToolExecution(tool_one)
|
||||
// 2: ToolExecution(tool_two)
|
||||
|
||||
expect(ep.steps).toHaveLength(3);
|
||||
|
||||
expect(ep.steps[0].type).toBe('AGENT_THOUGHT');
|
||||
expect((ep.steps[0] as AgentThought).text).toBe(
|
||||
'I will call them concurrently.',
|
||||
);
|
||||
|
||||
expect(ep.steps[1].type).toBe('TOOL_EXECUTION');
|
||||
expect((ep.steps[1] as ToolExecution).toolName).toBe('tool_one');
|
||||
expect((ep.steps[1] as ToolExecution).intent).toEqual({ p: 1 });
|
||||
expect((ep.steps[1] as ToolExecution).observation).toEqual({ r: 1 });
|
||||
|
||||
expect(ep.steps[2].type).toBe('TOOL_EXECUTION');
|
||||
expect((ep.steps[2] as ToolExecution).toolName).toBe('tool_two');
|
||||
expect((ep.steps[2] as ToolExecution).intent).toEqual({ p: 2 });
|
||||
expect((ep.steps[2] as ToolExecution).observation).toEqual({ r: 2 });
|
||||
|
||||
// The final model turn should become the yield
|
||||
expect(ep.yield).toBeDefined();
|
||||
expect(ep.yield?.type).toBe('AGENT_YIELD');
|
||||
expect(ep.yield?.text).toBe('Both complete.');
|
||||
|
||||
// Now verify we can reconstitute it without dropping the multiple calls
|
||||
const reconstituted = IrMapper.fromIr(episodes);
|
||||
|
||||
// The reconstituted history should have exactly 4 turns, same as original
|
||||
expect(reconstituted).toHaveLength(4);
|
||||
|
||||
// Check that the Model turn has both function calls
|
||||
expect(reconstituted[1].role).toBe('model');
|
||||
expect(reconstituted[1].parts).toHaveLength(3); // text + call1 + call2
|
||||
expect(reconstituted[1].parts![1].functionCall?.name).toBe('tool_one');
|
||||
expect(reconstituted[1].parts![2].functionCall?.name).toBe('tool_two');
|
||||
|
||||
// Check that the User turn has both function responses
|
||||
expect(reconstituted[2].role).toBe('user');
|
||||
expect(reconstituted[2].parts).toHaveLength(2); // response1 + response2
|
||||
expect(reconstituted[2].parts![0].functionResponse?.name).toBe('tool_one');
|
||||
expect(reconstituted[2].parts![1].functionResponse?.name).toBe('tool_two');
|
||||
});
|
||||
|
||||
it('should guarantee WeakMap ID stability across continuous mapping', () => {
|
||||
// 1. Initial history
|
||||
const history: Content[] = [
|
||||
{ role: 'user', parts: [{ text: 'Hello' }] },
|
||||
{ role: 'model', parts: [{ text: 'Hi there' }] },
|
||||
];
|
||||
|
||||
const tokenCalculator = new ContextTokenCalculator(4);
|
||||
const initialIr = IrMapper.toIr(history, tokenCalculator);
|
||||
expect(initialIr).toHaveLength(1);
|
||||
|
||||
// Save the uniquely generated deterministic ID for the first episode
|
||||
const episodeId = initialIr[0].id;
|
||||
const triggerId = initialIr[0].trigger.id;
|
||||
|
||||
// 2. Push new history (simulating a continuing conversation)
|
||||
history.push({ role: 'user', parts: [{ text: 'How are you?' }] });
|
||||
history.push({ role: 'model', parts: [{ text: 'I am an AI.' }] });
|
||||
|
||||
const updatedIr = IrMapper.toIr(history, tokenCalculator);
|
||||
expect(updatedIr).toHaveLength(2);
|
||||
|
||||
// 3. Verify ID Stability
|
||||
// The exact same ID must be generated for the first episode because the underlying Content object reference hasn't changed.
|
||||
// This proves the WeakMap successfully pinned the reference!
|
||||
expect(updatedIr[0].id).toBe(episodeId);
|
||||
expect(updatedIr[0].trigger.id).toBe(triggerId);
|
||||
|
||||
// Ensure the new episode has a different ID
|
||||
expect(updatedIr[1].id).not.toBe(episodeId);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,31 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Content } from '@google/genai';
|
||||
import type { Episode } from './types.js';
|
||||
import { toIr } from './toIr.js';
|
||||
import { fromIr } from './fromIr.js';
|
||||
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
|
||||
|
||||
export class IrMapper {
|
||||
/**
|
||||
* Translates a flat Gemini Content[] array into our rich Episodic Intermediate Representation.
|
||||
* Groups adjacent function calls and responses into unified ToolExecution nodes.
|
||||
*/
|
||||
static toIr(
|
||||
history: readonly Content[],
|
||||
tokenCalculator: ContextTokenCalculator,
|
||||
): Episode[] {
|
||||
return toIr(history, tokenCalculator);
|
||||
}
|
||||
|
||||
/**
|
||||
* Re-serializes the Episodic IR back into a flat Gemini Content[] array.
|
||||
*/
|
||||
static fromIr(episodes: Episode[]): Content[] {
|
||||
return fromIr(episodes);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Content } from '@google/genai';
|
||||
import { IrMapper } from './mapper.js';
|
||||
import type { Episode } from './types.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import type {
|
||||
ContextEnvironment,
|
||||
ContextTracer,
|
||||
} from '../sidecar/environment.js';
|
||||
import type { PipelineOrchestrator } from '../sidecar/orchestrator.js';
|
||||
import type { SidecarConfig } from '../sidecar/types.js';
|
||||
|
||||
export class IrProjector {
|
||||
/**
|
||||
* Orchestrates the final projection: takes a working buffer view,
|
||||
* applies the Immediate Sanitization pipeline, and enforces token boundaries.
|
||||
*/
|
||||
static async project(
|
||||
workingBuffer: Episode[],
|
||||
orchestrator: PipelineOrchestrator,
|
||||
sidecar: SidecarConfig,
|
||||
tracer: ContextTracer,
|
||||
env: ContextEnvironment,
|
||||
protectedIds: Set<string>,
|
||||
): Promise<Content[]> {
|
||||
if (!sidecar.budget) {
|
||||
const contents = IrMapper.fromIr(workingBuffer);
|
||||
tracer.logEvent('IrProjector', 'Projected Context to LLM (No Budget)', {
|
||||
projectedContext: contents,
|
||||
});
|
||||
return contents;
|
||||
}
|
||||
|
||||
const maxTokens = sidecar.budget.maxTokens;
|
||||
const currentTokens =
|
||||
env.tokenCalculator.calculateEpisodeListTokens(workingBuffer);
|
||||
|
||||
if (currentTokens <= maxTokens) {
|
||||
tracer.logEvent(
|
||||
'IrProjector',
|
||||
`View is within maxTokens (${currentTokens} <= ${maxTokens}). Returning view.`,
|
||||
);
|
||||
const contents = IrMapper.fromIr(workingBuffer);
|
||||
tracer.logEvent('IrProjector', 'Projected Context to LLM', {
|
||||
projectedContext: contents,
|
||||
});
|
||||
return contents;
|
||||
}
|
||||
|
||||
tracer.logEvent(
|
||||
'IrProjector',
|
||||
`View exceeds maxTokens (${currentTokens} > ${maxTokens}). Hitting Synchronous Pressure Barrier.`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`Context Manager Synchronous Barrier triggered: View at ${currentTokens} tokens (limit: ${maxTokens}).`,
|
||||
);
|
||||
|
||||
const processedEpisodes = await orchestrator.executePipeline(
|
||||
'Immediate Sanitization',
|
||||
workingBuffer,
|
||||
{
|
||||
currentTokens,
|
||||
maxTokens: sidecar.budget.maxTokens,
|
||||
retainedTokens: sidecar.budget.retainedTokens,
|
||||
deficitTokens: Math.max(0, currentTokens - sidecar.budget.maxTokens),
|
||||
protectedEpisodeIds: protectedIds,
|
||||
isBudgetSatisfied: currentTokens <= sidecar.budget.maxTokens,
|
||||
},
|
||||
);
|
||||
|
||||
const finalTokens =
|
||||
env.tokenCalculator.calculateEpisodeListTokens(processedEpisodes);
|
||||
tracer.logEvent(
|
||||
'IrProjector',
|
||||
`Finished projection. Final token count: ${finalTokens}.`,
|
||||
);
|
||||
debugLogger.log(
|
||||
`Context Manager finished. Final actual token count: ${finalTokens}.`,
|
||||
);
|
||||
|
||||
const contents = IrMapper.fromIr(processedEpisodes);
|
||||
tracer.logEvent('IrProjector', 'Projected Sanitized Context to LLM', {
|
||||
projectedContextSanitized: contents,
|
||||
});
|
||||
return contents;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,261 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Content, Part } from '@google/genai';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import type {
|
||||
Episode,
|
||||
IrMetadata,
|
||||
SemanticPart,
|
||||
ToolExecution,
|
||||
AgentThought,
|
||||
AgentYield,
|
||||
UserPrompt,
|
||||
SystemEvent,
|
||||
} from './types.js';
|
||||
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
|
||||
|
||||
// WeakMap to provide stable, deterministic identity across parses for the exact same Content/Part references
|
||||
const nodeIdentityMap = new WeakMap<object, string>();
|
||||
|
||||
export function getStableId(obj: object): string {
|
||||
let id = nodeIdentityMap.get(obj);
|
||||
if (!id) {
|
||||
id = randomUUID();
|
||||
nodeIdentityMap.set(obj, id);
|
||||
}
|
||||
return id;
|
||||
}
|
||||
|
||||
function isRecord(v: unknown): v is Record<string, unknown> {
|
||||
return typeof v === 'object' && v !== null && !Array.isArray(v);
|
||||
}
|
||||
|
||||
function isCompleteEpisode(ep: Partial<Episode>): ep is Episode {
|
||||
return (
|
||||
typeof ep.id === 'string' &&
|
||||
typeof ep.timestamp === 'number' &&
|
||||
!!ep.trigger &&
|
||||
Array.isArray(ep.steps)
|
||||
);
|
||||
}
|
||||
|
||||
export function toIr(
|
||||
history: readonly Content[],
|
||||
tokenCalculator: ContextTokenCalculator,
|
||||
): Episode[] {
|
||||
const episodes: Episode[] = [];
|
||||
let currentEpisode: Partial<Episode> | null = null;
|
||||
const pendingCallParts: Map<string, Part> = new Map();
|
||||
|
||||
const createMetadata = (parts: Part[]): IrMetadata => {
|
||||
const tokens = tokenCalculator.estimateTokensForParts(parts, 0);
|
||||
return {
|
||||
originalTokens: tokens,
|
||||
currentTokens: tokens,
|
||||
transformations: [],
|
||||
};
|
||||
};
|
||||
|
||||
const finalizeEpisode = () => {
|
||||
if (currentEpisode && isCompleteEpisode(currentEpisode)) {
|
||||
episodes.push(currentEpisode);
|
||||
}
|
||||
currentEpisode = null;
|
||||
};
|
||||
|
||||
for (const msg of history) {
|
||||
if (!msg.parts) continue;
|
||||
|
||||
if (msg.role === 'user') {
|
||||
const hasToolResponses = msg.parts.some((p) => !!p.functionResponse);
|
||||
const hasUserParts = msg.parts.some(
|
||||
(p) => !!p.text || !!p.inlineData || !!p.fileData,
|
||||
);
|
||||
|
||||
if (hasToolResponses) {
|
||||
currentEpisode = parseToolResponses(
|
||||
msg,
|
||||
currentEpisode,
|
||||
pendingCallParts,
|
||||
tokenCalculator,
|
||||
createMetadata,
|
||||
);
|
||||
}
|
||||
|
||||
if (hasUserParts) {
|
||||
finalizeEpisode();
|
||||
currentEpisode = parseUserParts(msg, createMetadata);
|
||||
}
|
||||
} else if (msg.role === 'model') {
|
||||
currentEpisode = parseModelParts(
|
||||
msg,
|
||||
currentEpisode,
|
||||
pendingCallParts,
|
||||
createMetadata,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if (currentEpisode) {
|
||||
finalizeYield(currentEpisode);
|
||||
finalizeEpisode();
|
||||
}
|
||||
|
||||
return episodes;
|
||||
}
|
||||
|
||||
function parseToolResponses(
|
||||
msg: Content,
|
||||
currentEpisode: Partial<Episode> | null,
|
||||
pendingCallParts: Map<string, Part>,
|
||||
tokenCalculator: ContextTokenCalculator,
|
||||
createMetadata: (parts: Part[]) => IrMetadata,
|
||||
): Partial<Episode> {
|
||||
if (!currentEpisode) {
|
||||
currentEpisode = {
|
||||
id: getStableId(msg),
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: getStableId(msg.parts![0] || msg),
|
||||
type: 'SYSTEM_EVENT',
|
||||
name: 'history_resume',
|
||||
payload: {},
|
||||
metadata: createMetadata([]),
|
||||
} as SystemEvent,
|
||||
steps: [],
|
||||
};
|
||||
}
|
||||
|
||||
for (const part of msg.parts!) {
|
||||
if (part.functionResponse) {
|
||||
const callId = part.functionResponse.id || '';
|
||||
const matchingCall = pendingCallParts.get(callId);
|
||||
|
||||
const intentTokens = matchingCall
|
||||
? tokenCalculator.estimateTokensForParts([matchingCall])
|
||||
: 0;
|
||||
const obsTokens = tokenCalculator.estimateTokensForParts([part]);
|
||||
|
||||
const step: ToolExecution = {
|
||||
id: getStableId(part),
|
||||
type: 'TOOL_EXECUTION',
|
||||
toolName: part.functionResponse.name || 'unknown',
|
||||
intent: isRecord(matchingCall?.functionCall?.args)
|
||||
? matchingCall.functionCall.args
|
||||
: {},
|
||||
observation: isRecord(part.functionResponse.response)
|
||||
? part.functionResponse.response
|
||||
: {},
|
||||
tokens: {
|
||||
intent: intentTokens,
|
||||
observation: obsTokens,
|
||||
},
|
||||
metadata: {
|
||||
originalTokens: intentTokens + obsTokens,
|
||||
currentTokens: intentTokens + obsTokens,
|
||||
transformations: [],
|
||||
},
|
||||
};
|
||||
currentEpisode.steps!.push(step);
|
||||
if (callId) pendingCallParts.delete(callId);
|
||||
}
|
||||
}
|
||||
return currentEpisode;
|
||||
}
|
||||
|
||||
function parseUserParts(
|
||||
msg: Content,
|
||||
createMetadata: (parts: Part[]) => IrMetadata,
|
||||
): Partial<Episode> {
|
||||
const semanticParts: SemanticPart[] = [];
|
||||
for (const p of msg.parts!) {
|
||||
if (p.text !== undefined)
|
||||
semanticParts.push({ type: 'text', text: p.text });
|
||||
else if (p.inlineData)
|
||||
semanticParts.push({
|
||||
type: 'inline_data',
|
||||
mimeType: p.inlineData.mimeType || '',
|
||||
data: p.inlineData.data || '',
|
||||
});
|
||||
else if (p.fileData)
|
||||
semanticParts.push({
|
||||
type: 'file_data',
|
||||
mimeType: p.fileData.mimeType || '',
|
||||
fileUri: p.fileData.fileUri || '',
|
||||
});
|
||||
else if (!p.functionResponse)
|
||||
semanticParts.push({ type: 'raw_part', part: p }); // Preserve unknowns
|
||||
}
|
||||
|
||||
const trigger: UserPrompt = {
|
||||
id: getStableId(msg.parts![0] || msg),
|
||||
type: 'USER_PROMPT',
|
||||
semanticParts,
|
||||
metadata: createMetadata(msg.parts!.filter((p) => !p.functionResponse)),
|
||||
};
|
||||
|
||||
return {
|
||||
id: getStableId(msg),
|
||||
timestamp: Date.now(),
|
||||
trigger,
|
||||
steps: [],
|
||||
};
|
||||
}
|
||||
|
||||
function parseModelParts(
|
||||
msg: Content,
|
||||
currentEpisode: Partial<Episode> | null,
|
||||
pendingCallParts: Map<string, Part>,
|
||||
createMetadata: (parts: Part[]) => IrMetadata,
|
||||
): Partial<Episode> {
|
||||
if (!currentEpisode) {
|
||||
currentEpisode = {
|
||||
id: getStableId(msg),
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: getStableId(msg.parts![0] || msg),
|
||||
type: 'SYSTEM_EVENT',
|
||||
name: 'model_init',
|
||||
payload: {},
|
||||
metadata: createMetadata([]),
|
||||
} as SystemEvent,
|
||||
steps: [],
|
||||
};
|
||||
}
|
||||
|
||||
for (const part of msg.parts!) {
|
||||
if (part.functionCall) {
|
||||
const callId = part.functionCall.id || '';
|
||||
if (callId) pendingCallParts.set(callId, part);
|
||||
} else if (part.text) {
|
||||
const thought: AgentThought = {
|
||||
id: getStableId(part),
|
||||
type: 'AGENT_THOUGHT',
|
||||
text: part.text,
|
||||
metadata: createMetadata([part]),
|
||||
};
|
||||
currentEpisode.steps!.push(thought);
|
||||
}
|
||||
}
|
||||
return currentEpisode;
|
||||
}
|
||||
|
||||
function finalizeYield(currentEpisode: Partial<Episode>) {
|
||||
if (currentEpisode.steps && currentEpisode.steps.length > 0) {
|
||||
const lastStep = currentEpisode.steps[currentEpisode.steps.length - 1];
|
||||
if (lastStep.type === 'AGENT_THOUGHT') {
|
||||
const yieldNode: AgentYield = {
|
||||
id: lastStep.id,
|
||||
type: 'AGENT_YIELD',
|
||||
text: lastStep.text,
|
||||
metadata: lastStep.metadata,
|
||||
};
|
||||
currentEpisode.steps.pop();
|
||||
currentEpisode.yield = yieldNode;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Part } from '@google/genai';
|
||||
|
||||
/**
|
||||
* Universal Audit Metadata
|
||||
* Tracks the lifecycle and transformations of a node or part within the IR.
|
||||
* This guarantees perfect reversibility and enables long-term memory offloading.
|
||||
*/
|
||||
export interface IrMetadata {
|
||||
/** The estimated number of tokens this entity originally consumed. */
|
||||
originalTokens: number;
|
||||
/** The current estimated number of tokens this entity consumes in its degraded state. */
|
||||
currentTokens: number;
|
||||
/** An audit trail of all transformations applied by ContextProcessors. */
|
||||
transformations: Array<{
|
||||
processorName: string;
|
||||
action:
|
||||
| 'MASKED'
|
||||
| 'TRUNCATED'
|
||||
| 'DEGRADED'
|
||||
| 'SUMMARIZED'
|
||||
| 'EVICTED'
|
||||
| 'SYNTHESIZED';
|
||||
timestamp: number;
|
||||
/** Pointer to where the original uncompressed payload was saved (if applicable) */
|
||||
diskPointer?: string;
|
||||
}>;
|
||||
}
|
||||
|
||||
export type IrNodeType =
|
||||
| 'USER_PROMPT'
|
||||
| 'SYSTEM_EVENT'
|
||||
| 'AGENT_THOUGHT'
|
||||
| 'TOOL_EXECUTION'
|
||||
| 'AGENT_YIELD';
|
||||
|
||||
/** Base interface for all nodes in the Episodic IR */
|
||||
export type VariantStatus = 'computing' | 'ready' | 'failed';
|
||||
|
||||
export interface BaseVariant {
|
||||
status: VariantStatus;
|
||||
recoveredTokens?: number;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
export interface SummaryVariant extends BaseVariant {
|
||||
type: 'summary';
|
||||
text: string;
|
||||
}
|
||||
|
||||
export interface MaskedVariant extends BaseVariant {
|
||||
type: 'masked';
|
||||
text: string;
|
||||
}
|
||||
|
||||
export interface SnapshotVariant extends BaseVariant {
|
||||
type: 'snapshot';
|
||||
episode: Episode;
|
||||
replacedEpisodeIds: string[];
|
||||
}
|
||||
|
||||
export type Variant = SummaryVariant | MaskedVariant | SnapshotVariant;
|
||||
|
||||
/** Base interface for all nodes in the Episodic IR */
|
||||
export interface IrNode {
|
||||
readonly id: string;
|
||||
readonly type: IrNodeType;
|
||||
metadata: IrMetadata;
|
||||
variants?: Record<string, Variant>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Semantic Parts for User Prompts
|
||||
* Ensures we can safely truncate text without deleting multi-modal parts (like images).
|
||||
*/
|
||||
export type SemanticPart =
|
||||
| {
|
||||
type: 'text';
|
||||
text: string;
|
||||
presentation?: { text: string; tokens: number };
|
||||
}
|
||||
| {
|
||||
type: 'inline_data';
|
||||
mimeType: string;
|
||||
data: string;
|
||||
presentation?: { text: string; tokens: number };
|
||||
}
|
||||
| {
|
||||
type: 'file_data';
|
||||
mimeType: string;
|
||||
fileUri: string;
|
||||
presentation?: { text: string; tokens: number };
|
||||
}
|
||||
| {
|
||||
type: 'raw_part';
|
||||
part: Part;
|
||||
presentation?: { text: string; tokens: number };
|
||||
};
|
||||
|
||||
/**
|
||||
* Trigger Nodes
|
||||
* Events that wake the agent up and initiate an Episode.
|
||||
*/
|
||||
export interface UserPrompt extends IrNode {
|
||||
readonly type: 'USER_PROMPT';
|
||||
/** The semantic breakdown of the user's multi-modal input */
|
||||
semanticParts: SemanticPart[];
|
||||
}
|
||||
|
||||
export interface SystemEvent extends IrNode {
|
||||
readonly type: 'SYSTEM_EVENT';
|
||||
name: string;
|
||||
payload: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export type EpisodeTrigger = UserPrompt | SystemEvent;
|
||||
|
||||
/**
|
||||
* Step Nodes
|
||||
* The internal autonomous actions taken by the agent during its loop.
|
||||
*/
|
||||
export interface AgentThought extends IrNode {
|
||||
readonly type: 'AGENT_THOUGHT';
|
||||
text: string;
|
||||
/** Overrides the rendered output for this thought */
|
||||
presentation?: {
|
||||
text: string;
|
||||
tokens: number;
|
||||
};
|
||||
}
|
||||
|
||||
export interface ToolExecution extends IrNode {
|
||||
readonly type: 'TOOL_EXECUTION';
|
||||
/** The name of the tool invoked */
|
||||
toolName: string;
|
||||
|
||||
/** The arguments passed to the tool (The 'FunctionCall') */
|
||||
intent: Record<string, unknown>;
|
||||
|
||||
/** The result returned by the tool (The 'FunctionResponse') */
|
||||
observation: string | Record<string, unknown>;
|
||||
|
||||
/** Granular token tracking for the different lifecycle phases of the tool */
|
||||
tokens: {
|
||||
intent: number;
|
||||
observation: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* The presentation layer. If defined, the IrMapper uses this instead of the
|
||||
* raw observation to build the functionResponse.
|
||||
* This preserves the immutable raw data for semantic queries while modifying the rendered output.
|
||||
*/
|
||||
presentation?: {
|
||||
intent?: Record<string, unknown>;
|
||||
observation?: string | Record<string, unknown>;
|
||||
tokens: {
|
||||
intent: number;
|
||||
observation: number;
|
||||
};
|
||||
};
|
||||
}
|
||||
|
||||
export type EpisodeStep = AgentThought | ToolExecution;
|
||||
|
||||
/**
|
||||
* Resolution Node
|
||||
* The final message where the agent yields control back to the user.
|
||||
*/
|
||||
export interface AgentYield extends IrNode {
|
||||
readonly type: 'AGENT_YIELD';
|
||||
text: string;
|
||||
presentation?: {
|
||||
text: string;
|
||||
tokens: number;
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* The Episode
|
||||
* A discrete, continuous run of the agent. Represents the full cycle from
|
||||
* taking control (Trigger) to returning control (Yield), encompassing all
|
||||
* internal reasoning and observations (Steps).
|
||||
*/
|
||||
export interface Episode {
|
||||
readonly id: string;
|
||||
/** When the episode began */
|
||||
readonly timestamp: number;
|
||||
variants?: Record<string, Variant>;
|
||||
|
||||
/** The event that initiated this run */
|
||||
trigger: EpisodeTrigger;
|
||||
|
||||
/** The sequence of autonomous actions and observations */
|
||||
steps: EpisodeStep[];
|
||||
|
||||
/** The final handover back to the user (can be undefined if the episode was aborted/errored) */
|
||||
yield?: AgentYield;
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { EpisodeEditor } from './ir/episodeEditor.js';
|
||||
|
||||
/**
|
||||
* State object passed through the processing pipeline.
|
||||
* Contains global accounting logic and semantic protection rules.
|
||||
*/
|
||||
export interface ContextAccountingState {
|
||||
readonly currentTokens: number;
|
||||
readonly maxTokens: number;
|
||||
readonly retainedTokens: number;
|
||||
|
||||
/** The exact number of tokens that need to be trimmed to reach the retainedTokens goal */
|
||||
readonly deficitTokens: number;
|
||||
|
||||
/**
|
||||
* Set of Episode IDs that the orchestrator has deemed highly protected.
|
||||
* Processors should generally skip mutating these episodes unless doing proactive/required transforms.
|
||||
*/
|
||||
readonly protectedEpisodeIds: Set<string>;
|
||||
|
||||
/**
|
||||
* True if currentTokens <= retainedTokens.
|
||||
*/
|
||||
readonly isBudgetSatisfied: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Interface for all context degradation strategies.
|
||||
*/
|
||||
export interface ContextProcessor {
|
||||
/** Unique name for telemetry and logging. */
|
||||
readonly name: string;
|
||||
|
||||
/**
|
||||
* Processes the episodic history payload via the provided EpisodeEditor, based on the current accounting state.
|
||||
* Processors should safely mutate or replace episodes using the editor's API.
|
||||
*/
|
||||
process(editor: EpisodeEditor, state: ContextAccountingState): Promise<void>;
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
createMockEnvironment,
|
||||
createDummyState,
|
||||
createDummyEpisode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { BlobDegradationProcessor } from './blobDegradationProcessor.js';
|
||||
import { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
import type { UserPrompt } from '../ir/types.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import type { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
|
||||
|
||||
describe('BlobDegradationProcessor', () => {
|
||||
let processor: BlobDegradationProcessor;
|
||||
let env: ContextEnvironment;
|
||||
let fileSystem: InMemoryFileSystem;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
env = createMockEnvironment();
|
||||
fileSystem = env.fileSystem as InMemoryFileSystem;
|
||||
processor = new BlobDegradationProcessor(env);
|
||||
});
|
||||
|
||||
it('degrades inline_data into a text reference and saves to disk', async () => {
|
||||
const dummyImageBase64 = Buffer.from('fake-image-data').toString('base64');
|
||||
|
||||
const ep = createDummyEpisode('ep-1', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'Look at this image:' },
|
||||
{
|
||||
type: 'inline_data',
|
||||
mimeType: 'image/png',
|
||||
data: dummyImageBase64,
|
||||
},
|
||||
]);
|
||||
|
||||
const state = createDummyState(false, 500);
|
||||
const editor = new EpisodeEditor([ep]);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
const parts = (result[0].trigger as UserPrompt).semanticParts;
|
||||
|
||||
// Text part should be untouched
|
||||
expect(parts[0].presentation).toBeUndefined();
|
||||
|
||||
// Inline data should be degraded
|
||||
expect(parts[1].presentation).toBeDefined();
|
||||
expect(parts[1].presentation!.text).toContain(
|
||||
'[Multi-Modal Blob (image/png',
|
||||
);
|
||||
expect(parts[1].presentation!.text).toContain(
|
||||
'degraded to text to preserve context window',
|
||||
);
|
||||
|
||||
// Verify it was written to fake FS
|
||||
expect(fileSystem.getFiles().size).toBeGreaterThan(0);
|
||||
const files = Array.from(fileSystem.getFiles().keys());
|
||||
expect(files[0]).toContain(
|
||||
'.gemini/tool-outputs/degraded-blobs/session-mock-session/blob_',
|
||||
);
|
||||
|
||||
expect(result[0].trigger.metadata.transformations.length).toBe(1);
|
||||
});
|
||||
|
||||
it('degrades file_data into a text reference without disk write', async () => {
|
||||
const ep = createDummyEpisode('ep-2', 'USER_PROMPT', [
|
||||
{
|
||||
type: 'file_data',
|
||||
mimeType: 'application/pdf',
|
||||
fileUri: 'gs://fake-bucket/doc.pdf',
|
||||
},
|
||||
]);
|
||||
|
||||
const state = createDummyState(false, 500);
|
||||
const editor = new EpisodeEditor([ep]);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
const parts = (result[0].trigger as UserPrompt).semanticParts;
|
||||
expect(parts[0].presentation).toBeDefined();
|
||||
expect(parts[0].presentation!.text).toContain(
|
||||
'[File Reference (application/pdf)',
|
||||
);
|
||||
expect(parts[0].presentation!.text).toContain(
|
||||
'Original URI: gs://fake-bucket/doc.pdf',
|
||||
);
|
||||
|
||||
expect(fileSystem.getFiles().size).toBe(0);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,142 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { sanitizeFilenamePart } from '../../utils/fileUtils.js';
|
||||
import type { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
|
||||
export type BlobDegradationProcessorOptions = Record<string, never>;
|
||||
|
||||
export class BlobDegradationProcessor implements ContextProcessor {
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
_options: BlobDegradationProcessorOptions,
|
||||
): BlobDegradationProcessor {
|
||||
return new BlobDegradationProcessor(env);
|
||||
}
|
||||
|
||||
readonly id = 'BlobDegradationProcessor';
|
||||
readonly name = 'BlobDegradationProcessor';
|
||||
readonly options = {};
|
||||
private env: ContextEnvironment;
|
||||
|
||||
constructor(env: ContextEnvironment) {
|
||||
this.env = env;
|
||||
}
|
||||
|
||||
async process(
|
||||
editor: EpisodeEditor,
|
||||
state: ContextAccountingState,
|
||||
): Promise<void> {
|
||||
if (state.isBudgetSatisfied) {
|
||||
return;
|
||||
}
|
||||
|
||||
let currentDeficit = state.deficitTokens;
|
||||
let directoryCreated = false;
|
||||
|
||||
let blobOutputsDir = this.env.fileSystem.join(
|
||||
this.env.projectTempDir,
|
||||
'degraded-blobs',
|
||||
);
|
||||
const sessionId = this.env.sessionId;
|
||||
if (sessionId) {
|
||||
blobOutputsDir = this.env.fileSystem.join(
|
||||
blobOutputsDir,
|
||||
`session-${sanitizeFilenamePart(sessionId)}`,
|
||||
);
|
||||
}
|
||||
|
||||
const ensureDir = async () => {
|
||||
if (!directoryCreated) {
|
||||
await this.env.fileSystem.mkdir(blobOutputsDir, { recursive: true });
|
||||
directoryCreated = true;
|
||||
}
|
||||
};
|
||||
|
||||
// Forward scan, looking for bloated non-text parts to degrade
|
||||
for (const ep of editor.episodes) {
|
||||
if (currentDeficit <= 0) break;
|
||||
if (state.protectedEpisodeIds.has(ep.id)) continue;
|
||||
|
||||
if (ep.trigger.type === 'USER_PROMPT') {
|
||||
for (let j = 0; j < ep.trigger.semanticParts.length; j++) {
|
||||
const part = ep.trigger.semanticParts[j];
|
||||
if (currentDeficit <= 0) break;
|
||||
// We only target non-text parts that haven't already been masked
|
||||
if (part.type === 'text' || part.presentation) continue;
|
||||
|
||||
let newText = '';
|
||||
let tokensSaved = 0;
|
||||
|
||||
if (part.type === 'inline_data') {
|
||||
await ensureDir();
|
||||
const ext = part.mimeType.split('/')[1] || 'bin';
|
||||
const fileName = `blob_${Date.now()}_${this.env.idGenerator.generateId()}.${ext}`;
|
||||
const filePath = this.env.fileSystem.join(blobOutputsDir, fileName);
|
||||
|
||||
// Base64 to buffer
|
||||
const buffer = Buffer.from(part.data, 'base64');
|
||||
await this.env.fileSystem.writeFile(filePath, buffer);
|
||||
|
||||
const mb = (buffer.byteLength / 1024 / 1024).toFixed(2);
|
||||
newText = `[Multi-Modal Blob (${part.mimeType}, ${mb}MB) degraded to text to preserve context window. Saved to: ${filePath}]`;
|
||||
|
||||
// Re-calculate tokens. Images are expensive (~258 tokens). The text is cheap (~20 tokens).
|
||||
const oldTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
{ inlineData: { mimeType: part.mimeType, data: part.data } },
|
||||
]);
|
||||
const newTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: newText },
|
||||
]);
|
||||
tokensSaved = oldTokens - newTokens;
|
||||
} else if (part.type === 'file_data') {
|
||||
newText = `[File Reference (${part.mimeType}) degraded to text to preserve context window. Original URI: ${part.fileUri}]`;
|
||||
const oldTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
{ fileData: { mimeType: part.mimeType, fileUri: part.fileUri } },
|
||||
]);
|
||||
const newTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: newText },
|
||||
]);
|
||||
tokensSaved = oldTokens - newTokens;
|
||||
} else if (part.type === 'raw_part') {
|
||||
newText = `[Unknown Part degraded to text to preserve context window.]`;
|
||||
const oldTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
part.part,
|
||||
]);
|
||||
const newTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: newText },
|
||||
]);
|
||||
tokensSaved = oldTokens - newTokens;
|
||||
}
|
||||
|
||||
if (newText && tokensSaved > 0) {
|
||||
const newTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: newText },
|
||||
]);
|
||||
|
||||
editor.editEpisode(ep.id, 'DEGRADE_BLOB', (draft) => {
|
||||
if (draft.trigger.type === 'USER_PROMPT') {
|
||||
draft.trigger.semanticParts[j].presentation = {
|
||||
text: newText,
|
||||
tokens: newTokens,
|
||||
};
|
||||
draft.trigger.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'DEGRADED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
currentDeficit -= tokensSaved;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
createMockEnvironment,
|
||||
createDummyState,
|
||||
createDummyEpisode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { EmergencyTruncationProcessor } from './emergencyTruncationProcessor.js';
|
||||
import { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
|
||||
describe('EmergencyTruncationProcessor', () => {
|
||||
let processor: EmergencyTruncationProcessor;
|
||||
let env: ContextEnvironment;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
env = createMockEnvironment();
|
||||
// Force token calculator to return exactly what we tell it for deterministic testing
|
||||
vi.spyOn(
|
||||
env.tokenCalculator,
|
||||
'calculateEpisodeListTokens',
|
||||
).mockImplementation((episodes) =>
|
||||
// Just sum up the metadata originalTokens for our dummy episodes
|
||||
episodes.reduce(
|
||||
(acc, ep) => acc + (ep.trigger.metadata.originalTokens || 100),
|
||||
0,
|
||||
),
|
||||
);
|
||||
|
||||
processor = new EmergencyTruncationProcessor(env, {});
|
||||
});
|
||||
|
||||
it('bypasses processing if currentTokens <= maxTokens', async () => {
|
||||
const episodes = [
|
||||
createDummyEpisode('ep-1', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'short' },
|
||||
]),
|
||||
];
|
||||
// State says we are under budget (5000 < 10000)
|
||||
const state = createDummyState(true, 0, new Set(), 5000, 10000);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
expect(result).toStrictEqual(episodes);
|
||||
expect(result.length).toBe(1);
|
||||
});
|
||||
|
||||
it('truncates episodes from the front (oldest) until targetTokens is met', async () => {
|
||||
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'oldest' },
|
||||
]);
|
||||
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'middle' },
|
||||
]);
|
||||
const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'newest' },
|
||||
]);
|
||||
|
||||
// Each is worth 100 tokens according to our mock
|
||||
const episodes = [ep1, ep2, ep3];
|
||||
|
||||
// We have 300 tokens, but max is 200. We need to drop 100 tokens.
|
||||
const state = createDummyState(false, 100, new Set(), 300, 200);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
// It should drop the FIRST episode (ep-1) and keep the rest.
|
||||
expect(result.length).toBe(2);
|
||||
expect(result[0].id).toBe('ep-2');
|
||||
expect(result[1].id).toBe('ep-3');
|
||||
});
|
||||
|
||||
it('never drops protected episodes (e.g. system instructions)', async () => {
|
||||
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'protected system prompt' },
|
||||
]);
|
||||
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'middle' },
|
||||
]);
|
||||
const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'newest' },
|
||||
]);
|
||||
|
||||
const episodes = [ep1, ep2, ep3];
|
||||
|
||||
// We have 300 tokens, max is 200. We need to drop 100 tokens.
|
||||
// However, ep-1 is protected!
|
||||
const state = createDummyState(false, 100, new Set(['ep-1']), 300, 200);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
// It should SKIP dropping ep-1 (protected) and drop ep-2 instead.
|
||||
expect(result.length).toBe(2);
|
||||
expect(result[0].id).toBe('ep-1'); // Protected, survived
|
||||
expect(result[1].id).toBe('ep-3'); // Survivor
|
||||
});
|
||||
|
||||
it('can drop multiple episodes if deficit is huge', async () => {
|
||||
const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', []);
|
||||
const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', []);
|
||||
const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', []);
|
||||
|
||||
const episodes = [ep1, ep2, ep3];
|
||||
|
||||
// We have 300 tokens, max is 50. We need to drop 250 tokens!
|
||||
const state = createDummyState(false, 250, new Set(), 300, 50);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
// It must drop ep1 (100t) and ep2 (100t).
|
||||
// Remaining is ep3 (100t).
|
||||
// Wait, if it drops ep1 (remaining=200) and ep2 (remaining=100),
|
||||
// when it looks at ep3, remaining (100) > max (50), so it will drop ep3 too!
|
||||
expect(result.length).toBe(0);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,60 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextProcessor, ContextAccountingState } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import type { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
|
||||
export type EmergencyTruncationProcessorOptions = Record<string, never>;
|
||||
|
||||
export class EmergencyTruncationProcessor implements ContextProcessor {
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: EmergencyTruncationProcessorOptions,
|
||||
): EmergencyTruncationProcessor {
|
||||
return new EmergencyTruncationProcessor(env, options);
|
||||
}
|
||||
|
||||
readonly id = 'EmergencyTruncationProcessor';
|
||||
readonly name = 'EmergencyTruncationProcessor';
|
||||
readonly options: EmergencyTruncationProcessorOptions;
|
||||
constructor(
|
||||
private readonly _env: ContextEnvironment,
|
||||
options: EmergencyTruncationProcessorOptions,
|
||||
) {
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
async process(
|
||||
editor: EpisodeEditor,
|
||||
state: ContextAccountingState,
|
||||
): Promise<void> {
|
||||
if (state.currentTokens <= state.maxTokens) return;
|
||||
|
||||
let remainingTokens = state.currentTokens;
|
||||
const targetTokens = state.maxTokens;
|
||||
const toRemove: string[] = [];
|
||||
|
||||
// We respect the global protected Episode IDs (like the system prompt at index 0)
|
||||
for (const ep of editor.episodes) {
|
||||
const epTokens = this._env.tokenCalculator.calculateEpisodeListTokens([
|
||||
ep,
|
||||
]);
|
||||
|
||||
if (
|
||||
remainingTokens > targetTokens &&
|
||||
!state.protectedEpisodeIds.has(ep.id)
|
||||
) {
|
||||
remainingTokens -= epTokens;
|
||||
toRemove.push(ep.id);
|
||||
}
|
||||
}
|
||||
|
||||
if (toRemove.length > 0) {
|
||||
editor.removeEpisodes(toRemove, 'TRUNCATED');
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
createMockEnvironment,
|
||||
createDummyState,
|
||||
createDummyEpisode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import { describe, it, expect, beforeEach } from 'vitest';
|
||||
import { HistorySquashingProcessor } from './historySquashingProcessor.js';
|
||||
import { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
import type { UserPrompt, AgentThought, AgentYield } from '../ir/types.js';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
|
||||
describe('HistorySquashingProcessor', () => {
|
||||
let processor: HistorySquashingProcessor;
|
||||
|
||||
beforeEach(() => {
|
||||
processor = new HistorySquashingProcessor(createMockEnvironment(), {
|
||||
maxTokensPerNode: 100,
|
||||
});
|
||||
});
|
||||
|
||||
const createThoughtEpisode = (
|
||||
id: string,
|
||||
userText: string,
|
||||
modelThought: string,
|
||||
) => {
|
||||
const ep = createDummyEpisode(id, 'USER_PROMPT', [
|
||||
{ type: 'text', text: userText },
|
||||
]);
|
||||
// Replace the tool steps with a thought step for this test
|
||||
ep.steps = [
|
||||
{
|
||||
id: randomUUID(),
|
||||
type: 'AGENT_THOUGHT',
|
||||
text: modelThought,
|
||||
metadata: {
|
||||
originalTokens: 1000,
|
||||
currentTokens: 1000,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
];
|
||||
return ep;
|
||||
};
|
||||
|
||||
it('bypasses processing if budget is satisfied', async () => {
|
||||
const episodes = [createThoughtEpisode('1', 'short text', 'short thought')];
|
||||
const state = createDummyState(true);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
expect(result).toStrictEqual(episodes);
|
||||
expect(
|
||||
(result[0].trigger as UserPrompt).semanticParts[0].presentation,
|
||||
).toBeUndefined();
|
||||
});
|
||||
|
||||
it('skips protected episodes', async () => {
|
||||
// 500 chars = ~125 tokens. Limit is 100 tokens, so it WOULD truncate if not protected.
|
||||
const longText = 'A'.repeat(500);
|
||||
const episodes = [createThoughtEpisode('ep-1', longText, 'short thought')];
|
||||
const state = createDummyState(false, 100, new Set(['ep-1']));
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
expect(
|
||||
(result[0].trigger as UserPrompt).semanticParts[0].presentation,
|
||||
).toBeUndefined();
|
||||
});
|
||||
|
||||
it('truncates both UserPrompts and AgentThoughts', async () => {
|
||||
const longUser = 'U'.repeat(1000); // ~250 tokens
|
||||
const longModel = 'M'.repeat(1000); // ~250 tokens
|
||||
const episodes = [createThoughtEpisode('ep-2', longUser, longModel)];
|
||||
const state = createDummyState(false, 500); // High deficit, force truncation
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
const userPart = (result[0].trigger as UserPrompt).semanticParts[0];
|
||||
const thoughtPart = result[0].steps[0] as AgentThought;
|
||||
|
||||
expect(userPart.presentation).toBeDefined();
|
||||
expect(userPart.presentation!.text).toContain(
|
||||
'[... OMITTED 600 chars ...]',
|
||||
);
|
||||
|
||||
expect(thoughtPart.presentation).toBeDefined();
|
||||
expect(thoughtPart.presentation!.text).toContain(
|
||||
'[... OMITTED 600 chars ...]',
|
||||
);
|
||||
|
||||
// Check audit trails
|
||||
expect(result[0].trigger.metadata.transformations.length).toBe(1);
|
||||
expect(thoughtPart.metadata.transformations.length).toBe(1);
|
||||
});
|
||||
|
||||
it('stops processing once deficit is resolved', async () => {
|
||||
const longUser1 = 'A'.repeat(1000);
|
||||
const longUser2 = 'B'.repeat(1000);
|
||||
const episodes = [
|
||||
createThoughtEpisode('ep-3', longUser1, 'short'),
|
||||
createThoughtEpisode('ep-4', longUser2, 'short'),
|
||||
];
|
||||
|
||||
// Set deficit to exactly what ONE truncation will save
|
||||
// Original = ~250 tokens. Limit = 100. Truncation saves ~150 tokens.
|
||||
const state = createDummyState(false, 150);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
// First episode should be truncated
|
||||
const ep1Part = (result[0].trigger as UserPrompt).semanticParts[0];
|
||||
expect(ep1Part.presentation).toBeDefined();
|
||||
|
||||
// Second episode should be untouched because the deficit hit 0
|
||||
const ep2Part = (result[1].trigger as UserPrompt).semanticParts[0];
|
||||
expect(ep2Part.presentation).toBeUndefined();
|
||||
});
|
||||
|
||||
it('truncates IrNodes', async () => {
|
||||
const longYield = 'Y'.repeat(1000); // ~250 tokens
|
||||
const ep = createThoughtEpisode('ep-5', 'short', 'short');
|
||||
ep.yield = {
|
||||
id: randomUUID(),
|
||||
type: 'AGENT_YIELD',
|
||||
text: longYield,
|
||||
metadata: {
|
||||
originalTokens: 250,
|
||||
currentTokens: 250,
|
||||
transformations: [],
|
||||
},
|
||||
};
|
||||
|
||||
const state = createDummyState(false, 500);
|
||||
const editor = new EpisodeEditor([ep]);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
const yieldPart = result[0].yield as AgentYield;
|
||||
const yieldPresentation = yieldPart.presentation as { text: string };
|
||||
expect(yieldPresentation).toBeDefined();
|
||||
expect(yieldPresentation.text).toContain('[... OMITTED 600 chars ...]');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,189 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { truncateProportionally } from '../truncation.js';
|
||||
import type { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
|
||||
export interface HistorySquashingProcessorOptions {
|
||||
maxTokensPerNode: number;
|
||||
}
|
||||
|
||||
export class HistorySquashingProcessor implements ContextProcessor {
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: HistorySquashingProcessorOptions,
|
||||
): HistorySquashingProcessor {
|
||||
return new HistorySquashingProcessor(env, options);
|
||||
}
|
||||
|
||||
static readonly schema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
maxTokensPerNode: {
|
||||
type: 'number',
|
||||
description:
|
||||
'The maximum tokens a node can have before being truncated.',
|
||||
},
|
||||
},
|
||||
required: ['maxTokensPerNode'],
|
||||
};
|
||||
|
||||
readonly id = 'HistorySquashingProcessor';
|
||||
readonly name = 'HistorySquashingProcessor';
|
||||
readonly options: HistorySquashingProcessorOptions;
|
||||
|
||||
constructor(
|
||||
env: ContextEnvironment,
|
||||
options: HistorySquashingProcessorOptions,
|
||||
) {
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
private tryApplySquash(
|
||||
text: string,
|
||||
limitChars: number,
|
||||
currentDeficit: number,
|
||||
setPresentation: (p: { text: string; tokens: number }) => void,
|
||||
recordAudit: () => void,
|
||||
): number {
|
||||
if (currentDeficit <= 0) return 0;
|
||||
const originalLength = text.length;
|
||||
if (originalLength <= limitChars) return 0;
|
||||
|
||||
const newText = truncateProportionally(
|
||||
text,
|
||||
limitChars,
|
||||
`\n\n[... OMITTED ${originalLength - limitChars} chars ...]\n\n`,
|
||||
);
|
||||
|
||||
if (newText !== text) {
|
||||
const newTokens = Math.floor(newText.length / 4);
|
||||
const oldTokens = Math.floor(originalLength / 4);
|
||||
const tokensSaved = oldTokens - newTokens;
|
||||
|
||||
setPresentation({ text: newText, tokens: newTokens });
|
||||
recordAudit();
|
||||
return tokensSaved;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
async process(
|
||||
editor: EpisodeEditor,
|
||||
state: ContextAccountingState,
|
||||
): Promise<void> {
|
||||
if (state.isBudgetSatisfied) {
|
||||
return;
|
||||
}
|
||||
|
||||
const { maxTokensPerNode } = this.options;
|
||||
// We estimate 4 chars per token for truncation logic
|
||||
const limitChars = maxTokensPerNode * 4;
|
||||
|
||||
// We track how many tokens we still need to cut. If we hit 0, we can stop early!
|
||||
let currentDeficit = state.deficitTokens;
|
||||
|
||||
for (const ep of editor.episodes) {
|
||||
if (currentDeficit <= 0) break;
|
||||
if (state.protectedEpisodeIds.has(ep.id)) continue;
|
||||
|
||||
// 1. Squash User Prompts
|
||||
if (ep.trigger.type === 'USER_PROMPT') {
|
||||
for (let j = 0; j < ep.trigger.semanticParts.length; j++) {
|
||||
const part = ep.trigger.semanticParts[j];
|
||||
if (part.type === 'text') {
|
||||
const saved = this.tryApplySquash(
|
||||
part.text,
|
||||
limitChars,
|
||||
currentDeficit,
|
||||
(p) => {
|
||||
editor.editEpisode(ep.id, 'SQUASH_PROMPT', (draft) => {
|
||||
if (draft.trigger.type === 'USER_PROMPT') {
|
||||
draft.trigger.semanticParts[j].presentation = p;
|
||||
}
|
||||
});
|
||||
},
|
||||
() => {
|
||||
editor.editEpisode(ep.id, 'SQUASH_PROMPT', (draft) => {
|
||||
draft.trigger.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'TRUNCATED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
});
|
||||
},
|
||||
);
|
||||
currentDeficit -= saved;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Squash Model Thoughts
|
||||
if (ep.steps) {
|
||||
for (let j = 0; j < ep.steps.length; j++) {
|
||||
const step = ep.steps[j];
|
||||
if (currentDeficit <= 0) break;
|
||||
if (step.type === 'AGENT_THOUGHT') {
|
||||
const saved = this.tryApplySquash(
|
||||
step.text,
|
||||
limitChars,
|
||||
currentDeficit,
|
||||
(p) => {
|
||||
editor.editEpisode(ep.id, 'SQUASH_THOUGHT', (draft) => {
|
||||
const draftStep = draft.steps[j];
|
||||
if (draftStep.type === 'AGENT_THOUGHT') {
|
||||
draftStep.presentation = p;
|
||||
}
|
||||
});
|
||||
},
|
||||
() => {
|
||||
editor.editEpisode(ep.id, 'SQUASH_THOUGHT', (draft) => {
|
||||
const draftStep = draft.steps[j];
|
||||
if (draftStep.type === 'AGENT_THOUGHT') {
|
||||
draftStep.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'TRUNCATED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
}
|
||||
});
|
||||
},
|
||||
);
|
||||
currentDeficit -= saved;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Squash Agent Yields
|
||||
if (currentDeficit > 0 && ep.yield) {
|
||||
const saved = this.tryApplySquash(
|
||||
ep.yield.text,
|
||||
limitChars,
|
||||
currentDeficit,
|
||||
(p) => {
|
||||
editor.editEpisode(ep.id, 'SQUASH_YIELD', (draft) => {
|
||||
if (draft.yield) draft.yield.presentation = p;
|
||||
});
|
||||
},
|
||||
() => {
|
||||
editor.editEpisode(ep.id, 'SQUASH_YIELD', (draft) => {
|
||||
if (draft.yield) {
|
||||
draft.yield.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'TRUNCATED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
}
|
||||
});
|
||||
},
|
||||
);
|
||||
currentDeficit -= saved;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
createMockEnvironment,
|
||||
createDummyState,
|
||||
createDummyEpisode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { SemanticCompressionProcessor } from './semanticCompressionProcessor.js';
|
||||
import { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
import type { UserPrompt, ToolExecution, AgentThought } from '../ir/types.js';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import type { BaseLlmClient } from 'src/core/baseLlmClient.js';
|
||||
|
||||
describe('SemanticCompressionProcessor', () => {
|
||||
let processor: SemanticCompressionProcessor;
|
||||
let generateContentMock: ReturnType<typeof vi.fn>;
|
||||
|
||||
beforeEach(() => {
|
||||
generateContentMock = vi.fn().mockResolvedValue({
|
||||
candidates: [{ content: { parts: [{ text: 'Mocked Summary!' }] } }],
|
||||
});
|
||||
|
||||
const env = createMockEnvironment();
|
||||
// Re-mock llmClient properly
|
||||
vi.spyOn(env, 'llmClient', 'get').mockReturnValue({
|
||||
generateContent: generateContentMock,
|
||||
} as unknown as BaseLlmClient);
|
||||
|
||||
processor = new SemanticCompressionProcessor(env, {
|
||||
nodeThresholdTokens: 2000,
|
||||
});
|
||||
});
|
||||
|
||||
const createEpisodeWithThoughtsAndTools = (
|
||||
id: string,
|
||||
userText: string,
|
||||
thoughtText: string,
|
||||
toolObs: string,
|
||||
) => {
|
||||
const ep = createDummyEpisode(id, 'USER_PROMPT', [
|
||||
{ type: 'text', text: userText },
|
||||
]);
|
||||
// We override metadata for threshold triggering
|
||||
ep.trigger.metadata.currentTokens = 3800;
|
||||
|
||||
ep.steps = [
|
||||
{
|
||||
id: randomUUID(),
|
||||
type: 'AGENT_THOUGHT',
|
||||
text: thoughtText,
|
||||
metadata: {
|
||||
originalTokens: 3800,
|
||||
currentTokens: 3800,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
{
|
||||
id: randomUUID(),
|
||||
type: 'TOOL_EXECUTION',
|
||||
toolName: 'test',
|
||||
intent: {},
|
||||
observation: toolObs,
|
||||
tokens: { intent: 10, observation: 3800 },
|
||||
metadata: {
|
||||
originalTokens: 3810,
|
||||
currentTokens: 3810,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
];
|
||||
return ep;
|
||||
};
|
||||
|
||||
it('bypasses processing if budget is satisfied', async () => {
|
||||
const episodes = [
|
||||
createEpisodeWithThoughtsAndTools('1', 'short', 'short', 'short'),
|
||||
];
|
||||
const state = createDummyState(true);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
|
||||
expect(generateContentMock).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('skips protected episodes even if over budget', async () => {
|
||||
const massiveStr = 'M'.repeat(15000);
|
||||
const episodes = [
|
||||
createEpisodeWithThoughtsAndTools(
|
||||
'ep-1',
|
||||
massiveStr,
|
||||
massiveStr,
|
||||
massiveStr,
|
||||
),
|
||||
];
|
||||
const state = createDummyState(false, 1000, new Set(['ep-1']));
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
|
||||
expect(generateContentMock).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('summarizes unprotected UserPrompts, Thoughts, and Tool observations until deficit is met', async () => {
|
||||
const massiveStr = 'M'.repeat(15000);
|
||||
const episodes = [
|
||||
createEpisodeWithThoughtsAndTools(
|
||||
'ep-1',
|
||||
massiveStr,
|
||||
massiveStr,
|
||||
massiveStr,
|
||||
),
|
||||
];
|
||||
const state = createDummyState(false, 50000); // Massive deficit, forces all 3 to summarize
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
|
||||
expect(generateContentMock).toHaveBeenCalledTimes(3);
|
||||
|
||||
// Verify presentation layers were injected
|
||||
const result = editor.getFinalEpisodes();
|
||||
const userPart = (result[0].trigger as UserPrompt).semanticParts[0];
|
||||
const thoughtPart = result[0].steps[0] as AgentThought;
|
||||
const toolPart = result[0].steps[1] as ToolExecution;
|
||||
|
||||
expect(userPart.presentation).toBeDefined();
|
||||
expect(userPart.presentation!.text).toContain('Mocked Summary!');
|
||||
|
||||
expect(thoughtPart.presentation).toBeDefined();
|
||||
expect(thoughtPart.presentation!.text).toContain('Mocked Summary!');
|
||||
|
||||
expect(toolPart.presentation).toBeDefined();
|
||||
expect(
|
||||
(toolPart.presentation!.observation as Record<string, string>)['summary'],
|
||||
).toContain('Mocked Summary!');
|
||||
});
|
||||
|
||||
it('stops calling LLM when deficit hits zero', async () => {
|
||||
const massiveStr = 'M'.repeat(15000);
|
||||
const episodes = [
|
||||
createEpisodeWithThoughtsAndTools(
|
||||
'ep-1',
|
||||
massiveStr,
|
||||
massiveStr,
|
||||
massiveStr,
|
||||
),
|
||||
];
|
||||
|
||||
// Set deficit low enough that ONE summary solves the problem
|
||||
const state = createDummyState(false, 5);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
|
||||
// It should only compress the UserPrompt and then stop
|
||||
expect(generateContentMock).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,261 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { LlmRole } from '../../telemetry/types.js';
|
||||
import { getResponseText } from '../../utils/partUtils.js';
|
||||
import type { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
|
||||
export interface SemanticCompressionProcessorOptions {
|
||||
nodeThresholdTokens: number;
|
||||
}
|
||||
|
||||
export class SemanticCompressionProcessor implements ContextProcessor {
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: SemanticCompressionProcessorOptions,
|
||||
): SemanticCompressionProcessor {
|
||||
return new SemanticCompressionProcessor(env, options);
|
||||
}
|
||||
|
||||
static readonly schema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
nodeThresholdTokens: {
|
||||
type: 'number',
|
||||
description: 'The token threshold above which nodes are summarized.',
|
||||
},
|
||||
},
|
||||
required: ['nodeThresholdTokens'],
|
||||
};
|
||||
|
||||
readonly id = 'SemanticCompressionProcessor';
|
||||
readonly name = 'SemanticCompressionProcessor';
|
||||
readonly options: SemanticCompressionProcessorOptions;
|
||||
private env: ContextEnvironment;
|
||||
private modelToUse: string = 'chat-compression-2.5-flash-lite';
|
||||
|
||||
constructor(
|
||||
env: ContextEnvironment,
|
||||
options: SemanticCompressionProcessorOptions,
|
||||
) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
async process(
|
||||
editor: EpisodeEditor,
|
||||
state: ContextAccountingState,
|
||||
): Promise<void> {
|
||||
// If the budget is satisfied, or semantic compression isn't enabled
|
||||
if (state.isBudgetSatisfied) {
|
||||
return;
|
||||
}
|
||||
|
||||
const semanticConfig = this.options;
|
||||
const limitTokens = semanticConfig.nodeThresholdTokens;
|
||||
const thresholdChars = this.env.tokenCalculator.tokensToChars(limitTokens);
|
||||
this.modelToUse = 'gemini-2.5-flash';
|
||||
|
||||
let currentDeficit = state.deficitTokens;
|
||||
|
||||
// We scan backwards (oldest to newest would also work, but older is safer to degrade first)
|
||||
for (const ep of editor.episodes) {
|
||||
if (currentDeficit <= 0) break;
|
||||
if (state.protectedEpisodeIds.has(ep.id)) continue;
|
||||
|
||||
// 1. Compress User Prompts
|
||||
if (ep.trigger.type === 'USER_PROMPT') {
|
||||
for (let j = 0; j < ep.trigger.semanticParts.length; j++) {
|
||||
const part = ep.trigger.semanticParts[j];
|
||||
if (currentDeficit <= 0) break;
|
||||
if (part.type !== 'text') continue;
|
||||
// If it's already got a presentation, we don't want to re-summarize a summary
|
||||
if (part.presentation) continue;
|
||||
|
||||
if (part.text.length > thresholdChars) {
|
||||
const summary = await this.generateSummary(
|
||||
part.text,
|
||||
'User Prompt',
|
||||
);
|
||||
const newTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: summary },
|
||||
]);
|
||||
const oldTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: part.text },
|
||||
]);
|
||||
|
||||
if (newTokens < oldTokens) {
|
||||
editor.editEpisode(ep.id, 'SUMMARIZE_PROMPT', (draft) => {
|
||||
if (draft.trigger.type === 'USER_PROMPT') {
|
||||
draft.trigger.semanticParts[j].presentation = {
|
||||
text: summary,
|
||||
tokens: newTokens,
|
||||
};
|
||||
draft.trigger.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'SUMMARIZED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
}
|
||||
});
|
||||
currentDeficit -= oldTokens - newTokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Compress Model Thoughts
|
||||
if (ep.steps) {
|
||||
for (let j = 0; j < ep.steps.length; j++) {
|
||||
const step = ep.steps[j];
|
||||
if (currentDeficit <= 0) break;
|
||||
if (step.type === 'AGENT_THOUGHT') {
|
||||
if (step.presentation) continue;
|
||||
if (step.text.length > thresholdChars) {
|
||||
const summary = await this.generateSummary(
|
||||
step.text,
|
||||
'Agent Thought',
|
||||
);
|
||||
const newTokens = this.env.tokenCalculator.estimateTokensForParts(
|
||||
[{ text: summary }],
|
||||
);
|
||||
const oldTokens = this.env.tokenCalculator.estimateTokensForParts(
|
||||
[{ text: step.text }],
|
||||
);
|
||||
|
||||
if (newTokens < oldTokens) {
|
||||
editor.editEpisode(ep.id, 'SUMMARIZE_THOUGHT', (draft) => {
|
||||
const draftStep = draft.steps[j];
|
||||
if (draftStep.type === 'AGENT_THOUGHT') {
|
||||
draftStep.presentation = {
|
||||
text: summary,
|
||||
tokens: newTokens,
|
||||
};
|
||||
draftStep.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'SUMMARIZED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
}
|
||||
});
|
||||
currentDeficit -= oldTokens - newTokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Compress Tool Observations
|
||||
if (step.type === 'TOOL_EXECUTION') {
|
||||
const rawObs = step.presentation?.observation ?? step.observation;
|
||||
|
||||
let stringifiedObs = '';
|
||||
if (typeof rawObs === 'string') {
|
||||
stringifiedObs = rawObs;
|
||||
} else {
|
||||
try {
|
||||
stringifiedObs = JSON.stringify(rawObs);
|
||||
} catch {
|
||||
stringifiedObs = String(rawObs);
|
||||
}
|
||||
}
|
||||
|
||||
if (
|
||||
stringifiedObs.length > thresholdChars &&
|
||||
!stringifiedObs.includes('<tool_output_masked>')
|
||||
) {
|
||||
const summary = await this.generateSummary(
|
||||
stringifiedObs,
|
||||
`Tool Output (${step.toolName})`,
|
||||
);
|
||||
|
||||
// Wrap the summary in an object so the Gemini API accepts it as a valid functionResponse.response
|
||||
const newObsObject = { summary };
|
||||
|
||||
const newObsTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([
|
||||
{
|
||||
functionResponse: {
|
||||
name: step.toolName,
|
||||
response: newObsObject,
|
||||
id: step.id,
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
const oldObsTokens =
|
||||
step.presentation?.tokens?.observation ??
|
||||
step.tokens?.observation ??
|
||||
step.tokens;
|
||||
const intentTokens =
|
||||
step.presentation?.tokens?.intent ?? step.tokens?.intent ?? 0;
|
||||
|
||||
if (newObsTokens < oldObsTokens) {
|
||||
editor.editEpisode(ep.id, 'SUMMARIZE_TOOL', (draft) => {
|
||||
const draftStep = draft.steps[j];
|
||||
if (draftStep.type === 'TOOL_EXECUTION') {
|
||||
draftStep.presentation = {
|
||||
intent:
|
||||
draftStep.presentation?.intent ?? draftStep.intent,
|
||||
observation: newObsObject,
|
||||
tokens: {
|
||||
intent: intentTokens,
|
||||
observation: newObsTokens,
|
||||
},
|
||||
};
|
||||
if (!draftStep.metadata) {
|
||||
draftStep.metadata = {
|
||||
transformations: [],
|
||||
currentTokens: 0,
|
||||
originalTokens: 0,
|
||||
};
|
||||
}
|
||||
if (!draftStep.metadata.transformations) {
|
||||
draftStep.metadata.transformations = [];
|
||||
}
|
||||
draftStep.metadata.transformations.push({
|
||||
processorName: this.name,
|
||||
action: 'SUMMARIZED',
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
}
|
||||
});
|
||||
currentDeficit -= oldObsTokens - newObsTokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async generateSummary(
|
||||
content: string,
|
||||
contentType: string,
|
||||
abortSignal?: AbortSignal,
|
||||
): Promise<string> {
|
||||
const promptMessage = `You are compressing an old episodic context buffer for an AI assistant.\nSummarize this ${contentType} block in 2-3 highly technical sentences. Keep all critical facts, file names, dependencies, and architectural decisions. Discard conversational filler and boilerplate.\n\nContent:\n${content.slice(0, 30000)}`;
|
||||
|
||||
const client = this.env.llmClient;
|
||||
try {
|
||||
const response = await client.generateContent({
|
||||
modelConfigKey: { model: this.modelToUse },
|
||||
contents: [{ role: 'user', parts: [{ text: promptMessage }] }],
|
||||
promptId: 'local-context-compression-summary',
|
||||
role: LlmRole.UTILITY_COMPRESSOR,
|
||||
abortSignal: abortSignal ?? new AbortController().signal,
|
||||
});
|
||||
const text = getResponseText(response) ?? '';
|
||||
return `[Semantic Summary of old ${contentType}]\n${text.trim()}`;
|
||||
} catch (e) {
|
||||
debugLogger.warn(`Semantic compression LLM call failed: ${e}`);
|
||||
// If we fail to summarize, we just return the original truncated by 50% as a fail-safe, or the original.
|
||||
// Returning original is safer to prevent data loss on API failure.
|
||||
return content;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
createMockEnvironment,
|
||||
createDummyState,
|
||||
createDummyEpisode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { StateSnapshotProcessor } from './stateSnapshotProcessor.js';
|
||||
import { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
|
||||
describe('StateSnapshotProcessor', () => {
|
||||
let processor: StateSnapshotProcessor;
|
||||
let env: ContextEnvironment;
|
||||
let generateContentMock: ReturnType<typeof vi.fn>;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
env = createMockEnvironment();
|
||||
|
||||
generateContentMock = vi.fn().mockResolvedValue({
|
||||
text: 'Mocked Compressed State Snapshot!',
|
||||
});
|
||||
vi.spyOn(env, 'llmClient', 'get').mockReturnValue({
|
||||
generateContent: generateContentMock,
|
||||
} as unknown as BaseLlmClient);
|
||||
|
||||
// Override token calc for testing
|
||||
vi.spyOn(env.tokenCalculator, 'estimateTokensForParts').mockReturnValue(
|
||||
100,
|
||||
);
|
||||
|
||||
processor = new StateSnapshotProcessor(env, {}, env.eventBus);
|
||||
});
|
||||
|
||||
it('bypasses processing if deficit is <= 0', async () => {
|
||||
const episodes = [
|
||||
createDummyEpisode('ep-1', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'hello' },
|
||||
]),
|
||||
];
|
||||
// current: 100, max: 1000, retained: 200 (deficit 0)
|
||||
const state = createDummyState(false, 0, new Set(), 100, 1000, 200);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
expect(result).toStrictEqual(episodes);
|
||||
expect(generateContentMock).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('bypasses processing if not enough episodes to summarize (needs at least 2 inner episodes)', async () => {
|
||||
const episodes = [
|
||||
createDummyEpisode('ep-sys', 'SYSTEM_EVENT', []),
|
||||
createDummyEpisode('ep-active', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'help' },
|
||||
]),
|
||||
];
|
||||
|
||||
// current: 1000, max: 10000, retained: 500. Target deficit = 500
|
||||
const state = createDummyState(false, 500, new Set(), 1000, 10000, 500);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
expect(result).toStrictEqual(episodes);
|
||||
expect(generateContentMock).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('summarizes intermediate episodes into a single snapshot episode', async () => {
|
||||
const episodes = [
|
||||
createDummyEpisode('ep-0', 'SYSTEM_EVENT', []),
|
||||
createDummyEpisode('ep-1', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'old 1' },
|
||||
]),
|
||||
createDummyEpisode('ep-2', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'old 2' },
|
||||
]),
|
||||
createDummyEpisode('ep-3', 'USER_PROMPT', [
|
||||
{ type: 'text', text: 'current' },
|
||||
]),
|
||||
];
|
||||
|
||||
// Target deficit = 200
|
||||
const state = createDummyState(false, 200, new Set(), 1000, 10000, 800);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
// We started with 4 episodes.
|
||||
// Episodes [1, 2] were synthesized into a single new Snapshot episode.
|
||||
// Final array should be: [0, SNAPSHOT, 3] = length 3.
|
||||
expect(result.length).toBe(3);
|
||||
expect(result[0].id).toBe('ep-0');
|
||||
|
||||
const snapshotEp = result[1];
|
||||
expect(snapshotEp.yield).toBeDefined();
|
||||
expect(snapshotEp.yield!.text).toContain('<CONTEXT_SNAPSHOT>');
|
||||
expect(snapshotEp.yield!.text).toContain(
|
||||
'Mocked Compressed State Snapshot!',
|
||||
);
|
||||
|
||||
expect(result[2].id).toBe('ep-3');
|
||||
|
||||
expect(generateContentMock).toHaveBeenCalledTimes(1);
|
||||
|
||||
const llmArgs = generateContentMock.mock.calls[0][0];
|
||||
expect(llmArgs.contents[0].parts[0].text).toContain('old 1');
|
||||
expect(llmArgs.contents[0].parts[0].text).toContain('old 2');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,182 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextProcessor, ContextAccountingState } from '../pipeline.js';
|
||||
import type { Episode } from '../ir/types.js';
|
||||
import type {
|
||||
ContextEnvironment,
|
||||
ContextEventBus,
|
||||
} from '../sidecar/environment.js';
|
||||
import { v4 as uuidv4 } from 'uuid';
|
||||
import { LlmRole } from '../../telemetry/llmRole.js';
|
||||
import { debugLogger } from 'src/utils/debugLogger.js';
|
||||
import type { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
|
||||
export interface StateSnapshotProcessorOptions {
|
||||
model?: string;
|
||||
systemInstruction?: string;
|
||||
triggerDeficitTokens?: number;
|
||||
}
|
||||
|
||||
export class StateSnapshotProcessor implements ContextProcessor {
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: StateSnapshotProcessorOptions,
|
||||
): StateSnapshotProcessor {
|
||||
return new StateSnapshotProcessor(env, options, env.eventBus);
|
||||
}
|
||||
readonly id = 'StateSnapshotProcessor';
|
||||
readonly name = 'StateSnapshotProcessor';
|
||||
readonly options: StateSnapshotProcessorOptions;
|
||||
private readonly env: ContextEnvironment;
|
||||
private isSynthesizing = false;
|
||||
|
||||
constructor(
|
||||
env: ContextEnvironment,
|
||||
options: StateSnapshotProcessorOptions,
|
||||
_eventBus: ContextEventBus,
|
||||
) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
async process(
|
||||
editor: EpisodeEditor,
|
||||
state: ContextAccountingState,
|
||||
): Promise<void> {
|
||||
const targetDeficit = Math.max(
|
||||
0,
|
||||
state.currentTokens - state.retainedTokens,
|
||||
);
|
||||
if (this.isSynthesizing || targetDeficit <= 0) return;
|
||||
|
||||
this.isSynthesizing = true;
|
||||
try {
|
||||
let deficitAccumulator = 0;
|
||||
const selectedEpisodes: Episode[] = [];
|
||||
|
||||
for (let i = 1; i < editor.episodes.length - 1; i++) {
|
||||
const ep = editor.episodes[i];
|
||||
selectedEpisodes.push(ep);
|
||||
let triggerText = '';
|
||||
if (ep.trigger?.type === 'USER_PROMPT') {
|
||||
const firstPart = ep.trigger.semanticParts?.[0];
|
||||
if (firstPart) {
|
||||
triggerText =
|
||||
firstPart.type === 'text'
|
||||
? firstPart.text
|
||||
: (firstPart.presentation?.text ?? '');
|
||||
}
|
||||
}
|
||||
deficitAccumulator += this.env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: triggerText },
|
||||
{ text: ep.yield?.text ?? '' },
|
||||
]);
|
||||
if (deficitAccumulator >= targetDeficit) break;
|
||||
}
|
||||
|
||||
if (selectedEpisodes.length < 2) return; // Not enough context to summarize
|
||||
|
||||
// Optimization: Do NOT emit VariantComputing, let the Orchestrator handle caching the final result.
|
||||
const snapshotEp: Episode =
|
||||
await this.synthesizeSnapshot(selectedEpisodes);
|
||||
|
||||
const oldIds = selectedEpisodes.map((ep) => ep.id);
|
||||
editor.replaceEpisodes(oldIds, snapshotEp, 'STATE_SNAPSHOT');
|
||||
} finally {
|
||||
this.isSynthesizing = false;
|
||||
}
|
||||
}
|
||||
|
||||
private async synthesizeSnapshot(episodes: Episode[]): Promise<Episode> {
|
||||
const client = this.env.llmClient;
|
||||
const systemPrompt =
|
||||
this.options.systemInstruction ??
|
||||
`You are an expert Context Memory Manager. You will be provided with a raw transcript of older conversation turns between a user and an AI assistant.
|
||||
Your task is to synthesize these turns into a single, dense, factual snapshot that preserves all critical context, preferences, active tasks, and factual knowledge, but discards conversational filler, pleasantries, and redundant back-and-forth iterations.
|
||||
|
||||
Output ONLY the raw factual snapshot, formatted compactly. Do not include markdown wrappers, prefixes like "Here is the snapshot", or conversational elements.`;
|
||||
|
||||
let userPromptText = 'TRANSCRIPT TO SNAPSHOT:\n\n';
|
||||
for (const ep of episodes) {
|
||||
if (ep.trigger?.type === 'USER_PROMPT') {
|
||||
const partsText = ep.trigger.semanticParts
|
||||
.map((p) => {
|
||||
if (p.type === 'text') return p.text;
|
||||
if (p.presentation) return p.presentation.text;
|
||||
return '';
|
||||
})
|
||||
.join('');
|
||||
userPromptText += `USER: ${partsText}\n`;
|
||||
} else if (ep.trigger?.type === 'SYSTEM_EVENT') {
|
||||
userPromptText += `[SYSTEM EVENT: ${ep.trigger.name}]\n`;
|
||||
}
|
||||
for (const step of ep.steps) {
|
||||
if (step.type === 'TOOL_EXECUTION') {
|
||||
userPromptText += `[Tool Called: ${step.toolName}]\n`;
|
||||
}
|
||||
}
|
||||
if (ep.yield) {
|
||||
userPromptText += `ASSISTANT: ${ep.yield.text}\n`;
|
||||
}
|
||||
userPromptText += '\n';
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await client.generateContent({
|
||||
modelConfigKey: { model: 'state-snapshot-processor' },
|
||||
contents: [{ role: 'user', parts: [{ text: userPromptText }] }],
|
||||
systemInstruction: { role: 'system', parts: [{ text: systemPrompt }] },
|
||||
promptId: this.env.promptId,
|
||||
role: LlmRole.UTILITY_STATE_SNAPSHOT_PROCESSOR,
|
||||
abortSignal: new AbortController().signal,
|
||||
});
|
||||
|
||||
const snapshotText = response.text;
|
||||
|
||||
// Synthesize a new "Episode" representing this compressed block
|
||||
const newId = uuidv4();
|
||||
const contentTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: snapshotText },
|
||||
]);
|
||||
|
||||
return {
|
||||
id: newId,
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: `${newId}-t`,
|
||||
type: 'USER_PROMPT',
|
||||
semanticParts: [],
|
||||
metadata: {
|
||||
originalTokens: 0,
|
||||
currentTokens: 0,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
steps: [],
|
||||
yield: {
|
||||
id: `${newId}-y`,
|
||||
type: 'AGENT_YIELD',
|
||||
text: `<CONTEXT_SNAPSHOT>\n${snapshotText}\n</CONTEXT_SNAPSHOT>`,
|
||||
metadata: {
|
||||
originalTokens: contentTokens,
|
||||
currentTokens: contentTokens,
|
||||
transformations: [
|
||||
{
|
||||
processorName: 'StateSnapshotProcessor',
|
||||
action: 'SYNTHESIZED',
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
};
|
||||
} catch (error) {
|
||||
debugLogger.error('Failed to synthesize snapshot:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { createMockEnvironment } from '../testing/contextTestUtils.js';
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { ToolMaskingProcessor } from './toolMaskingProcessor.js';
|
||||
import { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
import type { Episode, ToolExecution } from '../ir/types.js';
|
||||
import type { ContextAccountingState } from '../pipeline.js';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import type { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
|
||||
|
||||
describe('ToolMaskingProcessor', () => {
|
||||
let processor: ToolMaskingProcessor;
|
||||
let env: ContextEnvironment;
|
||||
let fileSystem: InMemoryFileSystem;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
env = createMockEnvironment();
|
||||
fileSystem = env.fileSystem as InMemoryFileSystem;
|
||||
|
||||
processor = new ToolMaskingProcessor(env, {
|
||||
stringLengthThresholdTokens: 100,
|
||||
});
|
||||
});
|
||||
|
||||
const getDummyState = (
|
||||
isSatisfied = false,
|
||||
deficit = 0,
|
||||
protectedIds = new Set<string>(),
|
||||
): ContextAccountingState => ({
|
||||
currentTokens: 5000,
|
||||
maxTokens: 10000,
|
||||
retainedTokens: 4000,
|
||||
deficitTokens: deficit,
|
||||
protectedEpisodeIds: protectedIds,
|
||||
isBudgetSatisfied: isSatisfied,
|
||||
});
|
||||
|
||||
const createDummyEpisode = (
|
||||
id: string,
|
||||
intent: Record<string, unknown>,
|
||||
observation: Record<string, unknown>,
|
||||
): Episode => ({
|
||||
id,
|
||||
timestamp: Date.now(),
|
||||
trigger: {
|
||||
id: randomUUID(),
|
||||
type: 'SYSTEM_EVENT',
|
||||
name: 'test',
|
||||
payload: {},
|
||||
metadata: { originalTokens: 10, currentTokens: 10, transformations: [] },
|
||||
},
|
||||
steps: [
|
||||
{
|
||||
id: randomUUID(),
|
||||
type: 'TOOL_EXECUTION',
|
||||
toolName: 'test_tool',
|
||||
intent,
|
||||
observation,
|
||||
tokens: { intent: 500, observation: 500 }, // Claim they are big enough to be masked
|
||||
metadata: {
|
||||
originalTokens: 1000,
|
||||
currentTokens: 1000,
|
||||
transformations: [],
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
it('bypasses processing if budget is satisfied', async () => {
|
||||
const episodes = [
|
||||
createDummyEpisode('1', { arg: 'short' }, { out: 'short' }),
|
||||
];
|
||||
const state = getDummyState(true);
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
expect(result).toStrictEqual(episodes);
|
||||
expect((result[0].steps[0] as ToolExecution).presentation).toBeUndefined();
|
||||
});
|
||||
|
||||
it('deep masks massive string intents and observations', async () => {
|
||||
// We need strings > limitChars (100 tokens * 4 chars = 400 chars)
|
||||
const massiveIntentString = 'I'.repeat(500);
|
||||
const massiveObsString = 'O'.repeat(500);
|
||||
|
||||
const intentPayload = { args: { nested: [massiveIntentString, 'short'] } };
|
||||
const obsPayload = { result: massiveObsString, error: null };
|
||||
|
||||
const episodes = [createDummyEpisode('ep-1', intentPayload, obsPayload)];
|
||||
const state = getDummyState(false, 1000, new Set()); // Huge deficit
|
||||
|
||||
const editor = new EpisodeEditor(episodes);
|
||||
await processor.process(editor, state);
|
||||
const result = editor.getFinalEpisodes();
|
||||
|
||||
const toolStep = result[0].steps[0] as ToolExecution;
|
||||
|
||||
expect(toolStep.presentation).toBeDefined();
|
||||
|
||||
// Check intent was deep masked
|
||||
const maskedIntent = toolStep.presentation!.intent as Record<
|
||||
string,
|
||||
unknown
|
||||
>;
|
||||
expect((maskedIntent['args'] as { nested: string }).nested[0]).toContain(
|
||||
'<tool_output_masked>',
|
||||
);
|
||||
expect((maskedIntent['args'] as { nested: string }).nested[1]).toBe(
|
||||
'short',
|
||||
); // Unchanged
|
||||
|
||||
// Check observation was deep masked
|
||||
const maskedObs = toolStep.presentation!.observation as Record<
|
||||
string,
|
||||
unknown
|
||||
>;
|
||||
expect((maskedObs as { result: string }).result).toContain(
|
||||
'<tool_output_masked>',
|
||||
);
|
||||
expect((maskedObs as { error: string }).error).toBeNull();
|
||||
|
||||
// Check disk writes occurred to fake FS
|
||||
expect(fileSystem.getFiles().size).toBe(2);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,319 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { sanitizeFilenamePart } from '../../utils/fileUtils.js';
|
||||
import {
|
||||
ACTIVATE_SKILL_TOOL_NAME,
|
||||
MEMORY_TOOL_NAME,
|
||||
ASK_USER_TOOL_NAME,
|
||||
ENTER_PLAN_MODE_TOOL_NAME,
|
||||
EXIT_PLAN_MODE_TOOL_NAME,
|
||||
} from '../../tools/tool-names.js';
|
||||
import type { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
|
||||
const UNMASKABLE_TOOLS = new Set([
|
||||
ACTIVATE_SKILL_TOOL_NAME,
|
||||
MEMORY_TOOL_NAME,
|
||||
ASK_USER_TOOL_NAME,
|
||||
ENTER_PLAN_MODE_TOOL_NAME,
|
||||
EXIT_PLAN_MODE_TOOL_NAME,
|
||||
]);
|
||||
|
||||
export interface ToolMaskingProcessorOptions {
|
||||
stringLengthThresholdTokens: number;
|
||||
}
|
||||
|
||||
type MaskableValue =
|
||||
| string
|
||||
| number
|
||||
| boolean
|
||||
| null
|
||||
| MaskableValue[]
|
||||
| { [key: string]: MaskableValue };
|
||||
|
||||
function isMaskableValue(val: unknown): val is MaskableValue {
|
||||
if (
|
||||
val === null ||
|
||||
typeof val === 'string' ||
|
||||
typeof val === 'number' ||
|
||||
typeof val === 'boolean'
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
if (Array.isArray(val)) {
|
||||
return val.every(isMaskableValue);
|
||||
}
|
||||
if (typeof val === 'object') {
|
||||
return Object.values(val).every(isMaskableValue);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
function isMaskableRecord(val: unknown): val is Record<string, MaskableValue> {
|
||||
return (
|
||||
typeof val === 'object' &&
|
||||
val !== null &&
|
||||
!Array.isArray(val) &&
|
||||
isMaskableValue(val)
|
||||
);
|
||||
}
|
||||
|
||||
export class ToolMaskingProcessor implements ContextProcessor {
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: ToolMaskingProcessorOptions,
|
||||
): ToolMaskingProcessor {
|
||||
return new ToolMaskingProcessor(env, options);
|
||||
}
|
||||
|
||||
static readonly schema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
stringLengthThresholdTokens: {
|
||||
type: 'number',
|
||||
description:
|
||||
'The token threshold above which tool intents/observations are masked.',
|
||||
},
|
||||
},
|
||||
required: ['stringLengthThresholdTokens'],
|
||||
};
|
||||
|
||||
readonly id = 'ToolMaskingProcessor';
|
||||
readonly name = 'ToolMaskingProcessor';
|
||||
readonly options: ToolMaskingProcessorOptions;
|
||||
private env: ContextEnvironment;
|
||||
|
||||
constructor(env: ContextEnvironment, options: ToolMaskingProcessorOptions) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
async process(
|
||||
editor: EpisodeEditor,
|
||||
state: ContextAccountingState,
|
||||
): Promise<void> {
|
||||
const maskingConfig = this.options;
|
||||
if (!maskingConfig) return;
|
||||
if (state.isBudgetSatisfied) return;
|
||||
|
||||
let currentDeficit = state.deficitTokens;
|
||||
const limitChars = this.env.tokenCalculator.tokensToChars(
|
||||
maskingConfig.stringLengthThresholdTokens,
|
||||
);
|
||||
|
||||
let toolOutputsDir = this.env.fileSystem.join(
|
||||
this.env.projectTempDir,
|
||||
'tool-outputs',
|
||||
);
|
||||
const sessionId = this.env.sessionId;
|
||||
if (sessionId) {
|
||||
toolOutputsDir = this.env.fileSystem.join(
|
||||
toolOutputsDir,
|
||||
`session-${sanitizeFilenamePart(sessionId)}`,
|
||||
);
|
||||
}
|
||||
|
||||
// We only create the directory if we actually mask something
|
||||
let directoryCreated = false;
|
||||
|
||||
// Helper to extract string and write to disk
|
||||
const handleMasking = async (
|
||||
content: string,
|
||||
toolName: string,
|
||||
callId: string,
|
||||
nodeType: string,
|
||||
): Promise<string> => {
|
||||
if (!directoryCreated) {
|
||||
await this.env.fileSystem.mkdir(toolOutputsDir, { recursive: true });
|
||||
directoryCreated = true;
|
||||
}
|
||||
|
||||
const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${this.env.idGenerator.generateId()}.txt`;
|
||||
const filePath = this.env.fileSystem.join(toolOutputsDir, fileName);
|
||||
|
||||
await this.env.fileSystem.writeFile(filePath, content);
|
||||
|
||||
const fileSizeMB = (
|
||||
Buffer.byteLength(content, 'utf8') /
|
||||
1024 /
|
||||
1024
|
||||
).toFixed(2);
|
||||
const totalLines = content.split('\n').length;
|
||||
return `<tool_output_masked>\n[Tool ${nodeType} string (${fileSizeMB}MB, ${totalLines} lines) masked to preserve context window. Full string saved to: ${filePath}]\n</tool_output_masked>`;
|
||||
};
|
||||
|
||||
// Forward scan, looking for massive intents or observations to mask
|
||||
for (const ep of editor.episodes) {
|
||||
if (currentDeficit <= 0) break;
|
||||
if (!ep || !ep.steps || state.protectedEpisodeIds.has(ep.id)) continue;
|
||||
|
||||
for (let j = 0; j < ep.steps.length; j++) {
|
||||
if (currentDeficit <= 0) break;
|
||||
const step = ep.steps[j];
|
||||
if (step.type !== 'TOOL_EXECUTION') continue;
|
||||
|
||||
const toolName = step.toolName;
|
||||
if (toolName && UNMASKABLE_TOOLS.has(toolName)) continue;
|
||||
|
||||
// Ensure presentation object exists
|
||||
if (!step.presentation) {
|
||||
step.presentation = {
|
||||
intent: step.intent,
|
||||
observation: step.observation,
|
||||
tokens: step.tokens, // Fallback to raw tokens initially
|
||||
};
|
||||
}
|
||||
|
||||
const callId = step.id || Date.now().toString();
|
||||
|
||||
const maskAsync = async (
|
||||
obj: MaskableValue,
|
||||
nodeType: string,
|
||||
): Promise<{ masked: MaskableValue; changed: boolean }> => {
|
||||
if (typeof obj === 'string') {
|
||||
if (obj.length > limitChars && !this.isAlreadyMasked(obj)) {
|
||||
const newString = await handleMasking(
|
||||
obj,
|
||||
toolName,
|
||||
callId,
|
||||
nodeType,
|
||||
);
|
||||
return { masked: newString, changed: true };
|
||||
}
|
||||
return { masked: obj, changed: false };
|
||||
}
|
||||
if (Array.isArray(obj)) {
|
||||
let changed = false;
|
||||
const masked: MaskableValue[] = [];
|
||||
for (const item of obj) {
|
||||
const res = await maskAsync(item, nodeType);
|
||||
if (res.changed) changed = true;
|
||||
masked.push(res.masked);
|
||||
}
|
||||
return { masked, changed };
|
||||
}
|
||||
if (typeof obj === 'object' && obj !== null) {
|
||||
let changed = false;
|
||||
const masked: Record<string, MaskableValue> = {};
|
||||
for (const [key, value] of Object.entries(obj)) {
|
||||
const res = await maskAsync(value, nodeType);
|
||||
if (res.changed) changed = true;
|
||||
masked[key] = res.masked;
|
||||
}
|
||||
return { masked, changed };
|
||||
}
|
||||
return { masked: obj, changed: false };
|
||||
};
|
||||
|
||||
const rawIntent = step.presentation?.intent ?? step.intent;
|
||||
const rawObs = step.presentation?.observation ?? step.observation;
|
||||
|
||||
if (!isMaskableRecord(rawIntent)) {
|
||||
throw new Error(
|
||||
`ToolMaskingProcessor: step intent is not a valid JSON record. CallID: ${callId}`,
|
||||
);
|
||||
}
|
||||
if (!isMaskableValue(rawObs)) {
|
||||
throw new Error(
|
||||
`ToolMaskingProcessor: step observation is not a valid JSON value. CallID: ${callId}`,
|
||||
);
|
||||
}
|
||||
|
||||
const intentRes = await maskAsync(rawIntent, 'intent');
|
||||
const obsRes = await maskAsync(rawObs, 'observation');
|
||||
|
||||
if (intentRes.changed || obsRes.changed) {
|
||||
const maskedIntent = isMaskableRecord(intentRes.masked)
|
||||
? (intentRes.masked as Record<string, unknown>)
|
||||
: undefined;
|
||||
const maskedObs = isMaskableRecord(obsRes.masked)
|
||||
? (obsRes.masked as Record<string, unknown>)
|
||||
: undefined;
|
||||
|
||||
// Recalculate tokens perfectly
|
||||
const newIntentTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([
|
||||
{
|
||||
functionCall: {
|
||||
name: toolName,
|
||||
args: maskedIntent,
|
||||
id: callId,
|
||||
},
|
||||
},
|
||||
]);
|
||||
const newObsTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
{
|
||||
functionResponse: {
|
||||
name: toolName,
|
||||
response:
|
||||
typeof obsRes.masked === 'string'
|
||||
? { message: obsRes.masked }
|
||||
: maskedObs,
|
||||
id: callId,
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
const oldTotal =
|
||||
step.presentation.tokens?.intent !== undefined
|
||||
? step.presentation.tokens.intent +
|
||||
step.presentation.tokens.observation
|
||||
: step.tokens.intent + step.tokens.observation;
|
||||
|
||||
const newTotal = newIntentTokens + newObsTokens;
|
||||
const savings = oldTotal - newTotal;
|
||||
|
||||
if (savings > 0) {
|
||||
currentDeficit -= savings;
|
||||
this.env.tracer.logEvent(
|
||||
'ToolMaskingProcessor',
|
||||
`Masked tool ${toolName}`,
|
||||
{ recoveredTokens: savings },
|
||||
);
|
||||
|
||||
editor.editEpisode(ep.id, 'MASK_TOOL', (draft) => {
|
||||
const draftStep = draft.steps[j];
|
||||
if (draftStep.type !== 'TOOL_EXECUTION') return;
|
||||
if (!draftStep.presentation) {
|
||||
draftStep.presentation = {
|
||||
intent: draftStep.intent,
|
||||
observation: draftStep.observation,
|
||||
tokens: draftStep.tokens,
|
||||
};
|
||||
}
|
||||
draftStep.presentation.intent = maskedIntent ?? {};
|
||||
draftStep.presentation.observation =
|
||||
typeof obsRes.masked === 'string'
|
||||
? { message: obsRes.masked }
|
||||
: (maskedObs ?? {});
|
||||
draftStep.presentation.tokens = {
|
||||
intent: newIntentTokens,
|
||||
observation: newObsTokens,
|
||||
};
|
||||
draftStep.metadata = {
|
||||
...draftStep.metadata,
|
||||
transformations: [
|
||||
...(draftStep.metadata?.transformations || []),
|
||||
{
|
||||
processorName: 'ToolMasking',
|
||||
action: 'MASKED',
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
};
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private isAlreadyMasked(content: string): boolean {
|
||||
return content.includes('<tool_output_masked>');
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { ProcessorRegistry } from './registry.js';
|
||||
import { registerBuiltInProcessors } from './builtins.js';
|
||||
import { describe, it, expect, beforeEach } from 'vitest';
|
||||
import { SidecarLoader } from './SidecarLoader.js';
|
||||
import { defaultSidecarProfile } from './profiles.js';
|
||||
import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
|
||||
import type { Config } from 'src/config/config.js';
|
||||
|
||||
describe('SidecarLoader (Fake FS)', () => {
|
||||
let fileSystem: InMemoryFileSystem;
|
||||
let registry: ProcessorRegistry;
|
||||
|
||||
beforeEach(() => {
|
||||
fileSystem = new InMemoryFileSystem();
|
||||
registry = new ProcessorRegistry();
|
||||
registerBuiltInProcessors(registry);
|
||||
});
|
||||
|
||||
const mockConfig = {
|
||||
getExperimentalContextSidecarConfig: () => '/path/to/sidecar.json',
|
||||
} as unknown as Config;
|
||||
|
||||
it('returns default profile if file does not exist', () => {
|
||||
const result = SidecarLoader.fromConfig(mockConfig, registry, fileSystem);
|
||||
expect(result).toBe(defaultSidecarProfile);
|
||||
});
|
||||
|
||||
it('returns default profile if file exists but is 0 bytes', () => {
|
||||
fileSystem.setFile('/path/to/sidecar.json', '');
|
||||
const result = SidecarLoader.fromConfig(mockConfig, registry, fileSystem);
|
||||
expect(result).toBe(defaultSidecarProfile);
|
||||
});
|
||||
|
||||
it('throws an error if file is empty whitespace', () => {
|
||||
fileSystem.setFile('/path/to/sidecar.json', ' \n ');
|
||||
expect(() =>
|
||||
SidecarLoader.fromConfig(mockConfig, registry, fileSystem),
|
||||
).toThrow('is empty');
|
||||
});
|
||||
|
||||
it('returns parsed config if file is valid', () => {
|
||||
const validConfig = {
|
||||
budget: { retainedTokens: 1000, maxTokens: 2000 },
|
||||
gcBackstop: { strategy: 'truncate', target: 'max' },
|
||||
pipelines: [],
|
||||
};
|
||||
fileSystem.setFile('/path/to/sidecar.json', JSON.stringify(validConfig));
|
||||
const result = SidecarLoader.fromConfig(mockConfig, registry, fileSystem);
|
||||
expect(result.budget.maxTokens).toBe(2000);
|
||||
});
|
||||
|
||||
it('throws validation error if file is invalid', () => {
|
||||
const invalidConfig = {
|
||||
budget: { retainedTokens: 1000 }, // missing maxTokens
|
||||
};
|
||||
fileSystem.setFile('/path/to/sidecar.json', JSON.stringify(invalidConfig));
|
||||
expect(() =>
|
||||
SidecarLoader.fromConfig(mockConfig, registry, fileSystem),
|
||||
).toThrow('Validation error:');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,87 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Config } from '../../config/config.js';
|
||||
import type { SidecarConfig } from './types.js';
|
||||
import { defaultSidecarProfile } from './profiles.js';
|
||||
import { SchemaValidator } from '../../utils/schemaValidator.js';
|
||||
import { getSidecarConfigSchema } from './schema.js';
|
||||
import type { IFileSystem } from '../system/IFileSystem.js';
|
||||
import { NodeFileSystem } from '../system/NodeFileSystem.js';
|
||||
import type { ProcessorRegistry } from './registry.js';
|
||||
|
||||
export class SidecarLoader {
|
||||
/**
|
||||
* Loads and validates a sidecar config from a specific file path.
|
||||
* Throws an error if the file cannot be read, parsed, or fails schema validation.
|
||||
*/
|
||||
static loadFromFile(
|
||||
sidecarPath: string,
|
||||
registry: ProcessorRegistry,
|
||||
fileSystem: IFileSystem = new NodeFileSystem(),
|
||||
): SidecarConfig {
|
||||
const fileContent = fileSystem.readFileSync(sidecarPath, 'utf8');
|
||||
|
||||
if (!fileContent.trim()) {
|
||||
throw new Error(`Sidecar configuration file at ${sidecarPath} is empty.`);
|
||||
}
|
||||
|
||||
let parsed: unknown;
|
||||
try {
|
||||
parsed = JSON.parse(fileContent);
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Failed to parse Sidecar configuration file at ${sidecarPath}: ${
|
||||
error instanceof Error ? error.message : String(error)
|
||||
}`,
|
||||
);
|
||||
}
|
||||
|
||||
const validationError = SchemaValidator.validate(
|
||||
getSidecarConfigSchema(registry),
|
||||
parsed,
|
||||
);
|
||||
if (validationError) {
|
||||
throw new Error(
|
||||
`Invalid sidecar configuration in ${sidecarPath}. Validation error: ${validationError}`,
|
||||
);
|
||||
}
|
||||
|
||||
// Schema has been validated.
|
||||
const isSidecarConfig = (val: unknown): val is SidecarConfig => true;
|
||||
if (isSidecarConfig(parsed)) {
|
||||
return parsed;
|
||||
}
|
||||
throw new Error(
|
||||
'Unreachable: schema validation passed but type predicate failed.',
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a Sidecar JSON graph from the experimental config file path or defaults.
|
||||
* If a config file is present but invalid, this will THROW to prevent silent misconfiguration.
|
||||
*/
|
||||
static fromConfig(
|
||||
config: Config,
|
||||
registry: ProcessorRegistry,
|
||||
fileSystem: IFileSystem = new NodeFileSystem(),
|
||||
): SidecarConfig {
|
||||
const sidecarPath = config.getExperimentalContextSidecarConfig();
|
||||
|
||||
if (sidecarPath && fileSystem.existsSync(sidecarPath)) {
|
||||
const size = fileSystem.statSyncSize(sidecarPath);
|
||||
// If the file exists but is completely empty (0 bytes), it's safe to fallback.
|
||||
if (size === 0) {
|
||||
return defaultSidecarProfile;
|
||||
}
|
||||
|
||||
// If the file has content, enforce strict validation and throw on failure.
|
||||
return this.loadFromFile(sidecarPath, registry, fileSystem);
|
||||
}
|
||||
|
||||
return defaultSidecarProfile;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ProcessorRegistry } from './registry.js';
|
||||
import {
|
||||
ToolMaskingProcessor,
|
||||
type ToolMaskingProcessorOptions,
|
||||
} from '../processors/toolMaskingProcessor.js';
|
||||
import { BlobDegradationProcessor } from '../processors/blobDegradationProcessor.js';
|
||||
import {
|
||||
SemanticCompressionProcessor,
|
||||
type SemanticCompressionProcessorOptions,
|
||||
} from '../processors/semanticCompressionProcessor.js';
|
||||
import {
|
||||
HistorySquashingProcessor,
|
||||
type HistorySquashingProcessorOptions,
|
||||
} from '../processors/historySquashingProcessor.js';
|
||||
import {
|
||||
StateSnapshotProcessor,
|
||||
type StateSnapshotProcessorOptions,
|
||||
} from '../processors/stateSnapshotProcessor.js';
|
||||
import {
|
||||
EmergencyTruncationProcessor,
|
||||
type EmergencyTruncationProcessorOptions,
|
||||
} from '../processors/emergencyTruncationProcessor.js';
|
||||
|
||||
export function registerBuiltInProcessors(registry: ProcessorRegistry) {
|
||||
registry.register<ToolMaskingProcessorOptions>({
|
||||
id: 'ToolMaskingProcessor',
|
||||
schema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
processorId: { const: 'ToolMaskingProcessor' },
|
||||
options: {
|
||||
type: 'object',
|
||||
properties: { stringLengthThresholdTokens: { type: 'number' } },
|
||||
required: ['stringLengthThresholdTokens'],
|
||||
},
|
||||
},
|
||||
required: ['processorId', 'options'],
|
||||
},
|
||||
create: (env, opts) => new ToolMaskingProcessor(env, opts),
|
||||
});
|
||||
|
||||
registry.register<Record<string, never>>({
|
||||
id: 'BlobDegradationProcessor',
|
||||
schema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
processorId: { const: 'BlobDegradationProcessor' },
|
||||
options: { type: 'object' },
|
||||
},
|
||||
required: ['processorId'],
|
||||
},
|
||||
create: (env) => new BlobDegradationProcessor(env),
|
||||
});
|
||||
|
||||
registry.register<SemanticCompressionProcessorOptions>({
|
||||
id: 'SemanticCompressionProcessor',
|
||||
schema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
processorId: { const: 'SemanticCompressionProcessor' },
|
||||
options: {
|
||||
type: 'object',
|
||||
properties: { nodeThresholdTokens: { type: 'number' } },
|
||||
required: ['nodeThresholdTokens'],
|
||||
},
|
||||
},
|
||||
required: ['processorId', 'options'],
|
||||
},
|
||||
create: (env, opts) => new SemanticCompressionProcessor(env, opts),
|
||||
});
|
||||
|
||||
registry.register<HistorySquashingProcessorOptions>({
|
||||
id: 'HistorySquashingProcessor',
|
||||
schema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
processorId: { const: 'HistorySquashingProcessor' },
|
||||
options: {
|
||||
type: 'object',
|
||||
properties: { maxTokensPerNode: { type: 'number' } },
|
||||
required: ['maxTokensPerNode'],
|
||||
},
|
||||
},
|
||||
required: ['processorId', 'options'],
|
||||
},
|
||||
create: (env, opts) => new HistorySquashingProcessor(env, opts),
|
||||
});
|
||||
|
||||
registry.register<StateSnapshotProcessorOptions>({
|
||||
id: 'StateSnapshotProcessor',
|
||||
schema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
processorId: { const: 'StateSnapshotProcessor' },
|
||||
options: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
model: { type: 'string' },
|
||||
systemInstruction: { type: 'string' },
|
||||
triggerDeficitTokens: { type: 'number' },
|
||||
},
|
||||
},
|
||||
},
|
||||
required: ['processorId'],
|
||||
},
|
||||
create: (env, opts) => StateSnapshotProcessor.create(env, opts),
|
||||
});
|
||||
|
||||
registry.register<EmergencyTruncationProcessorOptions>({
|
||||
id: 'EmergencyTruncationProcessor',
|
||||
schema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
processorId: { const: 'EmergencyTruncationProcessor' },
|
||||
options: { type: 'object' },
|
||||
},
|
||||
required: ['processorId'],
|
||||
},
|
||||
create: (env, opts) => EmergencyTruncationProcessor.create(env, opts),
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
import type { ContextEventBus } from '../eventBus.js';
|
||||
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
|
||||
import type { ContextTracer } from '../tracer.js';
|
||||
import type { IFileSystem } from '../system/IFileSystem.js';
|
||||
import type { IIdGenerator } from '../system/IIdGenerator.js';
|
||||
|
||||
export type { ContextTracer, ContextEventBus };
|
||||
|
||||
export interface ContextEnvironment {
|
||||
readonly llmClient: BaseLlmClient;
|
||||
readonly promptId: string;
|
||||
readonly sessionId: string;
|
||||
readonly traceDir: string;
|
||||
readonly projectTempDir: string;
|
||||
readonly tracer: ContextTracer;
|
||||
readonly charsPerToken: number;
|
||||
readonly tokenCalculator: ContextTokenCalculator;
|
||||
readonly fileSystem: IFileSystem;
|
||||
readonly idGenerator: IIdGenerator;
|
||||
readonly eventBus: ContextEventBus;
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
import type { ContextTracer } from '../tracer.js';
|
||||
import type { ContextEnvironment } from './environment.js';
|
||||
import type { ContextEventBus } from '../eventBus.js';
|
||||
import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
|
||||
import type { IFileSystem } from '../system/IFileSystem.js';
|
||||
import { NodeFileSystem } from '../system/NodeFileSystem.js';
|
||||
import type { IIdGenerator } from '../system/IIdGenerator.js';
|
||||
import { NodeIdGenerator } from '../system/NodeIdGenerator.js';
|
||||
|
||||
export class ContextEnvironmentImpl implements ContextEnvironment {
|
||||
readonly tokenCalculator: ContextTokenCalculator;
|
||||
readonly fileSystem: IFileSystem;
|
||||
readonly idGenerator: IIdGenerator;
|
||||
|
||||
constructor(
|
||||
readonly llmClient: BaseLlmClient,
|
||||
readonly sessionId: string,
|
||||
readonly promptId: string,
|
||||
readonly traceDir: string,
|
||||
readonly projectTempDir: string,
|
||||
readonly tracer: ContextTracer,
|
||||
readonly charsPerToken: number,
|
||||
readonly eventBus: ContextEventBus,
|
||||
fileSystem?: IFileSystem,
|
||||
idGenerator?: IIdGenerator,
|
||||
) {
|
||||
this.tokenCalculator = new ContextTokenCalculator(this.charsPerToken);
|
||||
this.fileSystem = fileSystem || new NodeFileSystem();
|
||||
this.idGenerator = idGenerator || new NodeIdGenerator();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,292 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
|
||||
import { PipelineOrchestrator } from './orchestrator.js';
|
||||
import { ProcessorRegistry } from './registry.js';
|
||||
import {
|
||||
createMockEnvironment,
|
||||
createDummyState,
|
||||
createDummyEpisode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import type { ContextEnvironment } from './environment.js';
|
||||
import type { ContextAccountingState, ContextProcessor } from '../pipeline.js';
|
||||
import type { PipelineDef, ProcessorConfig, SidecarConfig } from './types.js';
|
||||
import type { ContextEventBus } from '../eventBus.js';
|
||||
import type { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
|
||||
// Create a Dummy Processor for testing Orchestration routing
|
||||
class DummySyncProcessor implements ContextProcessor {
|
||||
static create() {
|
||||
return new DummySyncProcessor();
|
||||
}
|
||||
constructor() {}
|
||||
readonly name = 'DummySync';
|
||||
readonly id = 'DummySync';
|
||||
readonly options = {};
|
||||
async process(editor: EpisodeEditor, _state: ContextAccountingState) {
|
||||
editor.editEpisode(
|
||||
editor.episodes[0].id,
|
||||
'DUMMY_EDIT',
|
||||
(draft: unknown) => {
|
||||
(draft as Record<string, unknown>)['dummyModified'] = true;
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
class DummyAsyncProcessor implements ContextProcessor {
|
||||
static create() {
|
||||
return new DummyAsyncProcessor();
|
||||
}
|
||||
constructor() {}
|
||||
readonly name = 'DummyAsync';
|
||||
readonly id = 'DummyAsync';
|
||||
readonly options = {};
|
||||
async process(editor: EpisodeEditor, _state: ContextAccountingState) {
|
||||
editor.editEpisode(
|
||||
editor.episodes[0].id,
|
||||
'DUMMY_EDIT',
|
||||
(draft: unknown) => {
|
||||
(draft as Record<string, unknown>)['dummyAsyncModified'] = true;
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
class ThrowingProcessor implements ContextProcessor {
|
||||
static create() {
|
||||
return new ThrowingProcessor();
|
||||
}
|
||||
constructor() {}
|
||||
readonly name = 'Throwing';
|
||||
readonly id = 'Throwing';
|
||||
readonly options = {};
|
||||
async process(
|
||||
_editor: EpisodeEditor,
|
||||
_state: ContextAccountingState,
|
||||
): Promise<void> {
|
||||
throw new Error('Processor failed intentionally');
|
||||
}
|
||||
}
|
||||
|
||||
describe('PipelineOrchestrator (Component)', () => {
|
||||
let env: ContextEnvironment;
|
||||
let eventBus: ContextEventBus;
|
||||
let registry: ProcessorRegistry;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
env = createMockEnvironment();
|
||||
eventBus = env.eventBus;
|
||||
registry = new ProcessorRegistry();
|
||||
|
||||
// Register our test processors
|
||||
registry.register({
|
||||
id: 'DummySyncProcessor',
|
||||
schema: {},
|
||||
create: () => new DummySyncProcessor(),
|
||||
});
|
||||
registry.register({
|
||||
id: 'DummyAsyncProcessor',
|
||||
schema: {},
|
||||
create: () => new DummyAsyncProcessor(),
|
||||
});
|
||||
registry.register({
|
||||
id: 'ThrowingProcessor',
|
||||
schema: {},
|
||||
create: () => new ThrowingProcessor(),
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
// Cleanup registry to not pollute other tests
|
||||
registry.clear();
|
||||
});
|
||||
|
||||
const createConfig = (pipelines: PipelineDef[]): SidecarConfig => ({
|
||||
budget: { maxTokens: 100, retainedTokens: 50 },
|
||||
gcBackstop: { strategy: 'truncate', target: 'max' },
|
||||
pipelines,
|
||||
});
|
||||
|
||||
it('instantiates processors from the registry on initialization', () => {
|
||||
const config = createConfig([
|
||||
{
|
||||
name: 'Sync',
|
||||
execution: 'blocking',
|
||||
triggers: [],
|
||||
processors: [
|
||||
{ processorId: 'DummySyncProcessor' } as unknown as ProcessorConfig,
|
||||
],
|
||||
},
|
||||
]);
|
||||
|
||||
const orchestrator = new PipelineOrchestrator(
|
||||
config,
|
||||
env,
|
||||
eventBus,
|
||||
env.tracer,
|
||||
registry,
|
||||
);
|
||||
expect(
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
(orchestrator as any).instantiatedProcessors.has('DummySyncProcessor'),
|
||||
).toBe(true);
|
||||
});
|
||||
|
||||
it('throws an error if a config requests an unknown processor', () => {
|
||||
const config = createConfig([
|
||||
{
|
||||
name: 'Bad',
|
||||
execution: 'blocking',
|
||||
triggers: [],
|
||||
processors: [
|
||||
{ processorId: 'DoesNotExist' } as unknown as ProcessorConfig,
|
||||
],
|
||||
},
|
||||
]);
|
||||
|
||||
expect(
|
||||
() =>
|
||||
new PipelineOrchestrator(config, env, eventBus, env.tracer, registry),
|
||||
).toThrow('Context Processor [DoesNotExist] is not registered.');
|
||||
});
|
||||
|
||||
it('executes blocking pipelines synchronously and returns the modified array', async () => {
|
||||
const config = createConfig([
|
||||
{
|
||||
name: 'SyncPipe',
|
||||
execution: 'blocking',
|
||||
triggers: [],
|
||||
processors: [
|
||||
{ processorId: 'DummySyncProcessor' } as unknown as ProcessorConfig,
|
||||
],
|
||||
},
|
||||
]);
|
||||
const orchestrator = new PipelineOrchestrator(
|
||||
config,
|
||||
env,
|
||||
eventBus,
|
||||
env.tracer,
|
||||
registry,
|
||||
);
|
||||
|
||||
const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])];
|
||||
const state = createDummyState(false);
|
||||
|
||||
const result = await orchestrator.executePipeline(
|
||||
'SyncPipe',
|
||||
episodes,
|
||||
state,
|
||||
);
|
||||
|
||||
expect(result).toHaveLength(1);
|
||||
expect(
|
||||
(result[0] as unknown as { dummyModified: boolean }).dummyModified,
|
||||
).toBe(true);
|
||||
});
|
||||
|
||||
it('executes background pipelines asynchronously without blocking the return', async () => {
|
||||
const config = createConfig([
|
||||
{
|
||||
name: 'AsyncPipe',
|
||||
execution: 'background',
|
||||
triggers: [],
|
||||
processors: [
|
||||
{ processorId: 'DummyAsyncProcessor' } as unknown as ProcessorConfig,
|
||||
],
|
||||
},
|
||||
]);
|
||||
const orchestrator = new PipelineOrchestrator(
|
||||
config,
|
||||
env,
|
||||
eventBus,
|
||||
env.tracer,
|
||||
registry,
|
||||
);
|
||||
|
||||
const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])];
|
||||
const state = createDummyState(false);
|
||||
|
||||
// This should resolve immediately with the UNMODIFIED array because execution is background
|
||||
const result = await orchestrator.executePipeline(
|
||||
'AsyncPipe',
|
||||
episodes,
|
||||
state,
|
||||
);
|
||||
|
||||
expect(result).toHaveLength(1);
|
||||
expect(
|
||||
(result[0] as unknown as { asyncModified: unknown }).asyncModified,
|
||||
).toBeUndefined(); // Not modified yet!
|
||||
|
||||
// Wait for the background task to complete (50ms delay in DummyAsyncProcessor)
|
||||
await new Promise((resolve) => setTimeout(resolve, 60));
|
||||
});
|
||||
|
||||
it('gracefully handles and swallows processor errors in synchronous pipelines', async () => {
|
||||
const config = createConfig([
|
||||
{
|
||||
name: 'ThrowingPipe',
|
||||
execution: 'blocking',
|
||||
triggers: [],
|
||||
processors: [
|
||||
{ processorId: 'ThrowingProcessor' } as unknown as ProcessorConfig,
|
||||
],
|
||||
},
|
||||
]);
|
||||
const orchestrator = new PipelineOrchestrator(
|
||||
config,
|
||||
env,
|
||||
eventBus,
|
||||
env.tracer,
|
||||
registry,
|
||||
);
|
||||
|
||||
const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])];
|
||||
const state = createDummyState(false);
|
||||
|
||||
// It should not throw! It should swallow the error and return the unmodified array.
|
||||
const result = await orchestrator.executePipeline(
|
||||
'ThrowingPipe',
|
||||
episodes,
|
||||
state,
|
||||
);
|
||||
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result).toStrictEqual(episodes);
|
||||
});
|
||||
|
||||
it('automatically binds to budget_exceeded trigger via EventBus', () => {
|
||||
const config = createConfig([
|
||||
{
|
||||
name: 'PressureRelief',
|
||||
execution: 'background',
|
||||
triggers: ['budget_exceeded'],
|
||||
processors: [
|
||||
{ processorId: 'DummyAsyncProcessor' } as unknown as ProcessorConfig,
|
||||
],
|
||||
},
|
||||
]);
|
||||
|
||||
// Spy on the private method to see if the trigger fires it
|
||||
const executeSpy = vi.spyOn(
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
PipelineOrchestrator.prototype as any,
|
||||
'executePipelineAsync',
|
||||
);
|
||||
|
||||
new PipelineOrchestrator(config, env, eventBus, env.tracer, registry);
|
||||
|
||||
const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])];
|
||||
|
||||
// Emit the trigger
|
||||
eventBus.emitConsolidationNeeded({ episodes, targetDeficit: 100 });
|
||||
|
||||
expect(executeSpy).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,234 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Episode } from '../ir/types.js';
|
||||
import type { ContextProcessor, ContextAccountingState } from '../pipeline.js';
|
||||
import type { SidecarConfig, PipelineDef } from './types.js';
|
||||
import type {
|
||||
ContextEnvironment,
|
||||
ContextEventBus,
|
||||
ContextTracer,
|
||||
} from './environment.js';
|
||||
import type { ProcessorRegistry } from './registry.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { EpisodeEditor } from '../ir/episodeEditor.js';
|
||||
|
||||
export class PipelineOrchestrator {
|
||||
private activeTimers: NodeJS.Timeout[] = [];
|
||||
private readonly instantiatedProcessors = new Map<string, ContextProcessor>();
|
||||
|
||||
constructor(
|
||||
private readonly config: SidecarConfig,
|
||||
private readonly env: ContextEnvironment,
|
||||
private readonly eventBus: ContextEventBus,
|
||||
private readonly tracer: ContextTracer,
|
||||
private readonly registry: ProcessorRegistry,
|
||||
) {
|
||||
this.instantiateProcessors();
|
||||
this.registerTriggers();
|
||||
}
|
||||
|
||||
/**
|
||||
* Pre-loads and configures all processors defined in the sidecar config.
|
||||
*/
|
||||
private instantiateProcessors() {
|
||||
for (const pipeline of this.config.pipelines) {
|
||||
for (const procDef of pipeline.processors) {
|
||||
if (!this.instantiatedProcessors.has(procDef.processorId)) {
|
||||
const processorClass = this.registry.get(procDef.processorId);
|
||||
if (!processorClass) {
|
||||
throw new Error(
|
||||
`Context Processor [${procDef.processorId}] is not registered.`,
|
||||
);
|
||||
}
|
||||
// The Orchestrator injects standard dependencies required by processors
|
||||
// If a processor needs the eventBus (like Snapshot), it expects it via constructor.
|
||||
const instance = processorClass.create(
|
||||
this.env,
|
||||
procDef.options ?? {},
|
||||
);
|
||||
this.instantiatedProcessors.set(procDef.processorId, instance);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets up listeners for the triggers defined in the SidecarConfig.
|
||||
*/
|
||||
private registerTriggers() {
|
||||
for (const pipeline of this.config.pipelines) {
|
||||
for (const trigger of pipeline.triggers) {
|
||||
if (typeof trigger === 'object' && trigger.type === 'timer') {
|
||||
const timer = setInterval(() => {
|
||||
// For background timers, we need a way to get the latest state
|
||||
// But timers are generally disabled right now via the triggers config.
|
||||
// If needed, we will pass it via event bus.
|
||||
}, trigger.intervalMs);
|
||||
this.activeTimers.push(timer);
|
||||
} else if (trigger === 'budget_exceeded') {
|
||||
this.eventBus.onConsolidationNeeded((event) => {
|
||||
const state: ContextAccountingState = {
|
||||
currentTokens: 0,
|
||||
retainedTokens: this.config.budget.retainedTokens,
|
||||
maxTokens: this.config.budget.maxTokens,
|
||||
isBudgetSatisfied: false,
|
||||
deficitTokens: event.targetDeficit,
|
||||
protectedEpisodeIds: new Set(),
|
||||
};
|
||||
void this.executePipelineAsync(pipeline, event.episodes, state);
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
shutdown() {
|
||||
for (const timer of this.activeTimers) {
|
||||
clearInterval(timer);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Executes a pipeline asynchronously in the background. This is the "Eventual Consistency" path.
|
||||
* When the pipeline resolves, it emits a VariantReady event to cache the new graph.
|
||||
*/
|
||||
/**
|
||||
* Executes a pipeline based on its configured execution strategy ('blocking' or 'background').
|
||||
*/
|
||||
async executePipeline(
|
||||
pipelineName: string,
|
||||
episodes: Episode[],
|
||||
state: ContextAccountingState,
|
||||
): Promise<Episode[]> {
|
||||
const pipeline = this.config.pipelines.find((p) => p.name === pipelineName);
|
||||
if (!pipeline) return episodes;
|
||||
|
||||
if (pipeline.execution === 'background') {
|
||||
this.executePipelineAsync(pipeline, episodes, state).catch((e) => {
|
||||
debugLogger.error(`Background pipeline ${pipeline.name} failed:`, e);
|
||||
});
|
||||
return episodes; // Return immediately
|
||||
}
|
||||
|
||||
// Blocking execution
|
||||
this.tracer.logEvent(
|
||||
'Orchestrator',
|
||||
`Triggering synchronous pipeline: ${pipeline.name}`,
|
||||
);
|
||||
let currentEpisodes = [...episodes];
|
||||
for (let i = 0; i < pipeline.processors.length; i++) {
|
||||
const procDef = pipeline.processors[i];
|
||||
const processor = this.instantiatedProcessors.get(procDef.processorId);
|
||||
if (!processor) continue;
|
||||
|
||||
try {
|
||||
this.tracer.logEvent(
|
||||
'Orchestrator',
|
||||
`Executing processor: ${procDef.processorId}`,
|
||||
);
|
||||
const editor = new EpisodeEditor(currentEpisodes);
|
||||
await processor.process(editor, state);
|
||||
currentEpisodes = editor.getFinalEpisodes();
|
||||
} catch (error) {
|
||||
debugLogger.error(
|
||||
`Pipeline ${pipeline.name} failed synchronously at ${procDef.processorId}:`,
|
||||
error,
|
||||
);
|
||||
return currentEpisodes; // Return what we have so far
|
||||
}
|
||||
}
|
||||
|
||||
return currentEpisodes;
|
||||
}
|
||||
|
||||
/**
|
||||
* Internal method for running a pipeline entirely in the background.
|
||||
*/
|
||||
private async executePipelineAsync(
|
||||
pipeline: PipelineDef,
|
||||
currentState: Episode[],
|
||||
state: ContextAccountingState,
|
||||
) {
|
||||
this.tracer.logEvent(
|
||||
'Orchestrator',
|
||||
`Triggering async pipeline: ${pipeline.name}`,
|
||||
);
|
||||
if (!currentState || currentState.length === 0) return;
|
||||
|
||||
let currentEpisodes = [...currentState];
|
||||
|
||||
for (const procDef of pipeline.processors) {
|
||||
const processor = this.instantiatedProcessors.get(procDef.processorId);
|
||||
if (!processor) continue;
|
||||
|
||||
try {
|
||||
this.tracer.logEvent(
|
||||
'Orchestrator',
|
||||
`Executing processor: ${procDef.processorId} (async)`,
|
||||
);
|
||||
|
||||
const editor = new EpisodeEditor(currentEpisodes);
|
||||
await processor.process(editor, state);
|
||||
currentEpisodes = editor.getFinalEpisodes();
|
||||
|
||||
// Synthesize VariantReady events for anything that changed or was newly created
|
||||
for (const mutation of editor.getMutations()) {
|
||||
// We only broadcast modifications or replacements
|
||||
// (Insertions without replacement and deletions are not tracked as variants on an existing node)
|
||||
if (mutation.type === 'modified' || mutation.type === 'replaced') {
|
||||
const variantId = `v-${procDef.processorId.toLowerCase()}`;
|
||||
|
||||
let vType: 'snapshot' | 'summary' | 'masked' = 'masked';
|
||||
if (procDef.processorId.includes('Snapshot')) vType = 'snapshot';
|
||||
else if (procDef.processorId.includes('Semantic'))
|
||||
vType = 'summary';
|
||||
|
||||
const ep = mutation.episode!;
|
||||
let fallbackText = '';
|
||||
if (ep.yield?.text) fallbackText = ep.yield.text;
|
||||
else if (ep.trigger?.type === 'USER_PROMPT') {
|
||||
const firstPart = ep.trigger.semanticParts?.[0];
|
||||
if (firstPart) {
|
||||
fallbackText =
|
||||
firstPart.type === 'text'
|
||||
? firstPart.presentation?.text || firstPart.text
|
||||
: '';
|
||||
}
|
||||
}
|
||||
|
||||
this.eventBus.emitVariantReady({
|
||||
targetId:
|
||||
mutation.type === 'replaced' ? mutation.originalIds![0] : ep.id,
|
||||
variantId,
|
||||
variant:
|
||||
vType === 'snapshot'
|
||||
? {
|
||||
status: 'ready',
|
||||
type: 'snapshot',
|
||||
episode: ep,
|
||||
recoveredTokens: ep.yield?.metadata?.currentTokens || 10,
|
||||
replacedEpisodeIds: mutation.originalIds || [],
|
||||
}
|
||||
: {
|
||||
status: 'ready',
|
||||
type: vType,
|
||||
text: fallbackText,
|
||||
recoveredTokens: ep.yield?.metadata?.currentTokens || 10,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
debugLogger.error(
|
||||
`Pipeline ${pipeline.name} failed at ${procDef.processorId}:`,
|
||||
error,
|
||||
);
|
||||
return; // Halt pipeline
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { SidecarConfig } from './types.js';
|
||||
|
||||
/**
|
||||
* The standard default context management profile.
|
||||
* Optimized for safety, precision, and reliable summarization.
|
||||
*/
|
||||
export const defaultSidecarProfile: SidecarConfig = {
|
||||
budget: {
|
||||
retainedTokens: 65000,
|
||||
maxTokens: 150000,
|
||||
},
|
||||
gcBackstop: {
|
||||
strategy: 'truncate',
|
||||
target: 'incremental',
|
||||
freeTokensTarget: 10000,
|
||||
},
|
||||
pipelines: [
|
||||
{
|
||||
name: 'Immediate Sanitization',
|
||||
triggers: ['on_turn'],
|
||||
execution: 'blocking',
|
||||
processors: [
|
||||
{
|
||||
processorId: 'ToolMaskingProcessor',
|
||||
options: { stringLengthThresholdTokens: 8000 },
|
||||
},
|
||||
{ processorId: 'BlobDegradationProcessor', options: {} },
|
||||
{
|
||||
processorId: 'SemanticCompressionProcessor',
|
||||
options: { nodeThresholdTokens: 5000 },
|
||||
},
|
||||
{ processorId: 'EmergencyTruncationProcessor', options: {} },
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'Deep Background Compression',
|
||||
triggers: [{ type: 'timer', intervalMs: 5000 }, 'budget_exceeded'],
|
||||
execution: 'background',
|
||||
processors: [
|
||||
{
|
||||
processorId: 'HistorySquashingProcessor',
|
||||
options: { maxTokensPerNode: 3000 },
|
||||
},
|
||||
{ processorId: 'StateSnapshotProcessor', options: {} },
|
||||
],
|
||||
},
|
||||
],
|
||||
};
|
||||
@@ -0,0 +1,47 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextProcessor } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from './environment.js';
|
||||
|
||||
export interface ContextProcessorDef<TOptions = object> {
|
||||
readonly id: string;
|
||||
readonly schema: object;
|
||||
create(env: ContextEnvironment, options: TOptions): ContextProcessor;
|
||||
}
|
||||
|
||||
/**
|
||||
* Registry for mapping declarative sidecar configs to running Processor instances.
|
||||
*/
|
||||
export class ProcessorRegistry {
|
||||
private processors = new Map<string, ContextProcessorDef<unknown>>();
|
||||
|
||||
register<TOptions>(def: ContextProcessorDef<TOptions>) {
|
||||
this.processors.set(def.id, def);
|
||||
}
|
||||
|
||||
get(id: string): ContextProcessorDef {
|
||||
const def = this.processors.get(id);
|
||||
if (!def) {
|
||||
throw new Error(`Context Processor [${id}] is not registered.`);
|
||||
}
|
||||
return def;
|
||||
}
|
||||
|
||||
getSchemas(): object[] {
|
||||
const schemas: object[] = [];
|
||||
for (const def of this.processors.values()) {
|
||||
if (def.schema) {
|
||||
schemas.push(def.schema);
|
||||
}
|
||||
}
|
||||
return schemas;
|
||||
}
|
||||
|
||||
clear() {
|
||||
this.processors.clear();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ProcessorRegistry } from './registry.js';
|
||||
import './builtins.js';
|
||||
|
||||
export function getSidecarConfigSchema(registry: ProcessorRegistry) {
|
||||
return {
|
||||
$schema: 'http://json-schema.org/draft-07/schema#',
|
||||
title: 'SidecarConfig',
|
||||
description: 'The Data-Driven Schema for the Context Manager.',
|
||||
type: 'object',
|
||||
required: ['budget', 'gcBackstop', 'pipelines'],
|
||||
properties: {
|
||||
budget: {
|
||||
type: 'object',
|
||||
description: 'Defines the token ceilings and limits for the pipeline.',
|
||||
required: ['retainedTokens', 'maxTokens'],
|
||||
properties: {
|
||||
retainedTokens: {
|
||||
type: 'number',
|
||||
description:
|
||||
'The ideal token count the pipeline tries to shrink down to.',
|
||||
},
|
||||
maxTokens: {
|
||||
type: 'number',
|
||||
description:
|
||||
'The absolute maximum token count allowed before synchronous truncation kicks in.',
|
||||
},
|
||||
},
|
||||
},
|
||||
gcBackstop: {
|
||||
type: 'object',
|
||||
description:
|
||||
"Defines what happens when the pipeline fails to compress under 'maxTokens'",
|
||||
required: ['strategy', 'target'],
|
||||
properties: {
|
||||
strategy: {
|
||||
type: 'string',
|
||||
enum: ['truncate', 'compress', 'rollingSummarizer'],
|
||||
},
|
||||
target: {
|
||||
type: 'string',
|
||||
enum: ['incremental', 'freeNTokens', 'max'],
|
||||
},
|
||||
freeTokensTarget: {
|
||||
type: 'number',
|
||||
},
|
||||
},
|
||||
},
|
||||
pipelines: {
|
||||
type: 'array',
|
||||
description: 'The execution graphs for context manipulation.',
|
||||
items: {
|
||||
type: 'object',
|
||||
required: ['name', 'triggers', 'execution', 'processors'],
|
||||
properties: {
|
||||
name: {
|
||||
type: 'string',
|
||||
},
|
||||
triggers: {
|
||||
type: 'array',
|
||||
items: {
|
||||
anyOf: [
|
||||
{
|
||||
type: 'string',
|
||||
enum: ['on_turn', 'post_turn', 'budget_exceeded'],
|
||||
},
|
||||
{
|
||||
type: 'object',
|
||||
required: ['type', 'intervalMs'],
|
||||
properties: {
|
||||
type: {
|
||||
type: 'string',
|
||||
const: 'timer',
|
||||
},
|
||||
intervalMs: {
|
||||
type: 'number',
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
execution: {
|
||||
type: 'string',
|
||||
enum: ['blocking', 'background'],
|
||||
},
|
||||
processors: {
|
||||
type: 'array',
|
||||
items: {
|
||||
oneOf: registry.getSchemas(),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { StateSnapshotProcessorOptions } from '../processors/stateSnapshotProcessor.js';
|
||||
|
||||
/**
|
||||
* Definition of a processor or worker to be instantiated in the graph.
|
||||
*/
|
||||
export type ProcessorConfig =
|
||||
| {
|
||||
processorId: 'ToolMaskingProcessor';
|
||||
options: { stringLengthThresholdTokens: number };
|
||||
}
|
||||
| { processorId: 'BlobDegradationProcessor'; options?: object }
|
||||
| {
|
||||
processorId: 'SemanticCompressionProcessor';
|
||||
options: { nodeThresholdTokens: number };
|
||||
}
|
||||
| {
|
||||
processorId: 'HistorySquashingProcessor';
|
||||
options: { maxTokensPerNode: number };
|
||||
}
|
||||
| {
|
||||
processorId: 'StateSnapshotProcessor';
|
||||
options: StateSnapshotProcessorOptions;
|
||||
}
|
||||
| {
|
||||
processorId: 'EmergencyTruncationProcessor';
|
||||
options?: Record<string, unknown>;
|
||||
};
|
||||
|
||||
export type PipelineTrigger =
|
||||
| 'on_turn'
|
||||
| 'post_turn'
|
||||
| 'budget_exceeded'
|
||||
| { type: 'timer'; intervalMs: number };
|
||||
|
||||
export interface PipelineDef {
|
||||
name: string;
|
||||
triggers: PipelineTrigger[];
|
||||
execution: 'blocking' | 'background';
|
||||
processors: ProcessorConfig[];
|
||||
}
|
||||
|
||||
/**
|
||||
* The Data-Driven Schema for the Context Manager.
|
||||
*/
|
||||
export interface SidecarConfig {
|
||||
/** Defines the token ceilings and limits for the pipeline. */
|
||||
budget: {
|
||||
retainedTokens: number;
|
||||
maxTokens: number;
|
||||
};
|
||||
|
||||
/** Defines what happens when the pipeline fails to compress under 'maxTokens' */
|
||||
gcBackstop: {
|
||||
strategy: 'truncate' | 'compress' | 'rollingSummarizer';
|
||||
target: 'incremental' | 'freeNTokens' | 'max';
|
||||
freeTokensTarget?: number;
|
||||
};
|
||||
|
||||
/** The execution graphs for context manipulation */
|
||||
pipelines: PipelineDef[];
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { ContextManager } from '../contextManager.js';
|
||||
import { AgentChatHistory } from '../../core/agentChatHistory.js';
|
||||
import type { Content } from '@google/genai';
|
||||
import type { SidecarConfig } from '../sidecar/types.js';
|
||||
import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js';
|
||||
import { ContextTracer } from '../tracer.js';
|
||||
import { ContextEventBus } from '../eventBus.js';
|
||||
import { PipelineOrchestrator } from '../sidecar/orchestrator.js';
|
||||
import { registerBuiltInProcessors } from '../sidecar/builtins.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { ProcessorRegistry } from '../sidecar/registry.js';
|
||||
import { DeterministicIdGenerator } from '../system/DeterministicIdGenerator.js';
|
||||
import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
|
||||
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
|
||||
export interface TurnSummary {
|
||||
turnIndex: number;
|
||||
tokensBeforeBackground: number;
|
||||
tokensAfterBackground: number;
|
||||
}
|
||||
|
||||
export class SimulationHarness {
|
||||
readonly chatHistory: AgentChatHistory;
|
||||
contextManager!: ContextManager;
|
||||
env!: ContextEnvironmentImpl;
|
||||
orchestrator!: PipelineOrchestrator;
|
||||
readonly eventBus: ContextEventBus;
|
||||
config!: SidecarConfig;
|
||||
private tracer!: ContextTracer;
|
||||
private currentTurnIndex = 0;
|
||||
private tokenTrajectory: TurnSummary[] = [];
|
||||
|
||||
static async create(
|
||||
config: SidecarConfig,
|
||||
mockLlmClient: BaseLlmClient,
|
||||
mockTempDir = '/tmp/sim',
|
||||
): Promise<SimulationHarness> {
|
||||
const harness = new SimulationHarness();
|
||||
await harness.init(config, mockLlmClient, mockTempDir);
|
||||
return harness;
|
||||
}
|
||||
|
||||
private constructor() {
|
||||
this.chatHistory = new AgentChatHistory();
|
||||
this.eventBus = new ContextEventBus();
|
||||
}
|
||||
|
||||
private async init(
|
||||
config: SidecarConfig,
|
||||
mockLlmClient: BaseLlmClient,
|
||||
mockTempDir: string,
|
||||
) {
|
||||
this.config = config;
|
||||
const registry = new ProcessorRegistry();
|
||||
// Register all standard processors
|
||||
registerBuiltInProcessors(registry);
|
||||
|
||||
this.tracer = new ContextTracer({
|
||||
targetDir: mockTempDir,
|
||||
sessionId: 'sim-session',
|
||||
});
|
||||
this.env = new ContextEnvironmentImpl(
|
||||
mockLlmClient,
|
||||
'sim-prompt',
|
||||
'sim-session',
|
||||
mockTempDir,
|
||||
mockTempDir,
|
||||
this.tracer,
|
||||
4, // 4 chars per token average
|
||||
this.eventBus,
|
||||
new InMemoryFileSystem(),
|
||||
new DeterministicIdGenerator(),
|
||||
);
|
||||
|
||||
this.orchestrator = new PipelineOrchestrator(
|
||||
config,
|
||||
this.env,
|
||||
this.eventBus,
|
||||
this.tracer,
|
||||
registry,
|
||||
);
|
||||
this.contextManager = ContextManager.create(
|
||||
config,
|
||||
this.env,
|
||||
this.tracer,
|
||||
this.orchestrator,
|
||||
registry,
|
||||
);
|
||||
this.contextManager.subscribeToHistory(this.chatHistory);
|
||||
}
|
||||
|
||||
/**
|
||||
* Simulates a single "Turn" (User input + Model/Tool outputs)
|
||||
* A turn might consist of multiple Content messages (e.g. user prompt -> model call -> user response -> model answer)
|
||||
*/
|
||||
async simulateTurn(messages: Content[]) {
|
||||
// 1. Append the new messages
|
||||
const currentHistory = this.chatHistory.get();
|
||||
this.chatHistory.set([...currentHistory, ...messages]);
|
||||
|
||||
// 2. Measure tokens immediately after append (Before background processing)
|
||||
const tokensBefore = this.env.tokenCalculator.calculateEpisodeListTokens(
|
||||
this.contextManager.getWorkingBufferView(),
|
||||
);
|
||||
debugLogger.log(
|
||||
`[Turn ${this.currentTurnIndex}] Tokens BEFORE: ${tokensBefore}`,
|
||||
);
|
||||
|
||||
// 3. Yield to event loop to allow internal async subscribers and orchestrator to finish
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
|
||||
// 3.1 Simulate what projectCompressedHistory does with the sync handlers
|
||||
let currentView = this.contextManager.getWorkingBufferView();
|
||||
const currentTokens =
|
||||
this.env.tokenCalculator.calculateEpisodeListTokens(currentView);
|
||||
if (this.config.budget && currentTokens > this.config.budget.maxTokens) {
|
||||
debugLogger.log(
|
||||
`[Turn ${this.currentTurnIndex}] Sync panic triggered! ${currentTokens} > ${this.config.budget.maxTokens}`,
|
||||
);
|
||||
const syncPipelines = this.config.pipelines.filter(
|
||||
(p) => p.execution === 'blocking',
|
||||
);
|
||||
const orchestrator = this.orchestrator;
|
||||
for (const pipe of syncPipelines) {
|
||||
await orchestrator.executePipeline(pipe.name, currentView, {
|
||||
currentTokens,
|
||||
maxTokens: this.config.budget.maxTokens,
|
||||
retainedTokens: this.config.budget.retainedTokens,
|
||||
isBudgetSatisfied: false,
|
||||
deficitTokens: currentTokens - this.config.budget.maxTokens,
|
||||
protectedEpisodeIds: new Set(),
|
||||
});
|
||||
currentView = this.contextManager.getWorkingBufferView();
|
||||
}
|
||||
|
||||
// Inject the truncated view back into the graph
|
||||
for (let i = 0; i < currentView.length; i++) {
|
||||
const ep = currentView[i];
|
||||
if (
|
||||
!this.contextManager
|
||||
.getWorkingBufferView()
|
||||
.find((c) => c.id === ep.id)
|
||||
) {
|
||||
this.eventBus.emitVariantReady({
|
||||
targetId: ep.id,
|
||||
variantId: 'v-emergency',
|
||||
variant: {
|
||||
status: 'ready',
|
||||
type: 'masked', // Truncation is technically a mask
|
||||
text: ep.yield?.text || '',
|
||||
recoveredTokens: 0,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
// Wait for variant propagation
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
}
|
||||
|
||||
// 4. Measure tokens after background processors have (hopefully) emitted variants
|
||||
const tokensAfter = this.env.tokenCalculator.calculateEpisodeListTokens(
|
||||
this.contextManager.getWorkingBufferView(),
|
||||
);
|
||||
debugLogger.log(
|
||||
`[Turn ${this.currentTurnIndex}] Tokens AFTER: ${tokensAfter}`,
|
||||
);
|
||||
|
||||
this.tokenTrajectory.push({
|
||||
turnIndex: this.currentTurnIndex++,
|
||||
tokensBeforeBackground: tokensBefore,
|
||||
tokensAfterBackground: tokensAfter,
|
||||
});
|
||||
}
|
||||
|
||||
async getGoldenState() {
|
||||
const finalProjection =
|
||||
await this.contextManager.projectCompressedHistory();
|
||||
return {
|
||||
tokenTrajectory: this.tokenTrajectory,
|
||||
finalProjection,
|
||||
};
|
||||
}
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,146 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeAll, afterAll } from 'vitest';
|
||||
import { SimulationHarness } from './SimulationHarness.js';
|
||||
import type { SidecarConfig } from '../sidecar/types.js';
|
||||
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
|
||||
expect.addSnapshotSerializer({
|
||||
test: (val) =>
|
||||
typeof val === 'string' &&
|
||||
(/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i.test(
|
||||
val,
|
||||
) ||
|
||||
/^\/tmp\/sim/.test(val)), // Mask temp directories and UUIDs
|
||||
print: (val) =>
|
||||
typeof val === 'string' && /^\/tmp\/sim/.test(val)
|
||||
? '"<MOCKED_DIR>"'
|
||||
: '"<UUID>"',
|
||||
});
|
||||
|
||||
describe('System Lifecycle Golden Tests', () => {
|
||||
beforeAll(() => {
|
||||
vi.spyOn(Math, 'random').mockReturnValue(0.5);
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
const getAggressiveConfig = (): SidecarConfig => ({
|
||||
budget: { maxTokens: 4000, retainedTokens: 2000 }, // Extremely tight limits
|
||||
gcBackstop: { strategy: 'truncate', target: 'max' },
|
||||
pipelines: [
|
||||
{
|
||||
name: 'Pressure Relief', // Emits from eventBus 'budget_exceeded'
|
||||
execution: 'background',
|
||||
triggers: ['budget_exceeded'],
|
||||
processors: [
|
||||
{ processorId: 'BlobDegradationProcessor' },
|
||||
{
|
||||
processorId: 'ToolMaskingProcessor',
|
||||
options: { stringLengthThresholdTokens: 50 },
|
||||
}, // Mask any tool string > 200 chars
|
||||
{ processorId: 'StateSnapshotProcessor', options: {} }, // Squash old history
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'Immediate Sanitization', // The magic string the projector is hardcoded to use
|
||||
execution: 'blocking',
|
||||
triggers: ['budget_exceeded'],
|
||||
processors: [
|
||||
{ processorId: 'EmergencyTruncationProcessor', options: {} },
|
||||
],
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const mockLlmClient = {
|
||||
generateContent: vi.fn().mockResolvedValue({
|
||||
text: '<MOCKED_STATE_SNAPSHOT_SUMMARY>',
|
||||
}),
|
||||
} as unknown as BaseLlmClient;
|
||||
|
||||
it('Scenario 1: Organic Growth with Huge Tool Output & Images', async () => {
|
||||
const harness = await SimulationHarness.create(
|
||||
getAggressiveConfig(),
|
||||
mockLlmClient,
|
||||
);
|
||||
|
||||
// Turn 0: System Prompt
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'System Instructions' }] },
|
||||
{ role: 'model', parts: [{ text: 'Ack.' }] },
|
||||
]);
|
||||
|
||||
// Turn 1: Normal conversation
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'Hello!' }] },
|
||||
{ role: 'model', parts: [{ text: 'Hi, how can I help?' }] },
|
||||
]);
|
||||
|
||||
// Turn 2: Massive Tool Output (Should trigger ToolMaskingProcessor in background)
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'Read the logs.' }] },
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{
|
||||
functionCall: {
|
||||
name: 'run_shell_command',
|
||||
args: { cmd: 'cat server.log' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'run_shell_command',
|
||||
response: { output: 'LOG '.repeat(5000) },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{ role: 'model', parts: [{ text: 'The logs are very long.' }] },
|
||||
]);
|
||||
|
||||
// Turn 3: Multi-modal blob (Should trigger BlobDegradationProcessor)
|
||||
await harness.simulateTurn([
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{ text: 'Look at this architecture diagram:' },
|
||||
{
|
||||
inlineData: {
|
||||
mimeType: 'image/png',
|
||||
data: 'fake_base64_data_'.repeat(1000),
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{ role: 'model', parts: [{ text: 'Nice diagram.' }] },
|
||||
]);
|
||||
|
||||
// Turn 4: More conversation to trigger StateSnapshot
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'Can we refactor?' }] },
|
||||
{ role: 'model', parts: [{ text: 'Yes we can.' }] },
|
||||
]);
|
||||
|
||||
// Get final state
|
||||
const goldenState = await harness.getGoldenState();
|
||||
|
||||
// In a perfectly functioning opportunistic system, the token trajectory should show
|
||||
// the massive spikes in Turn 2 and 3 being immediately resolved by the background tasks.
|
||||
// The final projection should fit neatly under the Max Tokens limit.
|
||||
|
||||
expect(goldenState).toMatchSnapshot();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,18 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { IIdGenerator } from './IIdGenerator.js';
|
||||
|
||||
export class DeterministicIdGenerator implements IIdGenerator {
|
||||
private counter = 0;
|
||||
|
||||
constructor(private prefix: string = 'id-') {}
|
||||
|
||||
generateId(): string {
|
||||
this.counter++;
|
||||
return `${this.prefix}${this.counter}`;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
export interface IFileSystem {
|
||||
existsSync(path: string): boolean;
|
||||
statSyncSize(path: string): number;
|
||||
readFileSync(path: string, encoding: 'utf8'): string;
|
||||
writeFileSync(path: string, data: string | Buffer, encoding?: 'utf-8'): void;
|
||||
appendFileSync(path: string, data: string, encoding: 'utf-8'): void;
|
||||
mkdirSync(path: string, options?: { recursive?: boolean }): void;
|
||||
|
||||
writeFile(path: string, data: string | Buffer): Promise<void>;
|
||||
mkdir(path: string, options?: { recursive?: boolean }): Promise<void>;
|
||||
|
||||
join(...paths: string[]): string;
|
||||
dirname(path: string): string;
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
export interface IIdGenerator {
|
||||
generateId(): string;
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { IFileSystem } from './IFileSystem.js';
|
||||
|
||||
export class InMemoryFileSystem implements IFileSystem {
|
||||
private files = new Map<string, string | Buffer>();
|
||||
|
||||
getFiles(): ReadonlyMap<string, string | Buffer> {
|
||||
return this.files;
|
||||
}
|
||||
|
||||
setFile(path: string, content: string | Buffer) {
|
||||
this.files.set(this.normalize(path), content);
|
||||
}
|
||||
|
||||
private normalize(p: string): string {
|
||||
return p.replace(/\/+/g, '/');
|
||||
}
|
||||
|
||||
existsSync(p: string): boolean {
|
||||
return this.files.has(this.normalize(p));
|
||||
}
|
||||
|
||||
statSyncSize(p: string): number {
|
||||
const content = this.files.get(this.normalize(p));
|
||||
if (content === undefined) {
|
||||
throw new Error(`ENOENT: no such file or directory, stat '${p}'`);
|
||||
}
|
||||
return Buffer.isBuffer(content)
|
||||
? content.byteLength
|
||||
: Buffer.byteLength(content, 'utf8');
|
||||
}
|
||||
|
||||
readFileSync(p: string, encoding: 'utf8'): string {
|
||||
const content = this.files.get(this.normalize(p));
|
||||
if (content === undefined) {
|
||||
throw new Error(`ENOENT: no such file or directory, open '${p}'`);
|
||||
}
|
||||
if (Buffer.isBuffer(content)) {
|
||||
return content.toString(encoding);
|
||||
}
|
||||
return content;
|
||||
}
|
||||
|
||||
writeFileSync(p: string, data: string | Buffer, _encoding?: 'utf-8'): void {
|
||||
this.files.set(this.normalize(p), data);
|
||||
}
|
||||
|
||||
appendFileSync(p: string, data: string, _encoding: 'utf-8'): void {
|
||||
const norm = this.normalize(p);
|
||||
const existing = this.files.get(norm) || '';
|
||||
const existingStr = Buffer.isBuffer(existing)
|
||||
? existing.toString('utf8')
|
||||
: existing;
|
||||
this.files.set(norm, existingStr + data);
|
||||
}
|
||||
|
||||
mkdirSync(_p: string, _options?: { recursive?: boolean }): void {}
|
||||
|
||||
async writeFile(p: string, data: string | Buffer): Promise<void> {
|
||||
this.writeFileSync(p, data);
|
||||
}
|
||||
|
||||
async mkdir(_p: string, _options?: { recursive?: boolean }): Promise<void> {}
|
||||
|
||||
join(...paths: string[]): string {
|
||||
return this.normalize(paths.join('/'));
|
||||
}
|
||||
|
||||
dirname(p: string): string {
|
||||
const parts = this.normalize(p).split('/');
|
||||
parts.pop();
|
||||
return parts.length === 0 ? '.' : parts.join('/') || '/';
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import * as fs from 'node:fs';
|
||||
import * as fsPromises from 'node:fs/promises';
|
||||
import * as path from 'node:path';
|
||||
import type { IFileSystem } from './IFileSystem.js';
|
||||
|
||||
export class NodeFileSystem implements IFileSystem {
|
||||
existsSync(p: string): boolean {
|
||||
return fs.existsSync(p);
|
||||
}
|
||||
|
||||
statSyncSize(p: string): number {
|
||||
return fs.statSync(p).size;
|
||||
}
|
||||
|
||||
readFileSync(p: string, encoding: 'utf8'): string {
|
||||
return fs.readFileSync(p, encoding);
|
||||
}
|
||||
|
||||
writeFileSync(p: string, data: string | Buffer, encoding?: 'utf-8'): void {
|
||||
if (Buffer.isBuffer(data)) {
|
||||
fs.writeFileSync(p, data);
|
||||
} else {
|
||||
fs.writeFileSync(p, data, encoding);
|
||||
}
|
||||
}
|
||||
|
||||
appendFileSync(p: string, data: string, encoding: 'utf-8'): void {
|
||||
fs.appendFileSync(p, data, encoding);
|
||||
}
|
||||
|
||||
mkdirSync(p: string, options?: { recursive?: boolean }): void {
|
||||
fs.mkdirSync(p, options);
|
||||
}
|
||||
|
||||
async writeFile(p: string, data: string | Buffer): Promise<void> {
|
||||
await fsPromises.writeFile(p, data);
|
||||
}
|
||||
|
||||
async mkdir(p: string, options?: { recursive?: boolean }): Promise<void> {
|
||||
await fsPromises.mkdir(p, options);
|
||||
}
|
||||
|
||||
join(...paths: string[]): string {
|
||||
return path.join(...paths);
|
||||
}
|
||||
|
||||
dirname(p: string): string {
|
||||
return path.dirname(p);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import type { IIdGenerator } from './IIdGenerator.js';
|
||||
|
||||
export class NodeIdGenerator implements IIdGenerator {
|
||||
generateId(): string {
|
||||
return randomUUID();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { vi } from 'vitest';
|
||||
import type { Config } from '../../config/config.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import type { Content } from '@google/genai';
|
||||
import { AgentChatHistory } from '../../core/agentChatHistory.js';
|
||||
import { ContextManager } from '../contextManager.js';
|
||||
import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
|
||||
import { DeterministicIdGenerator } from '../system/DeterministicIdGenerator.js';
|
||||
import type {
|
||||
Episode,
|
||||
UserPrompt,
|
||||
SystemEvent,
|
||||
SemanticPart,
|
||||
} from '../ir/types.js';
|
||||
import type { ContextAccountingState } from '../pipeline.js';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
|
||||
export function createDummyState(
|
||||
isSatisfied = false,
|
||||
deficit = 0,
|
||||
protectedIds = new Set<string>(),
|
||||
currentTokens = 5000,
|
||||
maxTokens = 10000,
|
||||
retainedTokens = 4000,
|
||||
): ContextAccountingState {
|
||||
return {
|
||||
currentTokens,
|
||||
maxTokens,
|
||||
retainedTokens,
|
||||
deficitTokens: deficit,
|
||||
protectedEpisodeIds: protectedIds,
|
||||
isBudgetSatisfied: isSatisfied,
|
||||
};
|
||||
}
|
||||
|
||||
export function createDummyEpisode(
|
||||
id: string,
|
||||
type: 'USER_PROMPT' | 'SYSTEM_EVENT',
|
||||
parts: SemanticPart[] = [],
|
||||
toolSteps: Array<{
|
||||
intent: Record<string, unknown>;
|
||||
observation: Record<string, unknown>;
|
||||
toolName?: string;
|
||||
tokens?: { intent: number; observation: number };
|
||||
}> = [],
|
||||
): Episode {
|
||||
let trigger: UserPrompt | SystemEvent;
|
||||
|
||||
if (type === 'USER_PROMPT') {
|
||||
trigger = {
|
||||
id: randomUUID(),
|
||||
type: 'USER_PROMPT',
|
||||
semanticParts: parts,
|
||||
metadata: {
|
||||
originalTokens: 100,
|
||||
currentTokens: 100,
|
||||
transformations: [],
|
||||
},
|
||||
};
|
||||
} else {
|
||||
trigger = {
|
||||
id: randomUUID(),
|
||||
type: 'SYSTEM_EVENT',
|
||||
name: 'dummy_event',
|
||||
payload: {},
|
||||
metadata: {
|
||||
originalTokens: 100,
|
||||
currentTokens: 100,
|
||||
transformations: [],
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
id,
|
||||
timestamp: Date.now(),
|
||||
trigger,
|
||||
steps: toolSteps.map((step) => ({
|
||||
id: randomUUID(),
|
||||
type: 'TOOL_EXECUTION',
|
||||
toolName: step.toolName || 'test_tool',
|
||||
intent: step.intent,
|
||||
observation: step.observation,
|
||||
tokens: step.tokens || { intent: 50, observation: 50 },
|
||||
metadata: {
|
||||
originalTokens: 100,
|
||||
currentTokens: 100,
|
||||
transformations: [],
|
||||
},
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
export function createMockEnvironment(): ContextEnvironment {
|
||||
return {
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
llmClient: vi.fn().mockReturnValue({
|
||||
generateContent: vi.fn().mockResolvedValue({
|
||||
text: 'Mock LLM summary response',
|
||||
}),
|
||||
})() as unknown as BaseLlmClient,
|
||||
promptId: 'mock-prompt-id',
|
||||
sessionId: 'mock-session',
|
||||
traceDir: '/tmp/.gemini/trace',
|
||||
projectTempDir: '/tmp/.gemini/tool-outputs',
|
||||
eventBus: new ContextEventBus(),
|
||||
tracer: new ContextTracer({ targetDir: '/tmp', sessionId: 'mock-session' }),
|
||||
charsPerToken: 1,
|
||||
tokenCalculator: new ContextTokenCalculator(1),
|
||||
fileSystem: new InMemoryFileSystem(),
|
||||
idGenerator: new DeterministicIdGenerator('mock-uuid-'),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a block of synthetic conversation history designed to consume a specific number of tokens.
|
||||
* Assumes roughly 4 characters per token for standard English text.
|
||||
*/
|
||||
export function createSyntheticHistory(
|
||||
numTurns: number,
|
||||
tokensPerTurn: number,
|
||||
): Content[] {
|
||||
const history: Content[] = [];
|
||||
const charsPerTurn = tokensPerTurn * 1;
|
||||
|
||||
for (let i = 0; i < numTurns; i++) {
|
||||
history.push({
|
||||
role: 'user',
|
||||
parts: [{ text: `User turn ${i}. ` + 'A'.repeat(charsPerTurn) }],
|
||||
});
|
||||
history.push({
|
||||
role: 'model',
|
||||
parts: [{ text: `Model response ${i}. ` + 'B'.repeat(charsPerTurn) }],
|
||||
});
|
||||
}
|
||||
|
||||
return history;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a fully mocked Config object tailored for Context Component testing.
|
||||
*/
|
||||
export function createMockContextConfig(
|
||||
overrides?: Record<string, unknown>,
|
||||
llmClientOverride?: unknown,
|
||||
): Config {
|
||||
const defaultConfig = {
|
||||
isContextManagementEnabled: vi.fn().mockReturnValue(true),
|
||||
storage: {
|
||||
getProjectTempDir: vi.fn().mockReturnValue('/tmp/gemini-test'),
|
||||
},
|
||||
getBaseLlmClient: vi.fn().mockReturnValue(
|
||||
llmClientOverride || {
|
||||
generateContent: vi.fn().mockResolvedValue({
|
||||
text: '<mocked_snapshot>Synthesized state</mocked_snapshot>',
|
||||
}),
|
||||
},
|
||||
),
|
||||
getUsageStatisticsEnabled: vi.fn().mockReturnValue(false),
|
||||
getTargetDir: vi.fn().mockReturnValue('/tmp'),
|
||||
getSessionId: vi.fn().mockReturnValue('test-session'),
|
||||
getExperimentalContextSidecarConfig: vi.fn().mockReturnValue(undefined),
|
||||
};
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
return { ...defaultConfig, ...overrides } as unknown as Config;
|
||||
}
|
||||
|
||||
/**
|
||||
* Wires up a full ContextManager component with an AgentChatHistory and active background workers.
|
||||
*/
|
||||
import { ContextTracer } from '../tracer.js';
|
||||
import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js';
|
||||
import { SidecarLoader } from '../sidecar/SidecarLoader.js';
|
||||
import { ContextEventBus } from '../eventBus.js';
|
||||
import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
|
||||
import type { BaseLlmClient } from 'src/core/baseLlmClient.js';
|
||||
import { ProcessorRegistry } from '../sidecar/registry.js';
|
||||
import { registerBuiltInProcessors } from '../sidecar/builtins.js';
|
||||
|
||||
export function setupContextComponentTest(config: Config) {
|
||||
const chatHistory = new AgentChatHistory();
|
||||
const registry = new ProcessorRegistry();
|
||||
registerBuiltInProcessors(registry);
|
||||
const sidecar = SidecarLoader.fromConfig(config, registry);
|
||||
const tracer = new ContextTracer({
|
||||
targetDir: '/tmp',
|
||||
sessionId: 'test-session',
|
||||
});
|
||||
const eventBus = new ContextEventBus();
|
||||
const env = new ContextEnvironmentImpl(
|
||||
config.getBaseLlmClient(),
|
||||
'test prompt-id',
|
||||
'test-session',
|
||||
'/tmp',
|
||||
'/tmp/gemini-test',
|
||||
tracer,
|
||||
1,
|
||||
eventBus,
|
||||
);
|
||||
const contextManager = ContextManager.create(
|
||||
sidecar,
|
||||
env,
|
||||
tracer,
|
||||
undefined,
|
||||
registry,
|
||||
);
|
||||
|
||||
// The async worker is now internally managed by ContextManager
|
||||
|
||||
// Subscribe to history to enable the Eager/Opportunistic triggers
|
||||
contextManager.subscribeToHistory(chatHistory);
|
||||
|
||||
return { chatHistory, contextManager };
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { ContextTracer } from './tracer.js';
|
||||
import { InMemoryFileSystem } from './system/InMemoryFileSystem.js';
|
||||
import { DeterministicIdGenerator } from './system/DeterministicIdGenerator.js';
|
||||
|
||||
describe('ContextTracer (Fake FS & ID Gen)', () => {
|
||||
let fileSystem: InMemoryFileSystem;
|
||||
let idGenerator: DeterministicIdGenerator;
|
||||
|
||||
beforeEach(() => {
|
||||
fileSystem = new InMemoryFileSystem();
|
||||
idGenerator = new DeterministicIdGenerator('mock-uuid-');
|
||||
|
||||
// We must mock Date.now() to ensure asset file names are perfectly deterministic
|
||||
vi.useFakeTimers();
|
||||
vi.setSystemTime(new Date('2026-01-01T12:00:00Z'));
|
||||
});
|
||||
|
||||
it('initializes, logs events, and auto-saves large assets deterministically', () => {
|
||||
const tracer = new ContextTracer(
|
||||
{ enabled: true, targetDir: '/fake/target', sessionId: 'test-session' },
|
||||
fileSystem,
|
||||
idGenerator,
|
||||
);
|
||||
|
||||
// Verify Initialization
|
||||
const initTraceLog = fileSystem.readFileSync(
|
||||
'/fake/target/.gemini/context_trace/test-session/trace.log',
|
||||
'utf8',
|
||||
);
|
||||
expect(initTraceLog).toContain('[SYSTEM] Context Tracer Initialized');
|
||||
|
||||
// Small logging: shouldn't trigger saveAsset
|
||||
tracer.logEvent('TestComponent', 'TestAction', { key: 'value' });
|
||||
|
||||
const smallTraceLog = fileSystem.readFileSync(
|
||||
'/fake/target/.gemini/context_trace/test-session/trace.log',
|
||||
'utf8',
|
||||
);
|
||||
expect(smallTraceLog).toContain('[TestComponent] TestAction');
|
||||
expect(smallTraceLog).toContain('{"key":"value"}');
|
||||
|
||||
// Large logging: should trigger auto-asset save
|
||||
const hugeString = 'a'.repeat(2000);
|
||||
tracer.logEvent('TestComponent', 'LargeAction', { largeKey: hugeString });
|
||||
|
||||
// 1767268800000 is 2026-01-01T12:00:00Z
|
||||
const expectedAssetPath =
|
||||
'/fake/target/.gemini/context_trace/test-session/assets/1767268800000-mock-uuid-1-largeKey.json';
|
||||
|
||||
// Assert asset was written to FS
|
||||
expect(fileSystem.existsSync(expectedAssetPath)).toBe(true);
|
||||
|
||||
const largeTraceLog = fileSystem.readFileSync(
|
||||
'/fake/target/.gemini/context_trace/test-session/trace.log',
|
||||
'utf8',
|
||||
);
|
||||
expect(largeTraceLog).toContain('[TestComponent] LargeAction');
|
||||
expect(largeTraceLog).toContain(
|
||||
`{"largeKey":{"$asset":"1767268800000-mock-uuid-1-largeKey.json"}}`,
|
||||
);
|
||||
});
|
||||
|
||||
it('silently ignores logging when disabled', () => {
|
||||
const tracer = new ContextTracer(
|
||||
{ enabled: false, targetDir: '/fake/target', sessionId: 'test-session' },
|
||||
fileSystem,
|
||||
idGenerator,
|
||||
);
|
||||
|
||||
tracer.logEvent('TestComponent', 'TestAction');
|
||||
|
||||
const hugeString = 'a'.repeat(2000);
|
||||
tracer.logEvent('TestComponent', 'LargeAction', { largeKey: hugeString });
|
||||
|
||||
// FS should be completely empty
|
||||
expect(fileSystem.getFiles().size).toBe(0);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,120 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
import type { IFileSystem } from './system/IFileSystem.js';
|
||||
import { NodeFileSystem } from './system/NodeFileSystem.js';
|
||||
import type { IIdGenerator } from './system/IIdGenerator.js';
|
||||
import { NodeIdGenerator } from './system/NodeIdGenerator.js';
|
||||
|
||||
export interface ContextTracerOptions {
|
||||
enabled?: boolean;
|
||||
targetDir: string;
|
||||
sessionId: string;
|
||||
}
|
||||
|
||||
export class ContextTracer {
|
||||
private traceDir: string;
|
||||
private assetsDir: string;
|
||||
private enabled: boolean;
|
||||
private fileSystem: IFileSystem;
|
||||
private idGenerator: IIdGenerator;
|
||||
|
||||
private readonly MAX_INLINE_SIZE = 1000;
|
||||
|
||||
constructor(
|
||||
options: ContextTracerOptions,
|
||||
fileSystem: IFileSystem = new NodeFileSystem(),
|
||||
idGenerator: IIdGenerator = new NodeIdGenerator(),
|
||||
) {
|
||||
this.enabled = options.enabled ?? false;
|
||||
this.fileSystem = fileSystem;
|
||||
this.idGenerator = idGenerator;
|
||||
|
||||
this.traceDir = this.fileSystem.join(
|
||||
options.targetDir,
|
||||
'.gemini',
|
||||
'context_trace',
|
||||
options.sessionId,
|
||||
);
|
||||
this.assetsDir = this.fileSystem.join(this.traceDir, 'assets');
|
||||
|
||||
if (this.enabled) {
|
||||
try {
|
||||
this.fileSystem.mkdirSync(this.assetsDir, { recursive: true });
|
||||
this.logEvent('SYSTEM', 'Context Tracer Initialized', {
|
||||
sessionId: options.sessionId,
|
||||
});
|
||||
} catch (e) {
|
||||
debugLogger.error('Failed to initialize ContextTracer', e);
|
||||
this.enabled = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logEvent(
|
||||
component: string,
|
||||
action: string,
|
||||
details?: Record<string, unknown>,
|
||||
) {
|
||||
if (!this.enabled) return;
|
||||
try {
|
||||
let processedDetails: Record<string, unknown> | undefined;
|
||||
|
||||
if (details) {
|
||||
processedDetails = {};
|
||||
for (const [key, value] of Object.entries(details)) {
|
||||
const strValue =
|
||||
typeof value === 'string' ? value : JSON.stringify(value);
|
||||
if (strValue && strValue.length > this.MAX_INLINE_SIZE) {
|
||||
const assetId = this.saveAsset(component, key, value);
|
||||
processedDetails[key] = { $asset: assetId };
|
||||
} else {
|
||||
processedDetails[key] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const timestamp = new Date().toISOString();
|
||||
const detailsStr = processedDetails
|
||||
? ` | Details: ${JSON.stringify(processedDetails)}`
|
||||
: '';
|
||||
const logLine = `[${timestamp}] [${component}] ${action}${detailsStr}\n`;
|
||||
this.fileSystem.appendFileSync(
|
||||
this.fileSystem.join(this.traceDir, 'trace.log'),
|
||||
logLine,
|
||||
'utf-8',
|
||||
);
|
||||
} catch (e) {
|
||||
debugLogger.warn(`Tracing failed: ${e}`);
|
||||
}
|
||||
}
|
||||
|
||||
private saveAsset(
|
||||
component: string,
|
||||
assetName: string,
|
||||
data: unknown,
|
||||
): string {
|
||||
if (!this.enabled) return 'asset-recording-disabled';
|
||||
try {
|
||||
const assetId = `${Date.now()}-${this.idGenerator.generateId()}-${assetName}.json`;
|
||||
const assetPath = this.fileSystem.join(this.assetsDir, assetId);
|
||||
|
||||
this.fileSystem.writeFileSync(
|
||||
assetPath,
|
||||
JSON.stringify(data, null, 2),
|
||||
'utf-8',
|
||||
);
|
||||
this.logEvent(component, `Saved asset: ${assetName}`, { assetId });
|
||||
return assetId;
|
||||
} catch (e) {
|
||||
this.logEvent(component, `Failed to save asset: ${assetName}`, {
|
||||
error: String(e),
|
||||
});
|
||||
return 'asset-save-failed';
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Part } from '@google/genai';
|
||||
import { estimateTokenCountSync as baseEstimate } from '../../utils/tokenCalculation.js';
|
||||
import type { Episode } from '../ir/types.js';
|
||||
|
||||
/**
|
||||
* The flat token cost assigned to a single multi-modal asset (like an image tile)
|
||||
* by the Gemini API. We use this as a baseline heuristic for inlineData/fileData.
|
||||
*/
|
||||
const BASE_MULTIMODAL_TOKEN_COST = 258;
|
||||
|
||||
export class ContextTokenCalculator {
|
||||
constructor(private readonly charsPerToken: number) {}
|
||||
|
||||
/**
|
||||
* Fast, simple heuristic estimation for a raw string.
|
||||
*/
|
||||
estimateTokensForString(text: string): number {
|
||||
return Math.ceil(text.length / this.charsPerToken);
|
||||
}
|
||||
|
||||
/**
|
||||
* Fast, simple heuristic conversion from tokens to expected character length.
|
||||
* Useful for calculating truncation thresholds.
|
||||
*/
|
||||
tokensToChars(tokens: number): number {
|
||||
return tokens * this.charsPerToken;
|
||||
}
|
||||
|
||||
/**
|
||||
* Calculates the total token count for a complete Episodic IR graph.
|
||||
* This is fast because it relies on pre-computed metadata where available.
|
||||
*/
|
||||
calculateEpisodeListTokens(episodes: Episode[]): number {
|
||||
let tokens = 0;
|
||||
for (const ep of episodes) {
|
||||
if (ep.trigger) tokens += ep.trigger.metadata.currentTokens;
|
||||
for (const step of ep.steps) {
|
||||
tokens += step.metadata.currentTokens;
|
||||
}
|
||||
if (ep.yield) tokens += ep.yield.metadata.currentTokens;
|
||||
}
|
||||
return tokens;
|
||||
}
|
||||
|
||||
/**
|
||||
* Slower, precise estimation for a Gemini Content/Part graph.
|
||||
* Deeply inspects the nested structure and uses the base tokenization math.
|
||||
*/
|
||||
estimateTokensForParts(parts: Part[], depth: number = 0): number {
|
||||
let totalTokens = 0;
|
||||
for (const part of parts) {
|
||||
if (typeof part.text === 'string') {
|
||||
totalTokens += Math.ceil(part.text.length / this.charsPerToken);
|
||||
} else if (part.inlineData !== undefined || part.fileData !== undefined) {
|
||||
totalTokens += BASE_MULTIMODAL_TOKEN_COST;
|
||||
} else {
|
||||
totalTokens += Math.ceil(
|
||||
JSON.stringify(part).length / this.charsPerToken,
|
||||
);
|
||||
}
|
||||
}
|
||||
// Also include structural overhead
|
||||
return totalTokens + baseEstimate(parts, depth);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Content } from '@google/genai';
|
||||
|
||||
export type HistoryEventType = 'PUSH' | 'SYNC_FULL' | 'CLEAR';
|
||||
|
||||
export interface HistoryEvent {
|
||||
type: HistoryEventType;
|
||||
payload: readonly Content[];
|
||||
}
|
||||
|
||||
export type HistoryListener = (event: HistoryEvent) => void;
|
||||
|
||||
export class AgentChatHistory {
|
||||
private history: Content[];
|
||||
private listeners: Set<HistoryListener> = new Set();
|
||||
|
||||
constructor(initialHistory: Content[] = []) {
|
||||
this.history = [...initialHistory];
|
||||
}
|
||||
|
||||
subscribe(listener: HistoryListener): () => void {
|
||||
this.listeners.add(listener);
|
||||
// Emit initial state to new subscriber
|
||||
listener({ type: 'SYNC_FULL', payload: this.history });
|
||||
return () => this.listeners.delete(listener);
|
||||
}
|
||||
|
||||
private notify(type: HistoryEventType, payload: readonly Content[]) {
|
||||
const event: HistoryEvent = { type, payload };
|
||||
for (const listener of this.listeners) {
|
||||
listener(event);
|
||||
}
|
||||
}
|
||||
|
||||
push(content: Content) {
|
||||
this.history.push(content);
|
||||
this.notify('PUSH', [content]);
|
||||
}
|
||||
|
||||
set(history: readonly Content[]) {
|
||||
this.history = [...history];
|
||||
this.notify('SYNC_FULL', this.history);
|
||||
}
|
||||
|
||||
clear() {
|
||||
this.history = [];
|
||||
this.notify('CLEAR', []);
|
||||
}
|
||||
|
||||
get(): readonly Content[] {
|
||||
return this.history;
|
||||
}
|
||||
|
||||
map(callback: (value: Content, index: number, array: Content[]) => Content) {
|
||||
this.history = this.history.map(callback);
|
||||
this.notify('SYNC_FULL', this.history);
|
||||
}
|
||||
|
||||
flatMap<U>(
|
||||
callback: (
|
||||
value: Content,
|
||||
index: number,
|
||||
array: Content[],
|
||||
) => U | readonly U[],
|
||||
): U[] {
|
||||
return this.history.flatMap(callback);
|
||||
}
|
||||
|
||||
get length(): number {
|
||||
return this.history.length;
|
||||
}
|
||||
}
|
||||
@@ -16,4 +16,5 @@ export enum LlmRole {
|
||||
UTILITY_EDIT_CORRECTOR = 'utility_edit_corrector',
|
||||
UTILITY_AUTOCOMPLETE = 'utility_autocomplete',
|
||||
UTILITY_FAST_ACK_HELPER = 'utility_fast_ack_helper',
|
||||
UTILITY_STATE_SNAPSHOT_PROCESSOR = 'utility_state_snapshot_processr',
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user