diff --git a/packages/core/src/context/contextManager.ts b/packages/core/src/context/contextManager.ts index dc62fbcefd..9316f62362 100644 --- a/packages/core/src/context/contextManager.ts +++ b/packages/core/src/context/contextManager.ts @@ -405,9 +405,13 @@ export class ContextManager { : renderedHistory; const result = { - history: hardenHistory(combinedHistory, { - sentinels: this.sidecar.sentinels, - }), + history: hardenHistory( + combinedHistory, + { + sentinels: this.sidecar.sentinels, + }, + this.env.graphMapper.getIdService(), + ), didApplyManagement, baseUnits, processedNodes, diff --git a/packages/core/src/context/graph/fromGraph.ts b/packages/core/src/context/graph/fromGraph.ts index 32be710e58..e762049b02 100644 --- a/packages/core/src/context/graph/fromGraph.ts +++ b/packages/core/src/context/graph/fromGraph.ts @@ -7,13 +7,17 @@ import type { Content } from '@google/genai'; import type { ConcreteNode } from './types.js'; import { debugLogger } from '../../utils/debugLogger.js'; +import type { NodeIdService } from './nodeIdService.js'; /** * Reconstructs a valid Gemini Chat History 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. */ -export function fromGraph(nodes: readonly ConcreteNode[]): Content[] { +export function fromGraph( + nodes: readonly ConcreteNode[], + idService?: NodeIdService, +): Content[] { debugLogger.log( `[fromGraph] Reconstructing history from ${nodes.length} nodes`, ); @@ -23,7 +27,12 @@ export function fromGraph(nodes: readonly ConcreteNode[]): Content[] { for (const node of nodes) { const turnId = node.turnId; - const partWithId = { ...node.payload, _synthId: node.id }; + + // Register the payload in the identity service to ensure stability + // even if the turn content changes (e.g. after GC backstop). + if (idService) { + idService.set(node.payload, node.id); + } // We start a new turn if: // 1. We don't have a current turn. @@ -36,12 +45,12 @@ export function fromGraph(nodes: readonly ConcreteNode[]): Content[] { ) { currentTurn = { role: node.role, - parts: [partWithId], + parts: [node.payload], _turnId: turnId, }; history.push(currentTurn); } else { - currentTurn.parts = [...(currentTurn.parts || []), partWithId]; + currentTurn.parts = [...(currentTurn.parts || []), node.payload]; } } diff --git a/packages/core/src/context/graph/mapper.ts b/packages/core/src/context/graph/mapper.ts index d66928d58f..68d5498bc4 100644 --- a/packages/core/src/context/graph/mapper.ts +++ b/packages/core/src/context/graph/mapper.ts @@ -8,13 +8,14 @@ import { ContextGraphBuilder } from './toGraph.js'; import type { Content } from '@google/genai'; import type { HistoryEvent } from '../../core/agentChatHistory.js'; import { fromGraph } from './fromGraph.js'; +import { NodeIdService } from './nodeIdService.js'; export class ContextGraphMapper { - private readonly nodeIdentityMap = new WeakMap(); + private readonly idService = new NodeIdService(); private readonly builder: ContextGraphBuilder; constructor() { - this.builder = new ContextGraphBuilder(this.nodeIdentityMap); + this.builder = new ContextGraphBuilder(this.idService); } applyEvent(event: HistoryEvent): ConcreteNode[] { @@ -22,6 +23,10 @@ export class ContextGraphMapper { } fromGraph(nodes: readonly ConcreteNode[]): Content[] { - return fromGraph(nodes); + return fromGraph(nodes, this.idService); + } + + getIdService(): NodeIdService { + return this.idService; } } diff --git a/packages/core/src/context/graph/nodeIdService.ts b/packages/core/src/context/graph/nodeIdService.ts new file mode 100644 index 0000000000..a9e48e748e --- /dev/null +++ b/packages/core/src/context/graph/nodeIdService.ts @@ -0,0 +1,23 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Provides a durable mapping between history object references and their + * corresponding graph node IDs. This ensures that context management logic + * can track the identity of turns even after they are transformed (e.g. scrubbed + * or hardened) without polluting the raw JSON sent to the Gemini API. + */ +export class NodeIdService { + constructor(private readonly map: WeakMap = new WeakMap()) {} + + get(obj: object): string | undefined { + return this.map.get(obj); + } + + set(obj: object, id: string): void { + this.map.set(obj, id); + } +} diff --git a/packages/core/src/context/graph/toGraph.test.ts b/packages/core/src/context/graph/toGraph.test.ts index 9f221339bf..dce2257b23 100644 --- a/packages/core/src/context/graph/toGraph.test.ts +++ b/packages/core/src/context/graph/toGraph.test.ts @@ -8,6 +8,7 @@ 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'; describe('ContextGraphBuilder', () => { describe('toGraph', () => { @@ -26,7 +27,7 @@ describe('ContextGraphBuilder', () => { { role: 'user', parts: [{ text: 'Message 2' }] }, ]; - const builder = new ContextGraphBuilder(); + const builder = new ContextGraphBuilder(new NodeIdService()); const nodes = builder.processHistory(history); // We expect the first two messages and the last one to be present @@ -69,7 +70,7 @@ describe('ContextGraphBuilder', () => { ]; // 1. Initial Graph Generation - const builder1 = new ContextGraphBuilder(); + const builder1 = new ContextGraphBuilder(new NodeIdService()); const nodes1 = builder1.processHistory(complexHistory); // 2. Serialize and Deserialize (Simulating saving and loading from disk) @@ -77,7 +78,7 @@ describe('ContextGraphBuilder', () => { const parsedHistory = JSON.parse(serializedHistory) as Content[]; // 3. Second Graph Generation from parsed JSON - const builder2 = new ContextGraphBuilder(); + const builder2 = new ContextGraphBuilder(new NodeIdService()); const nodes2 = builder2.processHistory(parsedHistory); // Assertion: The arrays must be completely identical, including all generated UUIDs diff --git a/packages/core/src/context/graph/toGraph.ts b/packages/core/src/context/graph/toGraph.ts index b9c85b9f73..5a7c50a2fe 100644 --- a/packages/core/src/context/graph/toGraph.ts +++ b/packages/core/src/context/graph/toGraph.ts @@ -8,10 +8,7 @@ import type { Content, Part } from '@google/genai'; import { type ConcreteNode, NodeType } from './types.js'; import { randomUUID, createHash } from 'node:crypto'; import { debugLogger } from '../../utils/debugLogger.js'; - -interface PartWithSynthId extends Part { - _synthId?: string; -} +import type { NodeIdService } from './nodeIdService.js'; // Global WeakMap to cache hashes for Part objects. // This optimizes getStableId by avoiding redundant stringify/hash operations @@ -91,27 +88,24 @@ function isCodeExecutionResultPart( */ export function getStableId( obj: object, - nodeIdentityMap: WeakMap, + idService: NodeIdService, turnSalt: string = '', partIdx: number = 0, ): string { - let id = nodeIdentityMap.get(obj); + let id = idService.get(obj); if (id) return id; const cachedHash = PART_HASH_CACHE.get(obj); if (cachedHash) { id = `${cachedHash}_${turnSalt}_${partIdx}`; - nodeIdentityMap.set(obj, id); + idService.set(obj, id); return id; } - const part = obj as PartWithSynthId; + const part = obj as Part; let contentHash: string | undefined; - // If the object already has a synthetic ID property, use it. - if (typeof part._synthId === 'string') { - id = part._synthId; - } else if (isTextPart(part)) { + if (isTextPart(part)) { contentHash = createHash('sha256').update(part.text).digest('hex'); id = `text_${contentHash}_${turnSalt}_${partIdx}`; } else if (isInlineDataPart(part)) { @@ -167,7 +161,7 @@ export function getStableId( } } - nodeIdentityMap.set(obj, id); + idService.set(obj, id); return id; } @@ -176,9 +170,7 @@ export function getStableId( * Every Part in history is mapped to exactly one ConcreteNode. */ export class ContextGraphBuilder { - constructor( - private readonly nodeIdentityMap: WeakMap = new WeakMap(), - ) {} + constructor(private readonly idService: NodeIdService) {} processHistory(history: readonly Content[]): ConcreteNode[] { const nodes: ConcreteNode[] = []; @@ -213,7 +205,7 @@ export class ContextGraphBuilder { const occurrence = (seenHashes.get(h) || 0) + 1; seenHashes.set(h, occurrence); const turnSalt = `${h}_${occurrence}`; - const turnId = getStableId(msg, this.nodeIdentityMap, turnSalt, -1); + const turnId = getStableId(msg, this.idService, turnSalt, -1); if (msg.role === 'user') { for (let partIdx = 0; partIdx < msg.parts.length; partIdx++) { @@ -221,13 +213,13 @@ export class ContextGraphBuilder { const apiId = isFunctionResponsePart(part) && typeof part.functionResponse.id === 'string' - ? `resp_${part.functionResponse.id}_${turnSalt}_${partIdx}` + ? `resp_${part.functionResponse.id}` : isFunctionCallPart(part) && typeof part.functionCall.id === 'string' - ? `call_${part.functionCall.id}_${turnSalt}_${partIdx}` + ? `call_${part.functionCall.id}` : undefined; const id = - apiId || getStableId(part, this.nodeIdentityMap, turnSalt, partIdx); + apiId || getStableId(part, this.idService, turnSalt, partIdx); const node: ConcreteNode = { id, timestamp: Date.now(), @@ -245,10 +237,10 @@ export class ContextGraphBuilder { const part = msg.parts[partIdx]; const apiId = isFunctionCallPart(part) && typeof part.functionCall.id === 'string' - ? `call_${part.functionCall.id}_${turnSalt}_${partIdx}` + ? `call_${part.functionCall.id}` : undefined; const id = - apiId || getStableId(part, this.nodeIdentityMap, turnSalt, partIdx); + apiId || getStableId(part, this.idService, turnSalt, partIdx); const node: ConcreteNode = { id, timestamp: Date.now(), diff --git a/packages/core/src/utils/historyHardening.ts b/packages/core/src/utils/historyHardening.ts index 5ff071acd8..98122abf65 100644 --- a/packages/core/src/utils/historyHardening.ts +++ b/packages/core/src/utils/historyHardening.ts @@ -6,6 +6,7 @@ import type { Content, Part } from '@google/genai'; import { debugLogger } from './debugLogger.js'; +import type { NodeIdService } from '../context/graph/nodeIdService.js'; export const SYNTHETIC_THOUGHT_SIGNATURE = 'skip_thought_signature_validator'; @@ -37,6 +38,7 @@ const DEFAULT_SENTINELS = { export function hardenHistory( history: Content[], options: HardeningOptions = {}, + idService?: NodeIdService, ): Content[] { if (history.length === 0) return history; @@ -55,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); + final = scrubHistory(final, idService); return final; } @@ -293,10 +295,13 @@ function enforceRoleConstraints( * Deep-scrubs the history to remove any non-standard properties from Content and Part objects. * This ensures compatibility with strict APIs (like Vertex AI) that reject unknown fields. */ -export function scrubHistory(history: Content[]): Content[] { +export function scrubHistory( + history: Content[], + idService?: NodeIdService, +): Content[] { return history.map((content) => ({ role: content.role, - parts: (content.parts || []).map(scrubPart), + parts: (content.parts || []).map((p) => scrubPart(p, idService)), })); } @@ -308,7 +313,7 @@ function isThoughtPart(part: Part): part is ThoughtPart { return 'thoughtSignature' in part; } -function scrubPart(part: Part): Part { +function scrubPart(part: Part, idService?: NodeIdService): Part { const scrubbed: Record = {}; if ('text' in part && typeof part.text === 'string') { @@ -351,5 +356,17 @@ function scrubPart(part: Part): Part { } // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - return scrubbed as unknown as Part; + 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; }