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

This commit is contained in:
joshualitt
2026-05-18 09:44:59 -07:00
committed by GitHub
parent 9d01958cdb
commit 055e0f6452
57 changed files with 3143 additions and 751 deletions
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type { ProcessArgs, ContextProcessor } from '../pipeline.js';
import * as fs from 'node:fs/promises';
@@ -62,7 +62,8 @@ export function createBlobDegradationProcessor(
if (payload.inlineData?.data && payload.inlineData?.mimeType) {
await ensureDir();
const ext = payload.inlineData.mimeType.split('/')[1] || 'bin';
const fileName = `blob_${Date.now()}_${randomUUID()}.${ext}`;
// Use a stable filename based on the node ID
const fileName = `blob_${deriveStableId([node.id])}.${ext}`;
const filePath = path.join(blobOutputsDir, fileName);
const buffer = Buffer.from(payload.inlineData.data, 'base64');
@@ -92,7 +93,7 @@ export function createBlobDegradationProcessor(
if (newText && tokensSaved > 0) {
returnedNodes.push({
...node,
id: randomUUID(),
id: deriveStableId([node.id, 'degraded']),
payload: { text: newText },
replacesId: node.id,
turnId: node.turnId,
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import { type ConcreteNode, NodeType } from '../graph/types.js';
@@ -99,9 +99,10 @@ export function createNodeDistillationProcessor(
if (newTokens < oldTokens) {
const distilledPayload = updatePart(payload, { text: summary });
const newId = deriveStableId([node.id, 'distilled']);
returnedNodes.push({
...node,
id: randomUUID(),
id: newId,
payload: distilledPayload,
replacesId: node.id,
timestamp: node.timestamp,
@@ -158,9 +159,10 @@ export function createNodeDistillationProcessor(
functionResponse: newFR,
});
const newId = deriveStableId([node.id, 'distilled']);
returnedNodes.push({
...node,
id: randomUUID(),
id: newId,
payload: distilledPayload,
replacesId: node.id,
timestamp: node.timestamp,
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
@@ -79,9 +79,10 @@ export function createNodeTruncationProcessor(
if (text) {
const squashResult = tryApplySquash(text, limitChars);
if (squashResult) {
const newId = deriveStableId([node.id, 'truncated']);
returnedNodes.push({
...node,
id: randomUUID(),
id: newId,
payload: { ...payload, text: squashResult.text },
replacesId: node.id,
turnId: node.turnId,
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type {
ContextProcessor,
@@ -112,7 +112,8 @@ export function createRollingSummaryProcessor(
try {
// Synthesize the rolling summary synchronously
const snapshotText = await generateRollingSummary(nodesToSummarize);
const newId = randomUUID();
const consumedIds = nodesToSummarize.map((n) => n.id);
const newId = deriveStableId(consumedIds);
const summaryNode: RollingSummary = {
id: newId,
@@ -121,10 +122,9 @@ export function createRollingSummaryProcessor(
timestamp: nodesToSummarize[nodesToSummarize.length - 1].timestamp,
role: 'user',
payload: { text: snapshotText },
abstractsIds: nodesToSummarize.map((n) => n.id),
abstractsIds: consumedIds,
};
const consumedIds = nodesToSummarize.map((n) => n.id);
const returnedNodes = targets.filter(
(t) => !consumedIds.includes(t.id),
);
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type {
ContextProcessor,
@@ -90,8 +90,11 @@ export function createStateSnapshotProcessor(
const isValid = consumedIds.every((id) => targetIds.has(id));
if (isValid) {
debugLogger.log(
`[StateSnapshotProcessor] Successfully spliced PROPOSED_SNAPSHOT from Inbox into Graph. Consumed ${consumedIds.length} nodes.`,
);
// If valid, apply it!
const newId = randomUUID();
const newId = deriveStableId(consumedIds);
const snapshotNode: Snapshot = {
id: newId,
@@ -120,6 +123,10 @@ export function createStateSnapshotProcessor(
inbox.consume(proposed.id);
return returnedNodes;
} else {
debugLogger.log(
`[StateSnapshotProcessor] Rejected PROPOSED_SNAPSHOT from Inbox because one or more target IDs were missing from the current graph window.`,
);
}
}
}
@@ -179,11 +186,11 @@ export function createStateSnapshotProcessor(
maxStateTokens: options.maxStateTokens,
},
);
const newId = randomUUID();
const consumedIds = nodesToSummarize.map((n) => n.id);
if (baselineIdToConsume && !consumedIds.includes(baselineIdToConsume)) {
consumedIds.push(baselineIdToConsume);
}
const newId = deriveStableId(consumedIds);
const snapshotNode: Snapshot = {
id: newId,
@@ -3,7 +3,7 @@
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { deriveStableId } from '../../utils/cryptoUtils.js';
import type { JSONSchemaType } from 'ajv';
import type { ContextProcessor, ProcessArgs } from '../pipeline.js';
import * as fs from 'node:fs/promises';
@@ -120,7 +120,7 @@ export function createToolMaskingProcessor(
directoryCreated = true;
}
const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${randomUUID()}.txt`;
const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${deriveStableId([content])}.txt`;
const filePath = path.join(toolOutputsDir, fileName);
await fs.writeFile(filePath, content);
@@ -214,9 +214,10 @@ export function createToolMaskingProcessor(
functionCall: newFC,
});
const newId = deriveStableId([node.id, 'masked']);
returnedNodes.push({
...node,
id: randomUUID(),
id: newId,
payload: maskedPart,
replacesId: node.id,
turnId: node.turnId,
@@ -242,9 +243,10 @@ export function createToolMaskingProcessor(
functionResponse: newFR,
});
const newId = deriveStableId([node.id, 'masked']);
returnedNodes.push({
...node,
id: randomUUID(),
id: newId,
payload: maskedPart,
replacesId: node.id,
turnId: node.turnId,