From 2e2805d92f9093f6675983a7b29ca01b45b7f755 Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Thu, 3 Sep 2026 16:31:29 +0200 Subject: [PATCH] feat(agent): construct memory forks directly from live state Stream current values and lists through shared fork policy while preserving sequences and the source high-water mark. --- .../08-named-branch-streaming-forks.md | 12 ++ .../agent/src/harness/session/fork-policy.ts | 45 ++-- packages/agent/src/harness/session/fork.ts | 104 +++++---- .../session/in-memory-storage-state.ts | 201 +++++++++++++++--- packages/agent/src/harness/session/index.ts | 2 +- packages/agent/src/harness/session/memory.ts | 29 +-- .../test/harness/memory-session-repo.test.ts | 62 ------ 7 files changed, 266 insertions(+), 189 deletions(-) diff --git a/packages/agent/docs/work-packages/08-named-branch-streaming-forks.md b/packages/agent/docs/work-packages/08-named-branch-streaming-forks.md index faa09430c..4dd3aeb1b 100644 --- a/packages/agent/docs/work-packages/08-named-branch-streaming-forks.md +++ b/packages/agent/docs/work-packages/08-named-branch-streaming-forks.md @@ -238,3 +238,15 @@ WP08 is complete when: - shared conformance, focused backend tests, benchmarks, `npm run check`, and `./test.sh` pass; - normative docs and the roadmap reflect the new contract with historical documents untouched; - final Fable review reports no blocker. + +## 9. Open questions + +These do not change the current contract until resolved. + +### 9.1 Corruption checks + +Current: reject any malformed lane. Possible relaxation: trust supported atomic state and perform no fork-time corruption audit; branch forks only require the selected tip, config, and state. This removes global lane inventory validation. + +### 9.2 Sequence preservation + +Current: preserve source sequences, `nextSeq`, and list cursors. Possible relaxation: keep entry sequences but assign destination-local sequences to values, list elements, and reconstructed lane rows. Revisit if preservation materially complicates streaming. diff --git a/packages/agent/src/harness/session/fork-policy.ts b/packages/agent/src/harness/session/fork-policy.ts index 52523518e..13ac3a0cf 100644 --- a/packages/agent/src/harness/session/fork-policy.ts +++ b/packages/agent/src/harness/session/fork-policy.ts @@ -1,32 +1,35 @@ -/** - * Final action for current state at one address: - * - copy: emit the current value in the destination; - * - exclude: omit it from the destination; - * - reconstruct: do not copy the row because lane handling emits a coherent replacement. - */ -export type ForkDisposition = "copy" | "exclude" | "reconstruct"; +import type { CommittedListAppendWrite, CommittedValueSetWrite } from "./commit.ts"; -/** Decide the final fork action for one current scalar or list address. */ -export function classifyForkAddress( - address: { readonly namespace: string; readonly key: string }, - scope: "branch" | "tree", +export type ForkCurrentStatePlan = + | { scope: "branch"; branch: string; destinationTip: string | null } + | { scope: "tree" }; + +/** Project one current scalar row or surviving list element into destination state. */ +export function projectForkCurrentStateWrite( + write: CommittedValueSetWrite | CommittedListAppendWrite, + plan: ForkCurrentStatePlan, isEntryCopied: (entryId: string) => boolean, -): ForkDisposition { - switch (address.namespace) { +): CommittedValueSetWrite | CommittedListAppendWrite | undefined { + switch (write.namespace) { case "pi.session.name": - return "copy"; + return write; case "pi.entry.label": - return isEntryCopied(address.key) ? "copy" : "exclude"; + return isEntryCopied(write.key) ? write : undefined; case "pi.branch.tip": + if (plan.scope === "tree") return write; + return write.key === plan.branch ? { ...write, value: plan.destinationTip } : undefined; case "pi.lane.config": + return plan.scope === "tree" || write.key === plan.branch ? write : undefined; case "pi.lane.state": - return "reconstruct"; + return plan.scope === "tree" || write.key === plan.branch + ? { ...write, value: { currentOperationId: null, lastOperationId: null, inbox: [] } } + : undefined; case "pi.result": - return "exclude"; + return undefined; } - if (address.namespace.startsWith("pi.op.") || address.namespace.startsWith("pi.pending.")) return "exclude"; - if (address.namespace === "pi" || address.namespace.startsWith("pi.")) { - throw new Error(`Unknown reserved fork namespace: ${address.namespace}`); + if (write.namespace.startsWith("pi.op.") || write.namespace.startsWith("pi.pending.")) return undefined; + if (write.namespace === "pi" || write.namespace.startsWith("pi.")) { + throw new Error(`Unknown reserved fork namespace: ${write.namespace}`); } - return scope === "tree" ? "copy" : "exclude"; + return plan.scope === "tree" ? write : undefined; } diff --git a/packages/agent/src/harness/session/fork.ts b/packages/agent/src/harness/session/fork.ts index c7f6a79d7..7da189bfa 100644 --- a/packages/agent/src/harness/session/fork.ts +++ b/packages/agent/src/harness/session/fork.ts @@ -1,5 +1,5 @@ import type { CommittedWrite } from "./commit.ts"; -import { classifyForkAddress } from "./fork-policy.ts"; +import { type ForkCurrentStatePlan, projectForkCurrentStateWrite } from "./fork-policy.ts"; import type { Entry, ForkOptions } from "./types.ts"; import { branchTip, laneConfig, laneState, type StoredValue, type Value, value } from "./values.ts"; @@ -32,35 +32,31 @@ export function createForkSnapshot(source: ForkSourceSnapshot, options: ForkOpti const sourceTips = storedValuesInNamespace(source.scalarValues, branchTip("")); validateForkSourceSnapshot(source, sourceEntries, sourceTips, options); - const { entryIds, destinationTips } = selectForkContents(sourceEntries, sourceTips, options); + const { entryIds, plan } = selectForkContents(sourceEntries, sourceTips, options); const entries = new Map(); for (const id of entryIds) entries.set(id, sourceEntries.get(id)!); const scalarValues: StoredValue[] = []; let nextSeq = Math.max(0, ...[...entries.values()].map((entry) => entry.seq)) + 1; - const store = (address: Value, storedValue: T): void => { - scalarValues.push({ - address: value(address.namespace, address.key), - value: storedValue, - seq: nextSeq++, - }); - }; - for (const [branch, tipId] of destinationTips) { - const configuration = findStoredValue(source.scalarValues, laneConfig(branch)); - store(branchTip(branch), tipId); - if (configuration !== undefined) { - store(laneConfig(branch), configuration.value); - store(laneState(branch), { currentOperationId: null, lastOperationId: null, inbox: [] }); - } - } for (const stored of source.scalarValues) { - switch (classifyForkAddress(stored.address, options.scope, (entryId) => entryIds.has(entryId))) { - case "copy": - store(stored.address, stored.value); - break; - case "exclude": - case "reconstruct": - break; + const projected = projectForkCurrentStateWrite( + { + kind: "value", + op: "set", + seq: stored.seq, + namespace: stored.address.namespace, + key: stored.address.key, + value: stored.value, + }, + plan, + (entryId) => entryIds.has(entryId), + ); + if (projected !== undefined) { + scalarValues.push({ + address: value(projected.namespace, projected.key), + value: projected.value, + seq: nextSeq++, + }); } } @@ -87,41 +83,39 @@ function selectForkContents( sourceEntries: Map, sourceTips: StoredValue[], options: ForkOptions, -): { - entryIds: Set; - destinationTips: Map; -} { +): { entryIds: Set; plan: ForkCurrentStatePlan } { const entryIds = new Set(); - const destinationTips = new Map(); if (options.scope === "tree") { for (const id of sourceEntries.keys()) entryIds.add(id); - for (const stored of sourceTips) destinationTips.set(stored.address.key, stored.value); - } else { - const sourceTip = sourceTips.find((stored) => stored.address.key === options.branch); - if (sourceTip === undefined) throw new Error(`Unknown source branch: ${options.branch}`); - - const requested = options.entryId ?? sourceTip.value; - let found = requested === null; - let tipId: string | null = null; - let entryId = sourceTip.value; - while (entryId !== null) { - const entry = sourceEntries.get(entryId); - if (entry === undefined) throw new Error(`Corrupt source branch: missing parent ${entryId}`); - if (entry.id === requested) { - found = true; - tipId = options.position === "before" ? entry.parentId : entry.id; - if (options.position !== "before") entryIds.add(entry.id); - } else if (found) { - entryIds.add(entry.id); - } - entryId = entry.parentId; - } - if (!found) { - throw new Error(`Fork entry ${requested} is not on source branch ${JSON.stringify(options.branch)}`); - } - destinationTips.set(options.branch, tipId); + return { entryIds, plan: { scope: "tree" } }; } - return { entryIds, destinationTips }; + + const sourceTip = sourceTips.find((stored) => stored.address.key === options.branch); + if (sourceTip === undefined) throw new Error(`Unknown source branch: ${options.branch}`); + + const requested = options.entryId ?? sourceTip.value; + let found = requested === null; + let destinationTip: string | null = null; + let entryId = sourceTip.value; + while (entryId !== null) { + const entry = sourceEntries.get(entryId); + if (entry === undefined) throw new Error(`Corrupt source branch: missing parent ${entryId}`); + if (entry.id === requested) { + found = true; + destinationTip = options.position === "before" ? entry.parentId : entry.id; + if (options.position !== "before") entryIds.add(entry.id); + } else if (found) { + entryIds.add(entry.id); + } + entryId = entry.parentId; + } + if (!found) { + throw new Error(`Fork entry ${requested} is not on source branch ${JSON.stringify(options.branch)}`); + } + return { + entryIds, + plan: { scope: "branch", branch: options.branch, destinationTip }, + }; } function validateForkSourceSnapshot( diff --git a/packages/agent/src/harness/session/in-memory-storage-state.ts b/packages/agent/src/harness/session/in-memory-storage-state.ts index ae1b1737e..e1b4e58e0 100644 --- a/packages/agent/src/harness/session/in-memory-storage-state.ts +++ b/packages/agent/src/harness/session/in-memory-storage-state.ts @@ -1,5 +1,13 @@ import { addUsage, emptyUsage } from "../utils/usage.ts"; -import { type CommittedWrite, type PreparedCommit, prepareStorageCommit, validateCommittedWrites } from "./commit.ts"; +import { + type CommittedListAppendWrite, + type CommittedValueSetWrite, + type CommittedWrite, + type PreparedCommit, + prepareStorageCommit, + validateCommittedWrites, +} from "./commit.ts"; +import { projectForkCurrentStateWrite } from "./fork-policy.ts"; export type { CommittedEntryWrite, @@ -16,6 +24,7 @@ import type { Entry, EntryScan, EntryStructure, + ForkOptions, SessionStats, StorageBranchScan, UsageRow, @@ -23,6 +32,7 @@ import type { Write, } from "./types.ts"; import { + branchTip, type ListElement, type ListReadOptions, list, @@ -38,6 +48,16 @@ interface StoredListSnapshot { elements: ListElement[]; } +interface ForkLaneInventory { + tip: boolean; + configuration: boolean; + state: boolean; +} + +type MemoryForkPlan = + | { scope: "tree" } + | { scope: "branch"; branch: string; destinationTip: string | null; entryIds: Set }; + function physicalKey(namespace: string, key: string): string { return `${namespace}\u0000${key}`; } @@ -108,43 +128,164 @@ export class InMemoryStorageState { this.stats = { ...this.stats, usage: addUsage(this.stats.usage, row.usage) }; break; } - case "value": { - const key = physicalKey(write.namespace, write.key); - if (write.op === "delete") { - this.scalarValues.delete(key); - } else { - this.scalarValues.set(key, { - address: value(write.namespace, write.key), - value: write.value, - seq: write.seq, - }); - } + case "value": + if (write.op === "delete") this.scalarValues.delete(physicalKey(write.namespace, write.key)); + else this.applyValueSetOrListAppend(write); break; - } - case "list": { - const key = physicalKey(write.namespace, write.key); - if (write.op === "delete") { - this.listValues.delete(key); - } else { - const stored = this.listValues.get(key); - const element = { seq: write.seq, value: write.value }; - if (stored === undefined) { - this.listValues.set(key, { - address: list(write.namespace, write.key), - elements: [element], - }); - } else { - stored.elements.push(element); - } - } + case "list": + if (write.op === "delete") this.listValues.delete(physicalKey(write.namespace, write.key)); + else this.applyValueSetOrListAppend(write); break; - } } this.nextSeq = write.seq + 1; } return this.stats; } + createFork(options: ForkOptions): InMemoryStorageState { + const plan = this.selectForkPlan(options); + + const isEntryCopied = (entryId: string): boolean => { + if (plan.scope === "tree") return true; + return plan.entryIds.has(entryId); + }; + const destination = new InMemoryStorageState(); + let messageCount = 0; + for (const entry of this.entriesBySeq) { + if (!isEntryCopied(entry.id)) continue; + destination.entries.set(entry.id, entry); + destination.entriesBySeq.push(entry); + if (entry.type === "message") messageCount++; + } + destination.stats = { ...destination.stats, messageCount }; + + const lanes = new Map(); + const lane = (name: string): ForkLaneInventory => { + let inventory = lanes.get(name); + if (inventory === undefined) { + inventory = { tip: false, configuration: false, state: false }; + lanes.set(name, inventory); + } + return inventory; + }; + for (const stored of this.scalarValues.values()) { + switch (stored.address.namespace) { + case "pi.branch.tip": { + lane(stored.address.key).tip = true; + const tip = stored as StoredValue; + if (tip.value !== null && !this.entries.has(tip.value)) { + throw new Error(`Source session branch ${JSON.stringify(stored.address.key)} has an unknown tip`); + } + break; + } + case "pi.lane.config": + lane(stored.address.key).configuration = true; + break; + case "pi.lane.state": + lane(stored.address.key).state = true; + break; + } + + const projected = projectForkCurrentStateWrite( + { + kind: "value", + op: "set", + seq: stored.seq, + namespace: stored.address.namespace, + key: stored.address.key, + value: stored.value, + }, + plan, + isEntryCopied, + ); + if (projected !== undefined) destination.applyValueSetOrListAppend(projected); + } + + for (const [name, inventory] of lanes) { + if (!inventory.tip && (inventory.configuration || inventory.state)) { + throw new Error(`Source session branch ${JSON.stringify(name)} is missing branch.tip`); + } + if (inventory.configuration !== inventory.state) { + throw new Error(`Source session branch ${JSON.stringify(name)} has incomplete lane state`); + } + } + if (plan.scope === "branch" && !lanes.get(plan.branch)?.configuration) { + throw new Error(`Source branch ${JSON.stringify(plan.branch)} is not a configured AgentLane`); + } + + for (const stored of this.listValues.values()) { + for (const element of stored.elements) { + const projected = projectForkCurrentStateWrite( + { + kind: "list", + op: "append", + seq: element.seq, + namespace: stored.address.namespace, + key: stored.address.key, + value: element.value, + }, + plan, + isEntryCopied, + ); + if (projected !== undefined) destination.applyValueSetOrListAppend(projected); + } + } + destination.nextSeq = this.nextSeq; + return destination; + } + + private selectForkPlan(options: ForkOptions): MemoryForkPlan { + if (options.scope === "tree") return { scope: "tree" }; + + const sourceTip = this.getValue(branchTip(options.branch)); + if (sourceTip === undefined) throw new Error(`Unknown source branch: ${options.branch}`); + + const requested = options.entryId ?? sourceTip.value; + const entryIds = new Set(); + let found = requested === null; + let destinationTip: string | null = null; + let entryId = sourceTip.value; + while (entryId !== null) { + const entry = this.entries.get(entryId); + if (entry === undefined) throw new Error(`Corrupt source branch: missing parent ${entryId}`); + if (entry.id === requested) { + found = true; + destinationTip = options.position === "before" ? entry.parentId : entry.id; + if (options.position !== "before") entryIds.add(entry.id); + } else if (found) { + entryIds.add(entry.id); + } + entryId = entry.parentId; + } + if (!found) { + throw new Error(`Fork entry ${requested} is not on source branch ${JSON.stringify(options.branch)}`); + } + return { scope: "branch", branch: options.branch, destinationTip, entryIds }; + } + + private applyValueSetOrListAppend(write: CommittedValueSetWrite | CommittedListAppendWrite): void { + const key = physicalKey(write.namespace, write.key); + if (write.kind === "value") { + this.scalarValues.set(key, { + address: value(write.namespace, write.key), + value: write.value, + seq: write.seq, + }); + return; + } + + const element = { seq: write.seq, value: write.value }; + const stored = this.listValues.get(key); + if (stored === undefined) { + this.listValues.set(key, { + address: list(write.namespace, write.key), + elements: [element], + }); + } else { + stored.elements.push(element); + } + } + advanceNextSeq(nextSeq: number): void { if (!Number.isSafeInteger(nextSeq) || nextSeq < 1) { throw new Error(`Invalid storage sequence high-water mark: ${nextSeq}`); diff --git a/packages/agent/src/harness/session/index.ts b/packages/agent/src/harness/session/index.ts index 294403c87..5b9c99d58 100644 --- a/packages/agent/src/harness/session/index.ts +++ b/packages/agent/src/harness/session/index.ts @@ -11,7 +11,7 @@ export type { } from "./commit.ts"; export { commitWrite, insertEntry, insertUsage, prepareStorageCommit, validateCommittedWrites } from "./commit.ts"; export { createForkSnapshot, type ForkSourceSnapshot } from "./fork.ts"; -export { classifyForkAddress, type ForkDisposition } from "./fork-policy.ts"; +export { type ForkCurrentStatePlan, projectForkCurrentStateWrite } from "./fork-policy.ts"; export { JSONL_STORAGE_VERSION, type JsonlSessionCreateOptions, diff --git a/packages/agent/src/harness/session/memory.ts b/packages/agent/src/harness/session/memory.ts index ec836450a..119e70706 100644 --- a/packages/agent/src/harness/session/memory.ts +++ b/packages/agent/src/harness/session/memory.ts @@ -1,11 +1,5 @@ import { uuidv7 } from "@earendil-works/pi-ai/utils/uuid"; import type { Context } from "../context.ts"; -import { - createForkSnapshot, - type ForkDestinationSnapshot, - type ForkSourceSnapshot, - forkSnapshotWrites, -} from "./fork.ts"; import { InMemoryStorageState } from "./in-memory-storage-state.ts"; import { StorageBackedSession } from "./session.ts"; import type { @@ -114,10 +108,14 @@ export class MemoryStorage implements Storage { return Promise.resolve(this.storageState.getStats()); } - /** Capture the state needed to fork at one serialized boundary between commits. */ - captureForkSource(_context: Context): Promise { + /** Construct a destination storage at one serialized boundary between source commits. */ + fork(options: ForkOptions): Promise { if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed")); - const result = this.commitQueue.then(() => this.storageState.snapshotEntriesAndValues()); + const result = this.commitQueue.then(() => { + const destination = new MemoryStorage({ now: this.now }); + destination.storageState = this.storageState.createFork(options); + return destination; + }); this.commitQueue = result.then( () => undefined, () => undefined, @@ -133,14 +131,6 @@ export class MemoryStorage implements Storage { }); return this.closePromise; } - - static fromSnapshot(options: MemoryStorageOptions, snapshot: ForkDestinationSnapshot): MemoryStorage { - const storage = new MemoryStorage(options); - const writes = forkSnapshotWrites(snapshot); - storage.storageState.validateCommitted(writes); - storage.storageState.applyValidated(writes); - return storage; - } } const MEMORY_STORAGE_VERSION = 1; @@ -407,7 +397,7 @@ export class MemorySessionRepo implements SessionRepo { this.sessions.delete(metadata.id); } - async fork(source: SessionMetadata, options: ForkOptions, context: Context): Promise { + async fork(source: SessionMetadata, options: ForkOptions, _context: Context): Promise { this.assertOpen(); const sourceRecord = this.sessions.get(source.id); if (sourceRecord === undefined) throw new Error(`Unknown session: ${source.id}`); @@ -416,8 +406,7 @@ export class MemorySessionRepo implements SessionRepo { this.reserveId(id); try { - const snapshot = createForkSnapshot(await sourceRecord.storage.captureForkSource(context), options); - const storage = MemoryStorage.fromSnapshot({ now: this.now }, snapshot); + const storage = await sourceRecord.storage.fork(options); const metadata: SessionMetadata = { id, createdAt, diff --git a/packages/agent/test/harness/memory-session-repo.test.ts b/packages/agent/test/harness/memory-session-repo.test.ts index 1cd85a09c..ef465cf68 100644 --- a/packages/agent/test/harness/memory-session-repo.test.ts +++ b/packages/agent/test/harness/memory-session-repo.test.ts @@ -1,8 +1,6 @@ import { describe, expect, it } from "vitest"; import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts"; -import * as sessionWrites from "../../src/harness/session/commit.ts"; import { MemorySessionRepo } from "../../src/harness/session/index.ts"; -import * as storedValues from "../../src/harness/session/values.ts"; const NOW = 1_700_000_000_000; function uuidTimestamp(id: string): number { @@ -76,64 +74,4 @@ describe("MemorySessionRepo metadata", () => { const reopened = await repo.open(session.metadata, BACKGROUND_CONTEXT); await Promise.all([reopened.close(BACKGROUND_CONTEXT), repo.close(BACKGROUND_CONTEXT)]); }); - - it("captures fork options before waiting for its snapshot boundary", async () => { - const repo = new MemorySessionRepo({ now: () => NOW }); - const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT); - const rootId = "00000000-0000-7000-8000-000000000001"; - const childId = "00000000-0000-7000-8000-000000000002"; - const mutation = await source.beginMutation(BACKGROUND_CONTEXT); - const commit = mutation.commit( - [ - sessionWrites.insertEntry({ - id: rootId, - parentId: null, - type: "custom", - customType: "root", - }), - sessionWrites.insertEntry({ - id: childId, - parentId: rootId, - type: "custom", - customType: "child", - }), - storedValues.setValue(storedValues.branchTip("main"), childId), - storedValues.setValue(storedValues.laneConfig("main"), { - model: { provider: "test", modelId: "test" }, - thinkingLevel: "off", - activeToolNames: [], - }), - storedValues.setValue(storedValues.laneState("main"), { - currentOperationId: null, - lastOperationId: null, - inbox: [], - }), - ], - BACKGROUND_CONTEXT, - ); - const options = { - id: "fork", - scope: "branch" as const, - branch: "main", - entryId: childId, - position: "before" as "before" | "at", - }; - const fork = repo.fork(source.metadata, options, BACKGROUND_CONTEXT); - options.entryId = rootId; - options.position = "at"; - - await commit; - await mutation.end(BACKGROUND_CONTEXT); - await expect( - source.scanBranch({ start: childId, order: "oldestFirst" }, BACKGROUND_CONTEXT), - ).resolves.toMatchObject([{ id: rootId }, { id: childId }]); - const forked = await fork; - const forkedMain = await forked.branch("main", BACKGROUND_CONTEXT); - expect(await forkedMain?.getTipId(BACKGROUND_CONTEXT)).toBe(rootId); - await Promise.all([ - source.close(BACKGROUND_CONTEXT), - forked.close(BACKGROUND_CONTEXT), - repo.close(BACKGROUND_CONTEXT), - ]); - }); });