feat(agent): add JSONL session forking

Extract shared fork snapshot planning from the memory repository and use it to atomically publish JSONL branch and tree forks.

Split fork conformance into behavior and coordination cases so JSONL runs its currently supported behavior.

packages/agent/src/harness/session/fork.ts

@@ -0,0 +1,140 @@
+import type { Usage } from "@earendil-works/pi-ai";
+import type { StorageStateSnapshot } from "./storage-state.ts";
+import type { Entry, ForkOptions, Register, RegisterNamespace } from "./types.ts";
+
+export interface ForkSourceSnapshot {
+	entries: Entry[];
+	registers: Register[];
+}
+
+function isRegisterNamespace<TNamespace extends RegisterNamespace>(
+	register: Register,
+	namespace: TNamespace,
+): register is Register<TNamespace> {
+	return register.namespace === namespace;
+}
+
+function emptyUsage(): Usage {
+	return {
+		input: 0,
+		output: 0,
+		cacheRead: 0,
+		cacheWrite: 0,
+		totalTokens: 0,
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+	};
+}
+
+/** Build the complete logical state for a forked destination session. */
+export function createForkSnapshot(source: ForkSourceSnapshot, options: ForkOptions): StorageStateSnapshot {
+	const sourceEntries = new Map(source.entries.map((entry) => [entry.id, entry]));
+	const sourceLeaves = source.registers.filter((register) => isRegisterNamespace(register, "lane.leaf"));
+	validateForkSourceSnapshot(source, sourceEntries, sourceLeaves);
+
+	const { entryIds, laneToLeafId } = selectForkContents(sourceEntries, sourceLeaves, options);
+	const entries = new Map<string, Entry>();
+	for (const id of entryIds) entries.set(id, sourceEntries.get(id)!);
+
+	const registers: Register[] = [];
+	let nextSeq = Math.max(0, ...[...entries.values()].map((entry) => entry.seq)) + 1;
+	const setRegister = (namespace: RegisterNamespace, key: string, value: Register["value"]): void => {
+		registers.push({ namespace, key, value, seq: nextSeq++ } as Register);
+	};
+	for (const [lane, leaf] of laneToLeafId) {
+		const configuration = source.registers.find(
+			(register) => register.namespace === "lane.config" && register.key === lane,
+		);
+		if (configuration !== undefined) setRegister("lane.config", lane, configuration.value);
+		setRegister("lane.leaf", lane, leaf);
+		setRegister("lane.state", lane, { currentOperationId: null, pendingNextRun: [] });
+	}
+	for (const register of source.registers) {
+		if (
+			register.namespace === "fact.name" ||
+			register.namespace === "fact.custom" ||
+			(register.namespace === "fact.label" && entryIds.has(register.key))
+		) {
+			setRegister(register.namespace, register.key, register.value);
+		}
+	}
+
+	return {
+		entries,
+		registers,
+		usage: new Map(),
+		stats: {
+			messageCount: [...entries.values()].filter((entry) => entry.type === "message").length,
+			usage: emptyUsage(),
+		},
+		nextSeq,
+	};
+}
+
+function selectForkContents(
+	sourceEntries: Map<string, Entry>,
+	sourceLeaves: Register<"lane.leaf">[],
+	options: ForkOptions,
+): { entryIds: Set<string>; laneToLeafId: Map<string, string | null> } {
+	const entryIds = new Set<string>();
+	const laneToLeafId = new Map<string, string | null>();
+	if (options.scope === "tree") {
+		for (const id of sourceEntries.keys()) entryIds.add(id);
+		for (const register of sourceLeaves) laneToLeafId.set(register.key, register.value);
+	} else {
+		const mainLeaf = sourceLeaves.find((register) => register.key === "main");
+		if (mainLeaf === undefined) throw new Error("Source session is missing main lane");
+		const requested = options.entryId ?? mainLeaf.value;
+		let leaf = requested;
+		if (requested !== null) {
+			const target = sourceEntries.get(requested);
+			if (target === undefined) throw new Error(`Unknown fork entry: ${requested}`);
+			if (options.position === "before") leaf = target.parentId;
+		}
+
+		let entryId = leaf;
+		while (entryId !== null) {
+			const entry = sourceEntries.get(entryId);
+			if (entry === undefined) throw new Error(`Corrupt source branch: missing parent ${entryId}`);
+			entryIds.add(entryId);
+			entryId = entry.parentId;
+		}
+		laneToLeafId.set("main", leaf);
+	}
+	return { entryIds, laneToLeafId };
+}
+
+function validateForkSourceSnapshot(
+	source: ForkSourceSnapshot,
+	sourceEntries: Map<string, Entry>,
+	sourceLeaves: Register<"lane.leaf">[],
+): void {
+	const sourceLeafKeys = new Set(sourceLeaves.map((register) => register.key));
+
+	// TODO: do all these validations need to happen here? maybe somewhere else when the source session is loaded?
+	if (!sourceLeafKeys.has("main")) throw new Error("Source session is missing main lane");
+	for (const register of source.registers) {
+		if (
+			(register.namespace === "lane.config" ||
+				register.namespace === "lane.state" ||
+				register.namespace === "lane.lastResult") &&
+			!sourceLeafKeys.has(register.key)
+		) {
+			throw new Error(`Source session lane ${JSON.stringify(register.key)} is missing lane.leaf`);
+		}
+	}
+	for (const leaf of sourceLeaves) {
+		if (!hasRegister(source.registers, "lane.state", leaf.key)) {
+			throw new Error(`Source session lane ${JSON.stringify(leaf.key)} is missing lane.state`);
+		}
+		if (leaf.key !== "main" && !hasRegister(source.registers, "lane.config", leaf.key)) {
+			throw new Error(`Source session lane ${JSON.stringify(leaf.key)} is missing lane.config`);
+		}
+		if (leaf.value !== null && !sourceEntries.has(leaf.value)) {
+			throw new Error(`Source session lane ${JSON.stringify(leaf.key)} has an unknown leaf`);
+		}
+	}
+}
+
+function hasRegister(registers: Register[], namespace: RegisterNamespace, key: string): boolean {
+	return registers.some((register) => register.namespace === namespace && register.key === key);
+}

packages/agent/src/harness/session/jsonl/repo.ts

@@ -1,7 +1,8 @@
 import { uuidv7 } from "@earendil-works/pi-ai";
 import type { FileError, FileSystem, Result } from "../../types.ts";
+import { createForkSnapshot } from "../fork.ts";
 import { StorageBackedSession } from "../session.ts";
-import type { Session, SessionRepo } from "../types.ts";
+import type { ForkOptions, Session, SessionRepo } from "../types.ts";
 import { parseJsonlStorageHeader } from "./codec.ts";
 import { JsonlStorage } from "./storage.ts";
 import {
@@ -43,13 +44,10 @@ function sessionFileName(createdAt: number, id: string): string {
 	return `${timestamp}_${encodeURIComponent(id)}.jsonl`;
 }
 
-type JsonlSessionLifecycleRepo = Pick<
-	SessionRepo<JsonlSessionMetadata, JsonlSessionCreateOptions, JsonlSessionListOptions>,
-	"create" | "open" | "list" | "delete"
->;
-
 /** File-backed format-4 session repository lifecycle. */
-export class JsonlSessionRepo implements JsonlSessionLifecycleRepo {
+export class JsonlSessionRepo
+	implements SessionRepo<JsonlSessionMetadata, JsonlSessionCreateOptions, JsonlSessionListOptions>
+{
 	private readonly fileSystem: FileSystem;
 	private readonly sessionsRootInput: string;
 	private readonly now: () => number;
@@ -141,6 +139,41 @@ export class JsonlSessionRepo implements JsonlSessionLifecycleRepo {
 		fileValue(await this.fileSystem.remove(metadata.path), `Failed to delete session ${metadata.path}`);
 	}
 
+	async fork(source: JsonlSessionMetadata, options: ForkOptions): Promise<Session<JsonlSessionMetadata>> {
+		this.assertOpen();
+		const createdAt = this.now();
+		const id = options.id ?? uuidv7(createdAt);
+		if (await this.sessionIdExists(source.cwd, id)) throw new Error(`Session already exists: ${id}`);
+
+		const sourceStorage = await this.loadStorage(source);
+		const sourceSnapshot = await sourceStorage.snapshot().finally(() => sourceStorage.close());
+		const snapshot = createForkSnapshot(sourceSnapshot, options);
+		const path = await this.createPath(source.cwd, createdAt, id);
+		const header: JsonlStorageHeader = {
+			v: JSONL_FORMAT_VERSION,
+			kind: "header",
+			id,
+			storageVersion: JSONL_STORAGE_VERSION,
+			createdAt,
+			cwd: source.cwd,
+			parentSessionId: source.id,
+		};
+		let storage: JsonlStorage | undefined;
+		try {
+			storage = await JsonlStorage.createFromSnapshot(
+				{ fileSystem: this.fileSystem, path, now: this.now },
+				header,
+				snapshot,
+			);
+			const info = fileValue(await this.fileSystem.fileInfo(path), `Failed to read session ${path}`);
+			return new StorageBackedSession(metadataFromHeader(header, path, info.mtimeMs), storage);
+		} catch (error) {
+			await storage?.close().catch(() => undefined);
+			await this.fileSystem.remove(path, { force: true });
+			throw error;
+		}
+	}
+
 	close(): Promise<void> {
 		if (this.closePromise !== undefined) return this.closePromise;
 		this.closed = true;

packages/agent/src/harness/session/jsonl/storage.ts

@@ -6,6 +6,7 @@ import {
 	type CommittedUsageWrite,
 	type CommittedWrite,
 	StorageState,
+	type StorageStateSnapshot,
 } from "../storage-state.ts";
 import type {
 	CommitResult,
@@ -72,6 +73,15 @@ function splitCompleteLines(content: string): { lines: string[]; torn: boolean }
 	return { lines: content.slice(0, lastNewline).split("\n"), torn: true };
 }
 
+type JsonlSnapshotContents = Pick<StorageStateSnapshot, "entries" | "registers">;
+
+function snapshotWrites(snapshot: JsonlSnapshotContents): CommittedWrite[] {
+	const writes: CommittedWrite[] = [];
+	for (const entry of snapshot.entries.values()) writes.push({ kind: "entry", ...entry });
+	for (const register of snapshot.registers) writes.push({ kind: "register", op: "set", ...register });
+	return writes.sort((left, right) => left.seq - right.seq);
+}
+
 async function publishFileAtomically(fileSystem: FileSystem, destinationPath: string, content: string): Promise<void> {
 	const tempPath = `${destinationPath}.tmp`;
 	try {
@@ -112,6 +122,18 @@ export class JsonlStorage implements Storage {
 		return new JsonlStorage(options, header);
 	}
 
+	/** Atomically create storage from a complete prepared snapshot. */
+	static async createFromSnapshot(
+		options: JsonlStorageOptions,
+		header: JsonlStorageHeader,
+		snapshot: JsonlSnapshotContents,
+	): Promise<JsonlStorage> {
+		const writes = snapshotWrites(snapshot);
+		const content = `${[JSON.stringify(header), ...writes.map((write) => JSON.stringify(write))].join("\n")}\n`;
+		await publishFileAtomically(options.fileSystem, options.path, content);
+		return JsonlStorage.open(options);
+	}
+
 	static async open(options: JsonlStorageOptions): Promise<JsonlStorage> {
 		const content = fileValue(
 			await options.fileSystem.readTextFile(options.path),
@@ -206,6 +228,23 @@ export class JsonlStorage implements Storage {
 		return Promise.resolve(this.storageState.getStats());
 	}
 
+	/** Capture the current entries and registers at one serialized boundary between commits. */
+	snapshot(): Promise<{ entries: Entry[]; registers: Register[] }> {
+		if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed"));
+		const result = this.commitQueue.then(() => {
+			const snapshot = this.storageState.snapshot();
+			return {
+				entries: [...snapshot.entries.values()].sort((left, right) => left.seq - right.seq),
+				registers: snapshot.registers,
+			};
+		});
+		this.commitQueue = result.then(
+			() => undefined,
+			() => undefined,
+		);
+		return result;
+	}
+
 	close(): Promise<void> {
 		if (this.closePromise !== undefined) return this.closePromise;
 		this.state = "closing";

packages/agent/src/harness/session/memory.ts

@@ -1,5 +1,6 @@
-import { type Usage, uuidv7 } from "@earendil-works/pi-ai";
+import { uuidv7 } from "@earendil-works/pi-ai";
 import type { AgentMessage } from "../../types.ts";
+import { createForkSnapshot } from "./fork.ts";
 import { StorageBackedSession } from "./session.ts";
 import { StorageState, type StorageStateSnapshot } from "./storage-state.ts";
 import type {
@@ -37,24 +38,6 @@ export interface MemorySessionRepoOptions {
 	now?: () => number;
 }
 
-function isRegisterNamespace<TNamespace extends RegisterNamespace>(
-	register: Register,
-	namespace: TNamespace,
-): register is Register<TNamespace> {
-	return register.namespace === namespace;
-}
-
-function emptyUsage(): Usage {
-	return {
-		input: 0,
-		output: 0,
-		cacheRead: 0,
-		cacheWrite: 0,
-		totalTokens: 0,
-		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-	};
-}
-
 export class MemoryStorage implements Storage {
 	private readonly now: () => number;
 	private storageState = new StorageState();
@@ -406,8 +389,8 @@ export class MemorySessionRepo implements SessionRepo {
 		this.reserveId(id);
 
 		try {
-			const snapshot = await sourceRecord.storage.snapshot();
-			const storage = this.createForkStorage(snapshot, options);
+			const snapshot = createForkSnapshot(await sourceRecord.storage.snapshot(), options);
+			const storage = MemoryStorage.fromSnapshot({ now: this.now }, snapshot);
 			const metadata: SessionMetadata = {
 				id,
 				createdAt,
@@ -438,124 +421,6 @@ export class MemorySessionRepo implements SessionRepo {
 		});
 	}
 
-	private createForkStorage(
-		snapshot: { entries: Entry[]; registers: Register[] },
-		options: ForkOptions,
-	): MemoryStorage {
-		const sourceEntries = new Map(snapshot.entries.map((entry) => [entry.id, entry]));
-		const sourceLeaves = snapshot.registers.filter((register) => isRegisterNamespace(register, "lane.leaf"));
-		this.validateForkSourceSnapshot(snapshot, sourceEntries, sourceLeaves);
-
-		const { entryIds, laneToLeafId } = this.selectForkContents(sourceEntries, sourceLeaves, options);
-
-		const entries = new Map<string, Entry>();
-		for (const id of entryIds) entries.set(id, sourceEntries.get(id)!);
-		const registers: Register[] = [];
-		let nextSeq = Math.max(0, ...[...entries.values()].map((entry) => entry.seq)) + 1;
-		const setRegister = (namespace: RegisterNamespace, key: string, value: Register["value"]): void => {
-			registers.push({ namespace, key, value, seq: nextSeq++ } as Register);
-		};
-		for (const [lane, leaf] of laneToLeafId) {
-			const configuration = snapshot.registers.find(
-				(register) => register.namespace === "lane.config" && register.key === lane,
-			);
-			if (configuration !== undefined) setRegister("lane.config", lane, configuration.value);
-			setRegister("lane.leaf", lane, leaf);
-			setRegister("lane.state", lane, { currentOperationId: null, pendingNextRun: [] });
-		}
-		for (const register of snapshot.registers) {
-			if (
-				register.namespace === "fact.name" ||
-				register.namespace === "fact.custom" ||
-				(register.namespace === "fact.label" && entryIds.has(register.key))
-			) {
-				setRegister(register.namespace, register.key, register.value);
-			}
-		}
-		return MemoryStorage.fromSnapshot(
-			{ now: this.now },
-			{
-				entries,
-				registers,
-				usage: new Map(),
-				stats: {
-					messageCount: [...entries.values()].filter((entry) => entry.type === "message").length,
-					usage: emptyUsage(),
-				},
-				nextSeq,
-			},
-		);
-	}
-
-	private selectForkContents(
-		sourceEntries: Map<string, Entry>,
-		sourceLeaves: Register<"lane.leaf">[],
-		options: ForkOptions,
-	): { entryIds: Set<string>; laneToLeafId: Map<string, string | null> } {
-		const entryIds = new Set<string>();
-		const laneToLeafId = new Map<string, string | null>();
-		if (options.scope === "tree") {
-			for (const id of sourceEntries.keys()) entryIds.add(id);
-			for (const register of sourceLeaves) laneToLeafId.set(register.key, register.value);
-		} else {
-			const mainLeaf = sourceLeaves.find((register) => register.key === "main");
-			if (mainLeaf === undefined) throw new Error("Source session is missing main lane");
-			const requested = options.entryId ?? mainLeaf.value;
-			let leaf = requested;
-			if (requested !== null) {
-				const target = sourceEntries.get(requested);
-				if (target === undefined) throw new Error(`Unknown fork entry: ${requested}`);
-				if (options.position === "before") leaf = target.parentId;
-			}
-
-			let entryId = leaf;
-			while (entryId !== null) {
-				const entry = sourceEntries.get(entryId);
-				if (entry === undefined) throw new Error(`Corrupt source branch: missing parent ${entryId}`);
-				entryIds.add(entryId);
-				entryId = entry.parentId;
-			}
-			laneToLeafId.set("main", leaf);
-		}
-		return { entryIds, laneToLeafId };
-	}
-
-	private validateForkSourceSnapshot(
-		snapshot: { entries: Entry[]; registers: Register[] },
-		sourceEntries: Map<string, Entry>,
-		sourceLeaves: Register<"lane.leaf">[],
-	): void {
-		const sourceLeafKeys = new Set(sourceLeaves.map((register) => register.key));
-
-		// TODO: do all these validations need to happen here? maybe somewhere else when the source session is loaded?
-		if (!sourceLeafKeys.has("main")) throw new Error("Source session is missing main lane");
-		for (const register of snapshot.registers) {
-			if (
-				(register.namespace === "lane.config" ||
-					register.namespace === "lane.state" ||
-					register.namespace === "lane.lastResult") &&
-				!sourceLeafKeys.has(register.key)
-			) {
-				throw new Error(`Source session lane ${JSON.stringify(register.key)} is missing lane.leaf`);
-			}
-		}
-		for (const leaf of sourceLeaves) {
-			if (!this.hasRegister(snapshot.registers, "lane.state", leaf.key)) {
-				throw new Error(`Source session lane ${JSON.stringify(leaf.key)} is missing lane.state`);
-			}
-			if (leaf.key !== "main" && !this.hasRegister(snapshot.registers, "lane.config", leaf.key)) {
-				throw new Error(`Source session lane ${JSON.stringify(leaf.key)} is missing lane.config`);
-			}
-			if (leaf.value !== null && !sourceEntries.has(leaf.value)) {
-				throw new Error(`Source session lane ${JSON.stringify(leaf.key)} has an unknown leaf`);
-			}
-		}
-	}
-
-	private hasRegister(registers: Register[], namespace: RegisterNamespace, key: string): boolean {
-		return registers.some((register) => register.namespace === namespace && register.key === key);
-	}
-
 	private reserveId(id: string): void {
 		if (this.sessions.has(id) || this.pendingIds.has(id)) throw new Error(`Session already exists: ${id}`);
 		this.pendingIds.add(id);

packages/agent/src/harness/session/testing/conformance/session-repo.ts

@@ -231,45 +231,13 @@ export function createSessionRepoMessageConformance<TMetadata extends SessionMet
 	];
 }
 
-/** Creates fork cases for repositories that support creation, discovery, and forks. */
-export function createSessionRepoForkConformance<TMetadata extends SessionMetadata>(
+/** Creates fork-content cases that do not require concurrent repository coordination. */
+export function createSessionRepoForkBehaviorConformance<TMetadata extends SessionMetadata>(
 	backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create" | "list" | "fork">>,
 	onClose?: () => void | Promise<void>,
 ): readonly ConformanceCase[] {
 	const factory = prepareRepoCaseFactory(backendFactory, onClose);
 	return [
-		createCase(
-			factory,
-			"forks",
-			"publishes create when it reserves a shared destination id first",
-			async ({ repo }) => {
-				const source = await repo.create({ id: "source" });
-				const results = await Promise.allSettled([
-					repo.create({ id: "destination" }),
-					repo.fork(source.metadata, { id: "destination" }),
-				]);
-				strictEqual(results[0].status, "fulfilled");
-				strictEqual(results[1].status, "rejected");
-				if (results[0].status === "fulfilled") await results[0].value.close();
-				await source.close();
-			},
-		),
-		createCase(
-			factory,
-			"forks",
-			"publishes fork when it reserves a shared destination id first",
-			async ({ repo }) => {
-				const source = await repo.create({ id: "source" });
-				const results = await Promise.allSettled([
-					repo.fork(source.metadata, { id: "destination" }),
-					repo.create({ id: "destination" }),
-				]);
-				strictEqual(results[0].status, "fulfilled");
-				strictEqual(results[1].status, "rejected");
-				if (results[0].status === "fulfilled") await results[0].value.close();
-				await source.close();
-			},
-		),
 		createCase(factory, "forks", "forks a fresh session before first attachment", async ({ repo }) => {
 			const source = await repo.create({ id: "source" });
 			const fork = await repo.fork(source.metadata, { id: "fork" });
@@ -505,38 +473,96 @@ export function createSessionRepoForkConformance<TMetadata extends SessionMetada
 			deepStrictEqual((await fork.getRegister("lane.state", "review"))?.value, idleLaneState);
 			await Promise.all([source.close(), fork.close()]);
 		}),
-		createCase(factory, "forks", "captures one coherent boundary between source commits", async ({ repo }) => {
-			const source = await repo.create({ id: "source" });
-			const firstCommit = source.mutate("main", (mutator) =>
-				mutator.commit({
-					writes: [
-						{ kind: "entry", entry: { id: ROOT_ID, parentId: null, type: "custom", customType: "first" } },
-						{ kind: "register", op: "set", namespace: "lane.leaf", key: "main", value: ROOT_ID },
-					],
-				}),
-			);
-			const fork = repo.fork(source.metadata, { id: "fork" });
-			const secondCommit = source.mutate("main", (mutator) =>
-				mutator.commit({
-					writes: [
-						{
-							kind: "entry",
-							entry: { id: CHILD_ID, parentId: ROOT_ID, type: "custom", customType: "second" },
-						},
-						{ kind: "register", op: "set", namespace: "lane.leaf", key: "main", value: CHILD_ID },
-					],
-				}),
-			);
+	];
+}
 
-			const [, forked] = await Promise.all([firstCommit, fork]);
-			await secondCommit;
-			strictEqual(await forked.getLeafId(), ROOT_ID);
-			deepStrictEqual(
-				(await forked.findEntries({ order: "asc" })).map(({ id }) => id),
-				[ROOT_ID],
-			);
-			await Promise.all([source.close(), forked.close()]);
-		}),
+/** Creates fork cases that require destination reservation or active-source snapshot coordination. */
+export function createSessionRepoForkCoordinationConformance<TMetadata extends SessionMetadata>(
+	backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create" | "fork">>,
+	onClose?: () => void | Promise<void>,
+): readonly ConformanceCase[] {
+	const factory = prepareRepoCaseFactory(backendFactory, onClose);
+	return [
+		createCase(
+			factory,
+			"fork coordination",
+			"publishes create when it reserves a shared destination id first",
+			async ({ repo }) => {
+				const source = await repo.create({ id: "source" });
+				const results = await Promise.allSettled([
+					repo.create({ id: "destination" }),
+					repo.fork(source.metadata, { id: "destination" }),
+				]);
+				strictEqual(results[0].status, "fulfilled");
+				strictEqual(results[1].status, "rejected");
+				if (results[0].status === "fulfilled") await results[0].value.close();
+				await source.close();
+			},
+		),
+		createCase(
+			factory,
+			"fork coordination",
+			"publishes fork when it reserves a shared destination id first",
+			async ({ repo }) => {
+				const source = await repo.create({ id: "source" });
+				const results = await Promise.allSettled([
+					repo.fork(source.metadata, { id: "destination" }),
+					repo.create({ id: "destination" }),
+				]);
+				strictEqual(results[0].status, "fulfilled");
+				strictEqual(results[1].status, "rejected");
+				if (results[0].status === "fulfilled") await results[0].value.close();
+				await source.close();
+			},
+		),
+		createCase(
+			factory,
+			"fork coordination",
+			"captures one coherent boundary between source commits",
+			async ({ repo }) => {
+				const source = await repo.create({ id: "source" });
+				const firstCommit = source.mutate("main", (mutator) =>
+					mutator.commit({
+						writes: [
+							{ kind: "entry", entry: { id: ROOT_ID, parentId: null, type: "custom", customType: "first" } },
+							{ kind: "register", op: "set", namespace: "lane.leaf", key: "main", value: ROOT_ID },
+						],
+					}),
+				);
+				const fork = repo.fork(source.metadata, { id: "fork" });
+				const secondCommit = source.mutate("main", (mutator) =>
+					mutator.commit({
+						writes: [
+							{
+								kind: "entry",
+								entry: { id: CHILD_ID, parentId: ROOT_ID, type: "custom", customType: "second" },
+							},
+							{ kind: "register", op: "set", namespace: "lane.leaf", key: "main", value: CHILD_ID },
+						],
+					}),
+				);
+
+				const [, forked] = await Promise.all([firstCommit, fork]);
+				await secondCommit;
+				strictEqual(await forked.getLeafId(), ROOT_ID);
+				deepStrictEqual(
+					(await forked.findEntries({ order: "asc" })).map(({ id }) => id),
+					[ROOT_ID],
+				);
+				await Promise.all([source.close(), forked.close()]);
+			},
+		),
+	];
+}
+
+/** Creates every fork conformance case. */
+export function createSessionRepoForkConformance<TMetadata extends SessionMetadata>(
+	backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create" | "list" | "fork">>,
+	onClose?: () => void | Promise<void>,
+): readonly ConformanceCase[] {
+	return [
+		...createSessionRepoForkBehaviorConformance(backendFactory, onClose),
+		...createSessionRepoForkCoordinationConformance(backendFactory, onClose),
 	];
 }
 

packages/agent/src/harness/session/testing/index.ts

@@ -18,7 +18,9 @@ export {
 } from "./benchmark/storage.ts";
 export {
 	createSessionRepoConformance,
+	createSessionRepoForkBehaviorConformance,
 	createSessionRepoForkConformance,
+	createSessionRepoForkCoordinationConformance,
 	createSessionRepoLifecycleConformance,
 	createSessionRepoMessageConformance,
 	createSessionRepoOwnershipConformance,

packages/agent/test/harness/jsonl-session-repo-conformance.test.ts

@@ -3,9 +3,11 @@ import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
 import { type JsonlSessionMetadata, JsonlSessionRepo } from "../../src/harness/session/jsonl/index.ts";
 import {
 	type ConformanceCase,
+	createSessionRepoForkBehaviorConformance,
 	createSessionRepoLifecycleConformance,
 	createSessionRepoMessageConformance,
 } from "../../src/harness/session/testing/index.ts";
+import type { ForkOptions } from "../../src/harness/session/types.ts";
 import { createTempDir } from "./session-test-utils.ts";
 
 const NOW = 1_700_000_000_000;
@@ -33,10 +35,12 @@ async function createConformanceRepo() {
 		open: (metadata: JsonlSessionMetadata) => jsonlRepo.open(metadata),
 		list: () => jsonlRepo.list({ cwd: CONFORMANCE_CWD }),
 		delete: (metadata: JsonlSessionMetadata) => jsonlRepo.delete(metadata),
+		fork: (source: JsonlSessionMetadata, options: ForkOptions) => jsonlRepo.fork(source, options),
 	};
 }
 
 registerConformance("JsonlSessionRepo conformance", [
 	...createSessionRepoLifecycleConformance<JsonlSessionMetadata>(createConformanceRepo, () => jsonlRepo.close()),
 	...createSessionRepoMessageConformance<JsonlSessionMetadata>(createConformanceRepo, () => jsonlRepo.close()),
+	...createSessionRepoForkBehaviorConformance<JsonlSessionMetadata>(createConformanceRepo, () => jsonlRepo.close()),
 ]);

packages/agent/test/harness/jsonl-storage.test.ts

@@ -1,6 +1,7 @@
 import { describe, expect, it } from "vitest";
 import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
 import { JSONL_FORMAT_VERSION, JsonlStorage, type JsonlStorageHeader } from "../../src/harness/session/jsonl/index.ts";
+import type { StorageStateSnapshot } from "../../src/harness/session/storage-state.ts";
 import { getOrThrow } from "../../src/harness/types.ts";
 import { createTempDir } from "./session-test-utils.ts";
 
@@ -17,6 +18,84 @@ function header(id: string): JsonlStorageHeader {
 	};
 }
 
+function preparedSnapshot(): Pick<StorageStateSnapshot, "entries" | "registers"> {
+	return {
+		entries: new Map([
+			[
+				"root",
+				{
+					id: "root",
+					parentId: null,
+					type: "message",
+					message: { role: "user", content: "hello", timestamp: 1 },
+					seq: 1,
+					timestamp: NOW,
+				},
+			],
+		]),
+		registers: [
+			{ namespace: "lane.leaf", key: "main", value: "root", seq: 2 },
+			{ namespace: "lane.state", key: "main", value: { currentOperationId: null, pendingNextRun: [] }, seq: 3 },
+			{ namespace: "fact.name", key: "", value: "forked", seq: 4 },
+		],
+	};
+}
+
+describe("JsonlStorage snapshot creation", () => {
+	it("atomically publishes and opens a prepared snapshot", async () => {
+		const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() });
+		const options = { fileSystem, path: "fork.jsonl", now: () => NOW };
+		const storage = await JsonlStorage.createFromSnapshot(options, header("fork"), preparedSnapshot());
+
+		const lines = getOrThrow(await fileSystem.readTextFile("fork.jsonl"))
+			.trimEnd()
+			.split("\n");
+		expect(lines.slice(1).every((line) => !Array.isArray(JSON.parse(line)))).toBe(true);
+		expect((await storage.getEntries(["root"])).get("root")?.timestamp).toBe(NOW);
+		expect((await storage.getRegister("lane.leaf", "main"))?.value).toBe("root");
+		expect((await storage.getRegister("lane.state", "main"))?.value).toEqual({
+			currentOperationId: null,
+			pendingNextRun: [],
+		});
+		expect((await storage.getRegister("fact.name", ""))?.value).toBe("forked");
+		expect(await storage.getStats()).toEqual({
+			messageCount: 1,
+			usage: {
+				input: 0,
+				output: 0,
+				cacheRead: 0,
+				cacheWrite: 0,
+				totalTokens: 0,
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+			},
+		});
+		expect(await storage.scanUsage({ order: "asc" })).toEqual([]);
+		expect((await storage.commit({ writes: [] })).firstSeq).toBe(5);
+		expect(getOrThrow(await fileSystem.exists("fork.jsonl.tmp"))).toBe(false);
+		await storage.close();
+
+		const reopened = await JsonlStorage.open(options);
+		expect((await reopened.getEntries(["root"])).has("root")).toBe(true);
+		expect((await reopened.getRegister("fact.name", ""))?.value).toBe("forked");
+		await reopened.close();
+	});
+
+	it("removes the temporary file when publication fails", async () => {
+		const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() });
+		getOrThrow(await fileSystem.createDir("blocked.jsonl"));
+
+		await expect(
+			JsonlStorage.createFromSnapshot(
+				{ fileSystem, path: "blocked.jsonl", now: () => NOW },
+				header("blocked"),
+				preparedSnapshot(),
+			),
+		).rejects.toThrow("Failed to publish JSONL storage");
+		expect(getOrThrow(await fileSystem.exists("blocked.jsonl.tmp"))).toBe(false);
+		expect(getOrThrow(await fileSystem.fileInfo("blocked.jsonl"))).toMatchObject({ kind: "directory" });
+	});
+});
+
 describe("JsonlStorage persistence", () => {
 	it("writes one line per transaction and replays stamped state", async () => {
 		const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() });
@@ -99,6 +178,40 @@ describe("JsonlStorage persistence", () => {
 	});
 });
 
+describe("JsonlStorage snapshots", () => {
+	it("captures one serialized boundary between commits", async () => {
+		const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() });
+		const options = { fileSystem, path: "session.jsonl", now: () => NOW };
+		const storage = await JsonlStorage.create(options, header("snapshot"));
+
+		const firstCommit = storage.commit({
+			writes: [
+				{ kind: "entry", entry: { id: "root", parentId: null, type: "custom", customType: "root" } },
+				{ kind: "register", op: "set", namespace: "lane.leaf", key: "main", value: "root" },
+			],
+		});
+		const snapshot = storage.snapshot();
+		const secondCommit = storage.commit({
+			writes: [
+				{ kind: "entry", entry: { id: "child", parentId: "root", type: "custom", customType: "child" } },
+				{ kind: "register", op: "set", namespace: "lane.leaf", key: "main", value: "child" },
+			],
+		});
+
+		await firstCommit;
+		const captured = await snapshot;
+		await secondCommit;
+
+		expect(captured.entries.map(({ id }) => id)).toEqual(["root"]);
+		expect(captured.registers.find(({ namespace, key }) => namespace === "lane.leaf" && key === "main")?.value).toBe(
+			"root",
+		);
+		expect((await storage.getEntries(["child"])).has("child")).toBe(true);
+		expect((await storage.getRegister("lane.leaf", "main"))?.value).toBe("child");
+		await storage.close();
+	});
+});
+
 describe("JsonlStorage torn tail", () => {
 	function entryWrite(id: string) {
 		return {