diff --git a/packages/core/src/agents/browser/snapshotSuperseder.ts b/packages/core/src/agents/browser/snapshotSuperseder.ts index e8a5068dd9..935223c2e2 100644 --- a/packages/core/src/agents/browser/snapshotSuperseder.ts +++ b/packages/core/src/agents/browser/snapshotSuperseder.ts @@ -14,8 +14,8 @@ * model call so the model only ever sees the most recent snapshot in full. */ -import type { GeminiChat } from '../../core/geminiChat.js'; -import type { Content, Part } from '@google/genai'; +import type { GeminiChat, HistoryTurn } from '../../core/geminiChat.js'; +import type { Part } from '@google/genai'; import { debugLogger } from '../../utils/debugLogger.js'; const TAKE_SNAPSHOT_TOOL_NAME = 'take_snapshot'; @@ -39,7 +39,7 @@ export const SNAPSHOT_SUPERSEDED_PLACEHOLDER = * Uses {@link GeminiChat.setHistory} to apply the modified history. */ export function supersedeStaleSnapshots(chat: GeminiChat): void { - const history = chat.getHistory(); + const history = chat.getHistoryTurns(); // Locate all (contentIndex, partIndex) tuples for take_snapshot responses. const snapshotLocations: Array<{ @@ -48,7 +48,7 @@ export function supersedeStaleSnapshots(chat: GeminiChat): void { }> = []; for (let i = 0; i < history.length; i++) { - const parts = history[i].parts; + const parts = history[i].content.parts; if (!parts) continue; for (let j = 0; j < parts.length; j++) { const part = parts[j]; @@ -71,7 +71,7 @@ export function supersedeStaleSnapshots(chat: GeminiChat): void { const staleLocations = snapshotLocations.slice(0, -1); const needsUpdate = staleLocations.some(({ contentIdx, partIdx }) => { const output = getResponseOutput( - history[contentIdx].parts![partIdx].functionResponse?.response, + history[contentIdx].content.parts![partIdx].functionResponse?.response, ); return !output.includes(SNAPSHOT_SUPERSEDED_PLACEHOLDER); }); @@ -81,15 +81,18 @@ export function supersedeStaleSnapshots(chat: GeminiChat): void { } // Shallow-copy the history and replace stale snapshots. - const newHistory: Content[] = history.map((content) => ({ - ...content, - parts: content.parts ? [...content.parts] : undefined, + const newHistory: HistoryTurn[] = history.map((turn) => ({ + id: turn.id, + content: { + ...turn.content, + parts: turn.content.parts ? [...turn.content.parts] : undefined, + }, })); let replacedCount = 0; for (const { contentIdx, partIdx } of staleLocations) { - const originalPart = newHistory[contentIdx].parts![partIdx]; + const originalPart = newHistory[contentIdx].content.parts![partIdx]; if (!originalPart.functionResponse) continue; // Check if already superseded @@ -106,7 +109,7 @@ export function supersedeStaleSnapshots(chat: GeminiChat): void { }, }; - newHistory[contentIdx].parts![partIdx] = replacementPart; + newHistory[contentIdx].content.parts![partIdx] = replacementPart; replacedCount++; } diff --git a/packages/core/src/agents/local-executor.test.ts b/packages/core/src/agents/local-executor.test.ts index a1f3b72965..176371a808 100644 --- a/packages/core/src/agents/local-executor.test.ts +++ b/packages/core/src/agents/local-executor.test.ts @@ -756,12 +756,19 @@ describe('LocalAgentExecutor', () => { expect(startHistory).toBeDefined(); expect(startHistory).toHaveLength(2); + const history = startHistory!; // Perform checks on defined objects to satisfy TS - const firstPart = startHistory?.[0]?.parts?.[0]; + const firstPart = + 'content' in history[0] + ? history[0].content.parts?.[0] + : (history[0] as Content).parts?.[0]; expect(firstPart?.text).toBe('Goal: TestGoal'); - const secondPart = startHistory?.[1]?.parts?.[0]; + const secondPart = + 'content' in history[1] + ? history[1].content.parts?.[0] + : (history[1] as Content).parts?.[0]; expect(secondPart?.text).toBe('OK, starting on TestGoal.'); }); @@ -3601,7 +3608,14 @@ describe('LocalAgentExecutor', () => { expect(mockCompress).toHaveBeenCalledTimes(1); expect(mockSetHistory).toHaveBeenCalledTimes(1); - expect(mockSetHistory).toHaveBeenCalledWith(compressedHistory); + // History turns are now wrapped with IDs + expect(mockSetHistory).toHaveBeenCalledWith( + expect.arrayContaining([ + expect.objectContaining({ + content: expect.objectContaining({ role: 'user' }), + }), + ]), + ); }); it('should pass hasFailedCompressionAttempt=true to compression after a failure', async () => { @@ -3706,7 +3720,14 @@ describe('LocalAgentExecutor', () => { expect(mockCompress.mock.calls[2][5]).toBe(false); expect(mockSetHistory).toHaveBeenCalledTimes(1); - expect(mockSetHistory).toHaveBeenCalledWith(compressedHistory); + // History turns are now wrapped with IDs + expect(mockSetHistory).toHaveBeenCalledWith( + expect.arrayContaining([ + expect.objectContaining({ + content: expect.objectContaining({ role: 'user' }), + }), + ]), + ); }); }); diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index 266eb55a4c..7b320e7b18 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -15,6 +15,7 @@ import { type FunctionCall, type FunctionDeclaration, } from '@google/genai'; +import { randomUUID } from 'node:crypto'; import { ToolRegistry } from '../tools/tool-registry.js'; import { PromptRegistry } from '../prompts/prompt-registry.js'; import { ResourceRegistry } from '../resources/resource-registry.js'; @@ -919,12 +920,20 @@ export class LocalAgentExecutor { this.hasFailedCompressionAttempt = true; } else if (info.compressionStatus === CompressionStatus.COMPRESSED) { if (newHistory) { - chat.setHistory(newHistory); + const turns = newHistory.map((c) => ({ + id: randomUUID(), + content: c, + })); + chat.setHistory(turns); this.hasFailedCompressionAttempt = false; } } else if (info.compressionStatus === CompressionStatus.CONTENT_TRUNCATED) { if (newHistory) { - chat.setHistory(newHistory); + const turns = newHistory.map((c) => ({ + id: randomUUID(), + content: c, + })); + chat.setHistory(turns); // Do NOT reset hasFailedCompressionAttempt. // We only truncated content because summarization previously failed. // We want to keep avoiding expensive summarization calls. diff --git a/packages/core/src/context/contextManager.barrier.test.ts b/packages/core/src/context/contextManager.barrier.test.ts index e46637d7d8..9d869d42ab 100644 --- a/packages/core/src/context/contextManager.barrier.test.ts +++ b/packages/core/src/context/contextManager.barrier.test.ts @@ -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,29 @@ 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 +69,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.'); diff --git a/packages/core/src/context/contextManager.hotstart.test.ts b/packages/core/src/context/contextManager.hotstart.test.ts index 5d0d848267..2a391e453b 100644 --- a/packages/core/src/context/contextManager.hotstart.test.ts +++ b/packages/core/src/context/contextManager.hotstart.test.ts @@ -47,7 +47,7 @@ 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 +81,7 @@ 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(); diff --git a/packages/core/src/context/contextManager.ts b/packages/core/src/context/contextManager.ts index 9316f62362..90bf3a7f6e 100644 --- a/packages/core/src/context/contextManager.ts +++ b/packages/core/src/context/contextManager.ts @@ -5,7 +5,7 @@ */ import type { Content } from '@google/genai'; -import type { AgentChatHistory } from '../core/agentChatHistory.js'; +import type { AgentChatHistory, HistoryTurn } from '../core/agentChatHistory.js'; import { isToolExecution, type ConcreteNode } from './graph/types.js'; import type { ContextEventBus } from './eventBus.js'; import type { ContextTracer } from './tracer.js'; @@ -16,8 +16,6 @@ import { HistoryObserver } from './historyObserver.js'; import { render } from './graph/render.js'; import { ContextWorkingBufferImpl } from './pipeline/contextWorkingBuffer.js'; import { debugLogger } from '../utils/debugLogger.js'; -import { SnapshotStateHelper } from './utils/snapshotGenerator.js'; -import type { ContextEngineState } from '../services/chatRecordingTypes.js'; import { hardenHistory } from '../utils/historyHardening.js'; import { checkContextInvariants } from './utils/invariantChecker.js'; import type { AdvancedTokenCalculator } from './utils/contextTokenCalculator.js'; @@ -40,7 +38,8 @@ export class ContextManager { private lastRenderCache?: { nodesHash: string; result: { - history: Content[]; + history: HistoryTurn[]; + apiHistory: Content[]; didApplyManagement: boolean; baseUnits: number; processedNodes: readonly ConcreteNode[]; @@ -296,11 +295,12 @@ 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 = new Set(), abortSignal?: AbortSignal, ): Promise<{ - history: Content[]; + history: HistoryTurn[]; + apiHistory: Content[]; didApplyManagement: boolean; baseUnits: number; processedNodes: readonly ConcreteNode[]; @@ -400,18 +400,21 @@ export class ContextManager { this.tracer.logEvent('ContextManager', 'Finished rendering'); - const combinedHistory = header - ? [header, ...renderedHistory] - : renderedHistory; + const hardenedHistory = hardenHistory( + renderedHistory, + { + sentinels: this.sidecar.sentinels, + }, + ); + + const apiHistory = hardenedHistory.map((h) => h.content); + if (header) { + apiHistory.unshift(header); + } const result = { - history: hardenHistory( - combinedHistory, - { - sentinels: this.sidecar.sentinels, - }, - this.env.graphMapper.getIdService(), - ), + history: hardenedHistory, + apiHistory, didApplyManagement, baseUnits, processedNodes, @@ -433,10 +436,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) + @@ -468,13 +472,4 @@ export class ContextManager { ); } } - - exportState(): ContextEngineState { - return SnapshotStateHelper.exportState(this.buffer.nodes); - } - - restoreState(state: ContextEngineState): void { - if (!state) return; - SnapshotStateHelper.restoreState(state, this.env.inbox); - } } diff --git a/packages/core/src/context/graph/fromGraph.ts b/packages/core/src/context/graph/fromGraph.ts index e762049b02..92740cc02e 100644 --- a/packages/core/src/context/graph/fromGraph.ts +++ b/packages/core/src/context/graph/fromGraph.ts @@ -8,25 +8,27 @@ 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[], idService?: NodeIdService, -): Content[] { +): 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). @@ -40,26 +42,25 @@ export function fromGraph( // 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; } diff --git a/packages/core/src/context/graph/mapper.ts b/packages/core/src/context/graph/mapper.ts index 68d5498bc4..ffd9a4cac0 100644 --- a/packages/core/src/context/graph/mapper.ts +++ b/packages/core/src/context/graph/mapper.ts @@ -5,8 +5,7 @@ */ 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'; @@ -22,7 +21,7 @@ export class ContextGraphMapper { return this.builder.processHistory(event.payload); } - fromGraph(nodes: readonly ConcreteNode[]): Content[] { + fromGraph(nodes: readonly ConcreteNode[]): HistoryTurn[] { return fromGraph(nodes, this.idService); } diff --git a/packages/core/src/context/graph/render.ts b/packages/core/src/context/graph/render.ts index eae24641e7..df1ccea47e 100644 --- a/packages/core/src/context/graph/render.ts +++ b/packages/core/src/context/graph/render.ts @@ -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,7 +29,7 @@ export async function render( header?: Content, previewNodeIds: ReadonlySet = new Set(), ): Promise<{ - history: Content[]; + history: HistoryTurn[]; didApplyManagement: boolean; baseUnits: number; processedNodes: readonly ConcreteNode[]; @@ -98,7 +99,7 @@ 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, @@ -152,7 +153,7 @@ 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, diff --git a/packages/core/src/context/graph/toGraph.test.ts b/packages/core/src/context/graph/toGraph.test.ts index dce2257b23..d74adcb4d2 100644 --- a/packages/core/src/context/graph/toGraph.test.ts +++ b/packages/core/src/context/graph/toGraph.test.ts @@ -6,25 +6,37 @@ 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 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: '\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: '\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(new NodeIdService()); @@ -41,32 +53,44 @@ describe('ContextGraphBuilder', () => { it('should generate completely deterministic graph structure and UUIDs across JSON serialization cycles', () => { vi.spyOn(Date, 'now').mockReturnValue(0); - const complexHistory: Content[] = [ - { role: 'user', parts: [{ text: 'Step 1: complex analysis' }] }, + const complexHistory: HistoryTurn[] = [ { - role: 'model', - parts: [ - { text: 'Thinking about the tool to use.' }, - { - functionCall: { - name: 'fetch_data', - args: { query: 'test data' }, - }, - }, - ], + id: 'turn-1', + content: { role: 'user', parts: [{ text: 'Step 1: complex analysis' }] }, }, { - role: 'user', - parts: [ - { - functionResponse: { - name: 'fetch_data', - response: { status: 'success', data: [1, 2, 3] }, + 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.' }] }, }, - { role: 'model', parts: [{ text: 'Analysis complete.' }] }, ]; // 1. Initial Graph Generation @@ -75,7 +99,7 @@ describe('ContextGraphBuilder', () => { // 2. Serialize and Deserialize (Simulating saving and loading from disk) const serializedHistory = JSON.stringify(complexHistory); - const parsedHistory = JSON.parse(serializedHistory) as Content[]; + const parsedHistory = JSON.parse(serializedHistory) as HistoryTurn[]; // 3. Second Graph Generation from parsed JSON const builder2 = new ContextGraphBuilder(new NodeIdService()); diff --git a/packages/core/src/context/graph/toGraph.ts b/packages/core/src/context/graph/toGraph.ts index 5a7c50a2fe..51012fed04 100644 --- a/packages/core/src/context/graph/toGraph.ts +++ b/packages/core/src/context/graph/toGraph.ts @@ -4,11 +4,12 @@ * 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 { debugLogger } from '../../utils/debugLogger.js'; import type { NodeIdService } from './nodeIdService.js'; +import type { HistoryTurn } from '../../core/agentChatHistory.js'; // Global WeakMap to cache hashes for Part objects. // This optimizes getStableId by avoiding redundant stringify/hash operations @@ -82,7 +83,7 @@ function isCodeExecutionResultPart( } /** - * 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). */ @@ -154,7 +155,6 @@ export function getStableId( if (!id) { if (turnSalt && partIdx === -1) { - // Fallback for Turn objects (msg) since they don't have parts or content to hash directly here id = `turn_${turnSalt}`; } else { id = randomUUID(); @@ -172,14 +172,12 @@ export function getStableId( export class ContextGraphBuilder { 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(); - 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. @@ -197,15 +195,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.idService, turnSalt, -1); + const turnSalt = turn.id; + const turnId = `turn_${turnSalt}`; if (msg.role === 'user') { for (let partIdx = 0; partIdx < msg.parts.length; partIdx++) { @@ -213,13 +204,17 @@ export class ContextGraphBuilder { const apiId = isFunctionResponsePart(part) && typeof part.functionResponse.id === 'string' - ? `resp_${part.functionResponse.id}` + ? part.functionResponse.id : isFunctionCallPart(part) && typeof part.functionCall.id === 'string' - ? `call_${part.functionCall.id}` + ? part.functionCall.id : undefined; - const id = - apiId || getStableId(part, this.idService, turnSalt, partIdx); + + // 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(), @@ -231,16 +226,20 @@ export class ContextGraphBuilder { 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}` + ? part.functionCall.id : undefined; - const id = - apiId || getStableId(part, this.idService, turnSalt, partIdx); + + const id = apiId + ? `${apiId}_${turnSalt}_${partIdx}` + : `${turnSalt}_${partIdx}`; + const node: ConcreteNode = { id, timestamp: Date.now(), @@ -252,6 +251,7 @@ export class ContextGraphBuilder { turnId, }; nodes.push(node); + this.idService.set(part, id); } } } diff --git a/packages/core/src/context/processors/stateSnapshotProcessor.ts b/packages/core/src/context/processors/stateSnapshotProcessor.ts index 7473e0431e..52ab8f5338 100644 --- a/packages/core/src/context/processors/stateSnapshotProcessor.ts +++ b/packages/core/src/context/processors/stateSnapshotProcessor.ts @@ -3,7 +3,7 @@ * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ -import { randomUUID } from 'node:crypto'; +import { createHash } from 'node:crypto'; import type { JSONSchemaType } from 'ajv'; import type { ContextProcessor, @@ -50,6 +50,13 @@ export function createStateSnapshotProcessor( ): ContextProcessor { const generator = new SnapshotGenerator(env); + const generateStableId = (consumedIds: string[]) => { + return createHash('sha256') + .update(consumedIds.sort().join(',')) + .digest('hex') + .slice(0, 32); + }; + return { id, name: 'StateSnapshotProcessor', @@ -94,7 +101,7 @@ export function createStateSnapshotProcessor( `[StateSnapshotProcessor] Successfully spliced PROPOSED_SNAPSHOT from Inbox into Graph. Consumed ${consumedIds.length} nodes.`, ); // If valid, apply it! - const newId = randomUUID(); + const newId = generateStableId(consumedIds); const snapshotNode: Snapshot = { id: newId, @@ -186,11 +193,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 = generateStableId(consumedIds); const snapshotNode: Snapshot = { id: newId, diff --git a/packages/core/src/context/system-tests/__snapshots__/lifecycle.golden.test.ts.snap b/packages/core/src/context/system-tests/__snapshots__/lifecycle.golden.test.ts.snap index 3c36020dbc..e187c3d995 100644 --- a/packages/core/src/context/system-tests/__snapshots__/lifecycle.golden.test.ts.snap +++ b/packages/core/src/context/system-tests/__snapshots__/lifecycle.golden.test.ts.snap @@ -5,124 +5,163 @@ exports[`System Lifecycle Golden Tests > Scenario 1: Organic Growth with Huge To "baseUnits": 787, "finalProjection": [ { - "parts": [ - { - "text": "System Instructions", - }, - ], - "role": "user", - }, - { - "parts": [ - { - "text": "Ack.", - }, - ], - "role": "model", - }, - { - "parts": [ - { - "text": "Hello!", - }, - ], - "role": "user", - }, - { - "parts": [ - { - "text": "Hi, how can I help?", - }, - ], - "role": "model", - }, - { - "parts": [ - { - "text": "Read the logs.", - }, - ], - "role": "user", - }, - { - "parts": [ - { - "functionCall": { - "args": { - "cmd": "cat server.log", - }, - "name": "run_shell_command", + "content": { + "parts": [ + { + "text": "System Instructions", }, - "thoughtSignature": "skip_thought_signature_validator", - }, - ], - "role": "model", + ], + "role": "user", + }, + "id": "", }, { - "parts": [ - { - "functionResponse": { - "name": "run_shell_command", - "response": { - "output": " + "content": { + "parts": [ + { + "text": "Ack.", + }, + ], + "role": "model", + }, + "id": "", + }, + { + "content": { + "parts": [ + { + "text": "Hello!", + }, + ], + "role": "user", + }, + "id": "", + }, + { + "content": { + "parts": [ + { + "text": "Hi, how can I help?", + }, + ], + "role": "model", + }, + "id": "", + }, + { + "content": { + "parts": [ + { + "text": "Read the logs.", + }, + ], + "role": "user", + }, + "id": "", + }, + { + "content": { + "parts": [ + { + "functionCall": { + "args": { + "cmd": "cat server.log", + }, + "name": "run_shell_command", + }, + "thoughtSignature": "skip_thought_signature_validator", + }, + ], + "role": "model", + }, + "id": "", + }, + { + "content": { + "parts": [ + { + "functionResponse": { + "name": "run_shell_command", + "response": { + "output": " [Tool observation string (0.02MB, 1 lines) masked to preserve context window. Full string saved to: ] ", + }, }, }, - }, - ], - "role": "user", + ], + "role": "user", + }, + "id": "", }, { - "parts": [ - { - "text": "The logs are very long.", - }, - ], - "role": "model", + "content": { + "parts": [ + { + "text": "The logs are very long.", + }, + ], + "role": "model", + }, + "id": "", }, { - "parts": [ - { - "text": "Look at this architecture diagram:", - }, - { - "text": "[Multi-Modal Blob (image/png, 0.01MB) degraded to text to preserve context window. Saved to: ]", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "Look at this architecture diagram:", + }, + { + "text": "[Multi-Modal Blob (image/png, 0.01MB) degraded to text to preserve context window. Saved to: ]", + }, + ], + "role": "user", + }, + "id": "", }, { - "parts": [ - { - "text": "Nice diagram.", - }, - ], - "role": "model", + "content": { + "parts": [ + { + "text": "Nice diagram.", + }, + ], + "role": "model", + }, + "id": "", }, { - "parts": [ - { - "text": "Can we refactor?", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "Can we refactor?", + }, + ], + "role": "user", + }, + "id": "", }, { - "parts": [ - { - "text": "Yes we can.", - }, - ], - "role": "model", + "content": { + "parts": [ + { + "text": "Yes we can.", + }, + ], + "role": "model", + }, + "id": "", }, { - "parts": [ - { - "text": "Please continue.", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "Please continue.", + }, + ], + "role": "user", + }, + "id": "", }, ], "tokenTrajectory": [ @@ -160,44 +199,59 @@ exports[`System Lifecycle Golden Tests > Scenario 2: Under Budget (No Modificati "baseUnits": 68, "finalProjection": [ { - "parts": [ - { - "text": "System Instructions", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "System Instructions", + }, + ], + "role": "user", + }, + "id": "", }, { - "parts": [ - { - "text": "Ack.", - }, - ], - "role": "model", + "content": { + "parts": [ + { + "text": "Ack.", + }, + ], + "role": "model", + }, + "id": "", }, { - "parts": [ - { - "text": "Hello!", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "Hello!", + }, + ], + "role": "user", + }, + "id": "", }, { - "parts": [ - { - "text": "Hi, how can I help?", - }, - ], - "role": "model", + "content": { + "parts": [ + { + "text": "Hi, how can I help?", + }, + ], + "role": "model", + }, + "id": "", }, { - "parts": [ - { - "text": "Please continue.", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "Please continue.", + }, + ], + "role": "user", + }, + "id": "", }, ], "tokenTrajectory": [ @@ -220,60 +274,81 @@ exports[`System Lifecycle Golden Tests > Scenario 3: Node Distillation of Large "baseUnits": 5370, "finalProjection": [ { - "parts": [ - { - "text": "Mock response from: utility_compressor, for: {"text":"A...AAAAAAAA"}", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "Mock response from: utility_compressor, for: {"text":"A...AAAAAAAA"}", + }, + ], + "role": "user", + }, + "id": "", }, { - "parts": [ - { - "text": "Mock response from: utility_compressor, for: {"text":"B...BBBBBBBB"}", - }, - ], - "role": "model", + "content": { + "parts": [ + { + "text": "Mock response from: utility_compressor, for: {"text":"B...BBBBBBBB"}", + }, + ], + "role": "model", + }, + "id": "", }, { - "parts": [ - { - "text": "Mock response from: utility_compressor, for: {"text":"C...CCCCCCCC"}", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "Mock response from: utility_compressor, for: {"text":"C...CCCCCCCC"}", + }, + ], + "role": "user", + }, + "id": "", }, { - "parts": [ - { - "text": "Mock response from: utility_compressor, for: {"text":"D...DDDDDDDD"}", - }, - ], - "role": "model", + "content": { + "parts": [ + { + "text": "Mock response from: utility_compressor, for: {"text":"D...DDDDDDDD"}", + }, + ], + "role": "model", + }, + "id": "", }, { - "parts": [ - { - "text": "Mock response from: utility_compressor, for: {"text":"E...EEEEEEEE"}", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "Mock response from: utility_compressor, for: {"text":"E...EEEEEEEE"}", + }, + ], + "role": "user", + }, + "id": "", }, { - "parts": [ - { - "text": "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF", - }, - ], - "role": "model", + "content": { + "parts": [ + { + "text": "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF", + }, + ], + "role": "model", + }, + "id": "", }, { - "parts": [ - { - "text": "Please continue.", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "Please continue.", + }, + ], + "role": "user", + }, + "id": "", }, ], "tokenTrajectory": [ @@ -301,31 +376,40 @@ exports[`System Lifecycle Golden Tests > Scenario 4: Async-Driven Background GC "baseUnits": 505, "finalProjection": [ { - "parts": [ - { - "text": "{"active_tasks":[],"discovered_facts":[],"constraints_and_preferences":[],"recent_arc":[]}", - }, - { - "text": "Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 ..................................................", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "{"active_tasks":[],"discovered_facts":[],"constraints_and_preferences":[],"recent_arc":[]}", + }, + { + "text": "Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 Msg 8 ..................................................", + }, + ], + "role": "user", + }, + "id": "2371dc698715d731086209ad329ea7c9", }, { - "parts": [ - { - "text": "Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 ..................................................", - }, - ], - "role": "model", + "content": { + "parts": [ + { + "text": "Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 Msg 9 ..................................................", + }, + ], + "role": "model", + }, + "id": "", }, { - "parts": [ - { - "text": "Please continue.", - }, - ], - "role": "user", + "content": { + "parts": [ + { + "text": "Please continue.", + }, + ], + "role": "user", + }, + "id": "", }, ], "tokenTrajectory": [ diff --git a/packages/core/src/context/system-tests/hysteresis.test.ts b/packages/core/src/context/system-tests/hysteresis.test.ts index cf804a27f6..2858c27aba 100644 --- a/packages/core/src/context/system-tests/hysteresis.test.ts +++ b/packages/core/src/context/system-tests/hysteresis.test.ts @@ -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(['']); @@ -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,9 @@ 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 +58,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('')), + c.content.parts?.some((p) => p.text?.includes('')), ), ).toBe(false); @@ -79,7 +80,7 @@ describe('Context Manager Hysteresis Tests', () => { state = await harness.getGoldenState(); expect( state.finalProjection.some((c) => - c.parts?.some((p) => p.text?.includes('')), + c.content.parts?.some((p) => p.text?.includes('')), ), ).toBe(true); }); @@ -108,7 +109,7 @@ describe('Context Manager Hysteresis Tests', () => { let state = await harness.getGoldenState(); expect( state.finalProjection.some((c) => - c.parts?.some((p) => p.text?.includes('')), + c.content.parts?.some((p) => p.text?.includes('')), ), ).toBe(true); diff --git a/packages/core/src/context/system-tests/simulationHarness.ts b/packages/core/src/context/system-tests/simulationHarness.ts index c15c9fc26b..2955626ed3 100644 --- a/packages/core/src/context/system-tests/simulationHarness.ts +++ b/packages/core/src/context/system-tests/simulationHarness.ts @@ -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'; @@ -93,12 +94,13 @@ export class SimulationHarness { this.chatHistory, calculator, ); - } + } - async simulateTurn(messages: Content[]) { + 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( diff --git a/packages/core/src/context/utils/snapshotGenerator.ts b/packages/core/src/context/utils/snapshotGenerator.ts index b017ed3a0c..e715d7a336 100644 --- a/packages/core/src/context/utils/snapshotGenerator.ts +++ b/packages/core/src/context/utils/snapshotGenerator.ts @@ -10,9 +10,6 @@ import { LlmRole } from '../../telemetry/llmRole.js'; import { formatNodesForLlm } from './formatNodesForLlm.js'; import { randomUUID } from 'node:crypto'; import { isRecord } from '../../utils/markdownUtils.js'; -import type { LiveInbox } from '../pipeline/inbox.js'; -import type { ContextEngineState } from '../../services/chatRecordingTypes.js'; -import { debugLogger } from '../../utils/debugLogger.js'; function isStringArray(value: unknown): value is string[] { return ( @@ -83,80 +80,6 @@ export function findLatestSnapshotBaseline( return undefined; } -export const SnapshotStateHelper = { - exportState(nodes: readonly ConcreteNode[]): ContextEngineState { - const baseline = findLatestSnapshotBaseline(nodes); - if (!baseline) { - debugLogger.log( - '[SnapshotStateHelper] exportState: No snapshot baseline found in current nodes.', - ); - return {}; - } - - // Flatten abstractsIds to ensure only pristine/replayable IDs are persisted. - // This prevents deep nesting of synthetic snapshot IDs which cannot be reconstructed - // from saved chat messages during session resume. - const nodeMap = new Map(); - for (const n of nodes) nodeMap.set(n.id, n); - - const pristineIds = new Set(); - const toExpand = [...baseline.abstractsIds]; - const seen = new Set(); - - while (toExpand.length > 0) { - const id = toExpand.pop()!; - if (seen.has(id)) continue; - seen.add(id); - - const node = nodeMap.get(id); - if (node?.abstractsIds && node.abstractsIds.length > 0) { - toExpand.push(...node.abstractsIds); - } else { - pristineIds.add(id); - } - } - - debugLogger.log( - `[SnapshotStateHelper] exportState: Exporting snapshot ID ${baseline.id} representing ${pristineIds.size} pristine nodes.`, - ); - return { - snapshot: { - text: baseline.text, - consumedIds: Array.from(pristineIds), - timestamp: baseline.timestamp, - }, - }; - }, - - restoreState(state: ContextEngineState, inbox: LiveInbox): void { - if (!state.snapshot) { - debugLogger.log( - '[SnapshotStateHelper] restoreState: No snapshot found in provided ContextEngineState.', - ); - return; - } - - if ( - typeof state.snapshot.text === 'string' && - Array.isArray(state.snapshot.consumedIds) - ) { - debugLogger.log( - `[SnapshotStateHelper] restoreState: Publishing hydrated snapshot to LiveInbox with ${state.snapshot.consumedIds.length} consumed IDs.`, - ); - inbox.publish('PROPOSED_SNAPSHOT', { - newText: state.snapshot.text, - consumedIds: state.snapshot.consumedIds, - type: 'accumulate', - timestamp: state.snapshot.timestamp ?? Date.now(), - }); - } else { - debugLogger.log( - '[SnapshotStateHelper] restoreState: Invalid snapshot structural format.', - ); - } - }, -}; - export class SnapshotGenerator { constructor(private readonly env: ContextEnvironment) {} @@ -406,3 +329,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, + }, + }; + } +} diff --git a/packages/core/src/core/agentChatHistory.ts b/packages/core/src/core/agentChatHistory.ts index 7ef4b6a64d..0836bfc574 100644 --- a/packages/core/src/core/agentChatHistory.ts +++ b/packages/core/src/core/agentChatHistory.ts @@ -6,21 +6,35 @@ import type { Content } from '@google/genai'; +/** + * A durable wrapper for Gemini Content that carries a stable ID. + * This ID is preserved across all transformations and is used as the anchor + * for context graph node identity. + */ +export interface HistoryTurn { + readonly id: string; + readonly content: Content; +} + export type HistoryEventType = 'PUSH' | 'SYNC_FULL' | 'CLEAR' | 'SILENT_SYNC'; export interface HistoryEvent { type: HistoryEventType; - payload: readonly Content[]; + payload: readonly HistoryTurn[]; } export type HistoryListener = (event: HistoryEvent) => void; +/** + * The 'Strong Owner' of chat history turns. + * It ensures that every turn in the session is associated with a durable ID. + */ export class AgentChatHistory { - private history: Content[]; + private history: HistoryTurn[] = []; private listeners: Set = new Set(); - constructor(initialHistory: Content[] = []) { - this.history = [...initialHistory]; + constructor(initialTurns: HistoryTurn[] = []) { + this.history = [...initialTurns]; } subscribe(listener: HistoryListener): () => void { @@ -30,20 +44,27 @@ export class AgentChatHistory { return () => this.listeners.delete(listener); } - private notify(type: HistoryEventType, payload: readonly Content[]) { + private notify(type: HistoryEventType, payload: readonly HistoryTurn[]) { const event: HistoryEvent = { type, payload }; for (const listener of this.listeners) { listener(event); } } - push(content: Content) { - this.history.push(content); - this.notify('PUSH', [content]); + /** + * Adds a new turn to the history. + * Every turn must have a durable ID, usually provided by the ChatRecordingService. + */ + push(turn: HistoryTurn) { + this.history.push(turn); + this.notify('PUSH', [turn]); } - set(history: readonly Content[], options: { silent?: boolean } = {}) { - this.history = [...history]; + /** + * Overwrites the entire history with a new list of turns. + */ + set(turns: readonly HistoryTurn[], options: { silent?: boolean } = {}) { + this.history = [...turns]; this.notify(options.silent ? 'SILENT_SYNC' : 'SYNC_FULL', this.history); } @@ -52,20 +73,28 @@ export class AgentChatHistory { this.notify('CLEAR', []); } - get(): readonly Content[] { + get(): readonly HistoryTurn[] { return this.history; } - map(callback: (value: Content, index: number, array: Content[]) => Content) { - this.history = this.history.map(callback); - this.notify('SYNC_FULL', this.history); + /** + * Returns a copy of the raw Gemini Content[] for API consumption. + */ + getContents(): Content[] { + return this.history.map((h) => h.content); + } + + map( + callback: (value: HistoryTurn, index: number, array: HistoryTurn[]) => U, + ): U[] { + return this.history.map(callback); } flatMap( callback: ( - value: Content, + value: HistoryTurn, index: number, - array: Content[], + array: HistoryTurn[], ) => U | readonly U[], ): U[] { return this.history.flatMap(callback); @@ -75,3 +104,4 @@ export class AgentChatHistory { return this.history.length; } } + diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index da3335aa30..ce99fb37bc 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -46,6 +46,7 @@ import { LoopDetectionService } from '../services/loopDetectionService.js'; import { ChatCompressionService } from '../context/chatCompressionService.js'; import { AgentHistoryProvider } from '../context/agentHistoryProvider.js'; import type { ContextManager } from '../context/contextManager.js'; +import type { HistoryTurn } from './agentChatHistory.js'; import { ideContextStore } from '../ide/ideContext.js'; import { logNextSpeakerCheck } from '../telemetry/loggers.js'; import type { @@ -67,6 +68,7 @@ import { } from '../availability/policyHelpers.js'; import { getDisplayString, resolveModel } from '../config/models.js'; import { partToString } from '../utils/partUtils.js'; +import { randomUUID } from 'node:crypto'; import { coreEvents, CoreEvent, @@ -293,8 +295,14 @@ export class GeminiClient { this.getChat().stripThoughtsFromHistory(); } - setHistory(history: readonly Content[]) { - this.getChat().setHistory(history); + setHistory(history: readonly (Content | HistoryTurn)[]) { + const turns = history.map((item) => { + if ('id' in item && 'content' in item) { + return item as HistoryTurn; + } + return { id: randomUUID(), content: item as Content }; + }); + this.getChat().setHistory(turns); this.updateTelemetryTokenCount(); this.forceFullIdeContext = true; } @@ -414,14 +422,6 @@ export class GeminiClient { chat, this.lastPromptId, ); - if ( - this.contextManager && - resumedSessionData?.conversation.contextState - ) { - this.contextManager.restoreState( - resumedSessionData.conversation.contextState, - ); - } return chat; } catch (error) { await reportError( @@ -649,7 +649,11 @@ export class GeminiClient { if (this.config.getContextManagementConfig().enabled) { if (this.contextManager) { - const pendingRequest = createUserContent(request); + const rawPendingRequest = createUserContent(request); + const pendingRequest = { + id: randomUUID(), + content: rawPendingRequest, + }; const { history: newHistory, didApplyManagement, @@ -674,7 +678,11 @@ export class GeminiClient { signal, ); if (newHistory.length !== this.getHistory().length) { - this.getChat().setHistory(newHistory); + const turns = newHistory.map((c) => ({ + id: randomUUID(), + content: c, + })); + this.getChat().setHistory(turns); } } } else { @@ -827,9 +835,6 @@ export class GeminiClient { promptBaseUnits: currentBaseUnits, }); } - this.chat - ?.getChatRecordingService() - ?.saveContextState(this.contextManager.exportState()); } this.updateTelemetryTokenCount(); if (event.type === GeminiEventType.Error) { @@ -1237,7 +1242,11 @@ export class GeminiClient { if (newHistory) { // We truncated content to save space, but summarization is still "failed". // We update the chat context directly without resetting the failure flag. - this.getChat().setHistory(newHistory); + const turns = newHistory.map((c) => ({ + id: randomUUID(), + content: c, + })); + this.getChat().setHistory(turns); this.updateTelemetryTokenCount(); // We don't reset the chat session fully like in COMPRESSED because // this is a lighter-weight intervention. @@ -1256,7 +1265,11 @@ export class GeminiClient { this.config, ); if (result.maskedCount > 0) { - this.getChat().setHistory(result.newHistory); + const turns = result.newHistory.map((c) => ({ + id: randomUUID(), + content: c, + })); + this.getChat().setHistory(turns); } } diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 05a27f8bbc..0ecaae8491 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -19,6 +19,7 @@ import { SYNTHETIC_THOUGHT_SIGNATURE, type StreamEvent, stripToolCallIdPrefixes, + type HistoryTurn, } from './geminiChat.js'; import { type CompletedToolCall, @@ -234,9 +235,9 @@ describe('GeminiChat', () => { describe('constructor', () => { it('should initialize lastPromptTokenCount based on history size', () => { - const history: Content[] = [ - { role: 'user', parts: [{ text: 'Hello' }] }, - { role: 'model', parts: [{ text: 'Hi there' }] }, + const history: HistoryTurn[] = [ + { id: '1', content: { role: 'user', parts: [{ text: 'Hello' }] } }, + { id: '2', content: { role: 'model', parts: [{ text: 'Hi there' }] } }, ]; const chatWithHistory = new GeminiChat(mockConfig, '', [], history); // 'Hello': 5 chars * 0.25 = 1.25 @@ -253,8 +254,8 @@ describe('GeminiChat', () => { describe('setHistory', () => { it('should recalculate lastPromptTokenCount when history is updated', () => { - const initialHistory: Content[] = [ - { role: 'user', parts: [{ text: 'Hello' }] }, + const initialHistory: HistoryTurn[] = [ + { id: '1', content: { role: 'user', parts: [{ text: 'Hello' }] } }, ]; const chatWithHistory = new GeminiChat( mockConfig, @@ -264,14 +265,17 @@ describe('GeminiChat', () => { ); const initialCount = chatWithHistory.getLastPromptTokenCount(); - const newHistory: Content[] = [ + const newHistory: HistoryTurn[] = [ { - role: 'user', - parts: [ - { - text: 'This is a much longer history item that should result in more tokens than just hello.', - }, - ], + id: '2', + content: { + role: 'user', + parts: [ + { + text: 'This is a much longer history item that should result in more tokens than just hello.', + }, + ], + }, }, ]; chatWithHistory.setHistory(newHistory); @@ -331,9 +335,9 @@ describe('GeminiChat', () => { ).resolves.not.toThrow(); // 3. Verify history was recorded correctly - const history = chat.getHistory(); + const history = chat.getHistoryTurns(); expect(history.length).toBe(2); // user turn + model turn - const modelTurn = history[1]; + const modelTurn = history[1].content; expect(modelTurn?.parts?.length).toBe(1); // The empty part is discarded expect(modelTurn?.parts![0].functionCall).toBeDefined(); }); @@ -433,9 +437,9 @@ describe('GeminiChat', () => { ).resolves.not.toThrow(); // 3. Verify history was recorded correctly with only the valid part. - const history = chat.getHistory(); + const history = chat.getHistoryTurns(); expect(history.length).toBe(2); // user turn + model turn - const modelTurn = history[1]; + const modelTurn = history[1].content; expect(modelTurn?.parts?.length).toBe(1); expect(modelTurn?.parts![0].text).toBe('Initial valid content...'); }); @@ -478,9 +482,9 @@ describe('GeminiChat', () => { } // 3. Assert: Check that the final history was correctly consolidated. - const history = chat.getHistory(); + const history = chat.getHistoryTurns(); expect(history.length).toBe(2); - const modelTurn = history[1]; + const modelTurn = history[1].content; expect(modelTurn?.parts?.length).toBe(1); expect(modelTurn?.parts![0].text).toBe('Hello World!'); }); @@ -538,12 +542,12 @@ describe('GeminiChat', () => { } // 3. Assert: Check that the final history was correctly consolidated. - const history = chat.getHistory(); + const history = chat.getHistoryTurns(); // The history should contain the user's turn and ONE consolidated model turn. expect(history.length).toBe(2); - const modelTurn = history[1]; + const modelTurn = history[1].content; expect(modelTurn.role).toBe('model'); // The model turn should have 3 distinct parts: the merged text, the function call, and the final text. @@ -599,10 +603,10 @@ describe('GeminiChat', () => { } // 3. Assert: Check that the final history contains both function calls. - const history = chat.getHistory(); + const history = chat.getHistoryTurns(); expect(history.length).toBe(2); - const modelTurn = history[1]; + const modelTurn = history[1].content; expect(modelTurn.role).toBe('model'); expect(modelTurn.parts?.length).toBe(2); expect(modelTurn.parts![0].functionCall?.name).toBe('tool_A'); @@ -647,8 +651,8 @@ describe('GeminiChat', () => { // Consume the stream to trigger history recording } - const history = chat.getHistory(); - const modelTurn = history[1]; + const history = chat.getHistoryTurns(); + const modelTurn = history[1].content; expect(modelTurn.parts?.length).toBe(2); expect(modelTurn.parts![0].functionCall?.name).toBe('tool_X'); expect(modelTurn.parts![0].functionCall?.args).toEqual({ id: 1 }); @@ -694,12 +698,12 @@ describe('GeminiChat', () => { } // 3. Assert: Check the final state of the history. - const history = chat.getHistory(); + const history = chat.getHistoryTurns(); // The history should contain two turns: the user's message and the model's response. expect(history.length).toBe(2); - const modelTurn = history[1]; + const modelTurn = history[1].content; expect(modelTurn.role).toBe('model'); // CRUCIAL ASSERTION: @@ -713,21 +717,27 @@ describe('GeminiChat', () => { it('should throw an error when a tool call is followed by an empty stream response', async () => { // 1. Setup: A history where the model has just made a function call. - const initialHistory: Content[] = [ + const initialHistory: HistoryTurn[] = [ { - role: 'user', - parts: [{ text: 'Find a good Italian restaurant for me.' }], + id: '1', + content: { + role: 'user', + parts: [{ text: 'Find a good Italian restaurant for me.' }], + }, }, { - role: 'model', - parts: [ - { - functionCall: { - name: 'find_restaurant', - args: { cuisine: 'Italian' }, + id: '2', + content: { + role: 'model', + parts: [ + { + functionCall: { + name: 'find_restaurant', + args: { cuisine: 'Italian' }, + }, }, - }, - ], + ], + }, }, ]; chat.setHistory(initialHistory); @@ -1251,31 +1261,40 @@ describe('GeminiChat', () => { describe('addHistory', () => { it('should add a new content item to the history', () => { - const newContent: Content = { - role: 'user', - parts: [{ text: 'A new message' }], + const newTurn: HistoryTurn = { + id: '1', + content: { + role: 'user', + parts: [{ text: 'A new message' }], + }, }; - chat.addHistory(newContent); - const history = chat.getHistory(); + chat.addHistory(newTurn); + const history = chat.getHistoryTurns(); expect(history.length).toBe(1); - expect(history[0]).toEqual(newContent); + expect(history[0]).toEqual(newTurn); }); it('should add multiple items correctly', () => { - const content1: Content = { - role: 'user', - parts: [{ text: 'Message 1' }], + const turn1: HistoryTurn = { + id: '1', + content: { + role: 'user', + parts: [{ text: 'Message 1' }], + }, }; - const content2: Content = { - role: 'model', - parts: [{ text: 'Message 2' }], + const turn2: HistoryTurn = { + id: '2', + content: { + role: 'model', + parts: [{ text: 'Message 2' }], + }, }; - chat.addHistory(content1); - chat.addHistory(content2); - const history = chat.getHistory(); + chat.addHistory(turn1); + chat.addHistory(turn2); + const history = chat.getHistoryTurns(); expect(history.length).toBe(2); - expect(history[0]).toEqual(content1); - expect(history[1]).toEqual(content2); + expect(history[0]).toEqual(turn1); + expect(history[1]).toEqual(turn2); }); }); diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 6a728884a5..5f6b8f6a69 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -19,7 +19,10 @@ import { type GenerateContentParameters, type FunctionCall, } from '@google/genai'; -import { AgentChatHistory } from './agentChatHistory.js'; +export { AgentChatHistory, type HistoryTurn } from './agentChatHistory.js'; +import { AgentChatHistory, type HistoryTurn } from './agentChatHistory.js'; + +import { randomUUID } from 'node:crypto'; import { toParts } from '../code_assist/converter.js'; import { retryWithBackoff, @@ -159,8 +162,9 @@ function isValidContent(content: Content): boolean { * @throws Error if the history does not start with a user turn. * @throws Error if the history contains an invalid role. */ -function validateHistory(history: Content[]) { - for (const content of history) { +function validateHistory(history: (Content | HistoryTurn)[]) { + for (const item of history) { + const content = 'content' in item ? item.content : item; if (content.role !== 'user' && content.role !== 'model') { throw new Error(`Role must be user or model, but got ${content.role}.`); } @@ -175,23 +179,31 @@ function validateHistory(history: Content[]) { * filters or recitation). Extracting valid turns from the history * ensures that subsequent requests could be accepted by the model. */ -function extractCuratedHistory(comprehensiveHistory: Content[]): Content[] { - if (comprehensiveHistory === undefined || comprehensiveHistory.length === 0) { +function extractCuratedHistory( + comprehensiveHistory: readonly HistoryTurn[], +): HistoryTurn[] { + if ( + comprehensiveHistory === undefined || + comprehensiveHistory.length === 0 + ) { return []; } - const curatedHistory: Content[] = []; + const curatedHistory: HistoryTurn[] = []; const length = comprehensiveHistory.length; let i = 0; while (i < length) { - if (comprehensiveHistory[i].role === 'user') { + if (comprehensiveHistory[i].content.role === 'user') { curatedHistory.push(comprehensiveHistory[i]); i++; } else { - const modelOutput: Content[] = []; + const modelOutput: HistoryTurn[] = []; let isValid = true; - while (i < length && comprehensiveHistory[i].role === 'model') { + while ( + i < length && + comprehensiveHistory[i].content.role === 'model' + ) { modelOutput.push(comprehensiveHistory[i]); - if (isValid && !isValidContent(comprehensiveHistory[i])) { + if (isValid && !isValidContent(comprehensiveHistory[i].content)) { isValid = false; } i++; @@ -272,15 +284,33 @@ export class GeminiChat { readonly context: AgentLoopContext, private systemInstruction: string = '', private tools: Tool[] = [], - history: Content[] = [], + history: (Content | HistoryTurn)[] = [], resumedSessionData?: ResumedSessionData, private readonly onModelChanged?: (modelId: string) => Promise, ) { validateHistory(history); - this.agentHistory = new AgentChatHistory(history); + const initialHistory: HistoryTurn[] = resumedSessionData + ? resumedSessionData.conversation.messages + .filter((m) => m.type === 'user' || m.type === 'gemini') + .map((m) => ({ + id: m.id, + content: { + role: m.type === 'user' ? 'user' : 'model', + parts: Array.isArray(m.content) + ? (m.content as Part[]) + : [{ text: m.content as string }], + }, + })) + : history.map((item) => + 'id' in item && 'content' in item + ? item + : { id: randomUUID(), content: item }, + ); + + this.agentHistory = new AgentChatHistory(initialHistory); this.chatRecordingService = new ChatRecordingService(context); this.lastPromptTokenCount = estimateTokenCountSync( - this.agentHistory.flatMap((c) => c.parts || []), + this.agentHistory.flatMap((c) => c.content.parts || []), ); } @@ -362,41 +392,67 @@ export class GeminiChat { } } - this.chatRecordingService.recordMessage({ + const id = this.chatRecordingService.recordMessage({ model, type: 'user', content: userMessageParts, displayContent: finalDisplayContent, }); - } + this.agentHistory.push({ id, content: userContent }); + } else { + // Record tool response as a message to ensure durable ID and linear history for resume. + const id = this.chatRecordingService.recordSyntheticMessage( + 'user', + userContent.parts || [], + ); - // Add user content to history ONCE before any attempts. - const binaryInjections = this.extractBinaryInjections(userContent.parts); - if (binaryInjections) { - // Turn 1: The original tool response (now cleaned) - this.agentHistory.push(userContent); + // Binary injections: If the tool output contains binary data, we expand the history. + const binaryParts = this.extractBinaryInjections(userContent.parts); + if (binaryParts) { + // Turn 1: The original tool response (now cleaned) + this.agentHistory.push({ id, content: userContent }); - // Turn 2: Synthetic Model Acknowledgment - this.agentHistory.push({ - role: 'model', - parts: [ - { - text: 'Binary content received. Proceeding with analysis.', - thought: true, - thoughtSignature: SYNTHETIC_THOUGHT_SIGNATURE, + // Turn 2: Synthetic Model Acknowledgment + const modelId = this.chatRecordingService.recordSyntheticMessage( + 'gemini', + [ + { + text: 'Binary content received. Proceeding with analysis.', + thought: true, + thoughtSignature: SYNTHETIC_THOUGHT_SIGNATURE, + }, + ], + ); + this.agentHistory.push({ + id: modelId, + content: { + role: 'model', + parts: [ + { + text: 'Binary content received. Proceeding with analysis.', + thought: true, + thoughtSignature: SYNTHETIC_THOUGHT_SIGNATURE, + }, + ], }, - ], - }); + }); - // Turn 3: The actual binary data (becomes the current request message) - userContent = { - role: 'user', - parts: binaryInjections, - }; + // Turn 3: The actual binary data (becomes the current request message) + const binaryId = this.chatRecordingService.recordSyntheticMessage( + 'info', + binaryParts, + ); + userContent = { + role: 'user', + parts: binaryParts, + }; + this.agentHistory.push({ id: binaryId, content: userContent }); + } else { + this.agentHistory.push({ id, content: userContent }); + } } - this.agentHistory.push(userContent); - const requestContents = this.getHistory(true); + const requestHistory = this.getHistoryTurns(true); const streamWithRetries = async function* ( this: GeminiChat, @@ -420,7 +476,7 @@ export class GeminiChat { isConnectionPhase = true; const stream = await this.makeApiCallAndProcessStream( currentConfigKey, - requestContents, + requestHistory, prompt_id, signal, role, @@ -539,48 +595,45 @@ export class GeminiChat { return streamWithRetries.call(this); } - private extractBinaryInjections( - parts: Part[] | undefined, - ): Part[] | undefined { - if (!parts) { - return undefined; - } - - const binaryInjections: Part[] = []; - - for (const part of parts) { - const response = part.functionResponse?.response; - - if (response && BINARY_INJECTION_KEY in response) { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const binaryParts = response[BINARY_INJECTION_KEY] as Part[]; - delete response[BINARY_INJECTION_KEY]; - - if (Array.isArray(binaryParts)) { - binaryInjections.push(...binaryParts); + private extractBinaryInjections(parts: Part[] | undefined): Part[] | undefined { + const binaryParts: Part[] = []; + if (parts) { + for (const part of parts) { + const response = part.functionResponse?.response; + if (response && BINARY_INJECTION_KEY in response) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + const injected = response[BINARY_INJECTION_KEY] as Part[]; + delete response[BINARY_INJECTION_KEY]; + if (Array.isArray(injected)) { + binaryParts.push(...injected); + } } } } - return binaryInjections.length > 0 ? binaryInjections : undefined; + return binaryParts.length > 0 ? binaryParts : undefined; } private async makeApiCallAndProcessStream( modelConfigKey: ModelConfigKey, - requestContents: readonly Content[], + requestHistory: readonly HistoryTurn[], prompt_id: string, abortSignal: AbortSignal, role: LlmRole, ): Promise> { // Last mile scrubbing to remove internal tracking properties (e.g. callIndex) // before sending to the Gemini API. This whitelists only standard Gemini fields. - const scrubbedContents = this.context.config.isContextManagementEnabled() - ? scrubHistory([...requestContents]) - : [...requestContents]; + const scrubbedHistory = this.context.config.isContextManagementEnabled() + ? scrubHistory([...requestHistory]) + : [...requestHistory]; + + const scrubbedContents = scrubbedHistory.map((h) => h.content); const contentsForPreviewModel = this.ensureActiveLoopHasThoughtSignatures(scrubbedContents); + const requestContents = scrubbedContents; + // Track final request parameters for AfterModel hooks const { model: availabilityFinalModel, @@ -829,14 +882,21 @@ export class GeminiChat { * @return History contents alternating between user and model for the entire * chat session. */ - getHistory(curated: boolean = false): readonly Content[] { + getHistory(curated: boolean = false): Content[] { + return this.getHistoryTurns(curated).map((h) => h.content); + } + + /** + * Returns the chat history as HistoryTurns. + */ + getHistoryTurns(curated: boolean = false): HistoryTurn[] { const history = curated - ? extractCuratedHistory([...this.agentHistory.get()]) - : this.agentHistory.get(); + ? extractCuratedHistory(this.agentHistory.get()) + : [...this.agentHistory.get()]; return this.context.config.isContextManagementEnabled() - ? scrubHistory([...history]) - : [...history]; + ? scrubHistory(history) + : history; } /** @@ -849,24 +909,33 @@ export class GeminiChat { /** * Adds a new entry to the chat history. */ - addHistory(content: Content): void { - this.agentHistory.push(content); + addHistory(content: Content | HistoryTurn): void { + if ('id' in content && 'content' in content) { + this.agentHistory.push(content); + } else { + this.agentHistory.push({ id: randomUUID(), content }); + } } setHistory( - history: readonly Content[], + history: readonly (Content | HistoryTurn)[], options: { silent?: boolean } = {}, ): void { - this.agentHistory.set(history, options); - this.lastPromptTokenCount = estimateTokenCountSync( - this.agentHistory.flatMap((c) => c.parts || []), + const wrappedHistory: HistoryTurn[] = history.map((item) => + 'id' in item && 'content' in item + ? item + : { id: randomUUID(), content: item }, ); - this.chatRecordingService.updateMessagesFromHistory(history); + this.agentHistory.set(wrappedHistory, options); + this.lastPromptTokenCount = estimateTokenCountSync( + this.agentHistory.flatMap((c) => c.content.parts || []), + ); + this.chatRecordingService.updateMessagesFromHistory(this.agentHistory.get()); } stripThoughtsFromHistory(): void { - this.agentHistory.map((content) => { - const newContent = { ...content }; + const newHistory = this.agentHistory.map((turn) => { + const newContent = { ...turn.content }; if (newContent.parts) { newContent.parts = newContent.parts.map((part) => { if (part && typeof part === 'object' && 'thoughtSignature' in part) { @@ -877,8 +946,9 @@ export class GeminiChat { return part; }); } - return newContent; + return { id: turn.id, content: newContent }; }); + this.agentHistory.set(newHistory); } // To ensure our requests validate, the first function call in every model @@ -1162,15 +1232,22 @@ export class GeminiChat { .join('') .trim(); + let id: string; // Record model response text from the collected parts. // Also flush when there are thoughts or a tool call (even with no text) // so that BeforeTool hooks always see the latest transcript state. if (responseText || hasThoughts || hasToolCall) { - this.chatRecordingService.recordMessage({ + id = this.chatRecordingService.recordMessage({ model, type: 'gemini', content: responseText, }); + } else { + // Still need a durable ID even if response is empty (e.g. only tool calls) + id = this.chatRecordingService.recordSyntheticMessage( + 'gemini', + consolidatedParts, + ); } // Stream validation logic: A stream is considered successful if: @@ -1208,7 +1285,10 @@ export class GeminiChat { } } - this.agentHistory.push({ role: 'model', parts: consolidatedParts }); + this.agentHistory.push({ + id, + content: { role: 'model', parts: consolidatedParts }, + }); } getLastPromptTokenCount(): number { diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index 7af8380a5a..879570766e 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -47,9 +47,10 @@ import { } from './chatRecordingService.js'; import type { WorkspaceContext } from '../utils/workspaceContext.js'; import { CoreToolCallStatus } from '../scheduler/types.js'; -import type { Content, Part } from '@google/genai'; +import type { Part } from '@google/genai'; import type { Config } from '../config/config.js'; import { getProjectHash } from '../utils/paths.js'; +import type { HistoryTurn } from '../core/agentChatHistory.js'; vi.mock('../utils/paths.js'); vi.mock('node:crypto', async (importOriginal) => { @@ -1065,7 +1066,7 @@ describe('ChatRecordingService', () => { it('should update tool results from API history (masking sync)', async () => { // 1. Record an initial message and tool call - chatRecordingService.recordMessage({ + const modelMsgId = chatRecordingService.recordMessage({ type: 'gemini', content: 'I will list the files.', model: 'gemini-pro', @@ -1087,24 +1088,30 @@ describe('ChatRecordingService', () => { // 2. Prepare mock history with masked content const maskedSnippet = 'short preview'; - const history: Content[] = [ + const history: HistoryTurn[] = [ { - role: 'model', - parts: [ - { functionCall: { name: 'list_files', args: { path: '.' } } }, - ], + id: modelMsgId, + content: { + role: 'model', + parts: [ + { functionCall: { name: 'list_files', args: { path: '.' } } }, + ], + }, }, { - role: 'user', - parts: [ - { - functionResponse: { - name: 'list_files', - id: callId, - response: { output: maskedSnippet }, + id: 'user-id', + content: { + role: 'user', + parts: [ + { + functionResponse: { + name: 'list_files', + id: callId, + response: { output: maskedSnippet }, + }, }, - }, - ], + ], + }, }, ]; @@ -1132,8 +1139,15 @@ describe('ChatRecordingService', () => { output: maskedSnippet, }); }); + it('should preserve multi-modal sibling parts during sync', async () => { await chatRecordingService.initialize(); + const modelMsgId = chatRecordingService.recordMessage({ + type: 'gemini', + content: '', + model: 'gemini-pro', + }); + const callId = 'multi-modal-call'; const originalResult: Part[] = [ { @@ -1146,12 +1160,6 @@ describe('ChatRecordingService', () => { { inlineData: { mimeType: 'image/png', data: 'base64...' } }, ]; - chatRecordingService.recordMessage({ - type: 'gemini', - content: '', - model: 'gemini-pro', - }); - chatRecordingService.recordToolCalls('gemini-pro', [ { id: callId, @@ -1164,19 +1172,26 @@ describe('ChatRecordingService', () => { ]); const maskedSnippet = ''; - const history: Content[] = [ + const history: HistoryTurn[] = [ { - role: 'user', - parts: [ - { - functionResponse: { - name: 'read_file', - id: callId, - response: { output: maskedSnippet }, + id: modelMsgId, + content: { role: 'model', parts: [] }, + }, + { + id: 'user-id', + content: { + role: 'user', + parts: [ + { + functionResponse: { + name: 'read_file', + id: callId, + response: { output: maskedSnippet }, + }, }, - }, - { inlineData: { mimeType: 'image/png', data: 'base64...' } }, - ], + { inlineData: { mimeType: 'image/png', data: 'base64...' } }, + ], + }, }, ]; @@ -1201,14 +1216,14 @@ describe('ChatRecordingService', () => { it('should handle parts appearing BEFORE the functionResponse in a content block', async () => { await chatRecordingService.initialize(); - const callId = 'prefix-part-call'; - - chatRecordingService.recordMessage({ + const modelMsgId = chatRecordingService.recordMessage({ type: 'gemini', content: '', model: 'gemini-pro', }); + const callId = 'prefix-part-call'; + chatRecordingService.recordToolCalls('gemini-pro', [ { id: callId, @@ -1220,19 +1235,26 @@ describe('ChatRecordingService', () => { }, ]); - const history: Content[] = [ + const history: HistoryTurn[] = [ { - role: 'user', - parts: [ - { text: 'Prefix metadata or text' }, - { - functionResponse: { - name: 'read_file', - id: callId, - response: { output: 'file content' }, + id: modelMsgId, + content: { role: 'model', parts: [] }, + }, + { + id: 'user-id', + content: { + role: 'user', + parts: [ + { text: 'Prefix metadata or text' }, + { + functionResponse: { + name: 'read_file', + id: callId, + response: { output: 'file content' }, + }, }, - }, - ], + ], + }, }, ]; @@ -1263,25 +1285,30 @@ describe('ChatRecordingService', () => { appendFileSyncSpy.mockClear(); // History with a tool call ID that doesn't exist in the conversation - const history: Content[] = [ + const history: HistoryTurn[] = [ { - role: 'user', - parts: [ - { - functionResponse: { - name: 'read_file', - id: 'nonexistent-call-id', - response: { output: 'some content' }, + id: 'user-id', + content: { + role: 'user', + parts: [ + { + functionResponse: { + name: 'read_file', + id: 'nonexistent-call-id', + response: { output: 'some content' }, + }, }, - }, - ], + ], + }, }, ]; chatRecordingService.updateMessagesFromHistory(history); - // No tool calls matched, so writeFileSync should NOT have been called - expect(appendFileSyncSpy).not.toHaveBeenCalled(); + // In the new 'Strong Owner' architecture, updateMessagesFromHistory ensures that + // all turns in history (including new/synthetic ones) are recorded. + // Since 'user-id' was not in the original conversation, it is added. + expect(appendFileSyncSpy).toHaveBeenCalled(); }); }); diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 3bd4e618cd..bd1b0d4fd4 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -17,13 +17,12 @@ import { import readline from 'node:readline'; import { randomUUID } from 'node:crypto'; import type { - Content, - Part, PartListUnion, GenerateContentResponseUsageMetadata, } from '@google/genai'; import { debugLogger } from '../utils/debugLogger.js'; import type { AgentLoopContext } from '../config/agent-loop-context.js'; +import type { HistoryTurn } from '../core/agentChatHistory.js'; import { SESSION_FILE_PREFIX, type TokensSummary, @@ -36,7 +35,6 @@ import { type RewindRecord, type MetadataUpdateRecord, type PartialMetadataRecord, - type ContextEngineState, } from './chatRecordingTypes.js'; export * from './chatRecordingTypes.js'; @@ -276,7 +274,6 @@ export async function loadConversationRecord( memoryScratchpad: metadata.memoryScratchpad, directories: metadata.directories, kind: metadata.kind, - contextState: metadata.contextState, messages: options?.metadataOnly ? [] : loadedMessages, messageCount: options?.metadataOnly ? metadataMessages.length || messageIds.length @@ -499,9 +496,10 @@ export class ChatRecordingService { type: ConversationRecordExtra['type'], content: PartListUnion, displayContent?: PartListUnion, + id?: string, ): MessageRecord { return { - id: randomUUID(), + id: id || randomUUID(), timestamp: new Date().toISOString(), type, content, @@ -514,14 +512,16 @@ export class ChatRecordingService { type: ConversationRecordExtra['type']; content: PartListUnion; displayContent?: PartListUnion; - }): void { - if (!this.conversationFile || !this.cachedConversation) return; + id?: string; + }): string { + if (!this.conversationFile || !this.cachedConversation) return message.id || randomUUID(); try { const msg = this.newMessage( message.type, message.content, message.displayContent, + message.id, ); if (msg.type === 'gemini') { msg.thoughts = this.queuedThoughts; @@ -532,12 +532,30 @@ export class ChatRecordingService { } this.pushMessage(msg); this.updateMetadata({ lastUpdated: new Date().toISOString() }); + return msg.id; } catch (error) { debugLogger.error('Error saving message to chat history.', error); throw error; } } + /** + * Records a synthetic message (e.g. Binary Received, Snapshot/Summary) + * and returns its durable ID. + */ + recordSyntheticMessage( + type: ConversationRecordExtra['type'], + content: PartListUnion, + id?: string, + ): string { + return this.recordMessage({ + model: undefined, + type, + content, + id, + }); + } + recordThought(thought: ThoughtSummary): void { if (!this.conversationFile) return; this.queuedThoughts.push({ @@ -648,15 +666,6 @@ export class ChatRecordingService { } } - saveContextState(contextState: ContextEngineState): void { - if (!this.conversationFile) return; - try { - this.updateMetadata({ contextState } as Partial); - } catch (e: unknown) { - debugLogger.error('Error saving context state to chat history.', e); - } - } - saveSummary(summary: string): void { if (!this.conversationFile) return; try { @@ -880,48 +889,77 @@ export class ChatRecordingService { return this.cachedConversation; } - updateMessagesFromHistory(history: readonly Content[]): void { + updateMessagesFromHistory(history: readonly HistoryTurn[]): void { if (!this.conversationFile || !this.cachedConversation) return; try { - const partsMap = new Map(); - for (const content of history) { - if (content.role === 'user' && content.parts) { - const callIds = content.parts - .map((p) => p.functionResponse?.id) - .filter((id): id is string => !!id); + let updated = false; - if (callIds.length === 0) continue; + // 1. Sync content and IDs + const newMessages: MessageRecord[] = history.map((turn) => { + const existing = this.cachedConversation?.messages.find( + (m) => m.id === turn.id, + ); - let currentCallId = callIds[0]; - for (const part of content.parts) { - if (part.functionResponse?.id) { - currentCallId = part.functionResponse.id; + if (existing) { + // If content parts have changed (e.g. masking), update them + if ( + JSON.stringify(existing.content) !== + JSON.stringify(turn.content.parts) + ) { + updated = true; + } + return { + ...existing, + content: turn.content.parts || [], + }; + } + + // It's a new (possibly synthetic) turn like a summary + updated = true; + return this.newMessage( + turn.content.role === 'user' ? 'user' : 'gemini', + turn.content.parts || [], + undefined, + turn.id, + ); + }); + + // 2. Specialized 'Masking Sync' for tool call results + // If a user turn in history contains a functionResponse, we update the + // corresponding ToolCallRecord in the preceding gemini message. + for (const turn of history) { + if (turn.content.role !== 'user') continue; + for (const part of turn.content.parts || []) { + if (part.functionResponse) { + const callId = part.functionResponse.id; + // Find the gemini message that contains this tool call + const geminiMsg = newMessages.find( + (m) => + m.type === 'gemini' && + m.toolCalls?.some((tc) => tc.id === callId), + ) as MessageRecord & { type: 'gemini' }; + if (geminiMsg) { + const tc = geminiMsg.toolCalls!.find((tc) => tc.id === callId); + if (tc) { + // If the history version is different (e.g. masked), sync it into the record + // We sync the entire parts array of the user turn to ensure sibling parts are preserved + if (JSON.stringify(tc.result) !== JSON.stringify(turn.content.parts)) { + tc.result = turn.content.parts || []; + updated = true; + } + } } - - if (!partsMap.has(currentCallId)) { - partsMap.set(currentCallId, []); - } - partsMap.get(currentCallId)!.push(part); } } } - for (const message of this.cachedConversation.messages) { - let msgChanged = false; - if (message.type === 'gemini' && message.toolCalls) { - for (const toolCall of message.toolCalls) { - const newParts = partsMap.get(toolCall.id); - if (newParts !== undefined) { - toolCall.result = newParts; - msgChanged = true; - } - } - } - if (msgChanged) { - // Push updated message to log - this.pushMessage(message); - } + if (updated || newMessages.length !== this.cachedConversation.messages.length) { + this.cachedConversation.messages = newMessages; + this.updateMetadata({ + messages: newMessages, + lastUpdated: new Date().toISOString(), + }); } } catch (error) { debugLogger.error( diff --git a/packages/core/src/services/chatRecordingTypes.ts b/packages/core/src/services/chatRecordingTypes.ts index 926bd5b8a5..44c518a7e8 100644 --- a/packages/core/src/services/chatRecordingTypes.ts +++ b/packages/core/src/services/chatRecordingTypes.ts @@ -86,17 +86,6 @@ export type ConversationRecordExtra = */ export type MessageRecord = BaseMessageRecord & ConversationRecordExtra; -/** - * Complete conversation record stored in session files. - */ -export interface ContextEngineState { - snapshot?: { - text: string; - consumedIds: string[]; - timestamp?: number; - }; -} - export interface ConversationRecord { sessionId: string; projectHash: string; @@ -109,8 +98,6 @@ export interface ConversationRecord { directories?: string[]; /** The kind of conversation (main agent or subagent) */ kind?: 'main' | 'subagent'; - /** Opaque state object representing Context Engine state (e.g. snapshots) */ - contextState?: ContextEngineState; } /** * Data structure for resuming an existing session. @@ -146,5 +133,4 @@ export interface PartialMetadataRecord { memoryScratchpad?: MemoryScratchpad; directories?: string[]; kind?: 'main' | 'subagent'; - contextState?: ContextEngineState; } diff --git a/packages/core/src/utils/historyHardening.ts b/packages/core/src/utils/historyHardening.ts index 98122abf65..b635cbcb8f 100644 --- a/packages/core/src/utils/historyHardening.ts +++ b/packages/core/src/utils/historyHardening.ts @@ -4,9 +4,10 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { Content, Part } from '@google/genai'; +import { type Part } from '@google/genai'; import { debugLogger } from './debugLogger.js'; -import type { NodeIdService } from '../context/graph/nodeIdService.js'; +import { type HistoryTurn } from '../core/agentChatHistory.js'; +import { randomUUID } from 'node:crypto'; export const SYNTHETIC_THOUGHT_SIGNATURE = 'skip_thought_signature_validator'; @@ -36,10 +37,9 @@ const DEFAULT_SENTINELS = { * 5. Signatures: The first functionCall in a model turn must have a thoughtSignature. */ export function hardenHistory( - history: Content[], + history: HistoryTurn[], options: HardeningOptions = {}, - idService?: NodeIdService, -): Content[] { +): HistoryTurn[] { if (history.length === 0) return history; const sentinels = { ...DEFAULT_SENTINELS, ...options.sentinels }; @@ -57,7 +57,7 @@ export function hardenHistory( let final = enforceRoleConstraints(coalesced, sentinels); // Pass 5: Final Scrubbing (Remove custom/non-standard properties for API compatibility) - final = scrubHistory(final, idService); + final = scrubHistory(final); return final; } @@ -65,17 +65,20 @@ export function hardenHistory( /** * Combines adjacent turns with the same role and removes empty turns. */ -function coalesce(history: Content[]): Content[] { - const result: Content[] = []; +function coalesce(history: HistoryTurn[]): HistoryTurn[] { + const result: HistoryTurn[] = []; for (const turn of history) { - if (!turn.parts || turn.parts.length === 0) continue; + if (!turn.content.parts || turn.content.parts.length === 0) continue; const last = result[result.length - 1]; - if (last && last.role === turn.role) { - last.parts = [...(last.parts || []), ...(turn.parts || [])]; + if (last && last.content.role === turn.content.role) { + last.content.parts = [ + ...(last.content.parts || []), + ...(turn.content.parts || []), + ]; } else { - // Shallow clone the turn so we don't mutate the original history array structure - result.push({ ...turn }); + // Shallow clone the turn and content so we don't mutate the original history array structure + result.push({ id: turn.id, content: { ...turn.content } }); } } return result; @@ -85,10 +88,10 @@ function coalesce(history: Content[]): Content[] { * Ensures tool calls have matching responses and model turns have required signatures. */ function pairToolsAndEnforceSignatures( - history: Content[], + history: HistoryTurn[], sentinels: Required>, -): Content[] { - const result: Content[] = []; +): HistoryTurn[] { + const result: HistoryTurn[] = []; // We work on a copy to allow splicing in sentinel turns const work = [...history]; @@ -96,8 +99,8 @@ function pairToolsAndEnforceSignatures( for (let i = 0; i < work.length; i++) { const turn = work[i]; - if (turn.role === 'model') { - const parts = turn.parts || []; + if (turn.content.role === 'model') { + const parts = turn.content.parts || []; // A. Signatures let foundCall = false; @@ -125,8 +128,8 @@ function pairToolsAndEnforceSignatures( const name = call.functionCall!.name || 'unknown'; const hasResponse = - nextTurn?.role === 'user' && - nextTurn.parts?.some( + nextTurn?.content.role === 'user' && + nextTurn.content.parts?.some( (p) => p.functionResponse?.id === id && p.functionResponse?.name === name, @@ -145,17 +148,20 @@ function pairToolsAndEnforceSignatures( `[HistoryHardener] Detected ${missing.length} tool calls without responses. Injecting sentinel responses.`, ); - let targetUserTurn: Content; - if (nextTurn?.role === 'user') { + let targetUserTurn: HistoryTurn; + if (nextTurn?.content.role === 'user') { targetUserTurn = nextTurn; } else { - targetUserTurn = { role: 'user', parts: [] }; + targetUserTurn = { + id: randomUUID(), + content: { role: 'user', parts: [] }, + }; work.splice(i + 1, 0, targetUserTurn); } for (const m of missing) { - targetUserTurn.parts = targetUserTurn.parts || []; - targetUserTurn.parts.push({ + targetUserTurn.content.parts = targetUserTurn.content.parts || []; + targetUserTurn.content.parts.push({ functionResponse: { name: m.name, id: m.id, @@ -167,11 +173,11 @@ function pairToolsAndEnforceSignatures( } } } - } else if (turn.role === 'user') { + } else if (turn.content.role === 'user') { // C. Orphaned Responses // A user response MUST follow a model call. const prevTurn = result[result.length - 1]; - const parts = turn.parts || []; + const parts = turn.content.parts || []; const validParts: Part[] = []; for (const p of parts) { @@ -179,8 +185,8 @@ function pairToolsAndEnforceSignatures( const id = p.functionResponse.id; const name = p.functionResponse.name; const hasCall = - prevTurn?.role === 'model' && - prevTurn.parts?.some( + prevTurn?.content.role === 'model' && + prevTurn.content.parts?.some( (cp) => cp.functionCall?.id === id && cp.functionCall?.name === name, ); @@ -196,10 +202,10 @@ function pairToolsAndEnforceSignatures( validParts.push(p); } } - turn.parts = validParts; + turn.content.parts = validParts; } - if (turn.parts && turn.parts.length > 0) { + if (turn.content.parts && turn.content.parts.length > 0) { result.push(turn); } } @@ -210,21 +216,22 @@ function pairToolsAndEnforceSignatures( /** * Hoists and re-orders tool responses within user turns to match preceding model turns. */ -function refineToolResponses(history: Content[]): Content[] { +function refineToolResponses(history: HistoryTurn[]): HistoryTurn[] { for (let i = 1; i < history.length; i++) { const turn = history[i]; const prev = history[i - 1]; - if (turn.role === 'user' && prev.role === 'model') { + if (turn.content.role === 'user' && prev.content.role === 'model') { const callOrder = - prev.parts + prev.content.parts ?.filter((p) => !!p.functionCall) .map((p) => p.functionCall!.id) || []; if (callOrder.length > 0) { const responseParts = - turn.parts?.filter((p) => !!p.functionResponse) || []; - const otherParts = turn.parts?.filter((p) => !p.functionResponse) || []; + turn.content.parts?.filter((p) => !!p.functionResponse) || []; + const otherParts = + turn.content.parts?.filter((p) => !p.functionResponse) || []; if (responseParts.length > 0) { // 1. Re-order: Sort responses to match the model's call order @@ -242,7 +249,7 @@ function refineToolResponses(history: Content[]): Content[] { }); // 2. Hoisting: Place all sorted responses BEFORE text or other parts - turn.parts = [...responseParts, ...otherParts]; + turn.content.parts = [...responseParts, ...otherParts]; } } } @@ -254,36 +261,42 @@ function refineToolResponses(history: Content[]): Content[] { * Final pass to ensure start/end roles and alternation are correct. */ function enforceRoleConstraints( - history: Content[], + history: HistoryTurn[], sentinels: Required>, -): Content[] { +): HistoryTurn[] { if (history.length === 0) return []; // Re-coalesce first to catch any empty turns or adjacent roles introduced by pairing const base = coalesce(history); if (base.length === 0) return []; - const result: Content[] = [...base]; + const result: HistoryTurn[] = [...base]; // 1. Ensure starts with user - if (result[0].role === 'model') { + if (result[0].content.role === 'model') { debugLogger.log( '[HistoryHardener] Final history starts with model role. Prepending sentinel user turn.', ); result.unshift({ - role: 'user', - parts: [{ text: sentinels.continuation }], + id: randomUUID(), + content: { + role: 'user', + parts: [{ text: sentinels.continuation }], + }, }); } // 2. Ensure ends with user - if (result[result.length - 1].role === 'model') { + if (result[result.length - 1].content.role === 'model') { debugLogger.log( '[HistoryHardener] Final history ends with model role. Appending sentinel user turn.', ); result.push({ - role: 'user', - parts: [{ text: 'Please continue.' }], + id: randomUUID(), + content: { + role: 'user', + parts: [{ text: 'Please continue.' }], + }, }); } @@ -296,12 +309,14 @@ function enforceRoleConstraints( * This ensures compatibility with strict APIs (like Vertex AI) that reject unknown fields. */ export function scrubHistory( - history: Content[], - idService?: NodeIdService, -): Content[] { - return history.map((content) => ({ - role: content.role, - parts: (content.parts || []).map((p) => scrubPart(p, idService)), + history: HistoryTurn[], +): HistoryTurn[] { + return history.map((turn) => ({ + id: turn.id, + content: { + role: turn.content.role, + parts: (turn.content.parts || []).map((p) => scrubPart(p)), + }, })); } @@ -313,7 +328,7 @@ function isThoughtPart(part: Part): part is ThoughtPart { return 'thoughtSignature' in part; } -function scrubPart(part: Part, idService?: NodeIdService): Part { +function scrubPart(part: Part): Part { const scrubbed: Record = {}; if ('text' in part && typeof part.text === 'string') { @@ -356,17 +371,5 @@ function scrubPart(part: Part, idService?: NodeIdService): Part { } // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const result = scrubbed as unknown as Part; - - // Propagate durable identity to the scrubbed object. - // This allows the HistoryObserver to recognize nodes even after they've been - // projected into multiple history formats, without polluting the API JSON. - if (idService) { - const id = idService.get(part); - if (id) { - idService.set(result, id); - } - } - - return result; + return scrubbed as unknown as Part; }