sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fix(agent): create JSONL sessions atomically

Browse source 
Stage the header and initial main-lane transaction together, then publish the complete session with an atomic rename.
This commit is contained in:
David Brailovsky 2026-08-20 17:31:56 +02:00
parent 7ebf9087e7
commit 0cad8a04d6
5 changed files with 81 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

6
packages/agent/src/harness/session/jsonl/repo.ts
View file

@ -85,9 +85,9 @@ export class JsonlSessionRepo
cwd,
...(options.parentSessionId === undefined ? {} : { parentSessionId: options.parentSessionId }),
};
// TODO: do we want to make session creation atomic?
storage = await JsonlStorage.create({ fileSystem: this.fileSystem, path, now: this.now }, header, context);
await storage.commit(
storage = await JsonlStorage.create(
{ fileSystem: this.fileSystem, path, now: this.now },
header,
[
setValue(laneLeaf("main"), null),
setValue(laneState("main"), { currentOperationId: null, pendingNextRun: [] }),

23
packages/agent/src/harness/session/jsonl/storage.ts
View file

@ -75,6 +75,10 @@ function parseTransaction(line: string): CommittedWrite[] {
return (Array.isArray(value) ? value : [value]).map(parseCommittedWrite);
}
function serializeTransaction(writes: readonly CommittedWrite[]): string {
return JSON.stringify(writes.length === 1 ? writes[0] : writes);
}
function splitCompleteLines(content: string): { lines: string[]; torn: boolean } {
if (content.endsWith("\n")) return { lines: content.slice(0, -1).split("\n"), torn: false };
const lastNewline = content.lastIndexOf("\n");
@ -135,13 +139,16 @@ export class JsonlStorage implements Storage {
static async create(
options: JsonlStorageOptions,
header: JsonlStorageHeader,
initialWrites: Write[],
context: Context,
): Promise<JsonlStorage> {
fileValue(
await options.fileSystem.writeFile(options.path, `${JSON.stringify(header)}\n`, context),
`Failed to create JSONL storage ${options.path}`,
);
return new JsonlStorage(options, header, { kind: "v4" });
const storage = new JsonlStorage(options, header, { kind: "v4" });
const prepared = storage.storageState.prepareCommit(initialWrites, storage.now());
const lines = [JSON.stringify(header)];
if (prepared.writes.length !== 0) lines.push(serializeTransaction(prepared.writes));
await publishFileAtomically(options.fileSystem, options.path, `${lines.join("\n")}\n`, context);
storage.storageState.applyValidated(prepared.writes);
return storage;
}
/** Atomically create storage from a complete prepared snapshot. */
@ -232,11 +239,7 @@ export class JsonlStorage implements Storage {
const prepared = this.storageState.prepareCommit(writes, this.now());
if (prepared.writes.length !== 0) {
fileValue(
await this.fileSystem.appendFile(
this.path,
`${JSON.stringify(prepared.writes.length === 1 ? prepared.writes[0] : prepared.writes)}\n`,
context,
),
await this.fileSystem.appendFile(this.path, `${serializeTransaction(prepared.writes)}\n`, context),
`Failed to append JSONL storage ${this.path}`,
);
}

61
packages/agent/test/harness/jsonl-session-repo.test.ts
View file

@ -1,5 +1,5 @@
import { describe, expect, it } from "vitest";
import { BACKGROUND_CONTEXT } from "../../src/harness/context.ts";
import { BACKGROUND_CONTEXT, type Context } from "../../src/harness/context.ts";
import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
import { JSONL_STORAGE_VERSION, JsonlSessionRepo } from "../../src/harness/session/jsonl/index.ts";
import { getOrThrow } from "../../src/harness/types.ts";
@ -7,6 +7,26 @@ import { createTempDir } from "./session-test-utils.ts";
const NOW = 1_700_000_000_000;
class AtomicPublicationNodeExecutionEnv extends NodeExecutionEnv {
publication:
| { sourcePath: string; destinationPath: string; destinationExisted: boolean; stagedContent: string }
| undefined;
override async renameFile(sourcePath: string, destinationPath: string, context: Context) {
const destinationExists = await super.exists(destinationPath, context);
const stagedContent = await super.readTextFile(sourcePath, context);
if (destinationExists.ok && stagedContent.ok) {
this.publication = {
sourcePath,
destinationPath,
destinationExisted: destinationExists.value,
stagedContent: stagedContent.value,
};
}
return super.renameFile(sourcePath, destinationPath, context);
}
}
describe("JsonlSessionRepo cwd-scoped lifecycle", () => {
it("persists metadata and filters discovery by cwd", async () => {
const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() });
@ -46,6 +66,45 @@ describe("JsonlSessionRepo cwd-scoped lifecycle", () => {
await repo.close(BACKGROUND_CONTEXT);
});
it("atomically publishes the header and initialized main lane", async () => {
const fileSystem = new AtomicPublicationNodeExecutionEnv({ cwd: createTempDir() });
const repo = new JsonlSessionRepo({ fileSystem, sessionsRoot: "sessions", now: () => NOW });
const session = await repo.create({ id: "session", cwd: "/workspace" }, BACKGROUND_CONTEXT);
const publication = fileSystem.publication;
if (publication === undefined) throw new Error("Expected atomic session publication");
expect(publication.destinationPath).toBe(session.metadata.path);
expect(publication.destinationExisted).toBe(false);
const lines = publication.stagedContent.trimEnd().split("\n");
expect(lines).toHaveLength(2);
expect(JSON.parse(lines[0]!)).toMatchObject({ kind: "header", id: "session" });
expect(JSON.parse(lines[1]!)).toEqual([
{
kind: "value",
op: "set",
seq: 1,
namespace: "pi.lane.leaf",
key: "main",
value: null,
},
{
kind: "value",
op: "set",
seq: 2,
namespace: "pi.lane.state",
key: "main",
value: { currentOperationId: null, pendingNextRun: [] },
},
]);
expect(getOrThrow(await fileSystem.readTextFile(session.metadata.path, BACKGROUND_CONTEXT))).toBe(
publication.stagedContent,
);
expect(getOrThrow(await fileSystem.exists(publication.sourcePath, BACKGROUND_CONTEXT))).toBe(false);
await session.close(BACKGROUND_CONTEXT);
await repo.close(BACKGROUND_CONTEXT);
});
it("keeps fork destinations claimed until close and rejects deleting open sessions", async () => {
const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() });
const repo = new JsonlSessionRepo({ fileSystem, sessionsRoot: "sessions", now: () => NOW });

1
packages/agent/test/harness/jsonl-storage-conformance.test.ts
View file

@ -37,6 +37,7 @@ registerConformance(
createdAt: NOW,
cwd: "/workspace",
},
[],
BACKGROUND_CONTEXT,
);
return {

8
packages/agent/test/harness/jsonl-storage.test.ts
View file

@ -125,7 +125,7 @@ describe("JsonlStorage persistence", () => {
const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() });
const options = { fileSystem, path: "list-delete.jsonl", now: () => NOW };
const events = storedValues.list<string>("test.events");
const storage = await JsonlStorage.create(options, header("list-delete"), BACKGROUND_CONTEXT);
const storage = await JsonlStorage.create(options, header("list-delete"), [], BACKGROUND_CONTEXT);
await storage.commit(
[storedValues.appendList(events, "first"), storedValues.appendList(events, "second")],
BACKGROUND_CONTEXT,
@ -144,7 +144,7 @@ describe("JsonlStorage persistence", () => {
it("writes one line per transaction and replays stamped state", async () => {
const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() });
const options = { fileSystem, path: "session.jsonl", now: () => NOW };
const storage = await JsonlStorage.create(options, header("round-trip"), BACKGROUND_CONTEXT);
const storage = await JsonlStorage.create(options, header("round-trip"), [], BACKGROUND_CONTEXT);
const committed = await storage.commit(
[
sessionWrites.insertEntry({
@ -218,7 +218,7 @@ 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"), BACKGROUND_CONTEXT);
const storage = await JsonlStorage.create(options, header("snapshot"), [], BACKGROUND_CONTEXT);
const firstCommit = storage.commit(
[
@ -265,7 +265,7 @@ describe("JsonlStorage torn tail", () => {
async function seed() {
const fileSystem = new NodeExecutionEnv({ cwd: createTempDir() });
const options = { fileSystem, path: "session.jsonl" as const, now: () => NOW };
const storage = await JsonlStorage.create(options, header("torn"), BACKGROUND_CONTEXT);
const storage = await JsonlStorage.create(options, header("torn"), [], BACKGROUND_CONTEXT);
await storage.commit([entryWrite("kept")], BACKGROUND_CONTEXT);
await storage.close(BACKGROUND_CONTEXT);
const prefix = getOrThrow(await fileSystem.readTextFile("session.jsonl", BACKGROUND_CONTEXT));