diff --git a/packages/core/src/context/contextManager.ts b/packages/core/src/context/contextManager.ts index 7fe9e46a76..de967a1211 100644 --- a/packages/core/src/context/contextManager.ts +++ b/packages/core/src/context/contextManager.ts @@ -43,9 +43,12 @@ export class ContextManager { history: Content[]; didApplyManagement: boolean; baseUnits: number; + processedNodes: readonly ConcreteNode[]; }; }; + private lastAppliedNodes?: readonly ConcreteNode[]; + private hasPerformedHotStart = false; constructor( @@ -287,6 +290,7 @@ export class ContextManager { history: Content[]; didApplyManagement: boolean; baseUnits: number; + processedNodes: readonly ConcreteNode[]; }> { this.tracer.logEvent('ContextManager', 'Starting rendering of LLM context'); @@ -348,11 +352,7 @@ export class ContextManager { const protectionReasons = this.getProtectedNodeIds(nodes, activeTaskIds); // Apply final GC Backstop pressure barrier synchronously before mapping - const { - history: renderedHistory, - didApplyManagement, - baseUnits, - } = await render( + const renderResult = await render( nodes, this.orchestrator, this.sidecar, @@ -364,6 +364,15 @@ export class ContextManager { previewNodeIds, ); + const { + history: renderedHistory, + didApplyManagement, + baseUnits, + processedNodes, + } = renderResult; + + this.lastAppliedNodes = processedNodes; + // Structural validation in debug mode checkContextInvariants(this.buffer.nodes, 'RenderHistory'); @@ -379,6 +388,7 @@ export class ContextManager { }), didApplyManagement, baseUnits, + processedNodes, }; // Update cache @@ -434,7 +444,9 @@ export class ContextManager { } exportState(): ContextEngineState { - return SnapshotStateHelper.exportState(this.buffer.nodes); + return SnapshotStateHelper.exportState( + this.lastAppliedNodes || this.buffer.nodes, + ); } restoreState(state: ContextEngineState): void { diff --git a/packages/core/src/context/graph/render.ts b/packages/core/src/context/graph/render.ts index e16beb4f38..eae24641e7 100644 --- a/packages/core/src/context/graph/render.ts +++ b/packages/core/src/context/graph/render.ts @@ -31,6 +31,7 @@ export async function render( history: Content[]; didApplyManagement: boolean; baseUnits: number; + processedNodes: readonly ConcreteNode[]; }> { let headerTokens = 0; let headerBaseUnits = 0; @@ -52,7 +53,12 @@ export async function render( const baseUnits = advancedTokenCalculator.getRawBaseUnits(nodes) + headerBaseUnits; - return { history: contents, didApplyManagement: false, baseUnits }; + return { + history: contents, + didApplyManagement: false, + baseUnits, + processedNodes: nodes, + }; } const maxTokens = sidecar.config.budget.maxTokens; @@ -97,6 +103,7 @@ export async function render( history: contents, didApplyManagement: false, baseUnits: graphBaseUnits + headerBaseUnits, + processedNodes: nodes, }; } const targetDelta = currentTokens - sidecar.config.budget.retainedTokens; @@ -151,5 +158,6 @@ export async function render( didApplyManagement: true, baseUnits: advancedTokenCalculator.getRawBaseUnits(visibleNodes) + headerBaseUnits, + processedNodes, }; } diff --git a/packages/core/src/context/graph/toGraph.ts b/packages/core/src/context/graph/toGraph.ts index 8bf9459789..b9c85b9f73 100644 --- a/packages/core/src/context/graph/toGraph.ts +++ b/packages/core/src/context/graph/toGraph.ts @@ -64,11 +64,12 @@ function isFunctionResponsePart( function isExecutableCodePart( part: Part, -): part is Part & { executableCode: { code: string } } { +): part is Part & { executableCode: { code: string; language: string } } { return ( typeof part.executableCode === 'object' && part.executableCode !== null && - typeof part.executableCode.code === 'string' + typeof part.executableCode.code === 'string' && + typeof part.executableCode.language === 'string' ); } @@ -78,7 +79,8 @@ function isCodeExecutionResultPart( return ( typeof part.codeExecutionResult === 'object' && part.codeExecutionResult !== null && - typeof part.codeExecutionResult.output === 'string' + typeof part.codeExecutionResult.output === 'string' && + typeof part.codeExecutionResult.outcome === 'string' ); } @@ -138,12 +140,16 @@ export function getStableId( id = `resp_h_${contentHash}_${turnSalt}_${partIdx}`; } else if (isExecutableCodePart(part)) { contentHash = createHash('sha256') - .update(`exec:${part.executableCode.code}`) + .update( + `exec:${part.executableCode.language}:${part.executableCode.code}`, + ) .digest('hex'); id = `exec_${contentHash}_${turnSalt}_${partIdx}`; } else if (isCodeExecutionResultPart(part)) { contentHash = createHash('sha256') - .update(`result:${part.codeExecutionResult.output}`) + .update( + `result:${part.codeExecutionResult.outcome}:${part.codeExecutionResult.output}`, + ) .digest('hex'); id = `result_${contentHash}_${turnSalt}_${partIdx}`; } diff --git a/packages/core/src/context/utils/snapshotGenerator.ts b/packages/core/src/context/utils/snapshotGenerator.ts index 9cd5a1d27d..83cd00e66a 100644 --- a/packages/core/src/context/utils/snapshotGenerator.ts +++ b/packages/core/src/context/utils/snapshotGenerator.ts @@ -94,13 +94,36 @@ export const SnapshotStateHelper = { 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 ${baseline.abstractsIds.length} consumed nodes.`, + `[SnapshotStateHelper] exportState: Exporting snapshot ID ${baseline.id} representing ${pristineIds.size} pristine nodes.`, ); return { snapshot: { text: baseline.text, - consumedIds: baseline.abstractsIds, + consumedIds: Array.from(pristineIds), timestamp: baseline.timestamp, }, }; @@ -124,7 +147,7 @@ export const SnapshotStateHelper = { inbox.publish('PROPOSED_SNAPSHOT', { newText: state.snapshot.text, consumedIds: state.snapshot.consumedIds, - type: 'accumulate', + type: 'restored', timestamp: state.snapshot.timestamp ?? Date.now(), }); } else { diff --git a/packages/core/src/core/fakeContentGenerator.ts b/packages/core/src/core/fakeContentGenerator.ts index 0f5a1e35d0..4976604993 100644 --- a/packages/core/src/core/fakeContentGenerator.ts +++ b/packages/core/src/core/fakeContentGenerator.ts @@ -58,10 +58,14 @@ export class FakeContentGenerator implements ContentGenerator { userTierName?: string; paidTier?: GeminiUserTier; + private readonly responses: FakeResponse[]; + constructor( - private readonly responses: FakeResponse[], + responses: FakeResponse[], private readonly options: FakeContentGeneratorOptions = {}, - ) {} + ) { + this.responses = structuredClone(responses); + } static async fromFile( filePath: string, diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 4c1b47fdcf..3bd4e618cd 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -276,6 +276,7 @@ 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