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 {