refactor(agent): tighten harness session storage

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

@@ -1,7 +1,8 @@
-import { randomUUID } from "node:crypto";
+import { createReadStream } from "node:fs";
 import { appendFile, mkdir, readFile, writeFile } from "node:fs/promises";
 import { dirname, resolve } from "node:path";
-import type { CodingAgentSessionInfo, SessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js";
+import { createInterface } from "node:readline";
+import type { JsonlSessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js";
 
 interface SessionHeader {
 	type: "session";
@@ -12,116 +13,140 @@ interface SessionHeader {
 	parentSession?: string;
 }
 
-function headerToSessionInfo(header: SessionHeader, filePath?: string): CodingAgentSessionInfo {
+function headerToSessionInfo(header: SessionHeader, path: string): JsonlSessionInfo {
 	return {
 		id: header.id,
 		createdAt: header.timestamp,
-		parentSession: header.parentSession,
-		projectCwd: header.cwd,
-		filePath,
+		cwd: header.cwd,
+		path,
+		parentSessionPath: header.parentSession,
 	};
 }
 
-async function loadJsonlStorage(
-	filePath: string,
-): Promise<{ header?: SessionHeader; entries: SessionTreeEntry[]; leafId: string | null }> {
+export async function loadJsonlSessionInfo(filePath: string): Promise<JsonlSessionInfo> {
+	const stream = createReadStream(filePath, { encoding: "utf8" });
+	const lines = createInterface({ input: stream, crlfDelay: Infinity });
 	try {
-		const content = await readFile(filePath, "utf8");
-		const entries: SessionTreeEntry[] = [];
-		let header: SessionHeader | undefined;
-		let leafId: string | null = null;
-		for (const line of content.split("\n")) {
-			if (!line.trim()) continue;
+		for await (const line of lines) {
+			if (!line.trim()) break;
 			try {
-				const record = JSON.parse(line) as SessionHeader | SessionTreeEntry;
-				if (record.type === "session") {
-					header = record as SessionHeader;
-					continue;
-				}
-				entries.push(record as SessionTreeEntry);
-				leafId = (record as SessionTreeEntry).id;
+				const header = JSON.parse(line) as SessionHeader;
+				return headerToSessionInfo(header, resolve(filePath));
 			} catch {
-				// ignore malformed lines
+				throw new Error(`Invalid JSONL session file ${filePath}: first line is not a valid session header`);
 			}
 		}
-		return { header, entries, leafId };
-	} catch {
-		return { entries: [], leafId: null };
+		throw new Error(`Invalid JSONL session file ${filePath}: missing session header`);
+	} finally {
+		lines.close();
+		stream.destroy();
 	}
 }
 
-export class JsonlSessionTreeStorage implements SessionTreeStorage {
-	private filePath: string;
-	private cwd: string;
-	private headerInitialized = false;
-	private cacheLoaded = false;
-	private sessionInfo?: CodingAgentSessionInfo;
-	private entries: SessionTreeEntry[] = [];
-	private byId = new Map<string, SessionTreeEntry>();
-	private currentLeafId: string | null = null;
-	private requestedSessionId?: string;
-	private parentSession?: string;
+async function loadJsonlStorage(filePath: string): Promise<{
+	header: SessionHeader;
+	entries: SessionTreeEntry[];
+	leafId: string | null;
+}> {
+	const content = await readFile(filePath, "utf8");
+	const lines = content.split("\n").filter((line) => line.trim());
+	if (lines.length === 0) {
+		throw new Error(`Invalid JSONL session file ${filePath}: missing session header`);
+	}
 
-	constructor(filePath: string, options: { cwd: string; sessionId?: string; parentSession?: string }) {
+	let header: SessionHeader;
+	try {
+		header = JSON.parse(lines[0]!) as SessionHeader;
+	} catch {
+		throw new Error(`Invalid JSONL session file ${filePath}: first line is not a valid session header`);
+	}
+
+	const entries: SessionTreeEntry[] = [];
+	let leafId: string | null = null;
+	for (const line of lines.slice(1)) {
+		try {
+			const entry = JSON.parse(line) as SessionTreeEntry;
+			entries.push(entry);
+			leafId = entry.id;
+		} catch {
+			// ignore malformed entry lines
+		}
+	}
+	return { header, entries, leafId };
+}
+
+export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionInfo> {
+	private readonly filePath: string;
+	private readonly header: SessionHeader;
+	private readonly sessionInfo: JsonlSessionInfo;
+	private entries: SessionTreeEntry[];
+	private byId: Map<string, SessionTreeEntry>;
+	private currentLeafId: string | null;
+	private headerWritten: boolean;
+
+	private constructor(
+		filePath: string,
+		header: SessionHeader,
+		entries: SessionTreeEntry[],
+		leafId: string | null,
+		headerWritten: boolean,
+	) {
 		this.filePath = resolve(filePath);
-		this.cwd = options.cwd;
-		this.requestedSessionId = options.sessionId;
-		this.parentSession = options.parentSession;
+		this.header = header;
+		this.sessionInfo = headerToSessionInfo(header, this.filePath);
+		this.entries = entries;
+		this.byId = new Map(entries.map((entry) => [entry.id, entry]));
+		this.currentLeafId = leafId;
+		this.headerWritten = headerWritten;
 	}
 
-	private async ensureParentDir(): Promise<void> {
-		await mkdir(dirname(this.filePath), { recursive: true });
+	static async open(filePath: string): Promise<JsonlSessionTreeStorage> {
+		const resolvedPath = resolve(filePath);
+		const loaded = await loadJsonlStorage(resolvedPath);
+		return new JsonlSessionTreeStorage(resolvedPath, loaded.header, loaded.entries, loaded.leafId, true);
 	}
 
-	private async ensureLoaded(): Promise<void> {
-		if (this.cacheLoaded) {
-			return;
-		}
-		const loaded = await loadJsonlStorage(this.filePath);
-		this.entries = loaded.entries;
-		this.byId = new Map(loaded.entries.map((entry) => [entry.id, entry]));
-		this.currentLeafId = loaded.leafId;
-		this.headerInitialized = loaded.header !== undefined;
-		if (loaded.header) {
-			this.sessionInfo = headerToSessionInfo(loaded.header, this.filePath);
-		}
-		this.cacheLoaded = true;
-	}
-
-	private async ensureHeader(): Promise<void> {
-		await this.ensureLoaded();
-		if (this.headerInitialized) return;
-		await this.ensureParentDir();
+	static async create(
+		filePath: string,
+		options: {
+			cwd: string;
+			sessionId: string;
+			parentSessionPath?: string;
+		},
+	): Promise<JsonlSessionTreeStorage> {
+		const resolvedPath = resolve(filePath);
 		const header: SessionHeader = {
 			type: "session",
 			version: 3,
-			id: this.requestedSessionId ?? randomUUID(),
+			id: options.sessionId,
 			timestamp: new Date().toISOString(),
-			cwd: this.cwd,
-			parentSession: this.parentSession,
+			cwd: options.cwd,
+			parentSession: options.parentSessionPath,
 		};
-		await writeFile(this.filePath, `${JSON.stringify(header)}\n`);
-		this.sessionInfo = headerToSessionInfo(header, this.filePath);
-		this.headerInitialized = true;
+		return new JsonlSessionTreeStorage(resolvedPath, header, [], null, false);
 	}
 
-	async getSessionInfo(): Promise<SessionInfo> {
-		await this.ensureHeader();
-		return this.sessionInfo!;
+	async getSessionInfo(): Promise<JsonlSessionInfo> {
+		return this.sessionInfo;
 	}
 
 	async getLeafId(): Promise<string | null> {
-		await this.ensureLoaded();
 		return this.currentLeafId;
 	}
 
 	async setLeafId(leafId: string | null): Promise<void> {
-		await this.ensureLoaded();
+		if (leafId !== null && !this.byId.has(leafId)) {
+			throw new Error(`Entry ${leafId} not found`);
+		}
 		this.currentLeafId = leafId;
 	}
 
 	async appendEntry(entry: SessionTreeEntry): Promise<void> {
-		await this.ensureHeader();
+		if (!this.headerWritten) {
+			await mkdir(dirname(this.filePath), { recursive: true });
+			await writeFile(this.filePath, `${JSON.stringify(this.header)}\n`);
+			this.headerWritten = true;
+		}
 		await appendFile(this.filePath, `${JSON.stringify(entry)}\n`);
 		this.entries.push(entry);
 		this.byId.set(entry.id, entry);
@@ -129,12 +154,10 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage {
 	}
 
 	async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
-		await this.ensureLoaded();
 		return this.byId.get(id);
 	}
 
 	async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
-		await this.ensureLoaded();
 		if (leafId === null) return [];
 		const path: SessionTreeEntry[] = [];
 		let current = this.byId.get(leafId);
@@ -146,7 +169,6 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage {
 	}
 
 	async getEntries(): Promise<SessionTreeEntry[]> {
-		await this.ensureLoaded();
 		return [...this.entries];
 	}
 }

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

@@ -1,15 +1,20 @@
-import { randomUUID } from "crypto";
+import { v7 as uuidv7 } from "uuid";
 import type { SessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js";
 
 export class InMemorySessionTreeStorage implements SessionTreeStorage {
+	private readonly sessionInfo: SessionInfo;
 	private entries: SessionTreeEntry[];
+	private byId: Map<string, SessionTreeEntry>;
 	private leafId: string | null;
-	private sessionInfo: SessionInfo;
 
 	constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; sessionInfo?: SessionInfo }) {
 		this.entries = options?.entries ? [...options.entries] : [];
+		this.byId = new Map(this.entries.map((entry) => [entry.id, entry]));
 		this.leafId = options?.leafId ?? this.entries[this.entries.length - 1]?.id ?? null;
-		this.sessionInfo = options?.sessionInfo ?? { id: randomUUID(), createdAt: new Date().toISOString() };
+		if (this.leafId !== null && !this.byId.has(this.leafId)) {
+			throw new Error(`Entry ${this.leafId} not found`);
+		}
+		this.sessionInfo = options?.sessionInfo ?? { id: uuidv7(), createdAt: new Date().toISOString() };
 	}
 
 	async getSessionInfo(): Promise<SessionInfo> {
@@ -21,26 +26,29 @@ export class InMemorySessionTreeStorage implements SessionTreeStorage {
 	}
 
 	async setLeafId(leafId: string | null): Promise<void> {
+		if (leafId !== null && !this.byId.has(leafId)) {
+			throw new Error(`Entry ${leafId} not found`);
+		}
 		this.leafId = leafId;
 	}
 
 	async appendEntry(entry: SessionTreeEntry): Promise<void> {
 		this.entries.push(entry);
+		this.byId.set(entry.id, entry);
 		this.leafId = entry.id;
 	}
 
 	async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
-		return this.entries.find((entry) => entry.id === id);
+		return this.byId.get(id);
 	}
 
 	async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
 		if (leafId === null) return [];
-		const byId = new Map<string, SessionTreeEntry>(this.entries.map((entry) => [entry.id, entry]));
 		const path: SessionTreeEntry[] = [];
-		let current = byId.get(leafId);
+		let current = this.byId.get(leafId);
 		while (current) {
 			path.unshift(current);
-			current = current.parentId ? byId.get(current.parentId) : undefined;
+			current = current.parentId ? this.byId.get(current.parentId) : undefined;
 		}
 		return path;
 	}

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

@@ -1,13 +1,14 @@
-import { existsSync, mkdirSync, readdirSync, readFileSync, rmSync } from "node:fs";
+import { existsSync, mkdirSync, readdirSync, rmSync } from "node:fs";
 import { join, resolve } from "node:path";
 import { v7 as uuidv7 } from "uuid";
 import type {
-	CodingAgentSessionInfo,
-	CodingAgentSessionRepo,
+	JsonlSessionInfo,
+	JsonlSessionRepo,
 	Session,
 	SessionInfo,
 	SessionRepo,
 	SessionTreeEntry,
+	SessionTreeStorage,
 } from "../types.js";
 import { JsonlSessionTreeStorage } from "./jsonl-session-storage.js";
 import { InMemorySessionTreeStorage } from "./memory-session-storage.js";
@@ -21,8 +22,11 @@ function createTimestamp(): string {
 	return new Date().toISOString();
 }
 
-function toSession<TInfo extends SessionInfo>(info: TInfo, tree: DefaultSessionTree): Session<TInfo> {
-	return { info, tree };
+function toSession<TInfo extends SessionInfo>(
+	storage: SessionTreeStorage<TInfo>,
+	tree: DefaultSessionTree<TInfo>,
+): Session<TInfo> {
+	return { storage, tree };
 }
 
 function getPathEntriesToFork(
@@ -59,14 +63,13 @@ function getPathEntriesToFork(
 export class InMemorySessionRepo implements SessionRepo<string> {
 	private sessions = new Map<string, Session<SessionInfo>>();
 
-	async create(options?: { id?: string; parentSession?: string }): Promise<Session<SessionInfo>> {
+	async create(options?: { id?: string }): Promise<Session<SessionInfo>> {
 		const info: SessionInfo = {
 			id: options?.id ?? createSessionId(),
 			createdAt: createTimestamp(),
-			parentSession: options?.parentSession,
 		};
 		const storage = new InMemorySessionTreeStorage({ sessionInfo: info });
-		const session = toSession(info, new DefaultSessionTree(storage));
+		const session = toSession(storage, new DefaultSessionTree(storage));
 		this.sessions.set(info.id, session);
 		return session;
 	}
@@ -97,42 +100,16 @@ export class InMemorySessionRepo implements SessionRepo<string> {
 		const info: SessionInfo = {
 			id: options.id ?? createSessionId(),
 			createdAt: createTimestamp(),
-			parentSession: source.info.id,
 		};
 		const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null;
 		const storage = new InMemorySessionTreeStorage({ sessionInfo: info, entries: forkedEntries, leafId });
-		const session = toSession(info, new DefaultSessionTree(storage));
+		const session = toSession(storage, new DefaultSessionTree(storage));
 		this.sessions.set(info.id, session);
 		return session;
 	}
 }
 
-function readJsonlHeader(filePath: string): CodingAgentSessionInfo | undefined {
-	try {
-		const content = readFileSync(filePath, "utf8");
-		const firstLine = content.split("\n")[0];
-		if (!firstLine) return undefined;
-		const header = JSON.parse(firstLine) as {
-			type: string;
-			id: string;
-			timestamp: string;
-			cwd: string;
-			parentSession?: string;
-		};
-		if (header.type !== "session") return undefined;
-		return {
-			id: header.id,
-			createdAt: header.timestamp,
-			parentSession: header.parentSession,
-			projectCwd: header.cwd,
-			filePath,
-		};
-	} catch {
-		return undefined;
-	}
-}
-
-export class JsonlCodingAgentSessionRepo implements CodingAgentSessionRepo<string> {
+export class JsonlSessionFileRepo implements JsonlSessionRepo<string> {
 	private sessionDir: string;
 	private cwd: string;
 
@@ -146,55 +123,66 @@ export class JsonlCodingAgentSessionRepo implements CodingAgentSessionRepo<strin
 		return join(this.sessionDir, `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`);
 	}
 
-	async create(options?: { id?: string; parentSession?: string }): Promise<Session<CodingAgentSessionInfo>> {
+	async create(options?: { id?: string; parentSessionPath?: string }): Promise<Session<JsonlSessionInfo>> {
 		const id = options?.id ?? createSessionId();
 		const createdAt = createTimestamp();
 		const filePath = this.createSessionFilePath(id, createdAt);
-		const storage = new JsonlSessionTreeStorage(filePath, {
+		const storage = await JsonlSessionTreeStorage.create(filePath, {
 			cwd: this.cwd,
 			sessionId: id,
-			parentSession: options?.parentSession,
+			parentSessionPath: options?.parentSessionPath,
 		});
-		const info = (await storage.getSessionInfo()) as CodingAgentSessionInfo;
-		return toSession(info, new DefaultSessionTree(storage));
+		return toSession(storage, new DefaultSessionTree(storage));
 	}
 
-	async open(ref: string): Promise<Session<CodingAgentSessionInfo>> {
+	async open(ref: string): Promise<Session<JsonlSessionInfo>> {
 		const filePath = ref.includes("/") || ref.endsWith(".jsonl") ? resolve(ref) : join(this.sessionDir, ref);
 		if (!existsSync(filePath)) {
 			throw new Error(`Session not found: ${ref}`);
 		}
-		const storage = new JsonlSessionTreeStorage(filePath, { cwd: this.cwd });
-		const info = (await storage.getSessionInfo()) as CodingAgentSessionInfo;
-		return toSession(info, new DefaultSessionTree(storage));
+		const storage = await JsonlSessionTreeStorage.open(filePath);
+		return toSession(storage, new DefaultSessionTree(storage));
 	}
 
-	async list(): Promise<Array<Session<CodingAgentSessionInfo>>> {
+	async list(): Promise<Array<Session<JsonlSessionInfo>>> {
 		if (!existsSync(this.sessionDir)) {
 			return [];
 		}
 		const files = readdirSync(this.sessionDir)
 			.filter((file) => file.endsWith(".jsonl"))
 			.map((file) => join(this.sessionDir, file));
-		const sessions: Array<Session<CodingAgentSessionInfo>> = [];
+		const sessions: Array<Session<JsonlSessionInfo>> = [];
 		for (const filePath of files) {
-			const info = readJsonlHeader(filePath);
-			if (!info) continue;
-			sessions.push(
-				toSession(info, new DefaultSessionTree(new JsonlSessionTreeStorage(filePath, { cwd: info.projectCwd }))),
-			);
+			try {
+				const storage = await JsonlSessionTreeStorage.open(filePath);
+				sessions.push(toSession(storage, new DefaultSessionTree(storage)));
+			} catch {
+				// Ignore invalid session files when listing a directory.
+			}
 		}
 		return sessions;
 	}
 
-	async listByCwd(cwd: string): Promise<Array<Session<CodingAgentSessionInfo>>> {
-		return (await this.list()).filter((session) => session.info.projectCwd === cwd);
+	async listByCwd(cwd: string): Promise<Array<Session<JsonlSessionInfo>>> {
+		const sessions = await this.list();
+		const result: Array<Session<JsonlSessionInfo>> = [];
+		for (const session of sessions) {
+			if ((await session.storage.getSessionInfo()).cwd === cwd) {
+				result.push(session);
+			}
+		}
+		return result;
 	}
 
-	async getMostRecentByCwd(cwd: string): Promise<Session<CodingAgentSessionInfo> | undefined> {
-		const sessions = await this.listByCwd(cwd);
-		sessions.sort((a, b) => new Date(b.info.createdAt).getTime() - new Date(a.info.createdAt).getTime());
-		return sessions[0];
+	async getMostRecentByCwd(cwd: string): Promise<Session<JsonlSessionInfo> | undefined> {
+		const sessionsWithInfo = await Promise.all(
+			(await this.listByCwd(cwd)).map(async (session) => ({
+				session,
+				info: await session.storage.getSessionInfo(),
+			})),
+		);
+		sessionsWithInfo.sort((a, b) => new Date(b.info.createdAt).getTime() - new Date(a.info.createdAt).getTime());
+		return sessionsWithInfo[0]?.session;
 	}
 
 	async delete(ref: string): Promise<void> {
@@ -207,17 +195,18 @@ export class JsonlCodingAgentSessionRepo implements CodingAgentSessionRepo<strin
 	async fork(
 		ref: string,
 		options: { entryId: string; position?: "before" | "at"; id?: string },
-	): Promise<Session<CodingAgentSessionInfo>> {
+	): Promise<Session<JsonlSessionInfo>> {
 		const source = await this.open(ref);
 		const entries = await source.tree.getEntries();
 		const forkedEntries = getPathEntriesToFork(entries, options.entryId, options.position ?? "before");
 		const id = options.id ?? createSessionId();
 		const createdAt = createTimestamp();
 		const filePath = this.createSessionFilePath(id, createdAt);
-		const storage = new JsonlSessionTreeStorage(filePath, {
-			cwd: source.info.projectCwd,
+		const sourceInfo = await source.storage.getSessionInfo();
+		const storage = await JsonlSessionTreeStorage.create(filePath, {
+			cwd: sourceInfo.cwd,
 			sessionId: id,
-			parentSession: source.info.filePath ?? source.info.id,
+			parentSessionPath: sourceInfo.path,
 		});
 		for (const entry of forkedEntries) {
 			await storage.appendEntry(entry);
@@ -225,7 +214,6 @@ export class JsonlCodingAgentSessionRepo implements CodingAgentSessionRepo<strin
 		if (forkedEntries.length === 0) {
 			await storage.getSessionInfo();
 		}
-		const info = (await storage.getSessionInfo()) as CodingAgentSessionInfo;
-		return toSession(info, new DefaultSessionTree(storage));
+		return toSession(storage, new DefaultSessionTree(storage));
 	}
 }

packages/agent/src/harness/session/session-tree.ts

@@ -18,7 +18,6 @@ import type {
 	SessionTreeStorage,
 	ThinkingLevelChangeEntry,
 } from "../types.js";
-import { InMemorySessionTreeStorage } from "./memory-session-storage.js";
 
 function generateId(byId: { has(id: string): boolean }): string {
 	for (let i = 0; i < 100; i++) {
@@ -79,11 +78,11 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext
 	return { messages, thinkingLevel, model };
 }
 
-export class DefaultSessionTree implements SessionTree {
-	private storage: SessionTreeStorage;
+export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> implements SessionTree {
+	private storage: SessionTreeStorage<TInfo>;
 
-	constructor(storage?: SessionTreeStorage) {
-		this.storage = storage ?? new InMemorySessionTreeStorage();
+	constructor(storage: SessionTreeStorage<TInfo>) {
+		this.storage = storage;
 	}
 
 	getLeafId(): Promise<string | null> {
@@ -107,7 +106,7 @@ export class DefaultSessionTree implements SessionTree {
 		return buildSessionContext(await this.getBranch());
 	}
 
-	getSessionInfo(): Promise<SessionInfo> {
+	getSessionInfo(): Promise<TInfo> {
 		return this.storage.getSessionInfo();
 	}
 

packages/agent/src/harness/types.ts

@@ -160,16 +160,16 @@ export interface SessionContext {
 export interface SessionInfo {
 	id: string;
 	createdAt: string;
-	parentSession?: string;
 }
 
-export interface CodingAgentSessionInfo extends SessionInfo {
-	projectCwd: string;
-	filePath?: string;
+export interface JsonlSessionInfo extends SessionInfo {
+	cwd: string;
+	path: string;
+	parentSessionPath?: string;
 }
 
-export interface SessionTreeStorage {
-	getSessionInfo(): Promise<SessionInfo>;
+export interface SessionTreeStorage<TInfo extends SessionInfo = SessionInfo> {
+	getSessionInfo(): Promise<TInfo>;
 	getLeafId(): Promise<string | null>;
 	setLeafId(leafId: string | null): Promise<void>;
 	appendEntry(entry: SessionTreeEntry): Promise<void>;
@@ -179,7 +179,6 @@ export interface SessionTreeStorage {
 }
 
 export interface SessionTree {
-	getSessionInfo(): Promise<SessionInfo>;
 	getLeafId(): Promise<string | null>;
 	getEntry(id: string): Promise<SessionTreeEntry | undefined>;
 	getEntries(): Promise<SessionTreeEntry[]>;
@@ -213,21 +212,40 @@ export interface SessionTree {
 }
 
 export interface Session<TInfo extends SessionInfo = SessionInfo> {
-	info: TInfo;
+	storage: SessionTreeStorage<TInfo>;
 	tree: SessionTree;
 }
 
-export interface SessionRepo<TRef = string, TInfo extends SessionInfo = SessionInfo> {
-	create(options?: { id?: string; parentSession?: string }): Promise<Session<TInfo>>;
+export interface SessionCreateOptions {
+	id?: string;
+}
+
+export interface SessionForkOptions {
+	entryId: string;
+	position?: "before" | "at";
+	id?: string;
+}
+
+export interface SessionRepo<
+	TRef = string,
+	TInfo extends SessionInfo = SessionInfo,
+	TCreateOptions extends SessionCreateOptions = SessionCreateOptions,
+> {
+	create(options?: TCreateOptions): Promise<Session<TInfo>>;
 	open(ref: TRef): Promise<Session<TInfo>>;
 	list(): Promise<Array<Session<TInfo>>>;
 	delete(ref: TRef): Promise<void>;
-	fork(ref: TRef, options: { entryId: string; position?: "before" | "at"; id?: string }): Promise<Session<TInfo>>;
+	fork(ref: TRef, options: SessionForkOptions): Promise<Session<TInfo>>;
 }
 
-export interface CodingAgentSessionRepo<TRef = string> extends SessionRepo<TRef, CodingAgentSessionInfo> {
-	listByCwd(cwd: string): Promise<Array<Session<CodingAgentSessionInfo>>>;
-	getMostRecentByCwd(cwd: string): Promise<Session<CodingAgentSessionInfo> | undefined>;
+export interface JsonlSessionCreateOptions extends SessionCreateOptions {
+	parentSessionPath?: string;
+}
+
+export interface JsonlSessionRepo<TRef = string>
+	extends SessionRepo<TRef, JsonlSessionInfo, JsonlSessionCreateOptions> {
+	listByCwd(cwd: string): Promise<Array<Session<JsonlSessionInfo>>>;
+	getMostRecentByCwd(cwd: string): Promise<Session<JsonlSessionInfo> | undefined>;
 }
 
 export interface AgentHarnessPendingMutations {

packages/agent/test/harness/session-tree.test.ts

@@ -1,12 +1,12 @@
-import { existsSync, mkdirSync, readFileSync, rmSync } from "node:fs";
+import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import type { AgentMessage } from "@mariozechner/pi-agent-core";
 import { afterEach, describe, expect, it } from "vitest";
-import { JsonlSessionTreeStorage } from "../../src/harness/session/jsonl-session-storage.js";
+import { JsonlSessionTreeStorage, loadJsonlSessionInfo } from "../../src/harness/session/jsonl-session-storage.js";
 import { InMemorySessionTreeStorage } from "../../src/harness/session/memory-session-storage.js";
 import { DefaultSessionTree } from "../../src/harness/session/session-tree.js";
-import type { SessionTreeStorage } from "../../src/harness/types.js";
+import type { MessageEntry, SessionInfo, SessionTreeStorage } from "../../src/harness/types.js";
 
 function createUserMessage(text: string): AgentMessage {
 	return {
@@ -38,6 +38,13 @@ function createAssistantMessage(text: string): AgentMessage {
 
 const tempDirs: string[] = [];
 
+function createTempDir(): string {
+	const dir = join(tmpdir(), `pi-agent-session-tree-${Date.now()}-${Math.random().toString(36).slice(2)}`);
+	mkdirSync(dir, { recursive: true });
+	tempDirs.push(dir);
+	return dir;
+}
+
 afterEach(() => {
 	while (tempDirs.length > 0) {
 		const dir = tempDirs.pop()!;
@@ -47,10 +54,14 @@ afterEach(() => {
 	}
 });
 
-async function runSessionTreeSuite(name: string, createStorage: () => SessionTreeStorage, inspect?: () => void) {
+async function runSessionTreeSuite(
+	name: string,
+	createStorage: () => SessionTreeStorage | Promise<SessionTreeStorage>,
+	inspect?: () => void,
+) {
 	describe(name, () => {
 		it("appends messages and builds context in order", async () => {
-			const tree = new DefaultSessionTree(createStorage());
+			const tree = new DefaultSessionTree(await createStorage());
 			await tree.appendMessage(createUserMessage("one"));
 			await tree.appendMessage(createAssistantMessage("two"));
 			const context = await tree.buildContext();
@@ -58,7 +69,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
 		});
 
 		it("tracks model and thinking level changes", async () => {
-			const tree = new DefaultSessionTree(createStorage());
+			const tree = new DefaultSessionTree(await createStorage());
 			await tree.appendMessage(createUserMessage("one"));
 			await tree.appendModelChange("openai", "gpt-4.1");
 			await tree.appendThinkingLevelChange("high");
@@ -68,7 +79,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
 		});
 
 		it("supports branching by moving the leaf and appending a new branch", async () => {
-			const tree = new DefaultSessionTree(createStorage());
+			const tree = new DefaultSessionTree(await createStorage());
 			const user1 = await tree.appendMessage(createUserMessage("one"));
 			const assistant1 = await tree.appendMessage(createAssistantMessage("two"));
 			await tree.appendMessage(createUserMessage("three"));
@@ -82,7 +93,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
 		});
 
 		it("supports moving the leaf to root", async () => {
-			const tree = new DefaultSessionTree(createStorage());
+			const tree = new DefaultSessionTree(await createStorage());
 			await tree.appendMessage(createUserMessage("one"));
 			await tree.moveTo(null);
 			expect(await tree.getLeafId()).toBeNull();
@@ -90,7 +101,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
 		});
 
 		it("reconstructs compaction summaries in context", async () => {
-			const tree = new DefaultSessionTree(createStorage());
+			const tree = new DefaultSessionTree(await createStorage());
 			await tree.appendMessage(createUserMessage("one"));
 			await tree.appendMessage(createAssistantMessage("two"));
 			const user2 = await tree.appendMessage(createUserMessage("three"));
@@ -103,7 +114,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
 		});
 
 		it("supports branch summary entries in context", async () => {
-			const tree = new DefaultSessionTree(createStorage());
+			const tree = new DefaultSessionTree(await createStorage());
 			const user1 = await tree.appendMessage(createUserMessage("one"));
 			await tree.appendBranchSummary(user1, "summary text");
 			const context = await tree.buildContext();
@@ -111,7 +122,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
 		});
 
 		it("supports custom message entries in context", async () => {
-			const tree = new DefaultSessionTree(createStorage());
+			const tree = new DefaultSessionTree(await createStorage());
 			await tree.appendMessage(createUserMessage("one"));
 			await tree.appendCustomMessageEntry("custom", "hello", true, { ok: true });
 			const context = await tree.buildContext();
@@ -119,7 +130,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
 		});
 
 		it("supports labels and session info entries without affecting context", async () => {
-			const tree = new DefaultSessionTree(createStorage());
+			const tree = new DefaultSessionTree(await createStorage());
 			const user1 = await tree.appendMessage(createUserMessage("one"));
 			await tree.appendLabelChange(user1, "checkpoint");
 			await tree.appendSessionInfo("name");
@@ -132,7 +143,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
 		});
 
 		it("persists leaf changes and appended entries via storage", async () => {
-			const storage = createStorage();
+			const storage = await createStorage();
 			const tree = new DefaultSessionTree(storage);
 			const user1 = await tree.appendMessage(createUserMessage("one"));
 			await tree.appendMessage(createAssistantMessage("two"));
@@ -150,15 +161,194 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre
 	});
 }
 
+describe("InMemorySessionTreeStorage", () => {
+	it("returns configured session info", async () => {
+		const sessionInfo: SessionInfo = { id: "session-1", createdAt: "2026-01-01T00:00:00.000Z" };
+		const storage = new InMemorySessionTreeStorage({ sessionInfo });
+		expect(await storage.getSessionInfo()).toEqual(sessionInfo);
+	});
+
+	it("copies initial entries and tracks leaf independently", async () => {
+		const entry: MessageEntry = {
+			type: "message",
+			id: "entry-1",
+			parentId: null,
+			timestamp: "2026-01-01T00:00:00.000Z",
+			message: createUserMessage("one"),
+		};
+		const initialEntries = [entry];
+		const storage = new InMemorySessionTreeStorage({ entries: initialEntries });
+		initialEntries.push({ ...entry, id: "entry-2" });
+		expect((await storage.getEntries()).map((storedEntry) => storedEntry.id)).toEqual(["entry-1"]);
+		expect(await storage.getLeafId()).toBe("entry-1");
+		await storage.setLeafId(null);
+		expect(await storage.getLeafId()).toBeNull();
+	});
+
+	it("rejects invalid leaf ids", async () => {
+		const storage = new InMemorySessionTreeStorage();
+		await expect(storage.setLeafId("missing")).rejects.toThrow("Entry missing not found");
+		expect(() => new InMemorySessionTreeStorage({ leafId: "missing" })).toThrow("Entry missing not found");
+	});
+
+	it("walks paths to root", async () => {
+		const root: MessageEntry = {
+			type: "message",
+			id: "root",
+			parentId: null,
+			timestamp: "2026-01-01T00:00:00.000Z",
+			message: createUserMessage("root"),
+		};
+		const child: MessageEntry = {
+			...root,
+			id: "child",
+			parentId: "root",
+			message: createAssistantMessage("child"),
+		};
+		const storage = new InMemorySessionTreeStorage({ entries: [root, child] });
+		expect((await storage.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
+		expect(await storage.getPathToRoot(null)).toEqual([]);
+	});
+});
+
 runSessionTreeSuite("SessionTree with in-memory storage", () => new InMemorySessionTreeStorage());
 
+describe("JsonlSessionTreeStorage", () => {
+	it("throws for missing files when opening", async () => {
+		const dir = createTempDir();
+		const filePath = join(dir, "session.jsonl");
+		await expect(JsonlSessionTreeStorage.open(filePath)).rejects.toMatchObject({ code: "ENOENT" });
+	});
+
+	it("writes the header before the first appended entry", async () => {
+		const dir = createTempDir();
+		const filePath = join(dir, "session.jsonl");
+		const storage = await JsonlSessionTreeStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
+		expect(existsSync(filePath)).toBe(false);
+		expect(await storage.getLeafId()).toBeNull();
+		expect(await storage.getEntries()).toEqual([]);
+		await storage.appendEntry({
+			type: "message",
+			id: "user-1",
+			parentId: null,
+			timestamp: "2026-01-01T00:00:00.000Z",
+			message: createUserMessage("one"),
+		});
+		expect(existsSync(filePath)).toBe(true);
+		const lines = readFileSync(filePath, "utf8").trim().split("\n");
+		expect(JSON.parse(lines[0]!).type).toBe("session");
+		expect(JSON.parse(lines[1]!).id).toBe("user-1");
+		expect(lines).toHaveLength(2);
+	});
+
+	it("throws for malformed session headers", async () => {
+		const dir = createTempDir();
+		const filePath = join(dir, "session.jsonl");
+		writeFileSync(filePath, "not json\n");
+		await expect(JsonlSessionTreeStorage.open(filePath)).rejects.toThrow("first line is not a valid session header");
+	});
+
+	it("ignores malformed entry lines", async () => {
+		const dir = createTempDir();
+		const filePath = join(dir, "session.jsonl");
+		const header = {
+			type: "session",
+			version: 3,
+			id: "session-1",
+			timestamp: "2026-01-01T00:00:00.000Z",
+			cwd: dir,
+		};
+		const entry: MessageEntry = {
+			type: "message",
+			id: "entry-1",
+			parentId: null,
+			timestamp: "2026-01-01T00:00:00.000Z",
+			message: createUserMessage("one"),
+		};
+		writeFileSync(filePath, `${JSON.stringify(header)}\nnot json\n${JSON.stringify(entry)}\n`);
+		const storage = await JsonlSessionTreeStorage.open(filePath);
+		expect((await storage.getEntries()).map((loadedEntry) => loadedEntry.id)).toEqual(["entry-1"]);
+		expect(await storage.getLeafId()).toBe("entry-1");
+	});
+
+	it("creates and reads session info from the header", async () => {
+		const dir = createTempDir();
+		const filePath = join(dir, "session.jsonl");
+		const storage = await JsonlSessionTreeStorage.create(filePath, {
+			cwd: dir,
+			sessionId: "session-1",
+			parentSessionPath: "/tmp/parent.jsonl",
+		});
+		const info = await storage.getSessionInfo();
+		expect(info).toMatchObject({
+			id: "session-1",
+			cwd: dir,
+			path: filePath,
+			parentSessionPath: "/tmp/parent.jsonl",
+		});
+		expect(existsSync(filePath)).toBe(false);
+		await storage.appendEntry({
+			type: "message",
+			id: "user-1",
+			parentId: null,
+			timestamp: "2026-01-01T00:00:00.000Z",
+			message: createUserMessage("one"),
+		});
+		expect(await loadJsonlSessionInfo(filePath)).toEqual(info);
+	});
+
+	it("loads existing entries and reconstructs leaf", async () => {
+		const dir = createTempDir();
+		const filePath = join(dir, "session.jsonl");
+		const storage = await JsonlSessionTreeStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
+		const root: MessageEntry = {
+			type: "message",
+			id: "root",
+			parentId: null,
+			timestamp: "2026-01-01T00:00:00.000Z",
+			message: createUserMessage("root"),
+		};
+		const child: MessageEntry = {
+			...root,
+			id: "child",
+			parentId: "root",
+			message: createAssistantMessage("child"),
+		};
+		await storage.appendEntry(root);
+		await storage.appendEntry(child);
+		const loaded = await JsonlSessionTreeStorage.open(filePath);
+		expect(await loaded.getLeafId()).toBe("child");
+		expect((await loaded.getEntries()).map((entry) => entry.id)).toEqual(["root", "child"]);
+		expect((await loaded.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
+	});
+
+	it("reads session info from only the first JSONL line", async () => {
+		const dir = createTempDir();
+		const filePath = join(dir, "session.jsonl");
+		const header = {
+			type: "session",
+			version: 3,
+			id: "session-1",
+			timestamp: "2026-01-01T00:00:00.000Z",
+			cwd: dir,
+		};
+		const malformedSecondLine = "{".repeat(10000);
+		writeFileSync(filePath, `${JSON.stringify(header)}\n${malformedSecondLine}\n`);
+		expect(await loadJsonlSessionInfo(filePath)).toEqual({
+			id: "session-1",
+			createdAt: "2026-01-01T00:00:00.000Z",
+			cwd: dir,
+			path: filePath,
+			parentSessionPath: undefined,
+		});
+	});
+});
+
 runSessionTreeSuite(
 	"SessionTree with JSONL storage",
-	() => {
-		const dir = join(tmpdir(), `pi-agent-session-tree-${Date.now()}-${Math.random().toString(36).slice(2)}`);
-		mkdirSync(dir, { recursive: true });
-		tempDirs.push(dir);
-		return new JsonlSessionTreeStorage(join(dir, "session.jsonl"), { cwd: dir });
+	async () => {
+		const dir = createTempDir();
+		return await JsonlSessionTreeStorage.create(join(dir, "session.jsonl"), { cwd: dir, sessionId: "session-1" });
 	},
 	() => {
 		const dir = tempDirs[tempDirs.length - 1]!;