Restore cognition references during history rollback

graph/knowledge-state.js

@@ -64,6 +64,51 @@ function uniqueIds(values = []) {
   return result.slice(0, KNOWLEDGE_ENTRY_LIMIT);
 }
 
+function buildExistingGraphNodeIdSet(graph) {
+  const nodeIds = new Set();
+  for (const node of Array.isArray(graph?.nodes) ? graph.nodes : []) {
+    const nodeId = normalizeString(node?.id);
+    if (nodeId) nodeIds.add(nodeId);
+  }
+  return nodeIds;
+}
+
+function pruneKnowledgeOwnerNodeRefs(entry, graph = null) {
+  const normalizedEntry = createDefaultKnowledgeOwnerState(entry);
+  if (!graph || typeof graph !== "object") {
+    return normalizedEntry;
+  }
+
+  const existingNodeIds = buildExistingGraphNodeIdSet(graph);
+  const filterNodeIds = (values = []) =>
+    uniqueIds(values).filter((nodeId) => existingNodeIds.has(nodeId));
+
+  let ownerNodeId = normalizeString(normalizedEntry.nodeId);
+  if (ownerNodeId && !existingNodeIds.has(ownerNodeId)) {
+    const matches = findCharacterNodeByName(graph, normalizedEntry.ownerName);
+    ownerNodeId = matches.length === 1 ? normalizeString(matches[0]?.id) : "";
+  }
+
+  const visibilityScores = {};
+  for (const [nodeId, score] of Object.entries(
+    normalizedEntry.visibilityScores || {},
+  )) {
+    const normalizedNodeId = normalizeString(nodeId);
+    if (!normalizedNodeId || !existingNodeIds.has(normalizedNodeId)) continue;
+    visibilityScores[normalizedNodeId] = clampScore(score);
+  }
+
+  return createDefaultKnowledgeOwnerState({
+    ...normalizedEntry,
+    nodeId: ownerNodeId,
+    knownNodeIds: filterNodeIds(normalizedEntry.knownNodeIds),
+    mistakenNodeIds: filterNodeIds(normalizedEntry.mistakenNodeIds),
+    manualKnownNodeIds: filterNodeIds(normalizedEntry.manualKnownNodeIds),
+    manualHiddenNodeIds: filterNodeIds(normalizedEntry.manualHiddenNodeIds),
+    visibilityScores,
+  });
+}
+
 function buildOwnerAliasVariantSet(values = []) {
   const variants = new Set();
   for (const value of Array.isArray(values) ? values : [values]) {
@@ -122,14 +167,19 @@ function getKnowledgeOwnerEvidenceScore(owner = {}) {
   );
 }
 
-function findEquivalentCharacterOwnerEntry(ownerCollection, candidate = {}) {
-  if (normalizeOwnerType(candidate?.ownerType) !== OWNER_TYPE_CHARACTER) {
+function findEquivalentCharacterOwnerEntry(
+  ownerCollection,
+  candidate = {},
+  graph = null,
+) {
+  const normalizedCandidate = pruneKnowledgeOwnerNodeRefs(candidate, graph);
+  if (normalizeOwnerType(normalizedCandidate?.ownerType) !== OWNER_TYPE_CHARACTER) {
     return null;
   }
 
-  const candidateKey = normalizeString(candidate?.ownerKey);
-  const candidateNodeId = normalizeString(candidate?.nodeId);
-  const candidateAliasSet = getKnowledgeOwnerAliasVariantSet(candidate);
+  const candidateKey = normalizeString(normalizedCandidate?.ownerKey);
+  const candidateNodeId = normalizeString(normalizedCandidate?.nodeId);
+  const candidateAliasSet = getKnowledgeOwnerAliasVariantSet(normalizedCandidate);
   const matches = [];
   const values =
     ownerCollection instanceof Map
@@ -137,7 +187,7 @@ function findEquivalentCharacterOwnerEntry(ownerCollection, candidate = {}) {
       : Object.values(ownerCollection || {});
 
   for (const rawEntry of values) {
-    const entry = createDefaultKnowledgeOwnerState(rawEntry);
+    const entry = pruneKnowledgeOwnerNodeRefs(rawEntry, graph);
     if (!entry.ownerKey || normalizeOwnerType(entry.ownerType) !== OWNER_TYPE_CHARACTER) {
       continue;
     }
@@ -461,10 +511,10 @@ function resolveCanonicalKnowledgeEntry(
   entry,
   userAliasContext = null,
 ) {
-  const normalizedEntry = createDefaultKnowledgeOwnerState({
+  const normalizedEntry = pruneKnowledgeOwnerNodeRefs({
     ...entry,
     ownerKey,
-  });
+  }, graph);
   const resolvedOwner = resolveKnowledgeOwner(graph, {
     ownerType: normalizedEntry.ownerType,
     ownerName: normalizedEntry.ownerName,
@@ -501,7 +551,7 @@ export function normalizeKnowledgeState(state = {}, graph = null) {
     if (!canonicalEntry.ownerKey) continue;
     const equivalentEntry =
       owners[canonicalEntry.ownerKey] ||
-      findEquivalentCharacterOwnerEntry(owners, canonicalEntry);
+      findEquivalentCharacterOwnerEntry(owners, canonicalEntry, graph);
     const targetKey = equivalentEntry?.ownerKey || canonicalEntry.ownerKey;
     owners[targetKey] = owners[targetKey]
       ? mergeKnowledgeOwnerEntries(owners[targetKey], canonicalEntry)
@@ -625,6 +675,7 @@ export function resolveKnowledgeOwner(graph, input = {}) {
       nodeId,
       aliases,
     },
+    graph,
   );
   if (equivalentOwner?.ownerKey) {
     return {
@@ -1503,7 +1554,7 @@ export function listKnowledgeOwners(graph) {
     };
     const equivalentEntry =
       owners.get(normalizedEntry.ownerKey) ||
-      findEquivalentCharacterOwnerEntry(owners, displayEntry);
+      findEquivalentCharacterOwnerEntry(owners, displayEntry, graph);
     const targetKey = equivalentEntry?.ownerKey || normalizedEntry.ownerKey;
     owners.set(
       targetKey,

runtime/runtime-state.js

@@ -790,6 +790,18 @@ function buildJournalStateBefore(snapshotBefore, meta = {}) {
       ? snapshotBefore.historyState.historyDirtyFrom
       : null,
     vectorIndexState: clonePlain(snapshotBefore?.vectorIndexState || {}),
+    knowledgeState: clonePlain(
+      snapshotBefore?.knowledgeState || createDefaultKnowledgeState(),
+    ),
+    regionState: clonePlain(
+      snapshotBefore?.regionState || createDefaultRegionState(),
+    ),
+    timelineState: clonePlain(
+      snapshotBefore?.timelineState || createDefaultTimelineState(),
+    ),
+    summaryState: clonePlain(
+      snapshotBefore?.summaryState || createDefaultSummaryState(),
+    ),
     lastRecallResult: Array.isArray(snapshotBefore?.lastRecallResult)
       ? [...snapshotBefore.lastRecallResult]
       : null,
@@ -1157,6 +1169,21 @@ function applyJournalStateBefore(graph, stateBefore = {}) {
     ...createDefaultVectorIndexState(graph?.historyState?.chatId || ""),
     ...clonePlain(stateBefore.vectorIndexState || {}),
   };
+  graph.knowledgeState = createDefaultKnowledgeState(
+    clonePlain(stateBefore.knowledgeState || {}),
+  );
+  graph.regionState = createDefaultRegionState(
+    clonePlain(stateBefore.regionState || {}),
+  );
+  graph.timelineState = createDefaultTimelineState(
+    clonePlain(stateBefore.timelineState || {}),
+  );
+  graph.summaryState = createDefaultSummaryState(
+    clonePlain(stateBefore.summaryState || {}),
+  );
+  normalizeGraphCognitiveState(graph);
+  normalizeGraphStoryTimeline(graph);
+  normalizeGraphSummaryState(graph);
   graph.lastRecallResult = Array.isArray(stateBefore.lastRecallResult)
     ? [...stateBefore.lastRecallResult]
     : null;

tests/runtime-history.mjs

@@ -13,6 +13,7 @@ import {
   snapshotProcessedMessageHashes,
 } from "../runtime/runtime-state.js";
 import { createEmptyGraph } from "../graph/graph.js";
+import { normalizeKnowledgeState } from "../graph/knowledge-state.js";
 
 const chat = [
   { is_user: true, mes: "你好" },
@@ -129,6 +130,69 @@ assert.deepEqual(
   snapshotProcessedMessageHashes(chat, 3),
 );
 
+const danglingKnowledgeGraph = createEmptyGraph();
+danglingKnowledgeGraph.nodes.push({
+  id: "live-node",
+  type: "event",
+  fields: { title: "仍存在", summary: "仍存在的节点" },
+  seq: 1,
+  seqRange: [1, 1],
+  archived: false,
+  embedding: null,
+  importance: 5,
+  accessCount: 0,
+  lastAccessTime: Date.now(),
+  createdTime: Date.now(),
+  level: 0,
+  parentId: null,
+  childIds: [],
+  prevId: null,
+  nextId: null,
+  clusters: [],
+});
+danglingKnowledgeGraph.knowledgeState.owners["character:艾琳"] = {
+  ownerType: "character",
+  ownerKey: "character:艾琳",
+  ownerName: "艾琳",
+  nodeId: "ghost-owner-node",
+  knownNodeIds: ["ghost-node", "live-node"],
+  mistakenNodeIds: ["ghost-mistaken"],
+  manualKnownNodeIds: ["ghost-manual-known"],
+  manualHiddenNodeIds: ["ghost-manual-hidden"],
+  visibilityScores: {
+    "ghost-node": 1,
+    "live-node": 0.9,
+  },
+};
+const normalizedKnowledgeState = normalizeKnowledgeState(
+  danglingKnowledgeGraph.knowledgeState,
+  danglingKnowledgeGraph,
+);
+assert.deepEqual(
+  normalizedKnowledgeState.owners["character:艾琳"]?.knownNodeIds,
+  ["live-node"],
+);
+assert.deepEqual(
+  normalizedKnowledgeState.owners["character:艾琳"]?.mistakenNodeIds,
+  [],
+);
+assert.deepEqual(
+  normalizedKnowledgeState.owners["character:艾琳"]?.manualKnownNodeIds,
+  [],
+);
+assert.deepEqual(
+  normalizedKnowledgeState.owners["character:艾琳"]?.manualHiddenNodeIds,
+  [],
+);
+assert.deepEqual(
+  normalizedKnowledgeState.owners["character:艾琳"]?.visibilityScores,
+  { "live-node": 0.9 },
+);
+assert.equal(
+  normalizedKnowledgeState.owners["character:艾琳"]?.nodeId || "",
+  "",
+);
+
 const clearedGraph = normalizeGraphRuntimeState({
   historyState: {
     chatId: "chat-history-test",
@@ -177,6 +241,13 @@ graph.lastProcessedSeq = 3;
 graph.historyState.lastProcessedAssistantFloor = 3;
 graph.historyState.processedMessageHashes = hashes;
 graph.historyState.extractionCount = 4;
+graph.knowledgeState.owners["character:艾琳"] = {
+  ownerType: "character",
+  ownerKey: "character:艾琳",
+  ownerName: "艾琳",
+  knownNodeIds: ["node-1"],
+  visibilityScores: { "node-1": 1 },
+};
 const afterSnapshot = cloneGraphSnapshot(graph);
 appendBatchJournal(
   graph,
@@ -197,5 +268,9 @@ rollbackBatch(graph, recoveryPoint.affectedJournals[0]);
 assert.equal(graph.nodes.length, 0);
 assert.equal(graph.historyState.lastProcessedAssistantFloor, -1);
 assert.equal(graph.historyState.extractionCount, 0);
+assert.equal(
+  Object.keys(graph.knowledgeState?.owners || {}).length,
+  0,
+);
 
 console.log("runtime-history tests passed");