MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-07-12 19:10:55 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix(context): Fix snapshot recovery across sessions. (#26939)

Browse Source
This commit is contained in:
joshualitt
2026-05-18 09:44:59 -07:00
committed by GitHub
parent 9d01958cdb
commit 055e0f6452
57 changed files with 3143 additions and 751 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/context/chatCompressionService.test.ts
+1 -1
View File
@@ -196,7 +196,7 @@ describe('ChatCompressionService', () => {
} as unknown as Config;
vi.mocked(getInitialChatHistory).mockImplementation(
async (_config, extraHistory) => extraHistory || [],
async (_config, extraHistory) => (extraHistory ? [...extraHistory] : []),
);
});
packages/core/src/context/chatCompressionService.ts
+3 -1
View File
@@ -442,7 +442,9 @@ export class ChatCompressionService {
const fullNewHistory = await getInitialChatHistory(config, extraHistory);
const newTokenCount = await calculateRequestTokenCount(
fullNewHistory.flatMap((c) => c.parts || []),
fullNewHistory.flatMap(
(c) => ('content' in c ? c.content.parts : c.parts) || [],
),
config.getContentGenerator(),
model,
);
packages/core/src/context/contextManager.barrier.test.ts
+27 -10
View File
@@ -5,6 +5,7 @@
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { randomUUID } from 'node:crypto';
import { testTruncateProfile } from './testing/testProfile.js';
import {
createSyntheticHistory,
@@ -32,20 +33,35 @@ describe('ContextManager Sync Pressure Barrier Tests', () => {
// 2. Add System Prompt (Episode 0 - Protected)
chatHistory.set([
{ role: 'user', parts: [{ text: 'System prompt' }] },
{ role: 'model', parts: [{ text: 'Understood.' }] },
{
id: 'h1',
content: { role: 'user', parts: [{ text: 'System prompt' }] },
},
{
id: 'h2',
content: { role: 'model', parts: [{ text: 'Understood.' }] },
},
]);
// 3. Add massive history that blows past the 150k maxTokens limit
// 20 turns * ~20,000 tokens/turn (10k user + 10k model) = ~400,000 tokens
const massiveHistory = createSyntheticHistory(20, 10000);
const massiveHistory = createSyntheticHistory(20, 10000).map((c) => ({
id: randomUUID(),
content: c,
}));
chatHistory.set([...chatHistory.get(), ...massiveHistory]);
// 4. Add the Latest Turn (Protected)
chatHistory.set([
...chatHistory.get(),
{ role: 'user', parts: [{ text: 'Final question.' }] },
{ role: 'model', parts: [{ text: 'Final answer.' }] },
{
id: 'h-last-user',
content: { role: 'user', parts: [{ text: 'Final question.' }] },
},
{
id: 'h-last-model',
content: { role: 'model', parts: [{ text: 'Final answer.' }] },
},
]);
const rawHistoryLength = chatHistory.get().length;
@@ -59,21 +75,22 @@ describe('ContextManager Sync Pressure Barrier Tests', () => {
expect(projection.length).toBeLessThan(rawHistoryLength);
// Verify Episode 0 (System) was pruned, so we now start with a sentinel due to role alternation
expect(projection[0].role).toBe('user');
expect(projection[0].content.role).toBe('user');
const projectionString = JSON.stringify(projection);
expect(projectionString).toContain('User turn 17');
// Filter out synthetic Yield nodes (they are model responses without actual tool/text bodies)
const contentNodes = projection.filter(
(p) =>
p.parts && p.parts.some((part) => part.text && part.text !== 'Yield'),
p.content.parts &&
p.content.parts.some((part) => part.text && part.text !== 'Yield'),
);
// Verify the latest turn is perfectly preserved at the back
// Note: The HistoryHardener appends a "Please continue." user turn if we end on model,
// so we look at the turns before the sentinel.
const lastSentinel = contentNodes[contentNodes.length - 1];
const lastModel = contentNodes[contentNodes.length - 2];
const lastUser = contentNodes[contentNodes.length - 3];
const lastSentinel = contentNodes[contentNodes.length - 1].content;
const lastModel = contentNodes[contentNodes.length - 2].content;
const lastUser = contentNodes[contentNodes.length - 3].content;
expect(lastSentinel.role).toBe('user');
expect(lastSentinel.parts![0].text).toBe('Please continue.');
packages/core/src/context/contextManager.hotstart.test.ts
+6 -2
View File
@@ -47,7 +47,9 @@ describe('ContextManager - Hot Start Calibration', () => {
const emitGroundTruthSpy = vi.spyOn(env.eventBus, 'emitTokenGroundTruth');
// Add a node to make the buffer non-empty
chatHistory.set([{ role: 'user', parts: [{ text: 'Hello' }] }]);
chatHistory.set([
{ id: 'h1', content: { role: 'user', parts: [{ text: 'Hello' }] } },
]);
// First render should trigger calibration
await contextManager.renderHistory();
@@ -81,7 +83,9 @@ describe('ContextManager - Hot Start Calibration', () => {
);
// Add a node
chatHistory.set([{ role: 'user', parts: [{ text: 'Hello' }] }]);
chatHistory.set([
{ id: 'h1', content: { role: 'user', parts: [{ text: 'Hello' }] } },
]);
// Render should succeed without throwing
const result = await contextManager.renderHistory();
packages/core/src/context/contextManager.ts
+86 -36
View File
@@ -5,8 +5,11 @@
*/
import type { Content } from '@google/genai';
import type { AgentChatHistory } from '../core/agentChatHistory.js';
import { isToolExecution, type ConcreteNode } from './graph/types.js';
import type {
AgentChatHistory,
HistoryTurn,
} from '../core/agentChatHistory.js';
import type { ConcreteNode } from './graph/types.js';
import type { ContextEventBus } from './eventBus.js';
import type { ContextTracer } from './tracer.js';
import type { ContextEnvironment } from './pipeline/environment.js';
@@ -38,9 +41,11 @@ export class ContextManager {
private lastRenderCache?: {
nodesHash: string;
result: {
history: Content[];
history: HistoryTurn[];
apiHistory: Content[];
didApplyManagement: boolean;
baseUnits: number;
processedNodes: readonly ConcreteNode[];
};
};
@@ -75,6 +80,21 @@ export class ContextManager {
this.evaluateTriggers(event.newNodes);
});
this.eventBus.onProcessorResult((event) => {
// Defensive: Verify all targets are still present in the buffer.
// If a synchronous render or a previous async task already removed them,
// this result is stale and should be dropped.
const currentIds = new Set(this.buffer.nodes.map((n) => n.id));
const allTargetsPresent = event.targets.every((t) =>
currentIds.has(t.id),
);
if (!allTargetsPresent) {
debugLogger.log(
`[ContextManager] Dropping stale processor result from ${event.processorId}. One or more targets were already removed.`,
);
return;
}
this.buffer = this.buffer.applyProcessorResult(
event.processorId,
event.targets,
@@ -127,11 +147,11 @@ export class ContextManager {
const agedOutNodes = new Set<string>();
let rollingTokens = 0;
// Identify active tool calls that must NEVER be truncated
// Identify nodes that must NEVER be truncated
const protectedIds = this.getProtectedNodeIds(this.buffer.nodes);
if (protectedIds.size > 0) {
debugLogger.log(
`[ContextManager] Pinning ${protectedIds.size} active tool call nodes to prevent truncation.`,
`[ContextManager] Pinning ${protectedIds.size} nodes (recent_turn or external_active_task) to prevent truncation.`,
);
}
@@ -215,24 +235,7 @@ export class ContextManager {
}
}
// 2. Identify active tool calls that must NEVER be truncated
const calls = nodes.filter((n) => isToolExecution(n) && n.role === 'model');
const responses = new Set(
nodes
.filter((n) => isToolExecution(n) && n.role === 'user')
.map((n) => n.payload.functionResponse?.id)
.filter((id): id is string => !!id),
);
for (const call of calls) {
const id = call.payload.functionCall?.id;
// If we have a call but no response in the current graph, it's 'in flight'
if (id && !responses.has(id)) {
protectionMap.set(call.id, 'in_flight_tool_call');
}
}
// 3. Any externally requested protections
// 2. Any externally requested protections
for (const id of extraProtectedIds) {
protectionMap.set(id, 'external_active_task');
}
@@ -278,13 +281,15 @@ export class ContextManager {
* This is the primary method called by the agent framework before sending a request.
*/
async renderHistory(
pendingRequest?: Content,
pendingRequest?: HistoryTurn,
activeTaskIds: Set<string> = new Set(),
abortSignal?: AbortSignal,
): Promise<{
history: Content[];
history: HistoryTurn[];
apiHistory: Content[];
didApplyManagement: boolean;
baseUnits: number;
processedNodes: readonly ConcreteNode[];
}> {
this.tracer.logEvent('ContextManager', 'Starting rendering of LLM context');
@@ -302,6 +307,7 @@ export class ContextManager {
const hotStartPromise = (async () => {
if (!this.hasPerformedHotStart) {
this.hasPerformedHotStart = true;
if (this.buffer.nodes.length > 0) {
const nodesForHotStart = [...this.buffer.nodes, ...previewNodes];
await this.performHotStartCalibration(nodesForHotStart, abortSignal);
@@ -345,11 +351,7 @@ export class ContextManager {
const protectionReasons = this.getProtectedNodeIds(nodes, activeTaskIds);
// Apply final GC Backstop pressure barrier synchronously before mapping
const {
history: renderedHistory,
didApplyManagement,
baseUnits,
} = await render(
const renderResult = await render(
nodes,
this.orchestrator,
this.sidecar,
@@ -361,21 +363,68 @@ export class ContextManager {
previewNodeIds,
);
const {
history: renderedHistory,
didApplyManagement,
baseUnits,
processedNodes,
} = renderResult;
if (didApplyManagement) {
// Commit the GC backstop results back to the master buffer.
// We filter out preview nodes because they are ephemeral and will be
// added to history naturally by the client after the turn completes.
this.buffer = this.buffer.applyProcessorResult(
'sync_backstop',
this.buffer.nodes,
processedNodes.filter((n) => !previewNodeIds.has(n.id)),
);
}
// Structural validation in debug mode
checkContextInvariants(this.buffer.nodes, 'RenderHistory');
this.tracer.logEvent('ContextManager', 'Finished rendering');
const combinedHistory = header
? [header, ...renderedHistory]
// We must temporarily append the pendingRequest (if any) before hardening.
// Otherwise, the hardener will see dangling functionCalls and inject sentinels
// even though the pendingRequest provides the required functionResponses.
const fullHistoryToHarden = pendingRequest
? [...renderedHistory, pendingRequest]
: renderedHistory;
const hardenedHistory = hardenHistory(fullHistoryToHarden, {
sentinels: this.sidecar.sentinels,
});
if (pendingRequest) {
const last = hardenedHistory[hardenedHistory.length - 1];
if (last && last.content.parts) {
const numPartsToRemove = pendingRequest.content.parts?.length || 0;
if (
numPartsToRemove > 0 &&
last.content.parts.length > numPartsToRemove
) {
last.content.parts.splice(-numPartsToRemove);
} else {
hardenedHistory.pop();
}
} else {
hardenedHistory.pop();
}
}
const apiHistory = hardenedHistory.map((h) => h.content);
if (header) {
apiHistory.unshift(header);
}
const result = {
history: hardenHistory(combinedHistory, {
sentinels: this.sidecar.sentinels,
}),
history: hardenedHistory,
apiHistory,
didApplyManagement,
baseUnits,
processedNodes,
};
// Update cache
@@ -394,10 +443,11 @@ export class ContextManager {
);
const contents = this.env.graphMapper.fromGraph(nodes);
const rawContents = contents.map((h) => h.content);
const header = this.headerProvider
? await this.headerProvider()
: undefined;
const combinedHistory = header ? [header, ...contents] : contents;
const combinedHistory = header ? [header, ...rawContents] : rawContents;
const baseUnits =
this.advancedTokenCalculator.getRawBaseUnits(nodes) +
packages/core/src/context/graph/fromGraph.test.ts
+202
View File
@@ -0,0 +1,202 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { fromGraph } from './fromGraph.js';
import { NodeType, type ConcreteNode } from './types.js';
import { NodeIdService } from './nodeIdService.js';
describe('fromGraph', () => {
it('should reconstruct an empty history from empty nodes', () => {
expect(fromGraph([])).toEqual([]);
});
it('should reconstruct a single turn from a single node', () => {
const nodes: ConcreteNode[] = [
{
id: 'node_1',
turnId: 'turn_durable_1',
role: 'user',
type: NodeType.USER_PROMPT,
payload: { text: 'hello' },
timestamp: 100,
},
];
const history = fromGraph(nodes);
expect(history).toEqual([
{
id: 'durable_1',
content: {
role: 'user',
parts: [{ text: 'hello' }],
},
},
]);
});
it('should coalesce adjacent nodes with the same turnId into a single turn', () => {
const nodes: ConcreteNode[] = [
{
id: 'node_1',
turnId: 'turn_durable_1',
role: 'user',
type: NodeType.USER_PROMPT,
payload: { text: 'hello' },
timestamp: 100,
},
{
id: 'node_2',
turnId: 'turn_durable_1',
role: 'user',
type: NodeType.USER_PROMPT,
payload: { text: 'world' },
timestamp: 101,
},
];
const history = fromGraph(nodes);
expect(history).toEqual([
{
id: 'durable_1',
content: {
role: 'user',
parts: [{ text: 'hello' }, { text: 'world' }],
},
},
]);
});
it('should split turns when the role changes', () => {
const nodes: ConcreteNode[] = [
{
id: 'node_1',
turnId: 'turn_durable_1',
role: 'user',
type: NodeType.USER_PROMPT,
payload: { text: 'hello' },
timestamp: 100,
},
{
id: 'node_2',
turnId: 'turn_durable_2',
role: 'model',
type: NodeType.AGENT_THOUGHT,
payload: { text: 'hi' },
timestamp: 101,
},
];
const history = fromGraph(nodes);
expect(history).toEqual([
{
id: 'durable_1',
content: {
role: 'user',
parts: [{ text: 'hello' }],
},
},
{
id: 'durable_2',
content: {
role: 'model',
parts: [{ text: 'hi' }],
},
},
]);
});
it('should split turns when the turnId changes, even if role is the same', () => {
const nodes: ConcreteNode[] = [
{
id: 'node_1',
turnId: 'turn_durable_1',
role: 'user',
type: NodeType.USER_PROMPT,
payload: { text: 'hello' },
timestamp: 100,
},
{
id: 'node_2',
turnId: 'turn_durable_2',
role: 'user',
type: NodeType.USER_PROMPT,
payload: { text: 'world' },
timestamp: 101,
},
];
const history = fromGraph(nodes);
expect(history).toEqual([
{
id: 'durable_1',
content: {
role: 'user',
parts: [{ text: 'hello' }],
},
},
{
id: 'durable_2',
content: {
role: 'user',
parts: [{ text: 'world' }],
},
},
]);
});
it('should correctly strip the turn_ prefix from turnId', () => {
const nodes: ConcreteNode[] = [
{
id: 'node_1',
turnId: 'turn_my_stable_id_123',
role: 'user',
type: NodeType.USER_PROMPT,
payload: { text: 'hello' },
timestamp: 100,
},
];
const history = fromGraph(nodes);
expect(history[0].id).toBe('my_stable_id_123');
});
it('should handle orphan nodes gracefully', () => {
const nodes: ConcreteNode[] = [
{
id: 'node_1',
role: 'user',
type: NodeType.USER_PROMPT,
payload: { text: 'orphan part' },
timestamp: 100,
} as unknown as ConcreteNode,
];
const history = fromGraph(nodes);
expect(history[0].id).toBe('orphan');
expect(history[0].content.parts).toEqual([{ text: 'orphan part' }]);
});
it('should register identities with the NodeIdService if provided', () => {
const idService = new NodeIdService();
const payload = { text: 'hello' };
const nodes: ConcreteNode[] = [
{
id: 'node_1',
turnId: 'turn_1',
role: 'user',
type: NodeType.USER_PROMPT,
payload,
timestamp: 100,
},
];
fromGraph(nodes, idService);
// The payload object reference should map to the node ID
expect(idService.get(payload)).toBe('node_1');
});
});
packages/core/src/context/graph/fromGraph.ts
+29 -18
View File
@@ -7,22 +7,34 @@
import type { Content } from '@google/genai';
import type { ConcreteNode } from './types.js';
import { debugLogger } from '../../utils/debugLogger.js';
import type { NodeIdService } from './nodeIdService.js';
import type { HistoryTurn } from '../../core/agentChatHistory.js';
/**
* Reconstructs a valid Gemini Chat History from a list of Concrete Nodes.
* Reconstructs a list of HistoryTurns from a list of Concrete Nodes.
* This process is "role-alternation-aware" and uses turnId to
* preserve original turn boundaries even if multiple turns have the same role.
* preserve original turn boundaries and IDs.
*/
export function fromGraph(nodes: readonly ConcreteNode[]): Content[] {
export function fromGraph(
nodes: readonly ConcreteNode[],
idService?: NodeIdService,
): HistoryTurn[] {
debugLogger.log(
`[fromGraph] Reconstructing history from ${nodes.length} nodes`,
);
const history: Content[] = [];
let currentTurn: (Content & { _turnId?: string }) | null = null;
const history: HistoryTurn[] = [];
let currentTurn: { id: string; content: Content } | null = null;
for (const node of nodes) {
const turnId = node.turnId;
const turnId = node.turnId || 'orphan';
const durableId = turnId.startsWith('turn_') ? turnId.slice(5) : turnId;
// Register the payload in the identity service to ensure stability
// even if the turn content changes (e.g. after GC backstop).
if (idService) {
idService.set(node.payload, node.id);
}
// We start a new turn if:
// 1. We don't have a current turn.
@@ -30,26 +42,25 @@ export function fromGraph(nodes: readonly ConcreteNode[]): Content[] {
// 3. The turnId changes (Preserving distinct turns of the same role).
if (
!currentTurn ||
currentTurn.role !== node.role ||
currentTurn._turnId !== turnId
currentTurn.content.role !== node.role ||
currentTurn.id !== durableId
) {
currentTurn = {
role: node.role,
parts: [node.payload],
_turnId: turnId,
id: durableId,
content: {
role: node.role,
parts: [node.payload],
},
};
history.push(currentTurn);
} else {
currentTurn.parts = [...(currentTurn.parts || []), node.payload];
currentTurn.content.parts = [
...(currentTurn.content.parts || []),
node.payload,
];
}
}
// Final cleanup: remove our internal tracking field
for (const turn of history) {
const t = turn as Content & { _turnId?: string };
delete t._turnId;
}
debugLogger.log(`[fromGraph] Reconstructed ${history.length} turns`);
return history;
}
packages/core/src/context/graph/mapper.test.ts
+110
View File
@@ -0,0 +1,110 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { ContextGraphMapper } from './mapper.js';
import type { HistoryTurn } from '../../core/agentChatHistory.js';
import { hardenHistory } from '../../utils/historyHardening.js';
describe('ContextGraphMapper (Round-Trip Fidelity)', () => {
it('should flawlessly round-trip a complex history containing parallel tool calls and responses', () => {
// 1. Define a complex, worst-case scenario history
const originalHistory: HistoryTurn[] = [
{
id: 'system_prompt_id',
content: {
role: 'user',
parts: [{ text: '<session_context>\nSystem Prompt here' }],
},
},
{
id: 'user_turn_1',
content: {
role: 'user',
parts: [{ text: 'Please read file A and file B at the same time.' }],
},
},
{
id: 'model_turn_1',
content: {
role: 'model',
parts: [
{ text: 'I will read both files concurrently.' },
{
functionCall: {
id: 'call_A',
name: 'read_file',
args: { path: 'A.txt' },
},
thoughtSignature: 'synthetic_sig_xyz',
},
{
functionCall: {
id: 'call_B',
name: 'read_file',
args: { path: 'B.txt' },
},
},
],
},
},
// Note: GeminiChat records these as separate sequential user turns initially
{
id: 'tool_resp_B_id',
content: {
role: 'user',
parts: [
{
functionResponse: {
id: 'call_B',
name: 'read_file',
response: { content: 'File B' },
},
},
],
},
},
{
id: 'tool_resp_A_id',
content: {
role: 'user',
parts: [
{
functionResponse: {
id: 'call_A',
name: 'read_file',
response: { content: 'File A' },
},
},
],
},
},
];
// 2. We harden the original history first. The core agent loop feeds the hardener the pure history.
// We want our round-tripped history to match what the hardener WOULD have produced natively.
const hardenedOriginal = hardenHistory(originalHistory);
// 3. Translate History -> Graph
const mapper = new ContextGraphMapper();
// Simulate the HistoryObserver capturing the push
const nodes = mapper.applyEvent({
type: 'SYNC_FULL',
payload: originalHistory,
});
// 4. Translate Graph -> History
const reconstructedHistory = mapper.fromGraph(nodes);
// 5. Harden the reconstructed history (as the ContextManager does before sending to API)
const hardenedReconstructed = hardenHistory(reconstructedHistory);
// 6. Assert Absolute Equality
// The round-trip through the Context Graph and Hardener must exactly equal
// the original history put through the Hardener.
expect(hardenedReconstructed).toEqual(hardenedOriginal);
});
});
packages/core/src/context/graph/mapper.ts
+10 -6
View File
@@ -5,23 +5,27 @@
*/
import type { ConcreteNode } from './types.js';
import { ContextGraphBuilder } from './toGraph.js';
import type { Content } from '@google/genai';
import type { HistoryEvent } from '../../core/agentChatHistory.js';
import type { HistoryEvent, HistoryTurn } from '../../core/agentChatHistory.js';
import { fromGraph } from './fromGraph.js';
import { NodeIdService } from './nodeIdService.js';
export class ContextGraphMapper {
private readonly nodeIdentityMap = new WeakMap<object, string>();
private readonly idService = new NodeIdService();
private readonly builder: ContextGraphBuilder;
constructor() {
this.builder = new ContextGraphBuilder(this.nodeIdentityMap);
this.builder = new ContextGraphBuilder(this.idService);
}
applyEvent(event: HistoryEvent): ConcreteNode[] {
return this.builder.processHistory(event.payload);
}
fromGraph(nodes: readonly ConcreteNode[]): Content[] {
return fromGraph(nodes);
fromGraph(nodes: readonly ConcreteNode[]): HistoryTurn[] {
return fromGraph(nodes, this.idService);
}
getIdService(): NodeIdService {
return this.idService;
}
}
packages/core/src/context/graph/nodeIdService.ts
+23
View File
@@ -0,0 +1,23 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
/**
* Provides a durable mapping between history object references and their
* corresponding graph node IDs. This ensures that context management logic
* can track the identity of turns even after they are transformed (e.g. scrubbed
* or hardened) without polluting the raw JSON sent to the Gemini API.
*/
export class NodeIdService {
constructor(private readonly map: WeakMap<object, string> = new WeakMap()) {}
get(obj: object): string | undefined {
return this.map.get(obj);
}
set(obj: object, id: string): void {
this.map.set(obj, id);
}
}
packages/core/src/context/graph/render.ts
+22 -5
View File
@@ -12,9 +12,10 @@ import type { PipelineOrchestrator } from '../pipeline/orchestrator.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { performCalibration } from '../utils/tokenCalibration.js';
import type { AdvancedTokenCalculator } from '../utils/contextTokenCalculator.js';
import type { HistoryTurn } from '../../core/agentChatHistory.js';
/**
* Maps the Episodic Context Graph back into a raw Gemini Content[] array for transmission.
* Maps the Episodic Context Graph back into a list of HistoryTurns for transmission.
* It applies synchronous context management (GC backstop) if the budget is exceeded.
*/
export async function render(
@@ -28,9 +29,10 @@ export async function render(
header?: Content,
previewNodeIds: ReadonlySet<string> = new Set(),
): Promise<{
history: Content[];
history: HistoryTurn[];
didApplyManagement: boolean;
baseUnits: number;
processedNodes: readonly ConcreteNode[];
}> {
let headerTokens = 0;
let headerBaseUnits = 0;
@@ -52,7 +54,12 @@ export async function render(
const baseUnits =
advancedTokenCalculator.getRawBaseUnits(nodes) + headerBaseUnits;
return { history: contents, didApplyManagement: false, baseUnits };
return {
history: contents,
didApplyManagement: false,
baseUnits,
processedNodes: nodes,
};
}
const maxTokens = sidecar.config.budget.maxTokens;
@@ -92,11 +99,16 @@ export async function render(
tracer.logEvent('Render', 'Render Context for LLM', {
renderedContext: contents,
});
performCalibration(env, visibleNodes, contents);
performCalibration(
env,
visibleNodes,
contents.map((h) => h.content),
);
return {
history: contents,
didApplyManagement: false,
baseUnits: graphBaseUnits + headerBaseUnits,
processedNodes: nodes,
};
}
const targetDelta = currentTokens - sidecar.config.budget.retainedTokens;
@@ -145,11 +157,16 @@ export async function render(
tracer.logEvent('Render', 'Render Sanitized Context for LLM', {
renderedContextSanitized: contents,
});
performCalibration(env, visibleNodes, contents);
performCalibration(
env,
visibleNodes,
contents.map((h) => h.content),
);
return {
history: contents,
didApplyManagement: true,
baseUnits:
advancedTokenCalculator.getRawBaseUnits(visibleNodes) + headerBaseUnits,
processedNodes,
};
}
packages/core/src/context/graph/toGraph.test.ts
+102 -13
View File
@@ -4,29 +4,42 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { describe, it, expect, vi } from 'vitest';
import { ContextGraphBuilder } from './toGraph.js';
import type { Content } from '@google/genai';
import type { BaseConcreteNode } from './types.js';
import { NodeIdService } from './nodeIdService.js';
import type { HistoryTurn } from '../../core/agentChatHistory.js';
describe('ContextGraphBuilder', () => {
describe('toGraph', () => {
it('should skip legacy <session_context> headers even if they appear later in the history', () => {
const history: Content[] = [
{ role: 'user', parts: [{ text: 'Message 1' }] },
{ role: 'model', parts: [{ text: 'Reply 1' }] },
const history: HistoryTurn[] = [
{
role: 'user',
parts: [
{
text: '<session_context>\nThis is the Gemini CLI\nSome context...',
},
],
id: '1',
content: { role: 'user', parts: [{ text: 'Message 1' }] },
},
{
id: '2',
content: { role: 'model', parts: [{ text: 'Reply 1' }] },
},
{
id: '3',
content: {
role: 'user',
parts: [
{
text: '<session_context>\nThis is the Gemini CLI\nSome context...',
},
],
},
},
{
id: '4',
content: { role: 'user', parts: [{ text: 'Message 2' }] },
},
{ role: 'user', parts: [{ text: 'Message 2' }] },
];
const builder = new ContextGraphBuilder();
const builder = new ContextGraphBuilder(new NodeIdService());
const nodes = builder.processHistory(history);
// We expect the first two messages and the last one to be present
@@ -36,5 +49,81 @@ describe('ContextGraphBuilder', () => {
expect((nodes[1] as BaseConcreteNode).payload.text).toBe('Reply 1');
expect((nodes[2] as BaseConcreteNode).payload.text).toBe('Message 2');
});
it('should generate completely deterministic graph structure and UUIDs across JSON serialization cycles', () => {
vi.spyOn(Date, 'now').mockReturnValue(0);
const complexHistory: HistoryTurn[] = [
{
id: 'turn-1',
content: {
role: 'user',
parts: [{ text: 'Step 1: complex analysis' }],
},
},
{
id: 'turn-2',
content: {
role: 'model',
parts: [
{ text: 'Thinking about the tool to use.' },
{
functionCall: {
name: 'fetch_data',
args: { query: 'test data' },
},
},
],
},
},
{
id: 'turn-3',
content: {
role: 'user',
parts: [
{
functionResponse: {
name: 'fetch_data',
response: { status: 'success', data: [1, 2, 3] },
},
},
],
},
},
{
id: 'turn-4',
content: { role: 'model', parts: [{ text: 'Analysis complete.' }] },
},
];
// 1. Initial Graph Generation
const builder1 = new ContextGraphBuilder(new NodeIdService());
const nodes1 = builder1.processHistory(complexHistory);
// 2. Serialize and Deserialize (Simulating saving and loading from disk)
const serializedHistory = JSON.stringify(complexHistory);
const parsedHistory = JSON.parse(serializedHistory) as HistoryTurn[];
// 3. Second Graph Generation from parsed JSON
const builder2 = new ContextGraphBuilder(new NodeIdService());
const nodes2 = builder2.processHistory(parsedHistory);
// Assertion: The arrays must be completely identical, including all generated UUIDs
expect(nodes1).toEqual(nodes2);
// Sanity check to ensure IDs are actually populated and consistent
expect(nodes1.length).toBeGreaterThan(0);
nodes1.forEach((node, index) => {
expect(node.id).toBeDefined();
expect(node.id).toBe(nodes2[index].id);
expect(node.timestamp).toBe(0);
if ('turnId' in node) {
expect(node.turnId).toBeDefined();
expect(node.turnId).toBe((nodes2[index] as BaseConcreteNode).turnId);
}
});
vi.restoreAllMocks();
});
});
});
packages/core/src/context/graph/toGraph.ts
+80 -42
View File
@@ -4,14 +4,13 @@
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content, Part } from '@google/genai';
import type { Part } from '@google/genai';
import { type ConcreteNode, NodeType } from './types.js';
import { randomUUID, createHash } from 'node:crypto';
import { createHash } from 'node:crypto';
import { debugLogger } from '../../utils/debugLogger.js';
interface PartWithSynthId extends Part {
_synthId?: string;
}
import type { NodeIdService } from './nodeIdService.js';
import type { HistoryTurn } from '../../core/agentChatHistory.js';
import { isSnapshotState } from '../utils/snapshotGenerator.js';
// Global WeakMap to cache hashes for Part objects.
// This optimizes getStableId by avoiding redundant stringify/hash operations
@@ -62,34 +61,53 @@ function isFunctionResponsePart(
);
}
function isExecutableCodePart(
part: Part,
): part is Part & { executableCode: { code: string; language: string } } {
return (
typeof part.executableCode === 'object' &&
part.executableCode !== null &&
typeof part.executableCode.code === 'string' &&
typeof part.executableCode.language === 'string'
);
}
function isCodeExecutionResultPart(
part: Part,
): part is Part & { codeExecutionResult: { outcome: string; output: string } } {
return (
typeof part.codeExecutionResult === 'object' &&
part.codeExecutionResult !== null &&
typeof part.codeExecutionResult.output === 'string' &&
typeof part.codeExecutionResult.outcome === 'string'
);
}
/**
* Generates a stable ID for an object reference using a WeakMap.
* Generates a stable ID for an object reference using a NodeIdService.
* Falls back to content-based hashing for Part-like objects to ensure
* stability across object re-creations (e.g. during history mapping).
*/
export function getStableId(
obj: object,
nodeIdentityMap: WeakMap<object, string>,
idService: NodeIdService,
turnSalt: string = '',
partIdx: number = 0,
): string {
let id = nodeIdentityMap.get(obj);
let id = idService.get(obj);
if (id) return id;
const cachedHash = PART_HASH_CACHE.get(obj);
if (cachedHash) {
id = `${cachedHash}_${turnSalt}_${partIdx}`;
nodeIdentityMap.set(obj, id);
idService.set(obj, id);
return id;
}
const part = obj as PartWithSynthId;
const part = obj as Part;
let contentHash: string | undefined;
// If the object already has a synthetic ID property, use it.
if (typeof part._synthId === 'string') {
id = part._synthId;
} else if (isTextPart(part)) {
if (isTextPart(part)) {
contentHash = createHash('sha256').update(part.text).digest('hex');
id = `text_${contentHash}_${turnSalt}_${partIdx}`;
} else if (isInlineDataPart(part)) {
@@ -116,6 +134,20 @@ export function getStableId(
)
.digest('hex');
id = `resp_h_${contentHash}_${turnSalt}_${partIdx}`;
} else if (isExecutableCodePart(part)) {
contentHash = createHash('sha256')
.update(
`exec:${part.executableCode.language}:${part.executableCode.code}`,
)
.digest('hex');
id = `exec_${contentHash}_${turnSalt}_${partIdx}`;
} else if (isCodeExecutionResultPart(part)) {
contentHash = createHash('sha256')
.update(
`result:${part.codeExecutionResult.outcome}:${part.codeExecutionResult.output}`,
)
.digest('hex');
id = `result_${contentHash}_${turnSalt}_${partIdx}`;
}
if (contentHash) {
@@ -123,10 +155,14 @@ export function getStableId(
}
if (!id) {
id = randomUUID();
if (turnSalt && partIdx === -1) {
id = `turn_${turnSalt}`;
} else {
id = `${turnSalt}_f_${partIdx}`;
}
}
nodeIdentityMap.set(obj, id);
idService.set(obj, id);
return id;
}
@@ -135,18 +171,14 @@ export function getStableId(
* Every Part in history is mapped to exactly one ConcreteNode.
*/
export class ContextGraphBuilder {
constructor(
private readonly nodeIdentityMap: WeakMap<object, string> = new WeakMap(),
) {}
constructor(private readonly idService: NodeIdService) {}
processHistory(history: readonly Content[]): ConcreteNode[] {
processHistory(history: readonly HistoryTurn[]): ConcreteNode[] {
const nodes: ConcreteNode[] = [];
// Tracks occurrences of identical turn content to ensure unique stable IDs
const seenHashes = new Map<string, number>();
for (let turnIdx = 0; turnIdx < history.length; turnIdx++) {
const msg = history[turnIdx];
const turn = history[turnIdx];
const msg = turn.content;
if (!msg.parts) continue;
// Defensive: Skip legacy environment header regardless of where it appears.
@@ -164,15 +196,8 @@ export class ContextGraphBuilder {
}
}
// Generate a stable salt for this turn based on its role and content
const turnContent = JSON.stringify(msg.parts);
const h = createHash('md5')
.update(`${msg.role}:${turnContent}`)
.digest('hex');
const occurrence = (seenHashes.get(h) || 0) + 1;
seenHashes.set(h, occurrence);
const turnSalt = `${h}_${occurrence}`;
const turnId = getStableId(msg, this.nodeIdentityMap, turnSalt, -1);
const turnSalt = turn.id;
const turnId = `turn_${turnSalt}`;
if (msg.role === 'user') {
for (let partIdx = 0; partIdx < msg.parts.length; partIdx++) {
@@ -180,34 +205,46 @@ export class ContextGraphBuilder {
const apiId =
isFunctionResponsePart(part) &&
typeof part.functionResponse.id === 'string'
? `resp_${part.functionResponse.id}_${turnSalt}_${partIdx}`
? part.functionResponse.id
: isFunctionCallPart(part) &&
typeof part.functionCall.id === 'string'
? `call_${part.functionCall.id}_${turnSalt}_${partIdx}`
? part.functionCall.id
: undefined;
const id =
apiId || getStableId(part, this.nodeIdentityMap, turnSalt, partIdx);
const isSnapshot = isTextPart(part) && isSnapshotState(part.text);
// Use stable API ID if available, otherwise anchor to the turn and index.
const id = apiId
? `${apiId}_${turnSalt}_${partIdx}`
: `${turnSalt}_${partIdx}`;
const node: ConcreteNode = {
id,
timestamp: Date.now(),
type: isFunctionResponsePart(part)
? NodeType.TOOL_EXECUTION
: NodeType.USER_PROMPT,
: isSnapshot
? NodeType.SNAPSHOT
: NodeType.USER_PROMPT,
role: 'user',
payload: part,
turnId,
};
nodes.push(node);
this.idService.set(part, id);
}
} else if (msg.role === 'model') {
for (let partIdx = 0; partIdx < msg.parts.length; partIdx++) {
const part = msg.parts[partIdx];
const apiId =
isFunctionCallPart(part) && typeof part.functionCall.id === 'string'
? `call_${part.functionCall.id}_${turnSalt}_${partIdx}`
? part.functionCall.id
: undefined;
const id =
apiId || getStableId(part, this.nodeIdentityMap, turnSalt, partIdx);
const id = apiId
? `${apiId}_${turnSalt}_${partIdx}`
: `${turnSalt}_${partIdx}`;
const node: ConcreteNode = {
id,
timestamp: Date.now(),
@@ -219,6 +256,7 @@ export class ContextGraphBuilder {
turnId,
};
nodes.push(node);
this.idService.set(part, id);
}
}
}
packages/core/src/context/initializer.ts
+17
View File
@@ -24,6 +24,7 @@ import { StateSnapshotAsyncProcessorOptionsSchema } from './processors/stateSnap
import { RollingSummaryProcessorOptionsSchema } from './processors/rollingSummaryProcessor.js';
import { getEnvironmentContext } from '../utils/environmentContext.js';
import { AdaptiveTokenCalculator } from './utils/adaptiveTokenCalculator.js';
import { estimateContextBreakdown } from '../core/loggingContentGenerator.js';
import { NodeBehaviorRegistry } from './graph/behaviorRegistry.js';
import { registerBuiltInBehaviors } from './graph/builtinBehaviors.js';
@@ -92,10 +93,26 @@ export async function initializeContextManager(
const behaviorRegistry = new NodeBehaviorRegistry();
registerBuiltInBehaviors(behaviorRegistry);
const getOverheadTokens = () => {
const breakdown = estimateContextBreakdown([], {
systemInstruction: {
role: 'system',
parts: [{ text: chat.getSystemInstruction() }],
},
tools: chat.getTools(),
});
return (
breakdown.system_instructions +
breakdown.tool_definitions +
breakdown.mcp_servers
);
};
const calculator = new AdaptiveTokenCalculator(
charsPerToken,
behaviorRegistry,
eventBus,
getOverheadTokens,
);
const env = new ContextEnvironmentImpl(
packages/core/src/context/processors/blobDegradationProcessor.ts
+4 -3
View File
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type { ProcessArgs, ContextProcessor } from '../pipeline.js';
import * as fs from 'node:fs/promises';
@@ -62,7 +62,8 @@ export function createBlobDegradationProcessor(
if (payload.inlineData?.data && payload.inlineData?.mimeType) {
await ensureDir();
const ext = payload.inlineData.mimeType.split('/')[1] || 'bin';
const fileName = `blob_${Date.now()}_${randomUUID()}.${ext}`;
// Use a stable filename based on the node ID
const fileName = `blob_${deriveStableId([node.id])}.${ext}`;
const filePath = path.join(blobOutputsDir, fileName);
const buffer = Buffer.from(payload.inlineData.data, 'base64');
@@ -92,7 +93,7 @@ export function createBlobDegradationProcessor(
if (newText && tokensSaved > 0) {
returnedNodes.push({
...node,
id: randomUUID(),
id: deriveStableId([node.id, 'degraded']),
payload: { text: newText },
replacesId: node.id,
turnId: node.turnId,
packages/core/src/context/processors/nodeDistillationProcessor.ts
+5 -3
View File
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import { type ConcreteNode, NodeType } from '../graph/types.js';
@@ -99,9 +99,10 @@ export function createNodeDistillationProcessor(
if (newTokens < oldTokens) {
const distilledPayload = updatePart(payload, { text: summary });
const newId = deriveStableId([node.id, 'distilled']);
returnedNodes.push({
...node,
id: randomUUID(),
id: newId,
payload: distilledPayload,
replacesId: node.id,
timestamp: node.timestamp,
@@ -158,9 +159,10 @@ export function createNodeDistillationProcessor(
functionResponse: newFR,
});
const newId = deriveStableId([node.id, 'distilled']);
returnedNodes.push({
...node,
id: randomUUID(),
id: newId,
payload: distilledPayload,
replacesId: node.id,
timestamp: node.timestamp,
packages/core/src/context/processors/nodeTruncationProcessor.ts
+3 -2
View File
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
@@ -79,9 +79,10 @@ export function createNodeTruncationProcessor(
if (text) {
const squashResult = tryApplySquash(text, limitChars);
if (squashResult) {
const newId = deriveStableId([node.id, 'truncated']);
returnedNodes.push({
...node,
id: randomUUID(),
id: newId,
payload: { ...payload, text: squashResult.text },
replacesId: node.id,
turnId: node.turnId,
packages/core/src/context/processors/rollingSummaryProcessor.ts
+4 -4
View File
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type {
ContextProcessor,
@@ -112,7 +112,8 @@ export function createRollingSummaryProcessor(
try {
// Synthesize the rolling summary synchronously
const snapshotText = await generateRollingSummary(nodesToSummarize);
const newId = randomUUID();
const consumedIds = nodesToSummarize.map((n) => n.id);
const newId = deriveStableId(consumedIds);
const summaryNode: RollingSummary = {
id: newId,
@@ -121,10 +122,9 @@ export function createRollingSummaryProcessor(
timestamp: nodesToSummarize[nodesToSummarize.length - 1].timestamp,
role: 'user',
payload: { text: snapshotText },
abstractsIds: nodesToSummarize.map((n) => n.id),
abstractsIds: consumedIds,
};
const consumedIds = nodesToSummarize.map((n) => n.id);
const returnedNodes = targets.filter(
(t) => !consumedIds.includes(t.id),
);
packages/core/src/context/processors/stateSnapshotProcessor.ts
+10 -3
View File
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type {
ContextProcessor,
@@ -90,8 +90,11 @@ export function createStateSnapshotProcessor(
const isValid = consumedIds.every((id) => targetIds.has(id));
if (isValid) {
debugLogger.log(
`[StateSnapshotProcessor] Successfully spliced PROPOSED_SNAPSHOT from Inbox into Graph. Consumed ${consumedIds.length} nodes.`,
);
// If valid, apply it!
const newId = randomUUID();
const newId = deriveStableId(consumedIds);
const snapshotNode: Snapshot = {
id: newId,
@@ -120,6 +123,10 @@ export function createStateSnapshotProcessor(
inbox.consume(proposed.id);
return returnedNodes;
} else {
debugLogger.log(
`[StateSnapshotProcessor] Rejected PROPOSED_SNAPSHOT from Inbox because one or more target IDs were missing from the current graph window.`,
);
}
}
}
@@ -179,11 +186,11 @@ export function createStateSnapshotProcessor(
maxStateTokens: options.maxStateTokens,
},
);
const newId = randomUUID();
const consumedIds = nodesToSummarize.map((n) => n.id);
if (baselineIdToConsume && !consumedIds.includes(baselineIdToConsume)) {
consumedIds.push(baselineIdToConsume);
}
const newId = deriveStableId(consumedIds);
const snapshotNode: Snapshot = {
id: newId,
packages/core/src/context/processors/toolMaskingProcessor.ts
+6 -4
View File
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import * as fs from 'node:fs/promises';
@@ -120,7 +120,7 @@ export function createToolMaskingProcessor(
directoryCreated = true;
}
const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${randomUUID()}.txt`;
const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${deriveStableId([content])}.txt`;
const filePath = path.join(toolOutputsDir, fileName);
await fs.writeFile(filePath, content);
@@ -214,9 +214,10 @@ export function createToolMaskingProcessor(
functionCall: newFC,
});
const newId = deriveStableId([node.id, 'masked']);
returnedNodes.push({
...node,
id: randomUUID(),
id: newId,
payload: maskedPart,
replacesId: node.id,
turnId: node.turnId,
@@ -242,9 +243,10 @@ export function createToolMaskingProcessor(
functionResponse: newFR,
});
const newId = deriveStableId([node.id, 'masked']);
returnedNodes.push({
...node,
id: randomUUID(),
id: newId,
payload: maskedPart,
replacesId: node.id,
turnId: node.turnId,
packages/core/src/context/system-tests/__snapshots__/lifecycle.golden.test.ts.snap
+298 -198
View File
File diff suppressed because one or more lines are too long
packages/core/src/context/system-tests/hysteresis.test.ts
+11 -6
View File
@@ -9,7 +9,7 @@ import { SimulationHarness } from './simulationHarness.js';
import { createMockLlmClient } from '../testing/contextTestUtils.js';
import type { ContextProfile } from '../config/profiles.js';
import { generalistProfile } from '../config/profiles.js';
import type { Content } from '@google/genai';
import { type HistoryTurn } from '../../core/agentChatHistory.js';
describe('Context Manager Hysteresis Tests', () => {
const mockLlmClient = createMockLlmClient(['<SNAPSHOT>']);
@@ -18,6 +18,7 @@ describe('Context Manager Hysteresis Tests', () => {
...generalistProfile,
name: 'Hysteresis Stress Test',
config: {
...generalistProfile.config,
budget: {
maxTokens: 5000,
retainedTokens: 1000,
@@ -26,9 +27,13 @@ describe('Context Manager Hysteresis Tests', () => {
},
});
const getProjectionTokens = (proj: Content[], harness: SimulationHarness) =>
const getProjectionTokens = (
proj: HistoryTurn[],
harness: SimulationHarness,
) =>
proj.reduce(
(sum, c) => sum + harness.env.tokenCalculator.calculateContentTokens(c),
(sum, c) =>
sum + harness.env.tokenCalculator.calculateContentTokens(c.content),
0,
);
@@ -57,7 +62,7 @@ describe('Context Manager Hysteresis Tests', () => {
// No snapshot because maxTokens (5000) not exceeded, and deficit < threshold.
expect(
state.finalProjection.some((c) =>
c.parts?.some((p) => p.text?.includes('<SNAPSHOT>')),
c.content.parts?.some((p) => p.text?.includes('<SNAPSHOT>')),
),
).toBe(false);
@@ -79,7 +84,7 @@ describe('Context Manager Hysteresis Tests', () => {
state = await harness.getGoldenState();
expect(
state.finalProjection.some((c) =>
c.parts?.some((p) => p.text?.includes('<SNAPSHOT>')),
c.content.parts?.some((p) => p.text?.includes('<SNAPSHOT>')),
),
).toBe(true);
});
@@ -108,7 +113,7 @@ describe('Context Manager Hysteresis Tests', () => {
let state = await harness.getGoldenState();
expect(
state.finalProjection.some((c) =>
c.parts?.some((p) => p.text?.includes('<SNAPSHOT>')),
c.content.parts?.some((p) => p.text?.includes('<SNAPSHOT>')),
),
).toBe(true);
packages/core/src/context/system-tests/lifecycle.golden.test.ts
+2
View File
@@ -17,6 +17,7 @@ expect.addSnapshotSerializer({
(/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}/i.test(
val,
) ||
/^[0-9a-f]{32}$/i.test(val) ||
/[\\/]tmp[\\/]sim/.test(val)),
print: (val) => {
if (typeof val !== 'string') return `"${val}"`;
@@ -25,6 +26,7 @@ expect.addSnapshotSerializer({
/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}/gi,
'<UUID>',
)
.replace(/\b[0-9a-f]{32}\b/gi, '<UUID>')
.replace(/[\\/]tmp[\\/]sim[^\s"'\]]*/g, '<MOCKED_DIR>');
// Also scrub timestamps in filenames like blob_1234567890_...
packages/core/src/context/system-tests/simulationHarness.ts
+3 -1
View File
@@ -4,6 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { ContextManager } from '../contextManager.js';
import { AgentChatHistory } from '../../core/agentChatHistory.js';
import type { Content } from '@google/genai';
@@ -98,7 +99,8 @@ export class SimulationHarness {
async simulateTurn(messages: Content[]) {
// 1. Append the new messages
const currentHistory = this.chatHistory.get();
this.chatHistory.set([...currentHistory, ...messages]);
const turns = messages.map((m) => ({ id: randomUUID(), content: m }));
this.chatHistory.set([...currentHistory, ...turns]);
// 2. Measure tokens immediately after append
const tokensBefore = this.env.tokenCalculator.calculateConcreteListTokens(
packages/core/src/context/utils/adaptiveTokenCalculator.test.ts
+65
View File
@@ -122,4 +122,69 @@ describe('AdaptiveTokenCalculator', () => {
expect(calculator.getLearnedWeight()).toBe(1.0);
});
it('should subtract overhead tokens from actual tokens when determining target weight', () => {
const eventBus = new ContextEventBus();
const getOverheadTokens = () => 40;
const calculator = new AdaptiveTokenCalculator(
charsPerToken,
registry,
eventBus,
getOverheadTokens,
);
// Initial state: weight = 1.0
// Simulate an event where the API reported 100 tokens, and our base units were 100
// But overhead is 40.
// actualGraphTokens = 100 - 40 = 60
// rawTargetWeight = 60 / 100 = 0.6
// targetWeight = Math.max(0.5, 0.6) = 0.6
// newWeight = 1.0 * 0.8 + 0.6 * 0.2 = 0.8 + 0.12 = 0.92
eventBus.emitTokenGroundTruth({
actualTokens: 100,
promptBaseUnits: 100,
});
expect(calculator.getLearnedWeight()).toBeCloseTo(0.92, 5);
});
it('should enforce the maxStep limit to prevent violent oscillation from massive outliers', () => {
const eventBus = new ContextEventBus();
const maxStep = 0.05; // Tight limit
const calculator = new AdaptiveTokenCalculator(
charsPerToken,
registry,
eventBus,
undefined,
{ maxStep },
);
// Initial state: weight = 1.0
// Simulate a massive outlier where the API reports 10,000 tokens for 100 base units.
// rawTargetWeight = 100
// targetWeight = Math.min(100, 1.0 * 2.0) = 2.0
// emaWeight = 1.0 * 0.8 + 2.0 * 0.2 = 1.2
// BUT maxStep is 0.05, so the actual step is clamped.
// finalWeight = 1.0 + 0.05 = 1.05
eventBus.emitTokenGroundTruth({
actualTokens: 10000,
promptBaseUnits: 100,
});
expect(calculator.getLearnedWeight()).toBeCloseTo(1.05, 5);
// Simulate a massive under-estimation
// rawTargetWeight = 0
// targetWeight = Math.max(0, 1.05 * 0.5) = 0.525
// emaWeight = 1.05 * 0.8 + 0.525 * 0.2 = 0.84 + 0.105 = 0.945
// BUT maxStep is 0.05, so step is clamped: 1.05 - 0.05 = 1.0
eventBus.emitTokenGroundTruth({
actualTokens: 0,
promptBaseUnits: 100,
});
expect(calculator.getLearnedWeight()).toBeCloseTo(1.0, 5);
});
});
packages/core/src/context/utils/adaptiveTokenCalculator.ts
+43 -7
View File
@@ -14,6 +14,13 @@ import type { NodeBehaviorRegistry } from '../graph/behaviorRegistry.js';
import type { ContextEventBus, TokenGroundTruthEvent } from '../eventBus.js';
import { debugLogger } from '../../utils/debugLogger.js';
export interface AdaptiveLearningConfig {
/** The momentum factor for the Exponential Moving Average (EMA). Defaults to 0.2. */
learningRate?: number;
/** The absolute maximum change allowed to the weight in a single turn. Defaults to 0.15. */
maxStep?: number;
}
/**
* An Adaptive Token Calculator that dynamically learns the true token cost of the user's
* conversation by applying an Exponential Moving Average (EMA) gradient descent to
@@ -26,12 +33,18 @@ import { debugLogger } from '../../utils/debugLogger.js';
export class AdaptiveTokenCalculator implements AdvancedTokenCalculator {
private learnedWeight = 1.0;
private readonly baseCalculator: StaticTokenCalculator;
private readonly learningRate: number;
private readonly maxStep: number;
constructor(
charsPerToken: number,
registry: NodeBehaviorRegistry,
eventBus: ContextEventBus,
private readonly getOverheadTokens?: () => number,
config?: AdaptiveLearningConfig,
) {
this.learningRate = config?.learningRate ?? 0.2;
this.maxStep = config?.maxStep ?? 0.15;
this.baseCalculator = new StaticTokenCalculator(charsPerToken, registry);
eventBus.onTokenGroundTruth((event: TokenGroundTruthEvent) => {
this.handleGroundTruth(event.actualTokens, event.promptBaseUnits);
@@ -41,21 +54,44 @@ export class AdaptiveTokenCalculator implements AdvancedTokenCalculator {
private handleGroundTruth(actualTokens: number, promptBaseUnits: number) {
if (promptBaseUnits <= 0) return;
const overheadTokens = this.getOverheadTokens
? this.getOverheadTokens()
: 0;
// The Gemini API token count includes the static overhead (system instruction + tools)
// and the dynamic chat history (which we measure as promptBaseUnits).
// We subtract the overhead so the adaptive calculator is comparing "apples to apples"
// when learning the weight multiplier for the graph nodes.
const actualGraphTokens = Math.max(0, actualTokens - overheadTokens);
// Determine what ratio we should have used
const targetWeight = actualTokens / promptBaseUnits;
const rawTargetWeight = actualGraphTokens / promptBaseUnits;
const oldWeight = this.learnedWeight;
// Apply Momentum (Learning Rate)
const learningRate = 0.2;
const newWeight =
oldWeight * (1 - learningRate) + targetWeight * learningRate;
// Dampen extreme outliers *before* applying the EMA by capping the target weight
// to a reasonable multiple of the current weight. This prevents a single massive
// anomaly from destroying the running average.
const targetWeight = Math.max(
oldWeight * 0.5,
Math.min(rawTargetWeight, oldWeight * 2.0),
);
// Clamp to reasonable safety bounds to prevent rogue metadata poisoning the system
// Apply Momentum (Learning Rate)
let newWeight =
oldWeight * (1 - this.learningRate) + targetWeight * this.learningRate;
// Hard limit the maximum step size per turn to prevent violent oscillation
if (newWeight > oldWeight + this.maxStep)
newWeight = oldWeight + this.maxStep;
if (newWeight < oldWeight - this.maxStep)
newWeight = oldWeight - this.maxStep;
// Clamp to reasonable absolute safety bounds
this.learnedWeight = Math.max(0.5, Math.min(newWeight, 2.0));
debugLogger.log(
`[AdaptiveTokenCalculator] Learned weight updated to ${this.learnedWeight.toFixed(3)} ` +
`(API Tokens: ${actualTokens}, Base Units: ${promptBaseUnits}, Target Ratio: ${targetWeight.toFixed(3)})`,
`(API Tokens: ${actualTokens}, Overhead: ${overheadTokens}, Graph Tokens: ${actualGraphTokens}, Base Units: ${promptBaseUnits}, Target Ratio: ${targetWeight.toFixed(3)})`,
);
}
packages/core/src/context/utils/snapshotGenerator.test.ts
+121 -1
View File
@@ -5,12 +5,132 @@
*/
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { SnapshotGenerator, type SnapshotState } from './snapshotGenerator.js';
import {
SnapshotGenerator,
type SnapshotState,
SnapshotStateHelper,
} from './snapshotGenerator.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { NodeType, type ConcreteNode } from '../graph/types.js';
import type { Mock } from 'vitest';
describe('SnapshotStateHelper', () => {
describe('exportState', () => {
it('should flatten nested abstractsIds to pristine IDs', () => {
// Setup a graph with nested snapshots
// S3 abstracts [S2, N5]
// S2 abstracts [S1, N3, N4]
// S1 abstracts [N1, N2]
const nodes: ConcreteNode[] = [
{
id: 'N1',
type: NodeType.USER_PROMPT,
timestamp: 10,
role: 'user',
payload: { text: '1' },
turnId: 'T1',
},
{
id: 'N2',
type: NodeType.AGENT_THOUGHT,
timestamp: 20,
role: 'model',
payload: { text: '2' },
turnId: 'T1',
},
{
id: 'S1',
type: NodeType.SNAPSHOT,
timestamp: 30,
role: 'user',
payload: { text: 'State 1' },
turnId: 'S1',
abstractsIds: ['N1', 'N2'],
},
{
id: 'N3',
type: NodeType.USER_PROMPT,
timestamp: 40,
role: 'user',
payload: { text: '3' },
turnId: 'T2',
},
{
id: 'N4',
type: NodeType.AGENT_THOUGHT,
timestamp: 50,
role: 'model',
payload: { text: '4' },
turnId: 'T2',
},
{
id: 'S2',
type: NodeType.SNAPSHOT,
timestamp: 60,
role: 'user',
payload: { text: 'State 2' },
turnId: 'S2',
abstractsIds: ['S1', 'N3', 'N4'],
},
{
id: 'N5',
type: NodeType.USER_PROMPT,
timestamp: 70,
role: 'user',
payload: { text: '5' },
turnId: 'T3',
},
{
id: 'S3',
type: NodeType.SNAPSHOT,
timestamp: 80,
role: 'user',
payload: { text: 'State 3' },
turnId: 'S3',
abstractsIds: ['S2', 'N5'],
},
];
const state = SnapshotStateHelper.exportState(nodes);
expect(state.snapshot).toBeDefined();
expect(state.snapshot?.text).toBe('State 3');
// Should be flattened to only the "pristine" (non-snapshot) IDs
const consumedIds = state.snapshot?.consumedIds;
expect(consumedIds).toContain('N1');
expect(consumedIds).toContain('N2');
expect(consumedIds).toContain('N3');
expect(consumedIds).toContain('N4');
expect(consumedIds).toContain('N5');
// Should NOT contain the intermediate snapshot IDs
expect(consumedIds).not.toContain('S1');
expect(consumedIds).not.toContain('S2');
expect(consumedIds?.length).toBe(5);
});
it('should return empty state if no snapshot baseline is found', () => {
const nodes: ConcreteNode[] = [
{
id: 'N1',
type: NodeType.USER_PROMPT,
timestamp: 10,
role: 'user',
payload: { text: '1' },
turnId: 'T1',
},
];
const state = SnapshotStateHelper.exportState(nodes);
expect(state).toEqual({});
});
});
});
describe('SnapshotGenerator', () => {
let mockEnv: ContextEnvironment;
let mockGenerateJson: Mock;
packages/core/src/context/utils/snapshotGenerator.ts
+103
View File
@@ -48,10 +48,39 @@ export interface SnapshotState {
recent_arc: string[];
}
import { debugLogger } from '../../utils/debugLogger.js';
export function isSnapshotState(text: string): boolean {
const trimmed = text.trim();
if (!trimmed.startsWith('{') || !trimmed.endsWith('}')) {
return false;
}
try {
const parsed: unknown = JSON.parse(trimmed);
if (!isRecord(parsed)) return false;
const isSnap =
Array.isArray(parsed['active_tasks']) &&
Array.isArray(parsed['discovered_facts']) &&
Array.isArray(parsed['constraints_and_preferences']) &&
Array.isArray(parsed['recent_arc']);
if (!isSnap) {
debugLogger.log(
'[isSnapshotState] FAILED FOR JSON:',
JSON.stringify(parsed),
);
}
return isSnap;
} catch {
debugLogger.log('[isSnapshotState] PARSE FAILED FOR:', trimmed);
return false;
}
}
export interface BaselineSnapshotInfo {
text: string;
abstractsIds: string[];
id: string;
timestamp: number;
}
/**
@@ -61,6 +90,20 @@ export interface BaselineSnapshotInfo {
export function findLatestSnapshotBaseline(
targets: readonly ConcreteNode[],
): BaselineSnapshotInfo | undefined {
debugLogger.log(
'[findLatestSnapshotBaseline] Targets:',
targets.map((t) => ({
id: t.id,
type: t.type,
text:
t.payload &&
typeof t.payload === 'object' &&
'text' in t.payload &&
typeof t.payload.text === 'string'
? t.payload.text.substring(0, 20)
: '',
})),
);
const lastSnapshotNode = [...targets]
.reverse()
.find((n) => n.type === NodeType.SNAPSHOT && n.payload.text);
@@ -72,8 +115,10 @@ export function findLatestSnapshotBaseline(
? [...lastSnapshotNode.abstractsIds]
: [],
id: lastSnapshotNode.id,
timestamp: lastSnapshotNode.timestamp,
};
}
return undefined;
}
@@ -326,3 +371,61 @@ ${formatNodesForLlm(nodes)}`;
return JSON.stringify(newState);
}
}
/**
* Shared logic for working with Snapshot node state.
*/
export class SnapshotStateHelper {
/**
* Flatten nested abstract IDs to only the "pristine" (non-snapshot) IDs.
*/
static flattenAbstracts(
nodes: ConcreteNode[],
abstractsIds: readonly string[],
): string[] {
const pristineIds: string[] = [];
const nodeMap = new Map(nodes.map((n) => [n.id, n]));
const walk = (ids: readonly string[]) => {
for (const id of ids) {
const node = nodeMap.get(id);
if (!node) {
// Fallback: if node not in map, treat as pristine ID
pristineIds.push(id);
continue;
}
if (node.type === NodeType.SNAPSHOT && node.abstractsIds) {
walk(node.abstractsIds);
} else {
pristineIds.push(id);
}
}
};
walk(abstractsIds);
return Array.from(new Set(pristineIds)); // Dedupe
}
/**
* Helper to extract state from the most recent snapshot in a list of nodes.
*/
static exportState(nodes: ConcreteNode[]): {
snapshot?: { text: string; consumedIds: string[] };
} {
const baseline = findLatestSnapshotBaseline(nodes);
if (!baseline) return {};
const node = nodes.find((n) => n.id === baseline.id);
if (!node || node.type !== NodeType.SNAPSHOT) return {};
const consumedIds = this.flattenAbstracts(nodes, node.abstractsIds || []);
return {
snapshot: {
text: baseline.text,
consumedIds,
},
};
}
}