From fd8380f3239a9f50a21834faeb89e279b00e5e92 Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Fri, 14 Aug 2026 15:23:26 +0200 Subject: [PATCH] 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 | 140 ++++++++++++++++ .../agent/src/harness/session/jsonl/repo.ts | 47 +++++- .../src/harness/session/jsonl/storage.ts | 39 +++++ packages/agent/src/harness/session/memory.ts | 143 +--------------- .../testing/conformance/session-repo.ts | 156 ++++++++++-------- .../src/harness/session/testing/index.ts | 2 + .../jsonl-session-repo-conformance.test.ts | 4 + .../agent/test/harness/jsonl-storage.test.ts | 113 +++++++++++++ 8 files changed, 433 insertions(+), 211 deletions(-) create mode 100644 packages/agent/src/harness/session/fork.ts diff --git a/packages/agent/src/harness/session/fork.ts b/packages/agent/src/harness/session/fork.ts new file mode 100644 index 000000000..d6fa7ac0f --- /dev/null +++ b/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( + register: Register, + namespace: TNamespace, +): register is Register { + 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(); + 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, + sourceLeaves: Register<"lane.leaf">[], + options: ForkOptions, +): { entryIds: Set; laneToLeafId: Map } { + const entryIds = new Set(); + const laneToLeafId = new Map(); + 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, + 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); +} diff --git a/packages/agent/src/harness/session/jsonl/repo.ts b/packages/agent/src/harness/session/jsonl/repo.ts index 2ed56dbdb..143b23651 100644 --- a/packages/agent/src/harness/session/jsonl/repo.ts +++ b/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, - "create" | "open" | "list" | "delete" ->; - /** File-backed format-4 session repository lifecycle. */ -export class JsonlSessionRepo implements JsonlSessionLifecycleRepo { +export class JsonlSessionRepo + implements SessionRepo +{ 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> { + 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 { if (this.closePromise !== undefined) return this.closePromise; this.closed = true; diff --git a/packages/agent/src/harness/session/jsonl/storage.ts b/packages/agent/src/harness/session/jsonl/storage.ts index 0aa247afc..15a7864b6 100644 --- a/packages/agent/src/harness/session/jsonl/storage.ts +++ b/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; + +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 { 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 { + 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 { 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 { if (this.closePromise !== undefined) return this.closePromise; this.state = "closing"; diff --git a/packages/agent/src/harness/session/memory.ts b/packages/agent/src/harness/session/memory.ts index c2ad834bb..793f9ea5a 100644 --- a/packages/agent/src/harness/session/memory.ts +++ b/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( - register: Register, - namespace: TNamespace, -): register is Register { - 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(); - 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, - sourceLeaves: Register<"lane.leaf">[], - options: ForkOptions, - ): { entryIds: Set; laneToLeafId: Map } { - const entryIds = new Set(); - const laneToLeafId = new Map(); - 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, - 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); diff --git a/packages/agent/src/harness/session/testing/conformance/session-repo.ts b/packages/agent/src/harness/session/testing/conformance/session-repo.ts index c67f651c5..db89f5e04 100644 --- a/packages/agent/src/harness/session/testing/conformance/session-repo.ts +++ b/packages/agent/src/harness/session/testing/conformance/session-repo.ts @@ -231,45 +231,13 @@ export function createSessionRepoMessageConformance( +/** Creates fork-content cases that do not require concurrent repository coordination. */ +export function createSessionRepoForkBehaviorConformance( backendFactory: () => Promise, "create" | "list" | "fork">>, onClose?: () => void | Promise, ): 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 { - 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( + backendFactory: () => Promise, "create" | "fork">>, + onClose?: () => void | Promise, +): 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( + backendFactory: () => Promise, "create" | "list" | "fork">>, + onClose?: () => void | Promise, +): readonly ConformanceCase[] { + return [ + ...createSessionRepoForkBehaviorConformance(backendFactory, onClose), + ...createSessionRepoForkCoordinationConformance(backendFactory, onClose), ]; } diff --git a/packages/agent/src/harness/session/testing/index.ts b/packages/agent/src/harness/session/testing/index.ts index 8128bc879..839adc6fa 100644 --- a/packages/agent/src/harness/session/testing/index.ts +++ b/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, diff --git a/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts b/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts index 5b0e0d18a..7e54f89ed 100644 --- a/packages/agent/test/harness/jsonl-session-repo-conformance.test.ts +++ b/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(createConformanceRepo, () => jsonlRepo.close()), ...createSessionRepoMessageConformance(createConformanceRepo, () => jsonlRepo.close()), + ...createSessionRepoForkBehaviorConformance(createConformanceRepo, () => jsonlRepo.close()), ]); diff --git a/packages/agent/test/harness/jsonl-storage.test.ts b/packages/agent/test/harness/jsonl-storage.test.ts index 57e97222c..192f75f29 100644 --- a/packages/agent/test/harness/jsonl-storage.test.ts +++ b/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 { + 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 {