diff --git a/packages/memory-graph/src/__tests__/version-chain.test.ts b/packages/memory-graph/src/__tests__/version-chain.test.ts index 0503b771..64139aa4 100644 --- a/packages/memory-graph/src/__tests__/version-chain.test.ts +++ b/packages/memory-graph/src/__tests__/version-chain.test.ts @@ -71,7 +71,7 @@ describe("VersionChainIndex", () => { expect(chain!.map((e) => e.id)).toEqual(["m1", "m2", "m3"]) }) - it("getChain from middle element walks back to root", () => { + it("getChain from middle element returns full chain (backward + forward)", () => { const idx = new VersionChainIndex() const doc = makeDoc("d1", [ makeMem({ id: "m1", version: 1 }), @@ -90,12 +90,11 @@ describe("VersionChainIndex", () => { ]) idx.rebuild([doc]) - // Query from m2 (version 2) — walks back m2->m1, reverses to [m1,m2] - // Note: it doesn't walk forward to m3, only backward + // Query from m2 (version 2) — walks back to m1, forward to m3 const chain = idx.getChain("m2") expect(chain).not.toBeNull() - expect(chain!.length).toBe(2) - expect(chain!.map((e) => e.id)).toEqual(["m1", "m2"]) + expect(chain!.length).toBe(3) + expect(chain!.map((e) => e.id)).toEqual(["m1", "m2", "m3"]) }) it("caches chain results for all entries in the chain", () => { @@ -118,6 +117,32 @@ describe("VersionChainIndex", () => { expect(chain2).toBe(chain1) // same reference }) + it("getChain from v1 root with children returns full chain", () => { + const idx = new VersionChainIndex() + const doc = makeDoc("d1", [ + makeMem({ id: "m1", version: 1 }), + makeMem({ + id: "m2", + parentMemoryId: "m1", + rootMemoryId: "m1", + version: 2, + }), + makeMem({ + id: "m3", + parentMemoryId: "m2", + rootMemoryId: "m1", + version: 3, + }), + ]) + idx.rebuild([doc]) + + // Query from v1 root — walks forward to m2, m3 + const chain = idx.getChain("m1") + expect(chain).not.toBeNull() + expect(chain!.length).toBe(3) + expect(chain!.map((e) => e.id)).toEqual(["m1", "m2", "m3"]) + }) + it("getChain returns null for unknown ID", () => { const idx = new VersionChainIndex() idx.rebuild([makeDoc("d1", [makeMem({ id: "m1", version: 1 })])]) diff --git a/packages/memory-graph/src/canvas/version-chain.ts b/packages/memory-graph/src/canvas/version-chain.ts index 26de0bd0..a17bd97a 100644 --- a/packages/memory-graph/src/canvas/version-chain.ts +++ b/packages/memory-graph/src/canvas/version-chain.ts @@ -10,6 +10,7 @@ export interface ChainEntry { export class VersionChainIndex { private memoryMap = new Map() + private childrenMap = new Map() private cache = new Map() private lastDocs: GraphApiDocument[] | null = null @@ -17,11 +18,20 @@ export class VersionChainIndex { if (documents === this.lastDocs) return this.lastDocs = documents this.memoryMap.clear() + this.childrenMap.clear() this.cache.clear() for (const doc of documents) { for (const m of doc.memories) { this.memoryMap.set(m.id, m) + if (m.parentMemoryId) { + let children = this.childrenMap.get(m.parentMemoryId) + if (!children) { + children = [] + this.childrenMap.set(m.parentMemoryId, children) + } + children.push(m.id) + } } } } @@ -31,26 +41,47 @@ export class VersionChainIndex { if (cached) return cached const mem = this.memoryMap.get(memoryId) - if (!mem || mem.version <= 1) return null + if (!mem) return null - const chain: ChainEntry[] = [] + // Walk backward to root + const backward: GraphApiMemory[] = [] const visited = new Set() let current: GraphApiMemory | undefined = mem while (current && !visited.has(current.id)) { visited.add(current.id) - chain.push({ - id: current.id, - version: current.version, - memory: current.memory, - isForgotten: current.isForgotten, - isLatest: current.isLatest, - }) + backward.push(current) current = current.parentMemoryId ? this.memoryMap.get(current.parentMemoryId) : undefined } + backward.reverse() - chain.reverse() + // Walk forward from the selected node to find descendants + const forward: GraphApiMemory[] = [] + let tip: GraphApiMemory | undefined = mem + while (tip) { + const children = this.childrenMap.get(tip.id) + if (!children || children.length === 0) break + const child = this.memoryMap.get(children[0]) + if (!child || visited.has(child.id)) break + visited.add(child.id) + forward.push(child) + tip = child + } + + // Combine: backward (root..selected) + forward (selected+1..latest) + const all = [...backward, ...forward] + + // A single-entry chain (standalone v1 with no children) is not useful + if (all.length <= 1) return null + + const chain: ChainEntry[] = all.map((m) => ({ + id: m.id, + version: m.version, + memory: m.memory, + isForgotten: m.isForgotten, + isLatest: m.isLatest, + })) for (const entry of chain) { this.cache.set(entry.id, chain)