mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-07-12 02:50:41 -07:00
refactor(context): migrate context wiring to TS Profiles and convert Processors to pure functional HOFs
This commit acts on the design review feedback:
1. Eliminates dynamic JSON registry instantiation for pipelines.
2. Introduces TS ContextProfiles for strongly-typed declarative wiring.
3. Converts ContextProcessors to pure functions returned by factories (HOFs) holding local state via closures.
4. Converts ContextWorkers to simple {start, stop} lifecycle objects.
5. Removes now-obsolete registry and JSON parsing overhead from Orchestrator.
This commit is contained in:
@@ -0,0 +1,243 @@
|
||||
# Context Management Architecture: A Foundation for Scalable Context Exploration and Experimentation
|
||||
|
||||
## 1. Executive Summary & Motivation
|
||||
|
||||
As our agentic capabilities grow, the active context window becomes our most
|
||||
critical and constrained resource. Mismanaged context leads to broken caching,
|
||||
hallucination, and exorbitant token costs. Historically, context management has
|
||||
been decentralized, stateful, and highly coupled—making it dangerous to mutate
|
||||
and nearly impossible to safely experiment with.
|
||||
|
||||
Our primary goal with this architecture is to establish a rigorous, structured
|
||||
model for context computation that guarantees:
|
||||
|
||||
1. **Safety & Auditability:** Mutations happen in a predictable, auditable, and
|
||||
recoverable way.
|
||||
2. **Asynchronous Safety:** Long-running LLM-driven graph analysis can execute
|
||||
safely without blocking the user or creating race conditions.
|
||||
3. **Trivial Extensibility:** The system can be effortlessly augmented with new
|
||||
compression strategies to scale our experiments.
|
||||
4. **A Universal Data-Plane:** Beyond simple text compression, the architecture
|
||||
generalizes to safely structure _any_ computation in and around the context
|
||||
(e.g., continuous reflection, semantic routing, long-term memory extraction).
|
||||
|
||||
Because the ultimate "right answer" for context compression is unknown and
|
||||
constantly shifting, we have designed this architecture strictly around the
|
||||
**Open-Closed Principle**. We have "closed" the state-mutation engine to
|
||||
guarantee structural integrity, while leaving the behavioral logic entirely
|
||||
"open" for extension via decoupled primitives.
|
||||
|
||||
To be clear: this is not an exercise in over-engineering. Rather, we are
|
||||
applying proven, boring, industry-standard software paradigms—specifically
|
||||
Functional Reactive Programming (FRP), the Actor Model, and the Open-Closed
|
||||
Principle—to tame the inherent complexity of managing an agent's context and
|
||||
prevent the system from collapsing under its own state.
|
||||
|
||||
## 2. Embracing the Unknowns
|
||||
|
||||
Before defining the architecture, we must acknowledge the fundamental realities
|
||||
of context management:
|
||||
|
||||
- **The Trilemma (Quality vs. Cost vs. Caching):** There will never be one
|
||||
single correct way to manage history. Aggressive summarization saves tokens
|
||||
but breaks exact-string caching; retaining raw history maximizes caching but
|
||||
blows up token budgets.
|
||||
- **The Precision / Recall Tension:** Context compression inherently damages
|
||||
_recall_ (the ability to perfectly retrieve specific past details). However,
|
||||
it protects _precision_ (preventing the LLM from being distracted by
|
||||
irrelevant noise). We hypothesize that over time, active context management
|
||||
will focus primarily on protecting precision, while recall will be better
|
||||
protected in a targeted way via external memory, tasks, and planning systems.
|
||||
- **The Need for Rapid Experimentation:** Because these tradeoffs depend heavily
|
||||
on the specific model, task, and user budget, we need a system that lets us
|
||||
explore and express a wide range of options. We must be able to run a large
|
||||
number of experiments—including in production—without risking catastrophic
|
||||
state corruption.
|
||||
|
||||
## 3. The Core Architectural Primitives
|
||||
|
||||
To facilitate safe experimentation, we have separated the _execution_ of context
|
||||
mutations from the _logic_ of context mutations. This division reflects
|
||||
established, robust patterns.
|
||||
|
||||
### The "Closed" Foundation: Synchronous Pipelines (Functional Reactive Programming)
|
||||
|
||||
Drawing from the principles of **Functional Reactive Programming (FRP)**, the
|
||||
Context Working Buffer is treated as an immutable, ahead-of-time tracked graph.
|
||||
It can **only** be mutated synchronously, via an event-triggered **Pipeline**. A
|
||||
Pipeline is simply a linear list of functionally composed processors. By forcing
|
||||
all mutations through a synchronous, blocking pipeline of pure functions, we
|
||||
guarantee that the context is always modified in a sane, predictable, and
|
||||
mathematically sound sequence.
|
||||
|
||||
### The "Open" Extensions: Processors, Workers, and Inboxes (The Actor Model)
|
||||
|
||||
To extend the system, developers author two types of plugins:
|
||||
|
||||
1. **Context Processors:** Pure, fast, synchronous functions that take an input
|
||||
graph and return an immutable mutated graph. They run inside Pipelines.
|
||||
2. **Context Workers:** Inspired by the **Actor Model**, these are
|
||||
event-triggered background jobs designed for isolated, long-running async
|
||||
computations (e.g., asking an LLM to distill 50 turns of history).
|
||||
3. **Inboxes:** Because the graph can only be mutated synchronously, Workers
|
||||
cannot touch the graph directly (preventing race conditions). Instead, they
|
||||
drop their results via message-passing into point-in-time snapshots called
|
||||
_Inboxes_. Processors later read from these Inboxes during a synchronous
|
||||
pipeline run to safely apply the worker's findings.
|
||||
|
||||
## 4. Proofs of Construction
|
||||
|
||||
To demonstrate why these primitives are perfectly suited to the problem,
|
||||
consider the following structural pseudocode.
|
||||
|
||||
### Example A: Fast, Synchronous Sanitization (The Processor)
|
||||
|
||||
_Scenario: We need to immediately truncate massive tool outputs before they blow
|
||||
out the context window._
|
||||
|
||||
Because this requires no LLM calls, it is expressed as a simple **Processor**
|
||||
running in an ingestion pipeline.
|
||||
|
||||
```typescript
|
||||
// A pure function that guarantees safe mutation
|
||||
class ToolMaskingProcessor implements ContextProcessor {
|
||||
apply(graph: ContextGraph): ContextGraph {
|
||||
const mutatedGraph = graph.clone();
|
||||
|
||||
for (const node of mutatedGraph.getNodes('TOOL_OUTPUT')) {
|
||||
if (node.length > MAX_CHARS) {
|
||||
// Safely replace the node, retaining structural lineage
|
||||
mutatedGraph.replace(node, {
|
||||
type: 'MASKED_TOOL',
|
||||
text: `[Output truncated. Original size: ${node.length}]`,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return mutatedGraph;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Example B: Long-Running Summarization (Worker + Inbox + Processor)
|
||||
|
||||
_Scenario: The user has exceeded their token budget. We need to use an LLM to
|
||||
summarize the oldest 20 turns of conversation, but we cannot block the user from
|
||||
continuing to chat while the LLM generates the summary._
|
||||
|
||||
This requires our async-to-sync bridge.
|
||||
|
||||
**Step 1: The Worker (Async Analysis)**
|
||||
|
||||
```typescript
|
||||
class StateSnapshotWorker implements ContextWorker {
|
||||
// Triggers automatically in the background when the budget is exceeded
|
||||
async onBudgetExceeded(event: BudgetEvent, inbox: Inbox) {
|
||||
const agedOutNodes = event.getAgedOutNodes();
|
||||
|
||||
// Slow, async LLM call
|
||||
const summaryText = await llm.summarize(agedOutNodes);
|
||||
|
||||
// The worker CANNOT mutate the graph. It leaves a message in the Inbox.
|
||||
inbox.deliver('SUMMARY_READY', {
|
||||
targetNodes: agedOutNodes,
|
||||
summary: summaryText,
|
||||
});
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: The Processor (Sync Application)**
|
||||
|
||||
```typescript
|
||||
class StateSnapshotProcessor implements ContextProcessor {
|
||||
// Runs fast and synchronously during the next Pipeline execution
|
||||
apply(graph: ContextGraph, inbox: InboxSnapshot): ContextGraph {
|
||||
// Check if the background worker finished its job
|
||||
const messages = inbox.read('SUMMARY_READY');
|
||||
if (messages.isEmpty()) return graph;
|
||||
|
||||
const mutatedGraph = graph.clone();
|
||||
|
||||
for (const msg of messages) {
|
||||
// Safely swap the old nodes for the new summary
|
||||
mutatedGraph.collapseNodes(msg.targetNodes, {
|
||||
type: 'ROLLING_SUMMARY',
|
||||
text: msg.summary,
|
||||
});
|
||||
inbox.markConsumed(msg);
|
||||
}
|
||||
|
||||
return mutatedGraph;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Example C: Downstream Observation & Memory Extraction (The Ledger)
|
||||
|
||||
_Scenario: We want to extract long-term memories from conversation turns
|
||||
immediately before they are permanently deleted from the active context window
|
||||
by a garbage collector (GC)._
|
||||
|
||||
Because the Context Working Buffer immutably tracks its own mathematical deltas
|
||||
(the Audit Log) and computes structural lineage Ahead-Of-Time (AOT), downstream
|
||||
processors don't have to guess what happened; they can simply read the math.
|
||||
|
||||
```typescript
|
||||
class MemoryExtractionProcessor implements ContextProcessor {
|
||||
// Runs sequentially AFTER a Garbage Collection (GC) Processor
|
||||
apply(graph: ContextGraph): ContextGraph {
|
||||
// 1. Look at the immutable Audit Log to see what the previous step did
|
||||
const latestMutation = graph.getAuditLog().latest();
|
||||
if (latestMutation.processorId !== 'HistoryTruncationProcessor') {
|
||||
return graph; // Nothing was GC'd, do nothing
|
||||
}
|
||||
|
||||
// 2. Identify the exact pristine nodes that were permanently lost
|
||||
const lostPristineNodes = new Set<ContextNode>();
|
||||
|
||||
for (const removedId of latestMutation.removedIds) {
|
||||
// Because the graph tracks provenance Ahead-Of-Time (AOT),
|
||||
// we perfectly resolve what original thoughts this synthetic node represented.
|
||||
// (e.g. Deleting a ROLLING_SUMMARY implies losing the 3 original USER_PROMPTS it summarized).
|
||||
const roots = graph.getPristineNodes(removedId);
|
||||
roots.forEach((r) => lostPristineNodes.add(r));
|
||||
}
|
||||
|
||||
// 3. Compare against the currently surviving graph to find the TRUE delta
|
||||
// (Ensure the roots aren't still surviving inside some OTHER summary node)
|
||||
for (const survivingNode of graph.getNodes()) {
|
||||
const survivingRoots = graph.getPristineNodes(survivingNode.id);
|
||||
survivingRoots.forEach((r) => lostPristineNodes.delete(r));
|
||||
}
|
||||
|
||||
// 4. Dispatch the permanently lost nodes to the long-term memory subsystem
|
||||
if (lostPristineNodes.size > 0) {
|
||||
// Fire-and-forget async dispatch (Actor Model) to external DB
|
||||
LongTermMemorySystem.dispatchForEmbedding(Array.from(lostPristineNodes));
|
||||
}
|
||||
|
||||
// This processor is purely observational; it returns the graph unmutated
|
||||
return graph;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## 5. Conclusion & Future Evolution
|
||||
|
||||
By treating the Context Graph as an immutable ledger updated only via functional
|
||||
Pipelines, we have eliminated race conditions and untraceable graph corruption.
|
||||
By utilizing Workers and Inboxes, we have safely bridged the gap between slow
|
||||
LLM analysis and fast, synchronous terminal UI updates.
|
||||
|
||||
We recognize this is not the final form—future iterations may require strict
|
||||
simple priority to updates, or more advanced generational garbage collection.
|
||||
However, this architecture provides a rock-solid, extensible foundation.
|
||||
|
||||
More importantly, while this system was born from the need to manage tokens, its
|
||||
immutable ledger and reactive pipelines generalize to something far more
|
||||
profound. We have built a safe, predictable computational engine capable of
|
||||
driving the _entire_ agentic loop—from background reflection, to semantic
|
||||
routing, to long-term memory extraction. It empowers us to safely deploy,
|
||||
observe, and scale a wide array of strategies in pursuit of the optimal balance
|
||||
between token cost, caching efficiency, and agentic quality.
|
||||
@@ -10,7 +10,7 @@ import type { ConcreteNode } 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 type { ContextProfile } from './sidecar/profiles.js';
|
||||
import type { PipelineOrchestrator } from './sidecar/orchestrator.js';
|
||||
import { HistoryObserver } from './historyObserver.js';
|
||||
import { IrProjector } from './ir/projector.js';
|
||||
@@ -29,7 +29,7 @@ export class ContextManager {
|
||||
private readonly historyObserver: HistoryObserver;
|
||||
|
||||
constructor(
|
||||
private readonly sidecar: SidecarConfig,
|
||||
private readonly sidecar: ContextProfile,
|
||||
private readonly env: ContextEnvironment,
|
||||
private readonly tracer: ContextTracer,
|
||||
orchestrator: PipelineOrchestrator,
|
||||
@@ -74,7 +74,7 @@ export class ContextManager {
|
||||
* firing consolidation events if necessary.
|
||||
*/
|
||||
private evaluateTriggers(newNodes: Set<string>) {
|
||||
if (!this.sidecar.budget) return;
|
||||
if (!this.sidecar.config.budget) return;
|
||||
|
||||
if (newNodes.size > 0) {
|
||||
this.eventBus.emitChunkReceived({
|
||||
@@ -87,7 +87,7 @@ export class ContextManager {
|
||||
this.buffer.nodes,
|
||||
);
|
||||
|
||||
if (currentTokens > this.sidecar.budget.retainedTokens) {
|
||||
if (currentTokens > this.sidecar.config.budget.retainedTokens) {
|
||||
const agedOutNodes = new Set<string>();
|
||||
let rollingTokens = 0;
|
||||
// Walk backwards finding nodes that fall out of the retained budget
|
||||
@@ -96,7 +96,7 @@ export class ContextManager {
|
||||
rollingTokens += this.env.tokenCalculator.calculateConcreteListTokens([
|
||||
node,
|
||||
]);
|
||||
if (rollingTokens > this.sidecar.budget.retainedTokens) {
|
||||
if (rollingTokens > this.sidecar.config.budget.retainedTokens) {
|
||||
agedOutNodes.add(node.id);
|
||||
}
|
||||
}
|
||||
@@ -104,7 +104,7 @@ export class ContextManager {
|
||||
if (agedOutNodes.size > 0) {
|
||||
this.eventBus.emitConsolidationNeeded({
|
||||
nodes: this.buffer.nodes,
|
||||
targetDeficit: currentTokens - this.sidecar.budget.retainedTokens,
|
||||
targetDeficit: currentTokens - this.sidecar.config.budget.retainedTokens,
|
||||
targetNodeIds: agedOutNodes,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -12,7 +12,7 @@ import type {
|
||||
ContextTracer,
|
||||
} from '../sidecar/environment.js';
|
||||
import type { PipelineOrchestrator } from '../sidecar/orchestrator.js';
|
||||
import type { SidecarConfig } from '../sidecar/types.js';
|
||||
import type { ContextProfile } from '../sidecar/profiles.js';
|
||||
|
||||
export class IrProjector {
|
||||
/**
|
||||
@@ -22,12 +22,12 @@ export class IrProjector {
|
||||
static async project(
|
||||
nodes: readonly ConcreteNode[],
|
||||
orchestrator: PipelineOrchestrator,
|
||||
sidecar: SidecarConfig,
|
||||
sidecar: ContextProfile,
|
||||
tracer: ContextTracer,
|
||||
env: ContextEnvironment,
|
||||
protectedIds: Set<string>,
|
||||
): Promise<Content[]> {
|
||||
if (!sidecar.budget) {
|
||||
if (!sidecar.config.budget) {
|
||||
const contents = env.irMapper.fromIr(nodes);
|
||||
tracer.logEvent('IrProjector', 'Projected Context to LLM (No Budget)', {
|
||||
projectedContext: contents,
|
||||
@@ -35,7 +35,7 @@ export class IrProjector {
|
||||
return contents;
|
||||
}
|
||||
|
||||
const maxTokens = sidecar.budget.maxTokens;
|
||||
const maxTokens = sidecar.config.budget.maxTokens;
|
||||
const currentTokens =
|
||||
env.tokenCalculator.calculateConcreteListTokens(nodes);
|
||||
|
||||
@@ -79,7 +79,7 @@ export class IrProjector {
|
||||
node,
|
||||
]);
|
||||
rollingTokens += nodeTokens;
|
||||
if (rollingTokens > sidecar.budget.retainedTokens) {
|
||||
if (rollingTokens > sidecar.config.budget.retainedTokens) {
|
||||
agedOutNodes.add(node.id);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -38,10 +38,15 @@ export interface ProcessArgs {
|
||||
readonly inbox: InboxSnapshot;
|
||||
}
|
||||
|
||||
export interface ContextProcessor {
|
||||
/**
|
||||
* A ContextProcessor is now a pure function that returns a modified subset of nodes
|
||||
* (or the original targets if no changes are needed).
|
||||
* The Orchestrator will use this to generate a new graph delta.
|
||||
*/
|
||||
export interface ContextProcessorFn {
|
||||
readonly id: string;
|
||||
readonly name: string;
|
||||
process(args: ProcessArgs): Promise<readonly ConcreteNode[]>;
|
||||
(args: ProcessArgs): Promise<readonly ConcreteNode[]>;
|
||||
}
|
||||
|
||||
export interface ContextWorker {
|
||||
@@ -52,6 +57,8 @@ export interface ContextWorker {
|
||||
onNodesAgedOut?: boolean;
|
||||
onInboxTopics?: string[];
|
||||
};
|
||||
start(): void;
|
||||
stop(): void;
|
||||
execute(args: {
|
||||
targets: readonly ConcreteNode[];
|
||||
inbox: InboxSnapshot;
|
||||
|
||||
@@ -1,101 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import assert from 'node:assert';
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { BlobDegradationProcessor } from './blobDegradationProcessor.js';
|
||||
import {
|
||||
createMockProcessArgs,
|
||||
createMockEnvironment,
|
||||
createDummyNode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import type { UserPrompt, SemanticPart, ConcreteNode } from '../ir/types.js';
|
||||
|
||||
describe('BlobDegradationProcessor', () => {
|
||||
it('should ignore text parts and only target inline_data and file_data', async () => {
|
||||
const env = createMockEnvironment();
|
||||
// charsPerToken = 1
|
||||
// We want the degraded text to be cheaper than the original blob.
|
||||
// Degraded text is ~100 chars ("...degraded to text...").
|
||||
// So we make the blob data 200 chars.
|
||||
const fakeData = 'A'.repeat(200);
|
||||
|
||||
const processor = BlobDegradationProcessor.create(env, {});
|
||||
|
||||
const parts: SemanticPart[] = [
|
||||
{ type: 'text', text: 'Hello' },
|
||||
{ type: 'inline_data', mimeType: 'image/png', data: fakeData },
|
||||
{ type: 'text', text: 'World' },
|
||||
];
|
||||
|
||||
const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, {
|
||||
semanticParts: parts,
|
||||
}) as UserPrompt;
|
||||
|
||||
const targets = [prompt];
|
||||
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
|
||||
expect(result.length).toBe(1);
|
||||
const modifiedPrompt = result[0] as UserPrompt;
|
||||
|
||||
expect(modifiedPrompt.id).not.toBe(prompt.id);
|
||||
expect(modifiedPrompt.semanticParts.length).toBe(3);
|
||||
|
||||
// Text parts should be untouched
|
||||
expect(modifiedPrompt.semanticParts[0]).toEqual(parts[0]);
|
||||
expect(modifiedPrompt.semanticParts[2]).toEqual(parts[2]);
|
||||
|
||||
// The inline_data part should be replaced with text
|
||||
const degradedPart = modifiedPrompt.semanticParts[1];
|
||||
expect(degradedPart.type).toBe('text');
|
||||
assert(degradedPart.type === 'text');
|
||||
expect(degradedPart.text).toContain(
|
||||
'[Multi-Modal Blob (image/png, 0.00MB) degraded to text',
|
||||
);
|
||||
});
|
||||
|
||||
it('should degrade all blobs unconditionally', async () => {
|
||||
const env = createMockEnvironment();
|
||||
|
||||
const processor = BlobDegradationProcessor.create(env, {});
|
||||
|
||||
// Tokens for fileData = 258.
|
||||
// Degraded text = "[File Reference (video/mp4) degraded to text to preserve context window. Original URI: gs://test1]"
|
||||
// Degraded text length ~100 characters.
|
||||
// Since charsPerToken=1, degraded text = 100 tokens.
|
||||
// Tokens saved = 258 - 100 = 158. This is > 0, so it WILL degrade it!
|
||||
|
||||
const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, {
|
||||
semanticParts: [
|
||||
{ type: 'file_data', mimeType: 'video/mp4', fileUri: 'gs://test1' },
|
||||
{ type: 'file_data', mimeType: 'video/mp4', fileUri: 'gs://test2' },
|
||||
],
|
||||
}) as UserPrompt;
|
||||
|
||||
const targets = [prompt];
|
||||
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
|
||||
const modifiedPrompt = result[0] as UserPrompt;
|
||||
expect(modifiedPrompt.semanticParts.length).toBe(2);
|
||||
|
||||
// Both parts should be degraded
|
||||
expect(modifiedPrompt.semanticParts[0].type).toBe('text');
|
||||
expect(modifiedPrompt.semanticParts[1].type).toBe('text');
|
||||
});
|
||||
|
||||
it('should return exactly the targets array if targets are empty', async () => {
|
||||
const env = createMockEnvironment();
|
||||
|
||||
const processor = BlobDegradationProcessor.create(env, {});
|
||||
const targets: ConcreteNode[] = [];
|
||||
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
|
||||
expect(result).toBe(targets);
|
||||
});
|
||||
});
|
||||
@@ -3,45 +3,29 @@
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type { ProcessArgs, ContextProcessor } from '../pipeline.js';
|
||||
import type { ProcessArgs, ContextProcessorFn } from '../pipeline.js';
|
||||
import type { ConcreteNode, UserPrompt } from '../ir/types.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { sanitizeFilenamePart } from '../../utils/fileUtils.js';
|
||||
|
||||
export type BlobDegradationProcessorOptions = Record<string, never>;
|
||||
|
||||
export class BlobDegradationProcessor implements ContextProcessor {
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
_options: BlobDegradationProcessorOptions,
|
||||
): BlobDegradationProcessor {
|
||||
return new BlobDegradationProcessor(env);
|
||||
}
|
||||
|
||||
readonly componentType = 'processor';
|
||||
readonly id = 'BlobDegradationProcessor';
|
||||
readonly name = 'BlobDegradationProcessor';
|
||||
readonly options = {};
|
||||
private env: ContextEnvironment;
|
||||
|
||||
constructor(env: ContextEnvironment) {
|
||||
this.env = env;
|
||||
}
|
||||
|
||||
async process({ targets }: ProcessArgs): Promise<readonly ConcreteNode[]> {
|
||||
export function createBlobDegradationProcessor(
|
||||
id: string,
|
||||
env: ContextEnvironment,
|
||||
): ContextProcessorFn {
|
||||
const processor: any = async ({ targets }: ProcessArgs) => {
|
||||
if (targets.length === 0) {
|
||||
return targets;
|
||||
}
|
||||
|
||||
let directoryCreated = false;
|
||||
|
||||
let blobOutputsDir = this.env.fileSystem.join(
|
||||
this.env.projectTempDir,
|
||||
let blobOutputsDir = env.fileSystem.join(
|
||||
env.projectTempDir,
|
||||
'degraded-blobs',
|
||||
);
|
||||
const sessionId = this.env.sessionId;
|
||||
const sessionId = env.sessionId;
|
||||
if (sessionId) {
|
||||
blobOutputsDir = this.env.fileSystem.join(
|
||||
blobOutputsDir = env.fileSystem.join(
|
||||
blobOutputsDir,
|
||||
`session-${sanitizeFilenamePart(sessionId)}`,
|
||||
);
|
||||
@@ -49,7 +33,7 @@ export class BlobDegradationProcessor implements ContextProcessor {
|
||||
|
||||
const ensureDir = async () => {
|
||||
if (!directoryCreated) {
|
||||
await this.env.fileSystem.mkdir(blobOutputsDir, { recursive: true });
|
||||
await env.fileSystem.mkdir(blobOutputsDir, { recursive: true });
|
||||
directoryCreated = true;
|
||||
}
|
||||
};
|
||||
@@ -74,26 +58,26 @@ export class BlobDegradationProcessor implements ContextProcessor {
|
||||
case '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(
|
||||
const fileName = `blob_${Date.now()}_${env.idGenerator.generateId()}.${ext}`;
|
||||
const filePath = env.fileSystem.join(
|
||||
blobOutputsDir,
|
||||
fileName,
|
||||
);
|
||||
|
||||
const buffer = Buffer.from(part.data, 'base64');
|
||||
await this.env.fileSystem.writeFile(filePath, buffer);
|
||||
await 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}]`;
|
||||
|
||||
const oldTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([
|
||||
env.tokenCalculator.estimateTokensForParts([
|
||||
{
|
||||
inlineData: { mimeType: part.mimeType, data: part.data },
|
||||
},
|
||||
]);
|
||||
const newTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([
|
||||
env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: newText },
|
||||
]);
|
||||
tokensSaved = oldTokens - newTokens;
|
||||
@@ -102,7 +86,7 @@ export class BlobDegradationProcessor implements ContextProcessor {
|
||||
case 'file_data': {
|
||||
newText = `[File Reference (${part.mimeType}) degraded to text to preserve context window. Original URI: ${part.fileUri}]`;
|
||||
const oldTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([
|
||||
env.tokenCalculator.estimateTokensForParts([
|
||||
{
|
||||
fileData: {
|
||||
mimeType: part.mimeType,
|
||||
@@ -111,7 +95,7 @@ export class BlobDegradationProcessor implements ContextProcessor {
|
||||
},
|
||||
]);
|
||||
const newTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([
|
||||
env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: newText },
|
||||
]);
|
||||
tokensSaved = oldTokens - newTokens;
|
||||
@@ -120,9 +104,9 @@ export class BlobDegradationProcessor implements ContextProcessor {
|
||||
case 'raw_part': {
|
||||
newText = `[Unknown Part degraded to text to preserve context window.]`;
|
||||
const oldTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([part.part]);
|
||||
env.tokenCalculator.estimateTokensForParts([part.part]);
|
||||
const newTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([
|
||||
env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: newText },
|
||||
]);
|
||||
tokensSaved = oldTokens - newTokens;
|
||||
@@ -141,8 +125,9 @@ export class BlobDegradationProcessor implements ContextProcessor {
|
||||
if (modified) {
|
||||
const degradedNode: UserPrompt = {
|
||||
...node,
|
||||
id: this.env.idGenerator.generateId(),
|
||||
id: env.idGenerator.generateId(),
|
||||
semanticParts: newParts,
|
||||
replacesId: node.id,
|
||||
};
|
||||
returnedNodes.push(degradedNode);
|
||||
} else {
|
||||
@@ -157,5 +142,10 @@ export class BlobDegradationProcessor implements ContextProcessor {
|
||||
}
|
||||
|
||||
return returnedNodes;
|
||||
}
|
||||
};
|
||||
|
||||
processor.id = id;
|
||||
Object.defineProperty(processor, 'name', { value: 'BlobDegradationProcessor' });
|
||||
|
||||
return processor;
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
*/
|
||||
|
||||
import type {
|
||||
ContextProcessor,
|
||||
ContextProcessorFn,
|
||||
BackstopTargetOptions,
|
||||
ProcessArgs,
|
||||
} from '../pipeline.js';
|
||||
@@ -14,51 +14,20 @@ import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
|
||||
export type HistoryTruncationProcessorOptions = BackstopTargetOptions;
|
||||
|
||||
export class HistoryTruncationProcessor implements ContextProcessor {
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: HistoryTruncationProcessorOptions,
|
||||
): HistoryTruncationProcessor {
|
||||
return new HistoryTruncationProcessor(env, options);
|
||||
}
|
||||
|
||||
static readonly schema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
target: {
|
||||
type: 'string',
|
||||
enum: ['incremental', 'freeNTokens', 'max'],
|
||||
description: 'How much of the targeted history to truncate.',
|
||||
},
|
||||
freeTokensTarget: {
|
||||
type: 'number',
|
||||
description: 'The number of tokens to free if target is freeNTokens.',
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
readonly componentType = 'processor';
|
||||
readonly id = 'HistoryTruncationProcessor';
|
||||
readonly name = 'HistoryTruncationProcessor';
|
||||
private readonly env: ContextEnvironment;
|
||||
readonly options: HistoryTruncationProcessorOptions;
|
||||
constructor(
|
||||
env: ContextEnvironment,
|
||||
options: HistoryTruncationProcessorOptions,
|
||||
) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
async process({ targets }: ProcessArgs): Promise<readonly ConcreteNode[]> {
|
||||
export function createHistoryTruncationProcessor(
|
||||
id: string,
|
||||
env: ContextEnvironment,
|
||||
options: HistoryTruncationProcessorOptions,
|
||||
): ContextProcessorFn {
|
||||
const processor: any = async ({ targets }: ProcessArgs) => {
|
||||
// Calculate how many tokens we need to remove based on the configured knob
|
||||
let targetTokensToRemove = 0;
|
||||
const strategy = this.options.target ?? 'max';
|
||||
const strategy = options.target ?? 'max';
|
||||
|
||||
if (strategy === 'incremental') {
|
||||
targetTokensToRemove = Infinity;
|
||||
} else if (strategy === 'freeNTokens') {
|
||||
targetTokensToRemove = this.options.freeTokensTarget ?? 0;
|
||||
targetTokensToRemove = options.freeTokensTarget ?? 0;
|
||||
if (targetTokensToRemove <= 0) return targets;
|
||||
} else if (strategy === 'max') {
|
||||
// 'max' means we remove all targets without stopping early
|
||||
@@ -76,9 +45,14 @@ export class HistoryTruncationProcessor implements ContextProcessor {
|
||||
continue;
|
||||
}
|
||||
|
||||
removedTokens += this.env.tokenCalculator.getTokenCost(node);
|
||||
removedTokens += env.tokenCalculator.getTokenCost(node);
|
||||
}
|
||||
|
||||
return keptNodes;
|
||||
}
|
||||
};
|
||||
|
||||
processor.id = id;
|
||||
Object.defineProperty(processor, 'name', { value: 'HistoryTruncationProcessor' });
|
||||
|
||||
return processor;
|
||||
}
|
||||
|
||||
@@ -1,140 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import assert from 'node:assert';
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { NodeDistillationProcessor } from './nodeDistillationProcessor.js';
|
||||
import {
|
||||
createMockProcessArgs,
|
||||
createMockEnvironment,
|
||||
createDummyNode,
|
||||
createDummyToolNode,
|
||||
createMockLlmClient,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import type { UserPrompt, AgentThought, ToolExecution } from '../ir/types.js';
|
||||
|
||||
describe('NodeDistillationProcessor', () => {
|
||||
it('should trigger summarization via LLM for long text parts', async () => {
|
||||
const mockLlmClient = createMockLlmClient(['Mocked Summary!']);
|
||||
|
||||
// Use charsPerToken=1 naturally.
|
||||
const env = createMockEnvironment({
|
||||
llmClient: mockLlmClient,
|
||||
});
|
||||
|
||||
const processor = NodeDistillationProcessor.create(env, {
|
||||
nodeThresholdTokens: 10,
|
||||
});
|
||||
|
||||
const longText = 'A'.repeat(50); // 50 chars
|
||||
|
||||
const prompt = createDummyNode(
|
||||
'ep1',
|
||||
'USER_PROMPT',
|
||||
50,
|
||||
{
|
||||
semanticParts: [{ type: 'text', text: longText }],
|
||||
},
|
||||
'prompt-id',
|
||||
) as UserPrompt;
|
||||
|
||||
const thought = createDummyNode(
|
||||
'ep1',
|
||||
'AGENT_THOUGHT',
|
||||
50,
|
||||
{
|
||||
text: longText,
|
||||
},
|
||||
'thought-id',
|
||||
) as AgentThought;
|
||||
|
||||
const tool = createDummyToolNode(
|
||||
'ep1',
|
||||
5,
|
||||
500,
|
||||
{
|
||||
observation: { result: 'A'.repeat(500) },
|
||||
},
|
||||
'tool-id',
|
||||
);
|
||||
|
||||
const targets = [prompt, thought, tool];
|
||||
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
|
||||
expect(result.length).toBe(3);
|
||||
|
||||
// 1. User Prompt
|
||||
const compressedPrompt = result[0] as UserPrompt;
|
||||
expect(compressedPrompt.id).not.toBe(prompt.id);
|
||||
expect(compressedPrompt.semanticParts[0].type).toBe('text');
|
||||
assert(compressedPrompt.semanticParts[0].type === 'text');
|
||||
expect(compressedPrompt.semanticParts[0].text).toBe('Mocked Summary!');
|
||||
|
||||
// 2. Agent Thought
|
||||
const compressedThought = result[1] as AgentThought;
|
||||
expect(compressedThought.id).not.toBe(thought.id);
|
||||
expect(compressedThought.text).toBe('Mocked Summary!');
|
||||
|
||||
// 3. Tool Execution
|
||||
const compressedTool = result[2] as ToolExecution;
|
||||
expect(compressedTool.id).not.toBe(tool.id);
|
||||
expect(compressedTool.observation).toEqual({ summary: 'Mocked Summary!' });
|
||||
|
||||
expect(mockLlmClient.generateContent).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
|
||||
it('should ignore nodes that are below the threshold', async () => {
|
||||
const mockLlmClient = createMockLlmClient(['S']); // length = 1
|
||||
|
||||
const env = createMockEnvironment({
|
||||
llmClient: mockLlmClient,
|
||||
});
|
||||
|
||||
const processor = NodeDistillationProcessor.create(env, {
|
||||
nodeThresholdTokens: 100, // Very high threshold
|
||||
});
|
||||
|
||||
const shortText = 'Short text'; // 10 chars
|
||||
|
||||
const prompt = createDummyNode(
|
||||
'ep1',
|
||||
'USER_PROMPT',
|
||||
10,
|
||||
{
|
||||
semanticParts: [{ type: 'text', text: shortText }],
|
||||
},
|
||||
'prompt-id',
|
||||
) as UserPrompt;
|
||||
|
||||
const thought = createDummyNode(
|
||||
'ep1',
|
||||
'AGENT_THOUGHT',
|
||||
13,
|
||||
{
|
||||
text: 'Short thought',
|
||||
},
|
||||
'thought-id',
|
||||
) as AgentThought;
|
||||
|
||||
const targets = [prompt, thought];
|
||||
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
|
||||
expect(result.length).toBe(2);
|
||||
|
||||
// 1. User Prompt (NOT compressed)
|
||||
const untouchedPrompt = result[0] as UserPrompt;
|
||||
expect(untouchedPrompt.id).toBe(prompt.id);
|
||||
|
||||
// 2. Agent Thought (NOT compressed)
|
||||
const untouchedThought = result[1] as AgentThought;
|
||||
expect(untouchedThought.id).toBe(thought.id);
|
||||
|
||||
// LLM should not have been called
|
||||
expect(mockLlmClient.generateContent).toHaveBeenCalledTimes(0);
|
||||
});
|
||||
});
|
||||
@@ -3,7 +3,7 @@
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
|
||||
import type { ContextProcessorFn, ProcessArgs } from '../pipeline.js';
|
||||
import type { ConcreteNode } from '../ir/types.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
@@ -15,48 +15,20 @@ export interface NodeDistillationProcessorOptions {
|
||||
nodeThresholdTokens: number;
|
||||
}
|
||||
|
||||
export class NodeDistillationProcessor implements ContextProcessor {
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: NodeDistillationProcessorOptions,
|
||||
): NodeDistillationProcessor {
|
||||
return new NodeDistillationProcessor(env, options);
|
||||
}
|
||||
|
||||
static readonly schema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
nodeThresholdTokens: {
|
||||
type: 'number',
|
||||
description: 'The token threshold above which nodes are summarized.',
|
||||
},
|
||||
},
|
||||
required: ['nodeThresholdTokens'],
|
||||
};
|
||||
|
||||
readonly componentType = 'processor';
|
||||
readonly id = 'NodeDistillationProcessor';
|
||||
readonly name = 'NodeDistillationProcessor';
|
||||
readonly options: NodeDistillationProcessorOptions;
|
||||
private env: ContextEnvironment;
|
||||
|
||||
constructor(
|
||||
env: ContextEnvironment,
|
||||
options: NodeDistillationProcessorOptions,
|
||||
) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
private async generateSummary(
|
||||
export function createNodeDistillationProcessor(
|
||||
id: string,
|
||||
env: ContextEnvironment,
|
||||
options: NodeDistillationProcessorOptions,
|
||||
): ContextProcessorFn {
|
||||
const generateSummary = async (
|
||||
text: string,
|
||||
contextInfo: string,
|
||||
): Promise<string> {
|
||||
): Promise<string> => {
|
||||
try {
|
||||
const response = await this.env.llmClient.generateContent({
|
||||
const response = await env.llmClient.generateContent({
|
||||
role: LlmRole.UTILITY_COMPRESSOR,
|
||||
modelConfigKey: { model: 'gemini-3-flash-base' },
|
||||
promptId: this.env.promptId,
|
||||
promptId: env.promptId,
|
||||
abortSignal: new AbortController().signal,
|
||||
contents: [
|
||||
{
|
||||
@@ -81,12 +53,12 @@ export class NodeDistillationProcessor implements ContextProcessor {
|
||||
);
|
||||
return text; // Fallback to original text on API failure
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
async process({ targets }: ProcessArgs): Promise<readonly ConcreteNode[]> {
|
||||
const semanticConfig = this.options;
|
||||
const processor: any = async ({ targets }: ProcessArgs) => {
|
||||
const semanticConfig = options;
|
||||
const limitTokens = semanticConfig.nodeThresholdTokens;
|
||||
const thresholdChars = this.env.tokenCalculator.tokensToChars(limitTokens);
|
||||
const thresholdChars = env.tokenCalculator.tokensToChars(limitTokens);
|
||||
|
||||
const returnedNodes: ConcreteNode[] = [];
|
||||
|
||||
@@ -102,14 +74,14 @@ export class NodeDistillationProcessor implements ContextProcessor {
|
||||
if (part.type !== 'text') continue;
|
||||
|
||||
if (part.text.length > thresholdChars) {
|
||||
const summary = await this.generateSummary(
|
||||
const summary = await generateSummary(
|
||||
part.text,
|
||||
'User Prompt',
|
||||
);
|
||||
const newTokens = this.env.tokenCalculator.estimateTokensForParts(
|
||||
const newTokens = env.tokenCalculator.estimateTokensForParts(
|
||||
[{ text: summary }],
|
||||
);
|
||||
const oldTokens = this.env.tokenCalculator.estimateTokensForParts(
|
||||
const oldTokens = env.tokenCalculator.estimateTokensForParts(
|
||||
[{ text: part.text }],
|
||||
);
|
||||
|
||||
@@ -123,8 +95,9 @@ export class NodeDistillationProcessor implements ContextProcessor {
|
||||
if (modified) {
|
||||
returnedNodes.push({
|
||||
...node,
|
||||
id: this.env.idGenerator.generateId(),
|
||||
id: env.idGenerator.generateId(),
|
||||
semanticParts: newParts,
|
||||
replacesId: node.id,
|
||||
});
|
||||
} else {
|
||||
returnedNodes.push(node);
|
||||
@@ -134,20 +107,21 @@ export class NodeDistillationProcessor implements ContextProcessor {
|
||||
|
||||
case 'AGENT_THOUGHT': {
|
||||
if (node.text.length > thresholdChars) {
|
||||
const summary = await this.generateSummary(
|
||||
const summary = await generateSummary(
|
||||
node.text,
|
||||
'Agent Thought',
|
||||
);
|
||||
const newTokens = this.env.tokenCalculator.estimateTokensForParts([
|
||||
const newTokens = env.tokenCalculator.estimateTokensForParts([
|
||||
{ text: summary },
|
||||
]);
|
||||
const oldTokens = this.env.tokenCalculator.getTokenCost(node);
|
||||
const oldTokens = env.tokenCalculator.getTokenCost(node);
|
||||
|
||||
if (newTokens < oldTokens) {
|
||||
returnedNodes.push({
|
||||
...node,
|
||||
id: this.env.idGenerator.generateId(),
|
||||
id: env.idGenerator.generateId(),
|
||||
text: summary,
|
||||
replacesId: node.id,
|
||||
});
|
||||
break;
|
||||
}
|
||||
@@ -171,14 +145,14 @@ export class NodeDistillationProcessor implements ContextProcessor {
|
||||
}
|
||||
|
||||
if (stringifiedObs.length > thresholdChars) {
|
||||
const summary = await this.generateSummary(
|
||||
const summary = await generateSummary(
|
||||
stringifiedObs,
|
||||
node.toolName || 'unknown',
|
||||
);
|
||||
const newObsObject = { summary };
|
||||
|
||||
const newObsTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([
|
||||
env.tokenCalculator.estimateTokensForParts([
|
||||
{
|
||||
functionResponse: {
|
||||
name: node.toolName || 'unknown',
|
||||
@@ -190,18 +164,19 @@ export class NodeDistillationProcessor implements ContextProcessor {
|
||||
|
||||
const oldObsTokens =
|
||||
node.tokens?.observation ??
|
||||
this.env.tokenCalculator.getTokenCost(node);
|
||||
env.tokenCalculator.getTokenCost(node);
|
||||
const intentTokens = node.tokens?.intent ?? 0;
|
||||
|
||||
if (newObsTokens < oldObsTokens) {
|
||||
returnedNodes.push({
|
||||
...node,
|
||||
id: this.env.idGenerator.generateId(),
|
||||
id: env.idGenerator.generateId(),
|
||||
observation: newObsObject as Record<string, unknown>,
|
||||
tokens: {
|
||||
intent: intentTokens,
|
||||
observation: newObsTokens,
|
||||
},
|
||||
replacesId: node.id,
|
||||
});
|
||||
break;
|
||||
}
|
||||
@@ -217,5 +192,10 @@ export class NodeDistillationProcessor implements ContextProcessor {
|
||||
}
|
||||
|
||||
return returnedNodes;
|
||||
}
|
||||
};
|
||||
|
||||
processor.id = id;
|
||||
Object.defineProperty(processor, 'name', { value: 'NodeDistillationProcessor' });
|
||||
|
||||
return processor;
|
||||
}
|
||||
|
||||
@@ -1,128 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import assert from 'node:assert';
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { NodeTruncationProcessor } from './nodeTruncationProcessor.js';
|
||||
import {
|
||||
createMockProcessArgs,
|
||||
createMockEnvironment,
|
||||
createDummyNode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import type { UserPrompt, AgentThought, AgentYield } from '../ir/types.js';
|
||||
|
||||
describe('NodeTruncationProcessor', () => {
|
||||
it('should truncate nodes that exceed maxTokensPerNode', async () => {
|
||||
// env.tokenCalculator uses charsPerToken=1 natively.
|
||||
const env = createMockEnvironment();
|
||||
|
||||
const processor = NodeTruncationProcessor.create(env, {
|
||||
maxTokensPerNode: 10, // 10 chars limit
|
||||
});
|
||||
|
||||
const longText = 'A'.repeat(50); // 50 tokens
|
||||
|
||||
const prompt = createDummyNode(
|
||||
'ep1',
|
||||
'USER_PROMPT',
|
||||
50,
|
||||
{
|
||||
semanticParts: [{ type: 'text', text: longText }],
|
||||
},
|
||||
'prompt-id',
|
||||
) as UserPrompt;
|
||||
|
||||
const thought = createDummyNode(
|
||||
'ep1',
|
||||
'AGENT_THOUGHT',
|
||||
50,
|
||||
{
|
||||
text: longText,
|
||||
},
|
||||
'thought-id',
|
||||
) as AgentThought;
|
||||
|
||||
const yieldNode = createDummyNode(
|
||||
'ep1',
|
||||
'AGENT_YIELD',
|
||||
50,
|
||||
{
|
||||
text: longText,
|
||||
},
|
||||
'yield-id',
|
||||
) as AgentYield;
|
||||
|
||||
const targets = [prompt, thought, yieldNode];
|
||||
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
|
||||
expect(result.length).toBe(3);
|
||||
|
||||
// 1. User Prompt
|
||||
const squashedPrompt = result[0] as UserPrompt;
|
||||
expect(squashedPrompt.id).not.toBe(prompt.id);
|
||||
expect(squashedPrompt.semanticParts[0].type).toBe('text');
|
||||
assert(squashedPrompt.semanticParts[0].type === 'text');
|
||||
expect(squashedPrompt.semanticParts[0].text).toContain('[... OMITTED');
|
||||
|
||||
// 2. Agent Thought
|
||||
const squashedThought = result[1] as AgentThought;
|
||||
expect(squashedThought.id).not.toBe(thought.id);
|
||||
expect(squashedThought.text).toContain('[... OMITTED');
|
||||
|
||||
// 3. Agent Yield
|
||||
const squashedYield = result[2] as AgentYield;
|
||||
expect(squashedYield.id).not.toBe(yieldNode.id);
|
||||
expect(squashedYield.text).toContain('[... OMITTED');
|
||||
});
|
||||
|
||||
it('should ignore nodes that are below maxTokensPerNode', async () => {
|
||||
const env = createMockEnvironment();
|
||||
|
||||
const processor = NodeTruncationProcessor.create(env, {
|
||||
maxTokensPerNode: 100, // 100 chars limit
|
||||
});
|
||||
|
||||
const shortText = 'Short text'; // 10 chars
|
||||
|
||||
const prompt = createDummyNode(
|
||||
'ep1',
|
||||
'USER_PROMPT',
|
||||
10,
|
||||
{
|
||||
semanticParts: [{ type: 'text', text: shortText }],
|
||||
},
|
||||
'prompt-id',
|
||||
) as UserPrompt;
|
||||
|
||||
const thought = createDummyNode(
|
||||
'ep1',
|
||||
'AGENT_THOUGHT',
|
||||
13,
|
||||
{
|
||||
text: 'Short thought', // 13 chars
|
||||
},
|
||||
'thought-id',
|
||||
) as AgentThought;
|
||||
|
||||
const targets = [prompt, thought];
|
||||
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
|
||||
expect(result.length).toBe(2);
|
||||
|
||||
// 1. User Prompt (untouched)
|
||||
const squashedPrompt = result[0] as UserPrompt;
|
||||
expect(squashedPrompt.id).toBe(prompt.id);
|
||||
assert(squashedPrompt.semanticParts[0].type === 'text');
|
||||
expect(squashedPrompt.semanticParts[0].text).not.toContain('[... OMITTED');
|
||||
|
||||
// 2. Agent Thought (untouched)
|
||||
const untouchedThought = result[1] as AgentThought;
|
||||
expect(untouchedThought.id).toBe(thought.id);
|
||||
expect(untouchedThought.text).not.toContain('[... OMITTED');
|
||||
});
|
||||
});
|
||||
@@ -3,7 +3,7 @@
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
|
||||
import type { ContextProcessorFn, ProcessArgs } from '../pipeline.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { truncateProportionally } from '../truncation.js';
|
||||
import type { ConcreteNode } from '../ir/types.js';
|
||||
@@ -12,41 +12,12 @@ export interface NodeTruncationProcessorOptions {
|
||||
maxTokensPerNode: number;
|
||||
}
|
||||
|
||||
export class NodeTruncationProcessor implements ContextProcessor {
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: NodeTruncationProcessorOptions,
|
||||
): NodeTruncationProcessor {
|
||||
return new NodeTruncationProcessor(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 componentType = 'processor';
|
||||
readonly id = 'NodeTruncationProcessor';
|
||||
readonly name = 'NodeTruncationProcessor';
|
||||
readonly options: NodeTruncationProcessorOptions;
|
||||
private env: ContextEnvironment;
|
||||
|
||||
constructor(
|
||||
env: ContextEnvironment,
|
||||
options: NodeTruncationProcessorOptions,
|
||||
) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
private tryApplySquash(
|
||||
export function createNodeTruncationProcessor(
|
||||
id: string,
|
||||
env: ContextEnvironment,
|
||||
options: NodeTruncationProcessorOptions,
|
||||
): ContextProcessorFn {
|
||||
const tryApplySquash = (
|
||||
text: string,
|
||||
limitChars: number,
|
||||
): {
|
||||
@@ -54,7 +25,7 @@ export class NodeTruncationProcessor implements ContextProcessor {
|
||||
newTokens: number;
|
||||
oldTokens: number;
|
||||
tokensSaved: number;
|
||||
} | null {
|
||||
} | null => {
|
||||
const originalLength = text.length;
|
||||
if (originalLength <= limitChars) return null;
|
||||
|
||||
@@ -67,8 +38,8 @@ export class NodeTruncationProcessor implements ContextProcessor {
|
||||
if (newText !== text) {
|
||||
// Using accurate TokenCalculator instead of simple math
|
||||
const newTokens =
|
||||
this.env.tokenCalculator.estimateTokensForString(newText);
|
||||
const oldTokens = this.env.tokenCalculator.estimateTokensForString(text);
|
||||
env.tokenCalculator.estimateTokensForString(newText);
|
||||
const oldTokens = env.tokenCalculator.estimateTokensForString(text);
|
||||
const tokensSaved = oldTokens - newTokens;
|
||||
|
||||
if (tokensSaved > 0) {
|
||||
@@ -76,15 +47,15 @@ export class NodeTruncationProcessor implements ContextProcessor {
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
};
|
||||
|
||||
async process({ targets }: ProcessArgs): Promise<readonly ConcreteNode[]> {
|
||||
const processor: any = async ({ targets }: ProcessArgs) => {
|
||||
if (targets.length === 0) {
|
||||
return targets;
|
||||
}
|
||||
|
||||
const { maxTokensPerNode } = this.options;
|
||||
const limitChars = this.env.tokenCalculator.tokensToChars(maxTokensPerNode);
|
||||
const { maxTokensPerNode } = options;
|
||||
const limitChars = env.tokenCalculator.tokensToChars(maxTokensPerNode);
|
||||
|
||||
const returnedNodes: ConcreteNode[] = [];
|
||||
|
||||
@@ -97,7 +68,7 @@ export class NodeTruncationProcessor implements ContextProcessor {
|
||||
for (let j = 0; j < node.semanticParts.length; j++) {
|
||||
const part = node.semanticParts[j];
|
||||
if (part.type === 'text') {
|
||||
const squashResult = this.tryApplySquash(part.text, limitChars);
|
||||
const squashResult = tryApplySquash(part.text, limitChars);
|
||||
if (squashResult) {
|
||||
newParts[j] = { type: 'text', text: squashResult.text };
|
||||
modified = true;
|
||||
@@ -108,8 +79,9 @@ export class NodeTruncationProcessor implements ContextProcessor {
|
||||
if (modified) {
|
||||
returnedNodes.push({
|
||||
...node,
|
||||
id: this.env.idGenerator.generateId(),
|
||||
id: env.idGenerator.generateId(),
|
||||
semanticParts: newParts,
|
||||
replacesId: node.id,
|
||||
});
|
||||
} else {
|
||||
returnedNodes.push(node);
|
||||
@@ -118,12 +90,13 @@ export class NodeTruncationProcessor implements ContextProcessor {
|
||||
}
|
||||
|
||||
case 'AGENT_THOUGHT': {
|
||||
const squashResult = this.tryApplySquash(node.text, limitChars);
|
||||
const squashResult = tryApplySquash(node.text, limitChars);
|
||||
if (squashResult) {
|
||||
returnedNodes.push({
|
||||
...node,
|
||||
id: this.env.idGenerator.generateId(),
|
||||
id: env.idGenerator.generateId(),
|
||||
text: squashResult.text,
|
||||
replacesId: node.id,
|
||||
});
|
||||
} else {
|
||||
returnedNodes.push(node);
|
||||
@@ -132,12 +105,13 @@ export class NodeTruncationProcessor implements ContextProcessor {
|
||||
}
|
||||
|
||||
case 'AGENT_YIELD': {
|
||||
const squashResult = this.tryApplySquash(node.text, limitChars);
|
||||
const squashResult = tryApplySquash(node.text, limitChars);
|
||||
if (squashResult) {
|
||||
returnedNodes.push({
|
||||
...node,
|
||||
id: this.env.idGenerator.generateId(),
|
||||
id: env.idGenerator.generateId(),
|
||||
text: squashResult.text,
|
||||
replacesId: node.id,
|
||||
});
|
||||
} else {
|
||||
returnedNodes.push(node);
|
||||
@@ -152,5 +126,10 @@ export class NodeTruncationProcessor implements ContextProcessor {
|
||||
}
|
||||
|
||||
return returnedNodes;
|
||||
}
|
||||
};
|
||||
|
||||
processor.id = id;
|
||||
Object.defineProperty(processor, 'name', { value: 'NodeTruncationProcessor' });
|
||||
|
||||
return processor;
|
||||
}
|
||||
|
||||
@@ -1,86 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { RollingSummaryProcessor } from './rollingSummaryProcessor.js';
|
||||
import {
|
||||
createMockProcessArgs,
|
||||
createMockEnvironment,
|
||||
createDummyNode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
|
||||
describe('RollingSummaryProcessor', () => {
|
||||
it('should initialize with correct default options', () => {
|
||||
const env = createMockEnvironment();
|
||||
const processor = RollingSummaryProcessor.create(env, {
|
||||
target: 'incremental',
|
||||
});
|
||||
expect(processor.id).toBe('RollingSummaryProcessor');
|
||||
});
|
||||
|
||||
it('should summarize older nodes when the deficit exceeds the threshold', async () => {
|
||||
// env.tokenCalculator uses charsPerToken=1 based on createMockEnvironment
|
||||
const env = createMockEnvironment();
|
||||
|
||||
// We want to free exactly 100 tokens.
|
||||
// We will supply nodes that cost 50 tokens each.
|
||||
const processor = RollingSummaryProcessor.create(env, {
|
||||
target: 'freeNTokens',
|
||||
freeTokensTarget: 100,
|
||||
});
|
||||
|
||||
const text50 = 'A'.repeat(50);
|
||||
const targets = [
|
||||
createDummyNode(
|
||||
'ep1',
|
||||
'USER_PROMPT',
|
||||
50,
|
||||
{ semanticParts: [{ type: 'text', text: text50 }] },
|
||||
'id1',
|
||||
),
|
||||
createDummyNode('ep1', 'AGENT_THOUGHT', 50, { text: text50 }, 'id2'),
|
||||
createDummyNode('ep1', 'AGENT_YIELD', 50, { text: text50 }, 'id3'),
|
||||
];
|
||||
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
|
||||
// 3 nodes at 50 cost each.
|
||||
// The first node (id1) is the initial USER_PROMPT and is always skipped by RollingSummaryProcessor.
|
||||
// Node id2 adds 50 deficit. Node id3 adds 50 deficit. Total = 100 deficit, which hits the target break point.
|
||||
// Thus, id2 and id3 are summarized into a new ROLLING_SUMMARY node.
|
||||
expect(result.length).toBe(2);
|
||||
expect(result[0].type).toBe('USER_PROMPT');
|
||||
expect(result[1].type).toBe('ROLLING_SUMMARY');
|
||||
});
|
||||
|
||||
it('should preserve targets if deficit does not trigger summary', async () => {
|
||||
const env = createMockEnvironment();
|
||||
|
||||
// We want to free 100 tokens, but our nodes will only cost 10 tokens each.
|
||||
const processor = RollingSummaryProcessor.create(env, {
|
||||
target: 'freeNTokens',
|
||||
freeTokensTarget: 100,
|
||||
});
|
||||
|
||||
const text10 = 'A'.repeat(10);
|
||||
const targets = [
|
||||
createDummyNode(
|
||||
'ep1',
|
||||
'USER_PROMPT',
|
||||
10,
|
||||
{ semanticParts: [{ type: 'text', text: text10 }] },
|
||||
'id1',
|
||||
),
|
||||
createDummyNode('ep1', 'AGENT_THOUGHT', 10, { text: text10 }, 'id2'),
|
||||
];
|
||||
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
|
||||
// Deficit accumulator reaches 10. This is < 100 limit, and total summarizable nodes < 2 anyway.
|
||||
expect(result.length).toBe(2);
|
||||
expect(result[0].type).toBe('USER_PROMPT');
|
||||
expect(result[1].type).toBe('AGENT_THOUGHT');
|
||||
});
|
||||
});
|
||||
@@ -4,7 +4,7 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type {
|
||||
ContextProcessor,
|
||||
ContextProcessorFn,
|
||||
ProcessArgs,
|
||||
BackstopTargetOptions,
|
||||
} from '../pipeline.js';
|
||||
@@ -17,43 +17,17 @@ export interface RollingSummaryProcessorOptions extends BackstopTargetOptions {
|
||||
systemInstruction?: string;
|
||||
}
|
||||
|
||||
export class RollingSummaryProcessor implements ContextProcessor {
|
||||
static readonly schema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
target: { type: 'string', enum: ['incremental', 'freeNTokens', 'max'] },
|
||||
freeTokensTarget: { type: 'number' },
|
||||
systemInstruction: { type: 'string' },
|
||||
},
|
||||
};
|
||||
export function createRollingSummaryProcessor(
|
||||
id: string,
|
||||
env: ContextEnvironment,
|
||||
options: RollingSummaryProcessorOptions,
|
||||
): ContextProcessorFn {
|
||||
const generator = new SnapshotGenerator(env);
|
||||
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: RollingSummaryProcessorOptions,
|
||||
): RollingSummaryProcessor {
|
||||
return new RollingSummaryProcessor(env, options);
|
||||
}
|
||||
|
||||
readonly componentType = 'processor';
|
||||
readonly id = 'RollingSummaryProcessor';
|
||||
readonly name = 'RollingSummaryProcessor';
|
||||
readonly options: RollingSummaryProcessorOptions;
|
||||
private readonly env: ContextEnvironment;
|
||||
private readonly generator: SnapshotGenerator;
|
||||
|
||||
constructor(
|
||||
env: ContextEnvironment,
|
||||
options: RollingSummaryProcessorOptions,
|
||||
) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
this.generator = new SnapshotGenerator(env);
|
||||
}
|
||||
|
||||
async process({ targets }: ProcessArgs): Promise<readonly ConcreteNode[]> {
|
||||
const processor: any = async ({ targets }: ProcessArgs) => {
|
||||
if (targets.length === 0) return targets;
|
||||
|
||||
const strategy = this.options.target ?? 'max';
|
||||
const strategy = options.target ?? 'max';
|
||||
let targetTokensToRemove = 0;
|
||||
|
||||
if (strategy === 'incremental') {
|
||||
@@ -62,7 +36,7 @@ export class RollingSummaryProcessor implements ContextProcessor {
|
||||
// Ideally, the orchestrator should pass `tokensToRemove` explicitly.
|
||||
targetTokensToRemove = 10000;
|
||||
} else if (strategy === 'freeNTokens') {
|
||||
targetTokensToRemove = this.options.freeTokensTarget ?? Infinity;
|
||||
targetTokensToRemove = options.freeTokensTarget ?? Infinity;
|
||||
} else if (strategy === 'max') {
|
||||
targetTokensToRemove = Infinity;
|
||||
}
|
||||
@@ -80,7 +54,7 @@ export class RollingSummaryProcessor implements ContextProcessor {
|
||||
}
|
||||
|
||||
nodesToSummarize.push(node);
|
||||
deficitAccumulator += this.env.tokenCalculator.getTokenCost(node);
|
||||
deficitAccumulator += env.tokenCalculator.getTokenCost(node);
|
||||
|
||||
if (deficitAccumulator >= targetTokensToRemove) break;
|
||||
}
|
||||
@@ -89,11 +63,11 @@ export class RollingSummaryProcessor implements ContextProcessor {
|
||||
|
||||
try {
|
||||
// Synthesize the rolling summary synchronously
|
||||
const snapshotText = await this.generator.synthesizeSnapshot(
|
||||
const snapshotText = await generator.synthesizeSnapshot(
|
||||
nodesToSummarize,
|
||||
this.options.systemInstruction,
|
||||
options.systemInstruction,
|
||||
);
|
||||
const newId = this.env.idGenerator.generateId();
|
||||
const newId = env.idGenerator.generateId();
|
||||
|
||||
const summaryNode: RollingSummary = {
|
||||
id: newId,
|
||||
@@ -101,6 +75,7 @@ export class RollingSummaryProcessor implements ContextProcessor {
|
||||
type: 'ROLLING_SUMMARY',
|
||||
timestamp: Date.now(),
|
||||
text: snapshotText,
|
||||
abstractsIds: nodesToSummarize.map((n) => n.id),
|
||||
};
|
||||
|
||||
const consumedIds = nodesToSummarize.map((n) => n.id);
|
||||
@@ -121,5 +96,10 @@ export class RollingSummaryProcessor implements ContextProcessor {
|
||||
debugLogger.error('RollingSummaryProcessor failed sync backstop', e);
|
||||
return targets;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
processor.id = id;
|
||||
Object.defineProperty(processor, 'name', { value: 'RollingSummaryProcessor' });
|
||||
|
||||
return processor;
|
||||
}
|
||||
|
||||
@@ -1,115 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { StateSnapshotProcessor } from './stateSnapshotProcessor.js';
|
||||
import {
|
||||
createMockEnvironment,
|
||||
createDummyNode,
|
||||
createMockProcessArgs,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import type { InboxSnapshotImpl } from '../sidecar/inbox.js';
|
||||
|
||||
describe('StateSnapshotProcessor', () => {
|
||||
it('should ignore if budget is satisfied', async () => {
|
||||
const env = createMockEnvironment();
|
||||
const processor = StateSnapshotProcessor.create(env, {
|
||||
target: 'incremental',
|
||||
});
|
||||
const targets = [createDummyNode('ep1', 'USER_PROMPT')];
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
expect(result).toBe(targets); // Strict equality
|
||||
});
|
||||
|
||||
it('should apply a valid snapshot from the Inbox (Fast Path)', async () => {
|
||||
const env = createMockEnvironment();
|
||||
const processor = StateSnapshotProcessor.create(env, {
|
||||
target: 'incremental',
|
||||
});
|
||||
|
||||
const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A');
|
||||
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
|
||||
const nodeC = createDummyNode('ep2', 'USER_PROMPT', 50, {}, 'node-C');
|
||||
|
||||
const targets = [nodeA, nodeB, nodeC];
|
||||
|
||||
// The background worker created a snapshot of A and B
|
||||
const messages = [
|
||||
{
|
||||
id: 'msg-1',
|
||||
topic: 'PROPOSED_SNAPSHOT',
|
||||
timestamp: Date.now(),
|
||||
payload: {
|
||||
consumedIds: ['node-A', 'node-B'],
|
||||
newText: '<compressed A and B>',
|
||||
type: 'point-in-time',
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
const processArgs = createMockProcessArgs(targets, [], messages);
|
||||
const result = await processor.process(processArgs);
|
||||
|
||||
// Should remove A and B, insert Snapshot, keep C
|
||||
expect(result.length).toBe(2);
|
||||
expect(result[0].type).toBe('SNAPSHOT');
|
||||
expect(result[1].id).toBe('node-C');
|
||||
|
||||
// Should consume the message
|
||||
expect(
|
||||
(processArgs.inbox as InboxSnapshotImpl).getConsumedIds().has('msg-1'),
|
||||
).toBe(true);
|
||||
});
|
||||
|
||||
it('should reject a snapshot if the nodes were modified/deleted (Cache Invalidated)', async () => {
|
||||
const env = createMockEnvironment();
|
||||
const processor = StateSnapshotProcessor.create(env, {
|
||||
target: 'incremental',
|
||||
});
|
||||
// Make deficit 0 so we don't fall through to the sync backstop and fail the test that way
|
||||
|
||||
// node-A is MISSING (user deleted it)
|
||||
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
|
||||
const targets = [nodeB];
|
||||
|
||||
const messages = [
|
||||
{
|
||||
id: 'msg-1',
|
||||
topic: 'PROPOSED_SNAPSHOT',
|
||||
timestamp: Date.now(),
|
||||
payload: {
|
||||
consumedIds: ['node-A', 'node-B'],
|
||||
newText: '<compressed A and B>',
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
const processArgs = createMockProcessArgs(targets, [], messages);
|
||||
const result = await processor.process(processArgs);
|
||||
|
||||
// Because deficit is 0, and Inbox was rejected, nothing should change
|
||||
expect(result.length).toBe(1);
|
||||
expect(result[0].id).toBe('node-B');
|
||||
expect(
|
||||
(processArgs.inbox as InboxSnapshotImpl).getConsumedIds().has('msg-1'),
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it('should fall back to sync backstop if inbox is empty', async () => {
|
||||
const env = createMockEnvironment();
|
||||
const processor = StateSnapshotProcessor.create(env, { target: 'max' }); // Summarize all
|
||||
|
||||
const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A');
|
||||
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
|
||||
const nodeC = createDummyNode('ep2', 'USER_PROMPT', 50, {}, 'node-C');
|
||||
const targets = [nodeA, nodeB, nodeC];
|
||||
const result = await processor.process(createMockProcessArgs(targets));
|
||||
|
||||
// Should synthesize a new snapshot synchronously
|
||||
expect(env.llmClient.generateContent).toHaveBeenCalled();
|
||||
expect(result.length).toBe(2); // nodeA is skipped as "system prompt", snapshot + nodeA
|
||||
expect(result[1].type).toBe('SNAPSHOT');
|
||||
});
|
||||
});
|
||||
@@ -4,7 +4,7 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type {
|
||||
ContextProcessor,
|
||||
ContextProcessorFn,
|
||||
ProcessArgs,
|
||||
BackstopTargetOptions,
|
||||
} from '../pipeline.js';
|
||||
@@ -18,48 +18,23 @@ export interface StateSnapshotProcessorOptions extends BackstopTargetOptions {
|
||||
systemInstruction?: string;
|
||||
}
|
||||
|
||||
export class StateSnapshotProcessor implements ContextProcessor {
|
||||
static readonly schema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
target: { type: 'string', enum: ['incremental', 'freeNTokens', 'max'] },
|
||||
freeTokensTarget: { type: 'number' },
|
||||
model: { type: 'string' },
|
||||
systemInstruction: { type: 'string' },
|
||||
},
|
||||
};
|
||||
export function createStateSnapshotProcessor(
|
||||
id: string,
|
||||
env: ContextEnvironment,
|
||||
options: StateSnapshotProcessorOptions,
|
||||
): ContextProcessorFn {
|
||||
const generator = new SnapshotGenerator(env);
|
||||
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: StateSnapshotProcessorOptions,
|
||||
): StateSnapshotProcessor {
|
||||
return new StateSnapshotProcessor(env, options);
|
||||
}
|
||||
|
||||
readonly componentType = 'processor';
|
||||
readonly id = 'StateSnapshotProcessor';
|
||||
readonly name = 'StateSnapshotProcessor';
|
||||
readonly options: StateSnapshotProcessorOptions;
|
||||
private readonly env: ContextEnvironment;
|
||||
private readonly generator: SnapshotGenerator;
|
||||
|
||||
constructor(env: ContextEnvironment, options: StateSnapshotProcessorOptions) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
this.generator = new SnapshotGenerator(env);
|
||||
}
|
||||
|
||||
// --- ContextProcessor Interface (Sync Backstop / Cache Application) ---
|
||||
async process({
|
||||
const processor: any = async ({
|
||||
targets,
|
||||
inbox,
|
||||
}: ProcessArgs): Promise<readonly ConcreteNode[]> {
|
||||
}: ProcessArgs) => {
|
||||
if (targets.length === 0) {
|
||||
return targets;
|
||||
}
|
||||
|
||||
// Determine what mode we are looking for: 'incremental' -> 'point-in-time', 'max' -> 'accumulate'
|
||||
const strategy = this.options.target ?? 'max';
|
||||
const strategy = options.target ?? 'max';
|
||||
const expectedType =
|
||||
strategy === 'incremental' ? 'point-in-time' : 'accumulate';
|
||||
|
||||
@@ -90,7 +65,7 @@ export class StateSnapshotProcessor implements ContextProcessor {
|
||||
|
||||
if (isValid) {
|
||||
// If valid, apply it!
|
||||
const newId = this.env.idGenerator.generateId();
|
||||
const newId = env.idGenerator.generateId();
|
||||
|
||||
const snapshotNode: Snapshot = {
|
||||
id: newId,
|
||||
@@ -98,6 +73,7 @@ export class StateSnapshotProcessor implements ContextProcessor {
|
||||
type: 'SNAPSHOT',
|
||||
timestamp: Date.now(),
|
||||
text: newText,
|
||||
abstractsIds: consumedIds,
|
||||
};
|
||||
|
||||
// Remove the consumed nodes and insert the snapshot at the earliest index
|
||||
@@ -127,7 +103,7 @@ export class StateSnapshotProcessor implements ContextProcessor {
|
||||
if (strategy === 'incremental') {
|
||||
targetTokensToRemove = Infinity; // incremental implies removing as much as possible if no state is passed
|
||||
} else if (strategy === 'freeNTokens') {
|
||||
targetTokensToRemove = this.options.freeTokensTarget ?? Infinity;
|
||||
targetTokensToRemove = options.freeTokensTarget ?? Infinity;
|
||||
} else if (strategy === 'max') {
|
||||
targetTokensToRemove = Infinity;
|
||||
}
|
||||
@@ -144,7 +120,7 @@ export class StateSnapshotProcessor implements ContextProcessor {
|
||||
}
|
||||
|
||||
nodesToSummarize.push(node);
|
||||
deficitAccumulator += this.env.tokenCalculator.getTokenCost(node);
|
||||
deficitAccumulator += env.tokenCalculator.getTokenCost(node);
|
||||
|
||||
if (deficitAccumulator >= targetTokensToRemove) break;
|
||||
}
|
||||
@@ -152,17 +128,18 @@ export class StateSnapshotProcessor implements ContextProcessor {
|
||||
if (nodesToSummarize.length < 2) return targets; // Not enough context
|
||||
|
||||
try {
|
||||
const snapshotText = await this.generator.synthesizeSnapshot(
|
||||
const snapshotText = await generator.synthesizeSnapshot(
|
||||
nodesToSummarize,
|
||||
this.options.systemInstruction,
|
||||
options.systemInstruction,
|
||||
);
|
||||
const newId = this.env.idGenerator.generateId();
|
||||
const newId = env.idGenerator.generateId();
|
||||
const snapshotNode: Snapshot = {
|
||||
id: newId,
|
||||
logicalParentId: newId,
|
||||
type: 'SNAPSHOT',
|
||||
timestamp: Date.now(),
|
||||
text: snapshotText,
|
||||
abstractsIds: nodesToSummarize.map((n) => n.id),
|
||||
};
|
||||
|
||||
const consumedIds = nodesToSummarize.map((n) => n.id);
|
||||
@@ -183,5 +160,10 @@ export class StateSnapshotProcessor implements ContextProcessor {
|
||||
debugLogger.error('StateSnapshotProcessor failed sync backstop', e);
|
||||
return targets;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
processor.id = id;
|
||||
Object.defineProperty(processor, 'name', { value: 'StateSnapshotProcessor' });
|
||||
|
||||
return processor;
|
||||
}
|
||||
|
||||
@@ -1,112 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import { StateSnapshotWorker } from './stateSnapshotWorker.js';
|
||||
import {
|
||||
createMockEnvironment,
|
||||
createDummyNode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import { InboxSnapshotImpl } from '../sidecar/inbox.js';
|
||||
|
||||
describe('StateSnapshotWorker', () => {
|
||||
it('should generate a snapshot and publish it to the inbox', async () => {
|
||||
const env = createMockEnvironment();
|
||||
// Spy on the publish method
|
||||
const publishSpy = vi.spyOn(env.inbox, 'publish');
|
||||
|
||||
const worker = StateSnapshotWorker.create(env, { type: 'point-in-time' });
|
||||
|
||||
const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A');
|
||||
const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B');
|
||||
|
||||
const targets = [nodeA, nodeB];
|
||||
const inbox = new InboxSnapshotImpl([]);
|
||||
|
||||
await worker.execute({ targets, inbox });
|
||||
|
||||
// Ensure generateContent was called
|
||||
expect(env.llmClient.generateContent).toHaveBeenCalled();
|
||||
|
||||
// Verify it published to the inbox
|
||||
expect(publishSpy).toHaveBeenCalledWith(
|
||||
'PROPOSED_SNAPSHOT',
|
||||
expect.objectContaining({
|
||||
newText: 'Mock LLM summary response',
|
||||
consumedIds: ['node-A', 'node-B'],
|
||||
type: 'point-in-time',
|
||||
}),
|
||||
env.idGenerator,
|
||||
);
|
||||
});
|
||||
|
||||
it('should pull previous accumulate snapshot from inbox and append new targets', async () => {
|
||||
const env = createMockEnvironment();
|
||||
const publishSpy = vi.spyOn(env.inbox, 'publish');
|
||||
const drainSpy = vi.spyOn(env.inbox, 'drainConsumed');
|
||||
|
||||
const worker = StateSnapshotWorker.create(env, { type: 'accumulate' });
|
||||
|
||||
const nodeC = createDummyNode('ep2', 'USER_PROMPT', 50, {}, 'node-C');
|
||||
const targets = [nodeC];
|
||||
|
||||
// Simulate an existing accumulate draft in the inbox
|
||||
const inbox = new InboxSnapshotImpl([
|
||||
{
|
||||
id: 'draft-1',
|
||||
topic: 'PROPOSED_SNAPSHOT',
|
||||
timestamp: Date.now() - 1000,
|
||||
payload: {
|
||||
consumedIds: ['node-A', 'node-B'],
|
||||
newText: '<old snapshot>',
|
||||
type: 'accumulate',
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
await worker.execute({ targets, inbox });
|
||||
|
||||
// The old draft should be consumed
|
||||
expect(inbox.getConsumedIds().has('draft-1')).toBe(true);
|
||||
expect(drainSpy).toHaveBeenCalledWith(expect.any(Set));
|
||||
|
||||
// The new publish should contain ALL consumed IDs (old + new)
|
||||
expect(publishSpy).toHaveBeenCalledWith(
|
||||
'PROPOSED_SNAPSHOT',
|
||||
expect.objectContaining({
|
||||
newText: 'Mock LLM summary response',
|
||||
consumedIds: ['node-A', 'node-B', 'node-C'], // Aggregated!
|
||||
type: 'accumulate',
|
||||
}),
|
||||
env.idGenerator,
|
||||
);
|
||||
|
||||
// Verify the LLM was called with the old snapshot prepended
|
||||
expect(env.llmClient.generateContent).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
contents: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
parts: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
text: expect.stringContaining('<old snapshot>'),
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should ignore empty targets', async () => {
|
||||
const env = createMockEnvironment();
|
||||
const publishSpy = vi.spyOn(env.inbox, 'publish');
|
||||
const worker = StateSnapshotWorker.create(env, { type: 'accumulate' });
|
||||
|
||||
await worker.execute({ targets: [], inbox: new InboxSnapshotImpl([]) });
|
||||
|
||||
expect(env.llmClient.generateContent).not.toHaveBeenCalled();
|
||||
expect(publishSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -14,53 +14,29 @@ export interface StateSnapshotWorkerOptions {
|
||||
systemInstruction?: string;
|
||||
}
|
||||
|
||||
export class StateSnapshotWorker implements ContextWorker {
|
||||
static readonly schema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
type: { type: 'string', enum: ['accumulate', 'point-in-time'] },
|
||||
systemInstruction: { type: 'string' },
|
||||
},
|
||||
};
|
||||
export function createStateSnapshotWorker(
|
||||
id: string,
|
||||
env: ContextEnvironment,
|
||||
options: StateSnapshotWorkerOptions,
|
||||
): ContextWorker {
|
||||
const generator = new SnapshotGenerator(env);
|
||||
|
||||
static create(
|
||||
env: ContextEnvironment,
|
||||
options: StateSnapshotWorkerOptions,
|
||||
): StateSnapshotWorker {
|
||||
return new StateSnapshotWorker(env, options);
|
||||
}
|
||||
let isRunning = false;
|
||||
|
||||
readonly componentType = 'worker';
|
||||
readonly id = 'StateSnapshotWorker';
|
||||
readonly name = 'StateSnapshotWorker';
|
||||
readonly options: StateSnapshotWorkerOptions;
|
||||
private readonly env: ContextEnvironment;
|
||||
private readonly generator: SnapshotGenerator;
|
||||
|
||||
// Triggers when nodes exceed retained threshold (via retained_exceeded in Orchestrator)
|
||||
readonly triggers = {
|
||||
onNodesAgedOut: true,
|
||||
};
|
||||
|
||||
constructor(env: ContextEnvironment, options: StateSnapshotWorkerOptions) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
this.generator = new SnapshotGenerator(env);
|
||||
}
|
||||
|
||||
async execute({
|
||||
const execute = async ({
|
||||
targets,
|
||||
inbox,
|
||||
}: {
|
||||
targets: readonly ConcreteNode[];
|
||||
inbox: InboxSnapshot;
|
||||
}): Promise<void> {
|
||||
}): Promise<void> => {
|
||||
if (!isRunning) return;
|
||||
if (targets.length === 0) return;
|
||||
|
||||
try {
|
||||
let nodesToSummarize = [...targets];
|
||||
let previousConsumedIds: string[] = [];
|
||||
const workerType = this.options.type ?? 'point-in-time';
|
||||
const workerType = options.type ?? 'point-in-time';
|
||||
|
||||
if (workerType === 'accumulate') {
|
||||
// Look for the most recent unconsumed accumulate snapshot in the inbox
|
||||
@@ -83,13 +59,13 @@ export class StateSnapshotWorker implements ContextWorker {
|
||||
inbox.consume(latest.id);
|
||||
// And we must persist its consumption back to the live inbox immediately,
|
||||
// because we are effectively "taking" it from the shelf to modify.
|
||||
this.env.inbox.drainConsumed(new Set([latest.id]));
|
||||
env.inbox.drainConsumed(new Set([latest.id]));
|
||||
|
||||
previousConsumedIds = latest.payload.consumedIds;
|
||||
|
||||
// Prepend a synthetic node representing the previous rolling state
|
||||
const previousStateNode: ConcreteNode = {
|
||||
id: this.env.idGenerator.generateId(),
|
||||
id: env.idGenerator.generateId(),
|
||||
logicalParentId: '',
|
||||
type: 'SNAPSHOT',
|
||||
timestamp: latest.timestamp,
|
||||
@@ -100,9 +76,9 @@ export class StateSnapshotWorker implements ContextWorker {
|
||||
}
|
||||
}
|
||||
|
||||
const snapshotText = await this.generator.synthesizeSnapshot(
|
||||
const snapshotText = await generator.synthesizeSnapshot(
|
||||
nodesToSummarize,
|
||||
this.options.systemInstruction,
|
||||
options.systemInstruction,
|
||||
);
|
||||
|
||||
const newConsumedIds = [
|
||||
@@ -111,17 +87,32 @@ export class StateSnapshotWorker implements ContextWorker {
|
||||
];
|
||||
|
||||
// In V2, workers communicate their work to the inbox, and the processor picks it up.
|
||||
this.env.inbox.publish(
|
||||
env.inbox.publish(
|
||||
'PROPOSED_SNAPSHOT',
|
||||
{
|
||||
newText: snapshotText,
|
||||
consumedIds: newConsumedIds,
|
||||
type: workerType,
|
||||
},
|
||||
this.env.idGenerator,
|
||||
env.idGenerator,
|
||||
);
|
||||
} catch (e) {
|
||||
debugLogger.error('StateSnapshotWorker failed to generate snapshot', e);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
return {
|
||||
id,
|
||||
name: 'StateSnapshotWorker',
|
||||
triggers: {
|
||||
onNodesAgedOut: true,
|
||||
},
|
||||
start: () => {
|
||||
isRunning = true;
|
||||
},
|
||||
stop: () => {
|
||||
isRunning = false;
|
||||
},
|
||||
execute,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -1,68 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { ToolMaskingProcessor } from './toolMaskingProcessor.js';
|
||||
import {
|
||||
createMockProcessArgs,
|
||||
createMockEnvironment,
|
||||
createDummyToolNode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import type { ToolExecution } from '../ir/types.js';
|
||||
|
||||
describe('ToolMaskingProcessor', () => {
|
||||
it('should write large strings to disk and replace them with a masked pointer', async () => {
|
||||
const env = createMockEnvironment();
|
||||
// env uses charsPerToken=1 natively.
|
||||
// original string lengths > stringLengthThresholdTokens (which is 10) will be masked
|
||||
|
||||
const processor = ToolMaskingProcessor.create(env, {
|
||||
stringLengthThresholdTokens: 10,
|
||||
});
|
||||
|
||||
const longString = 'A'.repeat(500); // 500 chars
|
||||
|
||||
const toolStep = createDummyToolNode('ep1', 50, 500, {
|
||||
observation: {
|
||||
result: longString,
|
||||
metadata: 'short', // 5 chars, will not be masked
|
||||
},
|
||||
});
|
||||
|
||||
const result = await processor.process(createMockProcessArgs([toolStep]));
|
||||
|
||||
expect(result.length).toBe(1);
|
||||
const masked = result[0] as ToolExecution;
|
||||
|
||||
// It should have generated a new ID because it modified it
|
||||
expect(masked.id).not.toBe(toolStep.id);
|
||||
|
||||
// It should have masked the observation
|
||||
const obs = masked.observation as { result: string; metadata: string };
|
||||
expect(obs.result).toContain('<tool_output_masked>');
|
||||
expect(obs.metadata).toBe('short'); // Untouched
|
||||
});
|
||||
|
||||
it('should skip unmaskable tools', async () => {
|
||||
const env = createMockEnvironment();
|
||||
|
||||
const processor = ToolMaskingProcessor.create(env, {
|
||||
stringLengthThresholdTokens: 10,
|
||||
});
|
||||
|
||||
const toolStep = createDummyToolNode('ep1', 10, 10, {
|
||||
toolName: 'activate_skill',
|
||||
observation: {
|
||||
result:
|
||||
'this is a really long string that normally would get masked but wont because of the tool name',
|
||||
},
|
||||
});
|
||||
|
||||
const result = await processor.process(createMockProcessArgs([toolStep]));
|
||||
|
||||
// Returned the exact same object reference
|
||||
expect(result[0]).toBe(toolStep);
|
||||
});
|
||||
});
|
||||
@@ -3,7 +3,7 @@
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
|
||||
import type { ContextProcessorFn, ProcessArgs } from '../pipeline.js';
|
||||
import type { ConcreteNode, ToolExecution } from '../ir/types.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { sanitizeFilenamePart } from '../../utils/fileUtils.js';
|
||||
@@ -63,57 +63,31 @@ function isMaskableRecord(val: unknown): val is Record<string, MaskableValue> {
|
||||
);
|
||||
}
|
||||
|
||||
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'],
|
||||
export function createToolMaskingProcessor(
|
||||
id: string,
|
||||
env: ContextEnvironment,
|
||||
options: ToolMaskingProcessorOptions,
|
||||
): ContextProcessorFn {
|
||||
const isAlreadyMasked = (text: string): boolean => {
|
||||
return text.includes('<tool_output_masked>');
|
||||
};
|
||||
|
||||
readonly componentType = 'processor';
|
||||
readonly id = 'ToolMaskingProcessor';
|
||||
readonly name = 'ToolMaskingProcessor';
|
||||
readonly options: ToolMaskingProcessorOptions;
|
||||
private env: ContextEnvironment;
|
||||
|
||||
constructor(env: ContextEnvironment, options: ToolMaskingProcessorOptions) {
|
||||
this.env = env;
|
||||
this.options = options;
|
||||
}
|
||||
|
||||
private isAlreadyMasked(text: string): boolean {
|
||||
return text.includes('<tool_output_masked>');
|
||||
}
|
||||
|
||||
async process({ targets }: ProcessArgs): Promise<readonly ConcreteNode[]> {
|
||||
const maskingConfig = this.options;
|
||||
const processor: any = async ({ targets }: ProcessArgs) => {
|
||||
const maskingConfig = options;
|
||||
if (!maskingConfig) return targets;
|
||||
if (targets.length === 0) return targets;
|
||||
|
||||
const limitChars = this.env.tokenCalculator.tokensToChars(
|
||||
const limitChars = env.tokenCalculator.tokensToChars(
|
||||
maskingConfig.stringLengthThresholdTokens,
|
||||
);
|
||||
|
||||
let toolOutputsDir = this.env.fileSystem.join(
|
||||
this.env.projectTempDir,
|
||||
let toolOutputsDir = env.fileSystem.join(
|
||||
env.projectTempDir,
|
||||
'tool-outputs',
|
||||
);
|
||||
const sessionId = this.env.sessionId;
|
||||
const sessionId = env.sessionId;
|
||||
if (sessionId) {
|
||||
toolOutputsDir = this.env.fileSystem.join(
|
||||
toolOutputsDir = env.fileSystem.join(
|
||||
toolOutputsDir,
|
||||
`session-${sanitizeFilenamePart(sessionId)}`,
|
||||
);
|
||||
@@ -128,14 +102,14 @@ export class ToolMaskingProcessor implements ContextProcessor {
|
||||
nodeType: string,
|
||||
): Promise<string> => {
|
||||
if (!directoryCreated) {
|
||||
await this.env.fileSystem.mkdir(toolOutputsDir, { recursive: true });
|
||||
await 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);
|
||||
const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${env.idGenerator.generateId()}.txt`;
|
||||
const filePath = env.fileSystem.join(toolOutputsDir, fileName);
|
||||
|
||||
await this.env.fileSystem.writeFile(filePath, content);
|
||||
await env.fileSystem.writeFile(filePath, content);
|
||||
|
||||
const fileSizeMB = (
|
||||
Buffer.byteLength(content, 'utf8') /
|
||||
@@ -164,7 +138,7 @@ export class ToolMaskingProcessor implements ContextProcessor {
|
||||
nodeType: string,
|
||||
): Promise<{ masked: MaskableValue; changed: boolean }> => {
|
||||
if (typeof obj === 'string') {
|
||||
if (obj.length > limitChars && !this.isAlreadyMasked(obj)) {
|
||||
if (obj.length > limitChars && !isAlreadyMasked(obj)) {
|
||||
const newString = await handleMasking(
|
||||
obj,
|
||||
toolName || 'unknown',
|
||||
@@ -222,7 +196,7 @@ export class ToolMaskingProcessor implements ContextProcessor {
|
||||
: undefined;
|
||||
|
||||
const newIntentTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([
|
||||
env.tokenCalculator.estimateTokensForParts([
|
||||
{
|
||||
functionCall: {
|
||||
name: toolName || 'unknown',
|
||||
@@ -244,24 +218,25 @@ export class ToolMaskingProcessor implements ContextProcessor {
|
||||
}
|
||||
|
||||
const newObsTokens =
|
||||
this.env.tokenCalculator.estimateTokensForParts([
|
||||
env.tokenCalculator.estimateTokensForParts([
|
||||
obsPart as Part,
|
||||
]);
|
||||
|
||||
const tokensSaved =
|
||||
this.env.tokenCalculator.getTokenCost(node) -
|
||||
env.tokenCalculator.getTokenCost(node) -
|
||||
(newIntentTokens + newObsTokens);
|
||||
|
||||
if (tokensSaved > 0) {
|
||||
const maskedNode: ToolExecution = {
|
||||
...node,
|
||||
id: this.env.idGenerator.generateId(), // Modified, so generate new ID
|
||||
id: env.idGenerator.generateId(), // Modified, so generate new ID
|
||||
intent: maskedIntent ?? node.intent,
|
||||
observation: maskedObs ?? node.observation,
|
||||
tokens: {
|
||||
intent: newIntentTokens,
|
||||
observation: newObsTokens,
|
||||
},
|
||||
replacesId: node.id,
|
||||
};
|
||||
|
||||
returnedNodes.push(maskedNode);
|
||||
@@ -280,5 +255,10 @@ export class ToolMaskingProcessor implements ContextProcessor {
|
||||
}
|
||||
|
||||
return returnedNodes;
|
||||
}
|
||||
};
|
||||
|
||||
processor.id = id;
|
||||
Object.defineProperty(processor, 'name', { value: 'ToolMaskingProcessor' });
|
||||
|
||||
return processor;
|
||||
}
|
||||
|
||||
@@ -1,66 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { SidecarRegistry } 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: SidecarRegistry;
|
||||
|
||||
beforeEach(() => {
|
||||
fileSystem = new InMemoryFileSystem();
|
||||
registry = new SidecarRegistry();
|
||||
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 },
|
||||
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:');
|
||||
});
|
||||
});
|
||||
@@ -6,12 +6,9 @@
|
||||
|
||||
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 { defaultSidecarProfile, type ContextProfile } from './profiles.js';
|
||||
import type { IFileSystem } from '../system/IFileSystem.js';
|
||||
import { NodeFileSystem } from '../system/NodeFileSystem.js';
|
||||
import type { SidecarRegistry } from './registry.js';
|
||||
|
||||
export class SidecarLoader {
|
||||
/**
|
||||
@@ -20,9 +17,8 @@ export class SidecarLoader {
|
||||
*/
|
||||
static loadFromFile(
|
||||
sidecarPath: string,
|
||||
registry: SidecarRegistry,
|
||||
fileSystem: IFileSystem = new NodeFileSystem(),
|
||||
): SidecarConfig {
|
||||
): ContextProfile {
|
||||
const fileContent = fileSystem.readFileSync(sidecarPath, 'utf8');
|
||||
|
||||
if (!fileContent.trim()) {
|
||||
@@ -40,24 +36,15 @@ export class SidecarLoader {
|
||||
);
|
||||
}
|
||||
|
||||
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.',
|
||||
);
|
||||
const customConfig = parsed as Partial<SidecarConfig>;
|
||||
|
||||
return {
|
||||
...defaultSidecarProfile,
|
||||
config: {
|
||||
...defaultSidecarProfile.config,
|
||||
...(customConfig.budget ? { budget: customConfig.budget } : {})
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -66,9 +53,8 @@ export class SidecarLoader {
|
||||
*/
|
||||
static fromConfig(
|
||||
config: Config,
|
||||
registry: SidecarRegistry,
|
||||
fileSystem: IFileSystem = new NodeFileSystem(),
|
||||
): SidecarConfig {
|
||||
): ContextProfile {
|
||||
const sidecarPath = config.getExperimentalContextSidecarConfig();
|
||||
|
||||
if (sidecarPath && fileSystem.existsSync(sidecarPath)) {
|
||||
@@ -79,7 +65,7 @@ export class SidecarLoader {
|
||||
}
|
||||
|
||||
// If the file has content, enforce strict validation and throw on failure.
|
||||
return this.loadFromFile(sidecarPath, registry, fileSystem);
|
||||
return this.loadFromFile(sidecarPath, fileSystem);
|
||||
}
|
||||
|
||||
return defaultSidecarProfile;
|
||||
|
||||
@@ -1,85 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type { SidecarRegistry } from './registry.js';
|
||||
import {
|
||||
HistoryTruncationProcessor,
|
||||
type HistoryTruncationProcessorOptions,
|
||||
} from '../processors/historyTruncationProcessor.js';
|
||||
import { BlobDegradationProcessor } from '../processors/blobDegradationProcessor.js';
|
||||
import {
|
||||
NodeTruncationProcessor,
|
||||
type NodeTruncationProcessorOptions,
|
||||
} from '../processors/nodeTruncationProcessor.js';
|
||||
import {
|
||||
NodeDistillationProcessor,
|
||||
type NodeDistillationProcessorOptions,
|
||||
} from '../processors/nodeDistillationProcessor.js';
|
||||
import {
|
||||
ToolMaskingProcessor,
|
||||
type ToolMaskingProcessorOptions,
|
||||
} from '../processors/toolMaskingProcessor.js';
|
||||
import {
|
||||
StateSnapshotProcessor,
|
||||
type StateSnapshotProcessorOptions,
|
||||
} from '../processors/stateSnapshotProcessor.js';
|
||||
import {
|
||||
StateSnapshotWorker,
|
||||
type StateSnapshotWorkerOptions,
|
||||
} from '../processors/stateSnapshotWorker.js';
|
||||
import {
|
||||
RollingSummaryProcessor,
|
||||
type RollingSummaryProcessorOptions,
|
||||
} from '../processors/rollingSummaryProcessor.js';
|
||||
|
||||
export function registerBuiltInProcessors(registry: SidecarRegistry) {
|
||||
registry.registerProcessor<Record<string, never>>({
|
||||
id: 'BlobDegradationProcessor',
|
||||
schema: { type: 'object', properties: {} },
|
||||
create: (env) => new BlobDegradationProcessor(env),
|
||||
});
|
||||
|
||||
registry.registerProcessor<HistoryTruncationProcessorOptions>({
|
||||
id: 'HistoryTruncationProcessor',
|
||||
schema: HistoryTruncationProcessor.schema,
|
||||
create: (env, options) => HistoryTruncationProcessor.create(env, options),
|
||||
});
|
||||
|
||||
registry.registerProcessor<NodeTruncationProcessorOptions>({
|
||||
id: 'NodeTruncationProcessor',
|
||||
schema: NodeTruncationProcessor.schema,
|
||||
create: (env, options) => NodeTruncationProcessor.create(env, options),
|
||||
});
|
||||
|
||||
registry.registerProcessor<NodeDistillationProcessorOptions>({
|
||||
id: 'NodeDistillationProcessor',
|
||||
schema: NodeDistillationProcessor.schema,
|
||||
create: (env, options) => NodeDistillationProcessor.create(env, options),
|
||||
});
|
||||
|
||||
registry.registerProcessor<ToolMaskingProcessorOptions>({
|
||||
id: 'ToolMaskingProcessor',
|
||||
schema: ToolMaskingProcessor.schema,
|
||||
create: (env, options) => ToolMaskingProcessor.create(env, options),
|
||||
});
|
||||
|
||||
registry.registerProcessor<StateSnapshotProcessorOptions>({
|
||||
id: 'StateSnapshotProcessor',
|
||||
schema: StateSnapshotProcessor.schema,
|
||||
create: (env, options) => StateSnapshotProcessor.create(env, options),
|
||||
});
|
||||
|
||||
registry.registerWorker<StateSnapshotWorkerOptions>({
|
||||
id: 'StateSnapshotWorker',
|
||||
schema: StateSnapshotWorker.schema,
|
||||
create: (env, options) => StateSnapshotWorker.create(env, options),
|
||||
});
|
||||
|
||||
registry.registerProcessor<RollingSummaryProcessorOptions>({
|
||||
id: 'RollingSummaryProcessor',
|
||||
schema: RollingSummaryProcessor.schema,
|
||||
create: (env, options) => RollingSummaryProcessor.create(env, options),
|
||||
});
|
||||
}
|
||||
@@ -1,316 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import assert from 'node:assert';
|
||||
import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
|
||||
import { PipelineOrchestrator } from './orchestrator.js';
|
||||
import { SidecarRegistry } from './registry.js';
|
||||
import {
|
||||
createMockEnvironment,
|
||||
createDummyNode,
|
||||
} from '../testing/contextTestUtils.js';
|
||||
import type { ContextEnvironment } from './environment.js';
|
||||
import type {
|
||||
ContextProcessor,
|
||||
ContextWorker,
|
||||
InboxSnapshot,
|
||||
ProcessArgs,
|
||||
} from '../pipeline.js';
|
||||
import type { PipelineDef, ProcessorConfig, SidecarConfig } from './types.js';
|
||||
import type { ContextEventBus } from '../eventBus.js';
|
||||
import type { ConcreteNode, UserPrompt } from '../ir/types.js';
|
||||
|
||||
// A realistic mock processor that modifies the text of the first target node
|
||||
class ModifyingProcessor implements ContextProcessor {
|
||||
static create() {
|
||||
return new ModifyingProcessor();
|
||||
}
|
||||
constructor() {}
|
||||
readonly name = 'ModifyingProcessor';
|
||||
readonly id = 'ModifyingProcessor';
|
||||
readonly options = {};
|
||||
async process(args: ProcessArgs) {
|
||||
const newTargets = [...args.targets];
|
||||
if (newTargets.length > 0 && newTargets[0].type === 'USER_PROMPT') {
|
||||
const prompt = newTargets[0];
|
||||
const newParts = [...prompt.semanticParts];
|
||||
if (newParts.length > 0 && newParts[0].type === 'text') {
|
||||
newParts[0] = {
|
||||
...newParts[0],
|
||||
text: newParts[0].text + ' [modified]',
|
||||
};
|
||||
}
|
||||
newTargets[0] = {
|
||||
...prompt,
|
||||
id: prompt.id + '-modified',
|
||||
replacesId: prompt.id,
|
||||
semanticParts: newParts,
|
||||
};
|
||||
}
|
||||
return newTargets;
|
||||
}
|
||||
}
|
||||
|
||||
// A processor that just throws an error
|
||||
class ThrowingProcessor implements ContextProcessor {
|
||||
static create() {
|
||||
return new ThrowingProcessor();
|
||||
}
|
||||
constructor() {}
|
||||
readonly name = 'Throwing';
|
||||
readonly id = 'Throwing';
|
||||
readonly options = {};
|
||||
async process(): Promise<readonly ConcreteNode[]> {
|
||||
throw new Error('Processor failed intentionally');
|
||||
}
|
||||
}
|
||||
|
||||
// A mock worker that signals it ran
|
||||
class MockWorker implements ContextWorker {
|
||||
static create() {
|
||||
return new MockWorker();
|
||||
}
|
||||
constructor() {}
|
||||
readonly name = 'MockWorker';
|
||||
readonly id = 'MockWorker';
|
||||
readonly triggers = {
|
||||
onNodesAdded: true,
|
||||
};
|
||||
wasExecuted = false;
|
||||
|
||||
async execute(args: {
|
||||
targets: readonly ConcreteNode[];
|
||||
inbox: InboxSnapshot;
|
||||
}) {
|
||||
this.wasExecuted = true;
|
||||
if (args.targets.length > 0 && args.targets[0].type === 'USER_PROMPT') {
|
||||
const prompt = args.targets[0];
|
||||
if (prompt.semanticParts[0].type === 'text') {
|
||||
args.inbox.consume('test');
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
describe('PipelineOrchestrator (Component)', () => {
|
||||
let env: ContextEnvironment;
|
||||
let eventBus: ContextEventBus;
|
||||
let registry: SidecarRegistry;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.resetAllMocks();
|
||||
env = createMockEnvironment();
|
||||
eventBus = env.eventBus;
|
||||
registry = new SidecarRegistry();
|
||||
|
||||
registry.registerProcessor({
|
||||
id: 'ModifyingProcessor',
|
||||
schema: {},
|
||||
create: () => new ModifyingProcessor(),
|
||||
});
|
||||
registry.registerProcessor({
|
||||
id: 'ThrowingProcessor',
|
||||
schema: {},
|
||||
create: () => new ThrowingProcessor(),
|
||||
});
|
||||
registry.registerWorker({
|
||||
id: 'MockWorker',
|
||||
schema: {},
|
||||
create: () => new MockWorker(),
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
registry.clear();
|
||||
});
|
||||
|
||||
const createConfig = (
|
||||
pipelines: PipelineDef[],
|
||||
workers?: Array<{ workerId: string }>,
|
||||
): SidecarConfig => ({
|
||||
budget: { maxTokens: 100, retainedTokens: 50 },
|
||||
pipelines,
|
||||
workers,
|
||||
});
|
||||
|
||||
it('instantiates processors and workers from the registry on initialization', () => {
|
||||
const config = createConfig(
|
||||
[
|
||||
{
|
||||
name: 'SyncPipe',
|
||||
triggers: ['new_message'],
|
||||
processors: [
|
||||
{ processorId: 'ModifyingProcessor' } as unknown as ProcessorConfig,
|
||||
],
|
||||
},
|
||||
],
|
||||
[{ workerId: 'MockWorker' }],
|
||||
);
|
||||
|
||||
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('ModifyingProcessor'),
|
||||
).toBe(true);
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
expect((orchestrator as any).instantiatedWorkers.has('MockWorker')).toBe(
|
||||
true,
|
||||
);
|
||||
});
|
||||
|
||||
it('throws an error if a config requests an unknown processor', () => {
|
||||
const config = createConfig([
|
||||
{
|
||||
name: 'ThrowPipe',
|
||||
triggers: ['new_message'],
|
||||
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 synchronous routes (executeTriggerSync) and returns modified array', async () => {
|
||||
const config = createConfig([
|
||||
{
|
||||
name: 'SyncPipe',
|
||||
triggers: ['new_message'],
|
||||
processors: [
|
||||
{ processorId: 'ModifyingProcessor' } as unknown as ProcessorConfig,
|
||||
],
|
||||
},
|
||||
]);
|
||||
const orchestrator = new PipelineOrchestrator(
|
||||
config,
|
||||
env,
|
||||
eventBus,
|
||||
env.tracer,
|
||||
registry,
|
||||
);
|
||||
|
||||
const nodes = [
|
||||
createDummyNode(
|
||||
'not-protected-ep',
|
||||
'USER_PROMPT',
|
||||
100,
|
||||
{
|
||||
semanticParts: [{ type: 'text', text: 'original text' }],
|
||||
},
|
||||
'not-protected-id',
|
||||
),
|
||||
];
|
||||
|
||||
const result = await orchestrator.executeTriggerSync(
|
||||
'new_message',
|
||||
nodes,
|
||||
new Set(nodes.map((s) => s.id)),
|
||||
new Set(),
|
||||
);
|
||||
|
||||
expect(result).toHaveLength(1);
|
||||
const modifiedPrompt = result[0] as UserPrompt;
|
||||
assert(
|
||||
modifiedPrompt.semanticParts[0].type === 'text',
|
||||
'Expected a text part',
|
||||
);
|
||||
expect(modifiedPrompt.semanticParts[0].text).toBe(
|
||||
'original text [modified]',
|
||||
);
|
||||
});
|
||||
|
||||
it('gracefully handles and swallows processor errors in synchronous routes', async () => {
|
||||
const config = createConfig([
|
||||
{
|
||||
name: 'ThrowPipe',
|
||||
triggers: ['new_message'],
|
||||
processors: [
|
||||
{ processorId: 'ThrowingProcessor' } as unknown as ProcessorConfig,
|
||||
],
|
||||
},
|
||||
]);
|
||||
const orchestrator = new PipelineOrchestrator(
|
||||
config,
|
||||
env,
|
||||
eventBus,
|
||||
env.tracer,
|
||||
registry,
|
||||
);
|
||||
|
||||
const nodes = [
|
||||
createDummyNode(
|
||||
'not-protected-ep',
|
||||
'USER_PROMPT',
|
||||
100,
|
||||
undefined,
|
||||
'not-protected-id',
|
||||
),
|
||||
];
|
||||
|
||||
// It should not throw! It should swallow the error and return the unmodified array.
|
||||
const result = await orchestrator.executeTriggerSync(
|
||||
'new_message',
|
||||
nodes,
|
||||
new Set(nodes.map((s) => s.id)),
|
||||
new Set(),
|
||||
);
|
||||
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result).toStrictEqual(nodes);
|
||||
});
|
||||
|
||||
it('automatically dispatches workers when matching EventBus events occur', async () => {
|
||||
const config = createConfig([], [{ workerId: 'MockWorker' }]);
|
||||
|
||||
const orchestrator = new PipelineOrchestrator(
|
||||
config,
|
||||
env,
|
||||
eventBus,
|
||||
env.tracer,
|
||||
registry,
|
||||
);
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
const workerInstance = (orchestrator as any).instantiatedWorkers.get(
|
||||
'MockWorker',
|
||||
) as MockWorker;
|
||||
expect(workerInstance.wasExecuted).toBe(false);
|
||||
|
||||
const nodes = [
|
||||
createDummyNode(
|
||||
'not-protected-ep',
|
||||
'USER_PROMPT',
|
||||
100,
|
||||
{
|
||||
semanticParts: [{ type: 'text', text: 'worker trigger text' }],
|
||||
},
|
||||
'not-protected-id',
|
||||
),
|
||||
];
|
||||
|
||||
// Emit the new_message chunk which maps to onNodesAdded for workers
|
||||
eventBus.emitChunkReceived({
|
||||
nodes,
|
||||
targetNodeIds: new Set(nodes.map((n) => n.id)),
|
||||
});
|
||||
|
||||
// Worker execute is fire and forget, so we yield to the event loop
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
|
||||
expect(workerInstance.wasExecuted).toBe(true);
|
||||
});
|
||||
});
|
||||
@@ -5,33 +5,29 @@
|
||||
*/
|
||||
|
||||
import type { ConcreteNode } from '../ir/types.js';
|
||||
import type { ContextProcessor, ContextWorker } from '../pipeline.js';
|
||||
import type { SidecarConfig, PipelineDef, PipelineTrigger } from './types.js';
|
||||
import type { ContextWorker } from '../pipeline.js';
|
||||
import type { PipelineDef, PipelineTrigger } from './types.js';
|
||||
import type {
|
||||
ContextEnvironment,
|
||||
ContextEventBus,
|
||||
ContextTracer,
|
||||
} from './environment.js';
|
||||
import type { SidecarRegistry } from './registry.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import { InboxSnapshotImpl } from './inbox.js';
|
||||
import { ContextWorkingBufferImpl } from './contextWorkingBuffer.js';
|
||||
|
||||
export class PipelineOrchestrator {
|
||||
private activeTimers: NodeJS.Timeout[] = [];
|
||||
private readonly instantiatedProcessors = new Map<string, ContextProcessor>();
|
||||
private readonly instantiatedWorkers = new Map<string, ContextWorker>();
|
||||
|
||||
constructor(
|
||||
private readonly config: SidecarConfig,
|
||||
private readonly pipelines: PipelineDef[],
|
||||
private readonly workers: ContextWorker[],
|
||||
private readonly env: ContextEnvironment,
|
||||
private readonly eventBus: ContextEventBus,
|
||||
private readonly tracer: ContextTracer,
|
||||
private readonly registry: SidecarRegistry,
|
||||
) {
|
||||
this.instantiateProcessors();
|
||||
this.instantiateWorkers();
|
||||
this.setupTriggers();
|
||||
this.startWorkers();
|
||||
}
|
||||
|
||||
private isNodeAllowed(
|
||||
@@ -46,32 +42,19 @@ export class PipelineOrchestrator {
|
||||
);
|
||||
}
|
||||
|
||||
private instantiateProcessors() {
|
||||
for (const pipeline of this.config.pipelines) {
|
||||
for (const procDef of pipeline.processors) {
|
||||
if (!this.instantiatedProcessors.has(procDef.processorId)) {
|
||||
const factory = this.registry.getProcessor(procDef.processorId);
|
||||
const instance = factory.create(this.env, procDef.options || {});
|
||||
this.instantiatedProcessors.set(procDef.processorId, instance);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private instantiateWorkers() {
|
||||
if (!this.config.workers) return;
|
||||
for (const workerDef of this.config.workers) {
|
||||
if (!this.instantiatedWorkers.has(workerDef.workerId)) {
|
||||
const factory = this.registry.getWorker(workerDef.workerId);
|
||||
const instance = factory.create(this.env, workerDef.options || {});
|
||||
this.instantiatedWorkers.set(workerDef.workerId, instance);
|
||||
private startWorkers() {
|
||||
for (const worker of this.workers) {
|
||||
try {
|
||||
worker.start();
|
||||
} catch (e) {
|
||||
debugLogger.error(`Worker ${worker.name} failed to start:`, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private setupTriggers() {
|
||||
// 1. Pipeline Triggers
|
||||
for (const pipeline of this.config.pipelines) {
|
||||
for (const pipeline of this.pipelines) {
|
||||
for (const trigger of pipeline.triggers) {
|
||||
if (typeof trigger === 'object' && trigger.type === 'timer') {
|
||||
const timer = setInterval(() => {
|
||||
@@ -103,7 +86,7 @@ export class PipelineOrchestrator {
|
||||
// 2. Worker Triggers (onNodesAdded / onNodesAgedOut)
|
||||
this.eventBus.onChunkReceived((event) => {
|
||||
// Fire all workers that care about new nodes
|
||||
for (const worker of this.instantiatedWorkers.values()) {
|
||||
for (const worker of this.workers) {
|
||||
if (worker.triggers.onNodesAdded) {
|
||||
const inboxSnapshot = new InboxSnapshotImpl(
|
||||
this.env.inbox.getMessages() || [],
|
||||
@@ -121,7 +104,7 @@ export class PipelineOrchestrator {
|
||||
|
||||
this.eventBus.onConsolidationNeeded((event) => {
|
||||
// Fire all workers that care about aged out nodes
|
||||
for (const worker of this.instantiatedWorkers.values()) {
|
||||
for (const worker of this.workers) {
|
||||
if (worker.triggers.onNodesAgedOut) {
|
||||
const inboxSnapshot = new InboxSnapshotImpl(
|
||||
this.env.inbox.getMessages() || [],
|
||||
@@ -139,15 +122,19 @@ export class PipelineOrchestrator {
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// We don't have a formal event bus for inbox publish yet, but we will soon.
|
||||
// For now the workers are just registered.
|
||||
}
|
||||
|
||||
shutdown() {
|
||||
for (const timer of this.activeTimers) {
|
||||
clearInterval(timer);
|
||||
}
|
||||
for (const worker of this.workers) {
|
||||
try {
|
||||
worker.stop();
|
||||
} catch (e) {
|
||||
debugLogger.error(`Worker ${worker.name} failed to stop:`, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async executeTriggerSync(
|
||||
@@ -157,7 +144,7 @@ export class PipelineOrchestrator {
|
||||
protectedLogicalIds: ReadonlySet<string> = new Set(),
|
||||
): Promise<readonly ConcreteNode[]> {
|
||||
let currentBuffer = ContextWorkingBufferImpl.initialize(nodes);
|
||||
const pipelines = this.config.pipelines.filter((p) =>
|
||||
const triggerPipelines = this.pipelines.filter((p) =>
|
||||
p.triggers.includes(trigger),
|
||||
);
|
||||
|
||||
@@ -166,22 +153,19 @@ export class PipelineOrchestrator {
|
||||
this.env.inbox.getMessages() || [],
|
||||
);
|
||||
|
||||
for (const pipeline of pipelines) {
|
||||
for (const procDef of pipeline.processors) {
|
||||
const processor = this.instantiatedProcessors.get(procDef.processorId);
|
||||
if (!processor) continue;
|
||||
|
||||
for (const pipeline of triggerPipelines) {
|
||||
for (const processor of pipeline.processors) {
|
||||
try {
|
||||
this.tracer.logEvent(
|
||||
'Orchestrator',
|
||||
`Executing processor synchronously: ${procDef.processorId}`,
|
||||
`Executing processor synchronously: ${processor.id}`,
|
||||
);
|
||||
|
||||
const allowedTargets = currentBuffer.nodes.filter((n) =>
|
||||
this.isNodeAllowed(n, triggerTargets, protectedLogicalIds),
|
||||
);
|
||||
|
||||
const returnedNodes = await processor.process({
|
||||
const returnedNodes = await processor({
|
||||
buffer: currentBuffer,
|
||||
targets: allowedTargets,
|
||||
inbox: inboxSnapshot,
|
||||
@@ -194,7 +178,7 @@ export class PipelineOrchestrator {
|
||||
);
|
||||
} catch (error) {
|
||||
debugLogger.error(
|
||||
`Synchronous processor ${procDef.processorId} failed:`,
|
||||
`Synchronous processor ${processor.id} failed:`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
@@ -224,21 +208,18 @@ export class PipelineOrchestrator {
|
||||
this.env.inbox.getMessages() || [],
|
||||
);
|
||||
|
||||
for (const procDef of pipeline.processors) {
|
||||
const processor = this.instantiatedProcessors.get(procDef.processorId);
|
||||
if (!processor) continue;
|
||||
|
||||
for (const processor of pipeline.processors) {
|
||||
try {
|
||||
this.tracer.logEvent(
|
||||
'Orchestrator',
|
||||
`Executing processor: ${procDef.processorId} (async)`,
|
||||
`Executing processor: ${processor.id} (async)`,
|
||||
);
|
||||
|
||||
const allowedTargets = currentBuffer.nodes.filter((n) =>
|
||||
this.isNodeAllowed(n, triggerTargets, protectedLogicalIds),
|
||||
);
|
||||
|
||||
const returnedNodes = await processor.process({
|
||||
const returnedNodes = await processor({
|
||||
buffer: currentBuffer,
|
||||
targets: allowedTargets,
|
||||
inbox: inboxSnapshot,
|
||||
@@ -251,7 +232,7 @@ export class PipelineOrchestrator {
|
||||
);
|
||||
} catch (error) {
|
||||
debugLogger.error(
|
||||
`Pipeline ${pipeline.name} failed async at ${procDef.processorId}:`,
|
||||
`Pipeline ${pipeline.name} failed async at ${processor.id}:`,
|
||||
error,
|
||||
);
|
||||
return;
|
||||
|
||||
@@ -4,60 +4,63 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { SidecarConfig } from './types.js';
|
||||
import type { SidecarConfig, PipelineDef } from './types.js';
|
||||
import type { ContextEnvironment } from './environment.js';
|
||||
import type { ContextWorker } from '../pipeline.js';
|
||||
|
||||
// Import factories
|
||||
import { createToolMaskingProcessor } from '../processors/toolMaskingProcessor.js';
|
||||
import { createBlobDegradationProcessor } from '../processors/blobDegradationProcessor.js';
|
||||
import { createNodeTruncationProcessor } from '../processors/nodeTruncationProcessor.js';
|
||||
import { createNodeDistillationProcessor } from '../processors/nodeDistillationProcessor.js';
|
||||
import { createStateSnapshotProcessor } from '../processors/stateSnapshotProcessor.js';
|
||||
import { createStateSnapshotWorker } from '../processors/stateSnapshotWorker.js';
|
||||
|
||||
export interface ContextProfile {
|
||||
config: SidecarConfig;
|
||||
buildPipelines: (env: ContextEnvironment) => PipelineDef[];
|
||||
buildWorkers: (env: ContextEnvironment) => ContextWorker[];
|
||||
}
|
||||
|
||||
/**
|
||||
* The standard default context management profile.
|
||||
* Optimized for safety, precision, and reliable summarization.
|
||||
*/
|
||||
export const defaultSidecarProfile: SidecarConfig = {
|
||||
budget: {
|
||||
retainedTokens: 65000,
|
||||
maxTokens: 150000,
|
||||
},
|
||||
workers: [
|
||||
{
|
||||
workerId: 'StateSnapshotWorker',
|
||||
options: {
|
||||
type: 'accumulate',
|
||||
},
|
||||
export const defaultSidecarProfile: ContextProfile = {
|
||||
config: {
|
||||
budget: {
|
||||
retainedTokens: 65000,
|
||||
maxTokens: 150000,
|
||||
},
|
||||
],
|
||||
pipelines: [
|
||||
},
|
||||
|
||||
buildPipelines: (env: ContextEnvironment): PipelineDef[] => [
|
||||
{
|
||||
name: 'Immediate Sanitization',
|
||||
triggers: ['new_message'],
|
||||
processors: [
|
||||
{
|
||||
processorId: 'ToolMaskingProcessor',
|
||||
options: { stringLengthThresholdTokens: 8000 },
|
||||
},
|
||||
{ processorId: 'BlobDegradationProcessor', options: {} },
|
||||
createToolMaskingProcessor('ToolMasking', env, { stringLengthThresholdTokens: 8000 }),
|
||||
createBlobDegradationProcessor('BlobDegradation', env),
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'Normalization',
|
||||
triggers: ['retained_exceeded'],
|
||||
processors: [
|
||||
{
|
||||
processorId: 'NodeTruncationProcessor',
|
||||
options: { maxTokensPerNode: 3000 },
|
||||
},
|
||||
{
|
||||
processorId: 'NodeDistillationProcessor',
|
||||
options: { nodeThresholdTokens: 5000 },
|
||||
},
|
||||
createNodeTruncationProcessor('NodeTruncation', env, { maxTokensPerNode: 3000 }),
|
||||
createNodeDistillationProcessor('NodeDistillation', env, { nodeThresholdTokens: 5000 }),
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'Emergency Backstop',
|
||||
triggers: ['gc_backstop'],
|
||||
processors: [
|
||||
{
|
||||
processorId: 'StateSnapshotProcessor',
|
||||
options: { target: 'max' },
|
||||
},
|
||||
createStateSnapshotProcessor('StateSnapshotSync', env, { target: 'max' }),
|
||||
],
|
||||
},
|
||||
],
|
||||
|
||||
buildWorkers: (env: ContextEnvironment): ContextWorker[] => [
|
||||
createStateSnapshotWorker('StateSnapshotAsync', env, { type: 'accumulate' })
|
||||
]
|
||||
};
|
||||
|
||||
@@ -1,89 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { SidecarRegistry } from './registry.js';
|
||||
import type { ContextProcessorDef, ContextWorkerDef } from './registry.js';
|
||||
import type { ContextProcessor, ContextWorker } from '../pipeline.js';
|
||||
|
||||
describe('SidecarRegistry', () => {
|
||||
it('should register and retrieve processors correctly', () => {
|
||||
const registry = new SidecarRegistry();
|
||||
const processorDef: ContextProcessorDef = {
|
||||
id: 'TestProcessor',
|
||||
schema: { type: 'object' },
|
||||
create: () => ({}) as ContextProcessor,
|
||||
};
|
||||
|
||||
registry.registerProcessor(processorDef);
|
||||
const retrieved = registry.getProcessor('TestProcessor');
|
||||
expect(retrieved).toBe(processorDef);
|
||||
});
|
||||
|
||||
it('should register and retrieve workers correctly', () => {
|
||||
const registry = new SidecarRegistry();
|
||||
const workerDef: ContextWorkerDef = {
|
||||
id: 'TestWorker',
|
||||
schema: { type: 'object' },
|
||||
create: () => ({}) as ContextWorker,
|
||||
};
|
||||
|
||||
registry.registerWorker(workerDef);
|
||||
const retrieved = registry.getWorker('TestWorker');
|
||||
expect(retrieved).toBe(workerDef);
|
||||
});
|
||||
|
||||
it('should throw an error when retrieving unregistered processors', () => {
|
||||
const registry = new SidecarRegistry();
|
||||
expect(() => registry.getProcessor('Unknown')).toThrow(
|
||||
'Context Processor [Unknown] is not registered.',
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw an error when retrieving unregistered workers', () => {
|
||||
const registry = new SidecarRegistry();
|
||||
expect(() => registry.getWorker('Unknown')).toThrow(
|
||||
'Context Worker [Unknown] is not registered.',
|
||||
);
|
||||
});
|
||||
|
||||
it('should return combined schemas', () => {
|
||||
const registry = new SidecarRegistry();
|
||||
registry.registerProcessor({
|
||||
id: 'TestProcessor',
|
||||
schema: { title: 'processorSchema' },
|
||||
create: () => ({}) as ContextProcessor,
|
||||
});
|
||||
registry.registerWorker({
|
||||
id: 'TestWorker',
|
||||
schema: { title: 'workerSchema' },
|
||||
create: () => ({}) as ContextWorker,
|
||||
});
|
||||
|
||||
const schemas = registry.getSchemas() as Array<{ title?: string }>;
|
||||
expect(schemas.length).toBe(2);
|
||||
expect(schemas.find((s) => s.title === 'processorSchema')).toBeDefined();
|
||||
expect(schemas.find((s) => s.title === 'workerSchema')).toBeDefined();
|
||||
});
|
||||
|
||||
it('should safely clear the registry', () => {
|
||||
const registry = new SidecarRegistry();
|
||||
registry.registerProcessor({
|
||||
id: 'TestProcessor',
|
||||
schema: {},
|
||||
create: () => ({}) as ContextProcessor,
|
||||
});
|
||||
registry.registerWorker({
|
||||
id: 'TestWorker',
|
||||
schema: {},
|
||||
create: () => ({}) as ContextWorker,
|
||||
});
|
||||
|
||||
registry.clear();
|
||||
|
||||
expect(() => registry.getProcessor('TestProcessor')).toThrow();
|
||||
expect(() => registry.getWorker('TestWorker')).toThrow();
|
||||
});
|
||||
});
|
||||
@@ -1,68 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextProcessor, ContextWorker } 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;
|
||||
}
|
||||
|
||||
export interface ContextWorkerDef<TOptions = object> {
|
||||
readonly id: string;
|
||||
readonly schema: object;
|
||||
create(env: ContextEnvironment, options: TOptions): ContextWorker;
|
||||
}
|
||||
|
||||
/**
|
||||
* Registry for mapping declarative sidecar configs to running components.
|
||||
*/
|
||||
export class SidecarRegistry {
|
||||
private processors = new Map<string, ContextProcessorDef<unknown>>();
|
||||
private workers = new Map<string, ContextWorkerDef<unknown>>();
|
||||
|
||||
registerProcessor<TOptions>(def: ContextProcessorDef<TOptions>) {
|
||||
this.processors.set(def.id, def);
|
||||
}
|
||||
|
||||
registerWorker<TOptions>(def: ContextWorkerDef<TOptions>) {
|
||||
this.workers.set(def.id, def);
|
||||
}
|
||||
|
||||
getProcessor(id: string): ContextProcessorDef {
|
||||
const def = this.processors.get(id);
|
||||
if (!def) {
|
||||
throw new Error(`Context Processor [${id}] is not registered.`);
|
||||
}
|
||||
return def;
|
||||
}
|
||||
|
||||
getWorker(id: string): ContextWorkerDef {
|
||||
const def = this.workers.get(id);
|
||||
if (!def) {
|
||||
throw new Error(`Context Worker [${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);
|
||||
}
|
||||
for (const def of this.workers.values()) {
|
||||
if (def.schema) schemas.push(def.schema);
|
||||
}
|
||||
return schemas;
|
||||
}
|
||||
|
||||
clear() {
|
||||
this.processors.clear();
|
||||
this.workers.clear();
|
||||
}
|
||||
}
|
||||
@@ -1,92 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { SidecarRegistry } from './registry.js';
|
||||
import './builtins.js';
|
||||
|
||||
export function getSidecarConfigSchema(registry: SidecarRegistry) {
|
||||
return {
|
||||
$schema: 'http://json-schema.org/draft-07/schema#',
|
||||
title: 'SidecarConfig',
|
||||
description: 'The Data-Driven Schema for the Context Manager.',
|
||||
type: 'object',
|
||||
required: ['budget', '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.',
|
||||
},
|
||||
},
|
||||
},
|
||||
workers: {
|
||||
type: 'array',
|
||||
description: 'Background workers that proactively accumulate context.',
|
||||
items: {
|
||||
type: 'object',
|
||||
required: ['workerId'],
|
||||
properties: {
|
||||
workerId: { type: 'string' },
|
||||
options: { type: 'object' },
|
||||
},
|
||||
},
|
||||
},
|
||||
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: ['new_message', 'retained_exceeded', 'gc_backstop'],
|
||||
},
|
||||
{
|
||||
type: 'object',
|
||||
required: ['type', 'intervalMs'],
|
||||
properties: {
|
||||
type: {
|
||||
type: 'string',
|
||||
const: 'timer',
|
||||
},
|
||||
intervalMs: {
|
||||
type: 'number',
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
processors: {
|
||||
type: 'array',
|
||||
items: {
|
||||
oneOf: registry.getSchemas(),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -4,36 +4,7 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
/**
|
||||
* 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: 'NodeDistillationProcessor';
|
||||
options: { nodeThresholdTokens: number };
|
||||
}
|
||||
| {
|
||||
processorId: 'NodeTruncationProcessor';
|
||||
options: { maxTokensPerNode: number };
|
||||
}
|
||||
| {
|
||||
processorId: 'StateSnapshotProcessor';
|
||||
options?: Record<string, unknown>;
|
||||
}
|
||||
| {
|
||||
processorId: 'HistoryTruncationProcessor';
|
||||
options?: Record<string, unknown>;
|
||||
};
|
||||
|
||||
export interface WorkerConfig {
|
||||
workerId: string;
|
||||
options?: Record<string, unknown>;
|
||||
}
|
||||
import type { ContextProcessorFn } from '../pipeline.js';
|
||||
|
||||
export type PipelineTrigger =
|
||||
| 'new_message'
|
||||
@@ -44,7 +15,7 @@ export type PipelineTrigger =
|
||||
export interface PipelineDef {
|
||||
name: string;
|
||||
triggers: PipelineTrigger[];
|
||||
processors: ProcessorConfig[];
|
||||
processors: ContextProcessorFn[];
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -56,10 +27,4 @@ export interface SidecarConfig {
|
||||
retainedTokens: number;
|
||||
maxTokens: number;
|
||||
};
|
||||
|
||||
/** The execution graphs for context manipulation */
|
||||
pipelines: PipelineDef[];
|
||||
|
||||
/** Background actors that generate data for pipelines */
|
||||
workers?: WorkerConfig[];
|
||||
}
|
||||
|
||||
@@ -7,14 +7,12 @@
|
||||
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 type { ContextProfile } from '../sidecar/profiles.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 { SidecarRegistry } from '../sidecar/registry.js';
|
||||
import { DeterministicIdGenerator } from '../system/DeterministicIdGenerator.js';
|
||||
import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js';
|
||||
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
@@ -31,13 +29,13 @@ export class SimulationHarness {
|
||||
env!: ContextEnvironmentImpl;
|
||||
orchestrator!: PipelineOrchestrator;
|
||||
readonly eventBus: ContextEventBus;
|
||||
config!: SidecarConfig;
|
||||
config!: ContextProfile;
|
||||
private tracer!: ContextTracer;
|
||||
private currentTurnIndex = 0;
|
||||
private tokenTrajectory: TurnSummary[] = [];
|
||||
|
||||
static async create(
|
||||
config: SidecarConfig,
|
||||
config: ContextProfile,
|
||||
mockLlmClient: BaseLlmClient,
|
||||
mockTempDir = '/tmp/sim',
|
||||
): Promise<SimulationHarness> {
|
||||
@@ -52,14 +50,11 @@ export class SimulationHarness {
|
||||
}
|
||||
|
||||
private async init(
|
||||
config: SidecarConfig,
|
||||
config: ContextProfile,
|
||||
mockLlmClient: BaseLlmClient,
|
||||
mockTempDir: string,
|
||||
) {
|
||||
this.config = config;
|
||||
const registry = new SidecarRegistry();
|
||||
// Register all standard processors
|
||||
registerBuiltInProcessors(registry);
|
||||
|
||||
this.tracer = new ContextTracer({
|
||||
targetDir: mockTempDir,
|
||||
@@ -79,11 +74,11 @@ export class SimulationHarness {
|
||||
);
|
||||
|
||||
this.orchestrator = new PipelineOrchestrator(
|
||||
config,
|
||||
config.buildPipelines(this.env),
|
||||
config.buildWorkers(this.env),
|
||||
this.env,
|
||||
this.eventBus,
|
||||
this.tracer,
|
||||
registry,
|
||||
);
|
||||
this.contextManager = new ContextManager(
|
||||
config,
|
||||
@@ -118,9 +113,9 @@ export class SimulationHarness {
|
||||
let currentView = this.contextManager.getNodes();
|
||||
const currentTokens =
|
||||
this.env.tokenCalculator.calculateConcreteListTokens(currentView);
|
||||
if (this.config.budget && currentTokens > this.config.budget.maxTokens) {
|
||||
if (this.config.config.budget && currentTokens > this.config.config.budget.maxTokens) {
|
||||
debugLogger.log(
|
||||
`[Turn ${this.currentTurnIndex}] Sync panic triggered! ${currentTokens} > ${this.config.budget.maxTokens}`,
|
||||
`[Turn ${this.currentTurnIndex}] Sync panic triggered! ${currentTokens} > ${this.config.config.budget.maxTokens}`,
|
||||
);
|
||||
const orchestrator = this.orchestrator;
|
||||
// In the V2 simulation, we trigger the 'gc_backstop' to simulate emergency pressure.
|
||||
|
||||
@@ -7,7 +7,12 @@
|
||||
import { describe, it, expect, vi, beforeAll, afterAll } from 'vitest';
|
||||
import { SimulationHarness } from './SimulationHarness.js';
|
||||
import { createMockLlmClient } from '../testing/contextTestUtils.js';
|
||||
import type { SidecarConfig } from '../sidecar/types.js';
|
||||
import type { ContextProfile } from '../sidecar/profiles.js';
|
||||
import { createToolMaskingProcessor } from '../processors/toolMaskingProcessor.js';
|
||||
import { createBlobDegradationProcessor } from '../processors/blobDegradationProcessor.js';
|
||||
import { createStateSnapshotProcessor } from '../processors/stateSnapshotProcessor.js';
|
||||
import { createHistoryTruncationProcessor } from '../processors/historyTruncationProcessor.js';
|
||||
import { createStateSnapshotWorker } from '../processors/stateSnapshotWorker.js';
|
||||
|
||||
expect.addSnapshotSerializer({
|
||||
test: (val) =>
|
||||
@@ -31,30 +36,29 @@ describe('System Lifecycle Golden Tests', () => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
const getAggressiveConfig = (): SidecarConfig => ({
|
||||
budget: { maxTokens: 1000, retainedTokens: 500 }, // Extremely tight limits
|
||||
pipelines: [
|
||||
const getAggressiveConfig = (): ContextProfile => ({
|
||||
config: {
|
||||
budget: { maxTokens: 1000, retainedTokens: 500 }, // Extremely tight limits
|
||||
},
|
||||
buildPipelines: (env) => [
|
||||
{
|
||||
name: 'Pressure Relief', // Emits from eventBus 'retained_exceeded'
|
||||
triggers: ['retained_exceeded'],
|
||||
processors: [
|
||||
{ processorId: 'BlobDegradationProcessor' },
|
||||
{
|
||||
processorId: 'ToolMaskingProcessor',
|
||||
options: { stringLengthThresholdTokens: 50 },
|
||||
}, // Mask any tool string > 50 chars
|
||||
{ processorId: 'StateSnapshotProcessor', options: {} }, // Squash old history
|
||||
createBlobDegradationProcessor('BlobDegradationProcessor', env),
|
||||
createToolMaskingProcessor('ToolMaskingProcessor', env, { stringLengthThresholdTokens: 50 }),
|
||||
createStateSnapshotProcessor('StateSnapshotProcessor', env, {}),
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'Immediate Sanitization', // The magic string the projector is hardcoded to use
|
||||
triggers: ['retained_exceeded'],
|
||||
processors: [
|
||||
{ processorId: 'HistoryTruncationProcessor', options: {} },
|
||||
createHistoryTruncationProcessor('HistoryTruncationProcessor', env, {}),
|
||||
],
|
||||
},
|
||||
],
|
||||
workers: [{ workerId: 'StateSnapshotWorker' }],
|
||||
buildWorkers: (env) => [createStateSnapshotWorker('StateSnapshotWorker', env, {})],
|
||||
});
|
||||
|
||||
const mockLlmClient = createMockLlmClient([
|
||||
@@ -141,10 +145,12 @@ describe('System Lifecycle Golden Tests', () => {
|
||||
});
|
||||
|
||||
it('Scenario 2: Under Budget (No Modifications)', async () => {
|
||||
const generousConfig: SidecarConfig = {
|
||||
budget: { maxTokens: 100000, retainedTokens: 50000 },
|
||||
pipelines: [], // No triggers
|
||||
workers: [],
|
||||
const generousConfig: ContextProfile = {
|
||||
config: {
|
||||
budget: { maxTokens: 100000, retainedTokens: 50000 },
|
||||
},
|
||||
buildPipelines: () => [],
|
||||
buildWorkers: () => [],
|
||||
};
|
||||
|
||||
const harness = await SimulationHarness.create(
|
||||
@@ -171,12 +177,12 @@ describe('System Lifecycle Golden Tests', () => {
|
||||
});
|
||||
|
||||
it('Scenario 3: Worker-Driven Background GC', async () => {
|
||||
const gcConfig: SidecarConfig = {
|
||||
budget: { maxTokens: 200, retainedTokens: 100 },
|
||||
pipelines: [], // No standard pipelines
|
||||
workers: [
|
||||
{ workerId: 'StateSnapshotWorker' }, // This should fire on chunk events
|
||||
],
|
||||
const gcConfig: ContextProfile = {
|
||||
config: {
|
||||
budget: { maxTokens: 200, retainedTokens: 100 },
|
||||
},
|
||||
buildPipelines: () => [],
|
||||
buildWorkers: (env) => [createStateSnapshotWorker('StateSnapshotWorker', env, {})],
|
||||
};
|
||||
|
||||
const harness = await SimulationHarness.create(gcConfig, mockLlmClient);
|
||||
|
||||
@@ -14,8 +14,6 @@ import { ContextTracer } from '../tracer.js';
|
||||
import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js';
|
||||
import { SidecarLoader } from '../sidecar/SidecarLoader.js';
|
||||
import { ContextEventBus } from '../eventBus.js';
|
||||
import { SidecarRegistry } from '../sidecar/registry.js';
|
||||
import { registerBuiltInProcessors } from '../sidecar/builtins.js';
|
||||
import { PipelineOrchestrator } from '../sidecar/orchestrator.js';
|
||||
import type { ConcreteNode, ToolExecution } from '../ir/types.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
@@ -24,6 +22,7 @@ import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
import type { Content, GenerateContentResponse } from '@google/genai';
|
||||
import { InboxSnapshotImpl } from '../sidecar/inbox.js';
|
||||
import type { InboxMessage, ProcessArgs } from '../pipeline.js';
|
||||
import type { ContextProfile } from '../sidecar/profiles.js';
|
||||
|
||||
/**
|
||||
* Creates a valid mock GenerateContentResponse with the provided text.
|
||||
@@ -95,7 +94,7 @@ export function createDummyToolNode(
|
||||
}
|
||||
|
||||
import type { Mock } from 'vitest';
|
||||
import type { SidecarConfig } from '../sidecar/types.js';
|
||||
|
||||
|
||||
export interface MockLlmClient extends BaseLlmClient {
|
||||
generateContent: Mock;
|
||||
@@ -244,12 +243,10 @@ export function createMockContextConfig(
|
||||
|
||||
export function setupContextComponentTest(
|
||||
config: Config,
|
||||
sidecarOverride?: SidecarConfig,
|
||||
sidecarOverride?: ContextProfile,
|
||||
): { chatHistory: AgentChatHistory; contextManager: ContextManager } {
|
||||
const chatHistory = new AgentChatHistory();
|
||||
const registry = new SidecarRegistry();
|
||||
registerBuiltInProcessors(registry);
|
||||
const sidecar = sidecarOverride || SidecarLoader.fromConfig(config, registry);
|
||||
const sidecar = sidecarOverride || SidecarLoader.fromConfig(config);
|
||||
const tracer = new ContextTracer({
|
||||
targetDir: '/tmp',
|
||||
sessionId: 'test-session',
|
||||
@@ -267,11 +264,11 @@ export function setupContextComponentTest(
|
||||
);
|
||||
|
||||
const orchestrator = new PipelineOrchestrator(
|
||||
sidecar,
|
||||
sidecar.buildPipelines(env),
|
||||
sidecar.buildWorkers(env),
|
||||
env,
|
||||
eventBus,
|
||||
tracer,
|
||||
registry,
|
||||
);
|
||||
|
||||
const contextManager = new ContextManager(
|
||||
|
||||
@@ -3,18 +3,26 @@
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import type { SidecarConfig } from '../sidecar/types.js';
|
||||
import type { ContextProfile } from '../sidecar/profiles.js';
|
||||
import type { PipelineDef } from '../sidecar/types.js';
|
||||
import type { ContextEnvironment } from '../sidecar/environment.js';
|
||||
import { createHistoryTruncationProcessor } from '../processors/historyTruncationProcessor.js';
|
||||
|
||||
export const testTruncateProfile: SidecarConfig = {
|
||||
budget: {
|
||||
retainedTokens: 65000,
|
||||
maxTokens: 150000,
|
||||
export const testTruncateProfile: ContextProfile = {
|
||||
config: {
|
||||
budget: {
|
||||
retainedTokens: 65000,
|
||||
maxTokens: 150000,
|
||||
},
|
||||
},
|
||||
pipelines: [
|
||||
buildPipelines: (env: ContextEnvironment): PipelineDef[] => [
|
||||
{
|
||||
name: 'Emergency Backstop (Truncate Only)',
|
||||
triggers: ['gc_backstop', 'retained_exceeded'],
|
||||
processors: [{ processorId: 'HistoryTruncationProcessor', options: {} }],
|
||||
processors: [
|
||||
createHistoryTruncationProcessor('HistoryTruncation', env, {}),
|
||||
],
|
||||
},
|
||||
],
|
||||
buildWorkers: () => [],
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user