refactor(agent): consolidate harness session abstraction

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

@@ -0,0 +1,142 @@
+import { constants } from "node:fs";
+import { access, mkdir, readdir, rm } from "node:fs/promises";
+import { join, resolve } from "node:path";
+import type {
+	JsonlSessionCreateOptions,
+	JsonlSessionListQuery,
+	JsonlSessionMetadata,
+	JsonlSessionRef,
+	JsonlSessionRepoApi,
+	JsonlSessionResolveOptions,
+	Session,
+} from "../types.js";
+import { JsonlSessionStorage, loadJsonlSessionMetadata } from "./jsonl-session-storage.js";
+import { createSessionId, createTimestamp, getPathEntriesToFork, toSession } from "./session-repo.js";
+
+async function exists(path: string): Promise<boolean> {
+	try {
+		await access(path, constants.F_OK);
+		return true;
+	} catch {
+		return false;
+	}
+}
+
+function encodeCwd(cwd: string): string {
+	return `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`;
+}
+
+export class JsonlSessionRepo implements JsonlSessionRepoApi {
+	private sessionsRoot: string;
+
+	constructor(options: { sessionsRoot: string }) {
+		this.sessionsRoot = resolve(options.sessionsRoot);
+	}
+
+	private getSessionDir(cwd: string): string {
+		return join(this.sessionsRoot, encodeCwd(cwd));
+	}
+
+	private createSessionFilePath(cwd: string, sessionId: string, timestamp: string): string {
+		return join(this.getSessionDir(cwd), `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`);
+	}
+
+	private refPath(ref: JsonlSessionRef): string {
+		return resolve(ref.path);
+	}
+
+	async create(options: JsonlSessionCreateOptions): Promise<Session<JsonlSessionMetadata>> {
+		await mkdir(this.sessionsRoot, { recursive: true });
+		const id = options.id ?? createSessionId();
+		const createdAt = createTimestamp();
+		const filePath = this.createSessionFilePath(options.cwd, id, createdAt);
+		const storage = await JsonlSessionStorage.create(filePath, {
+			cwd: options.cwd,
+			sessionId: id,
+			parentSessionPath: options.parentSessionPath,
+		});
+		return toSession(storage);
+	}
+
+	async open(ref: JsonlSessionRef): Promise<Session<JsonlSessionMetadata>> {
+		const filePath = this.refPath(ref);
+		if (!(await exists(filePath))) {
+			throw new Error(`Session not found: ${filePath}`);
+		}
+		const storage = await JsonlSessionStorage.open(filePath);
+		return toSession(storage);
+	}
+
+	async list(query: JsonlSessionListQuery = {}): Promise<JsonlSessionMetadata[]> {
+		const dirs = query.cwd ? [this.getSessionDir(query.cwd)] : await this.listSessionDirs();
+		const sessions: JsonlSessionMetadata[] = [];
+		for (const dir of dirs) {
+			if (!(await exists(dir))) continue;
+			const files = (await readdir(dir)).filter((file) => file.endsWith(".jsonl")).map((file) => join(dir, file));
+			for (const filePath of files) {
+				try {
+					sessions.push(await loadJsonlSessionMetadata(filePath));
+				} catch {
+					// Ignore invalid session files when listing a directory.
+				}
+			}
+		}
+		sessions.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime());
+		return sessions;
+	}
+
+	async resolve(ref: string, options: JsonlSessionResolveOptions = {}): Promise<JsonlSessionMetadata[]> {
+		if (ref.includes("/") || ref.includes("\\") || ref.endsWith(".jsonl")) {
+			try {
+				return [await loadJsonlSessionMetadata(resolve(ref))];
+			} catch {
+				return [];
+			}
+		}
+
+		const local = options.cwd
+			? (await this.list({ cwd: options.cwd })).filter((session) => session.id.startsWith(ref))
+			: [];
+		if (local.length > 0 || !options.searchAll) return local;
+		return (await this.list()).filter((session) => session.id.startsWith(ref));
+	}
+
+	async getMostRecent(query: JsonlSessionListQuery = {}): Promise<JsonlSessionMetadata | undefined> {
+		return (await this.list(query))[0];
+	}
+
+	async delete(ref: JsonlSessionRef): Promise<void> {
+		const filePath = this.refPath(ref);
+		await rm(filePath, { force: true });
+	}
+
+	async fork(
+		ref: JsonlSessionRef,
+		options: JsonlSessionCreateOptions & { entryId: string; position?: "before" | "at"; id?: string },
+	): Promise<Session<JsonlSessionMetadata>> {
+		const source = await this.open(ref);
+		const forkedEntries = await getPathEntriesToFork(
+			source.getStorage(),
+			options.entryId,
+			options.position ?? "before",
+		);
+		const sourceInfo = await source.getMetadata();
+		const id = options.id ?? createSessionId();
+		const createdAt = createTimestamp();
+		const storage = await JsonlSessionStorage.create(this.createSessionFilePath(options.cwd, id, createdAt), {
+			cwd: options.cwd,
+			sessionId: id,
+			parentSessionPath: options.parentSessionPath ?? sourceInfo.path,
+		});
+		for (const entry of forkedEntries) {
+			await storage.appendEntry(entry);
+		}
+		return toSession(storage);
+	}
+
+	private async listSessionDirs(): Promise<string[]> {
+		if (!(await exists(this.sessionsRoot))) return [];
+		const entries = await readdir(this.sessionsRoot, { withFileTypes: true });
+		return entries.filter((entry) => entry.isDirectory()).map((entry) => join(this.sessionsRoot, entry.name));
+	}
+}

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

@@ -2,7 +2,7 @@ import { createReadStream } from "node:fs";
 import { appendFile, mkdir, readFile, writeFile } from "node:fs/promises";
 import { dirname, resolve } from "node:path";
 import { createInterface } from "node:readline";
-import type { JsonlSessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js";
+import type { JsonlSessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js";
 
 interface SessionHeader {
 	type: "session";
@@ -13,7 +13,25 @@ interface SessionHeader {
 	parentSession?: string;
 }
 
-function headerToSessionInfo(header: SessionHeader, path: string): JsonlSessionInfo {
+function updateLabelCache(labelsById: Map<string, string>, entry: SessionTreeEntry): void {
+	if (entry.type !== "label") return;
+	const label = entry.label?.trim();
+	if (label) {
+		labelsById.set(entry.targetId, label);
+	} else {
+		labelsById.delete(entry.targetId);
+	}
+}
+
+function buildLabelsById(entries: SessionTreeEntry[]): Map<string, string> {
+	const labelsById = new Map<string, string>();
+	for (const entry of entries) {
+		updateLabelCache(labelsById, entry);
+	}
+	return labelsById;
+}
+
+function headerToSessionMetadata(header: SessionHeader, path: string): JsonlSessionMetadata {
 	return {
 		id: header.id,
 		createdAt: header.timestamp,
@@ -23,7 +41,7 @@ function headerToSessionInfo(header: SessionHeader, path: string): JsonlSessionI
 	};
 }
 
-export async function loadJsonlSessionInfo(filePath: string): Promise<JsonlSessionInfo> {
+export async function loadJsonlSessionMetadata(filePath: string): Promise<JsonlSessionMetadata> {
 	const stream = createReadStream(filePath, { encoding: "utf8" });
 	const lines = createInterface({ input: stream, crlfDelay: Infinity });
 	try {
@@ -31,7 +49,7 @@ export async function loadJsonlSessionInfo(filePath: string): Promise<JsonlSessi
 			if (!line.trim()) break;
 			try {
 				const header = JSON.parse(line) as SessionHeader;
-				return headerToSessionInfo(header, resolve(filePath));
+				return headerToSessionMetadata(header, resolve(filePath));
 			} catch {
 				throw new Error(`Invalid JSONL session file ${filePath}: first line is not a valid session header`);
 			}
@@ -75,35 +93,27 @@ async function loadJsonlStorage(filePath: string): Promise<{
 	return { header, entries, leafId };
 }
 
-export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionInfo> {
+export class JsonlSessionStorage implements SessionStorage<JsonlSessionMetadata> {
 	private readonly filePath: string;
-	private readonly header: SessionHeader;
-	private readonly sessionInfo: JsonlSessionInfo;
+	private readonly metadata: JsonlSessionMetadata;
 	private entries: SessionTreeEntry[];
 	private byId: Map<string, SessionTreeEntry>;
+	private labelsById: Map<string, string>;
 	private currentLeafId: string | null;
-	private headerWritten: boolean;
 
-	private constructor(
-		filePath: string,
-		header: SessionHeader,
-		entries: SessionTreeEntry[],
-		leafId: string | null,
-		headerWritten: boolean,
-	) {
+	private constructor(filePath: string, header: SessionHeader, entries: SessionTreeEntry[], leafId: string | null) {
 		this.filePath = resolve(filePath);
-		this.header = header;
-		this.sessionInfo = headerToSessionInfo(header, this.filePath);
+		this.metadata = headerToSessionMetadata(header, this.filePath);
 		this.entries = entries;
 		this.byId = new Map(entries.map((entry) => [entry.id, entry]));
+		this.labelsById = buildLabelsById(entries);
 		this.currentLeafId = leafId;
-		this.headerWritten = headerWritten;
 	}
 
-	static async open(filePath: string): Promise<JsonlSessionTreeStorage> {
+	static async open(filePath: string): Promise<JsonlSessionStorage> {
 		const resolvedPath = resolve(filePath);
 		const loaded = await loadJsonlStorage(resolvedPath);
-		return new JsonlSessionTreeStorage(resolvedPath, loaded.header, loaded.entries, loaded.leafId, true);
+		return new JsonlSessionStorage(resolvedPath, loaded.header, loaded.entries, loaded.leafId);
 	}
 
 	static async create(
@@ -113,7 +123,7 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionI
 			sessionId: string;
 			parentSessionPath?: string;
 		},
-	): Promise<JsonlSessionTreeStorage> {
+	): Promise<JsonlSessionStorage> {
 		const resolvedPath = resolve(filePath);
 		const header: SessionHeader = {
 			type: "session",
@@ -123,11 +133,13 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionI
 			cwd: options.cwd,
 			parentSession: options.parentSessionPath,
 		};
-		return new JsonlSessionTreeStorage(resolvedPath, header, [], null, false);
+		await mkdir(dirname(resolvedPath), { recursive: true });
+		await writeFile(resolvedPath, `${JSON.stringify(header)}\n`);
+		return new JsonlSessionStorage(resolvedPath, header, [], null);
 	}
 
-	async getSessionInfo(): Promise<JsonlSessionInfo> {
-		return this.sessionInfo;
+	async getMetadata(): Promise<JsonlSessionMetadata> {
+		return this.metadata;
 	}
 
 	async getLeafId(): Promise<string | null> {
@@ -142,14 +154,10 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionI
 	}
 
 	async appendEntry(entry: SessionTreeEntry): Promise<void> {
-		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);
+		updateLabelCache(this.labelsById, entry);
 		this.currentLeafId = entry.id;
 	}
 
@@ -157,6 +165,10 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage<JsonlSessionI
 		return this.byId.get(id);
 	}
 
+	async getLabel(id: string): Promise<string | undefined> {
+		return this.labelsById.get(id);
+	}
+
 	async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
 		if (leafId === null) return [];
 		const path: SessionTreeEntry[] = [];

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

@@ -0,0 +1,55 @@
+import type { Session, SessionMetadata, SessionRepo } from "../types.js";
+import { InMemorySessionStorage } from "./memory-session-storage.js";
+import { createSessionId, createTimestamp, getPathEntriesToFork, toSession } from "./session-repo.js";
+
+export class InMemorySessionRepo implements SessionRepo<SessionMetadata, { id?: string }, string, void> {
+	private sessions = new Map<string, Session<SessionMetadata>>();
+
+	async create(options: { id?: string } = {}): Promise<Session<SessionMetadata>> {
+		const info: SessionMetadata = {
+			id: options.id ?? createSessionId(),
+			createdAt: createTimestamp(),
+		};
+		const storage = new InMemorySessionStorage({ metadata: info });
+		const session = toSession(storage);
+		this.sessions.set(info.id, session);
+		return session;
+	}
+
+	async open(ref: string): Promise<Session<SessionMetadata>> {
+		const session = this.sessions.get(ref);
+		if (!session) {
+			throw new Error(`Session not found: ${ref}`);
+		}
+		return session;
+	}
+
+	async list(): Promise<SessionMetadata[]> {
+		return Promise.all([...this.sessions.values()].map((session) => session.getMetadata()));
+	}
+
+	async delete(ref: string): Promise<void> {
+		this.sessions.delete(ref);
+	}
+
+	async fork(
+		ref: string,
+		options: { entryId: string; position?: "before" | "at"; id?: string },
+	): Promise<Session<SessionMetadata>> {
+		const source = await this.open(ref);
+		const forkedEntries = await getPathEntriesToFork(
+			source.getStorage(),
+			options.entryId,
+			options.position ?? "before",
+		);
+		const info: SessionMetadata = {
+			id: options.id ?? createSessionId(),
+			createdAt: createTimestamp(),
+		};
+		const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null;
+		const storage = new InMemorySessionStorage({ metadata: info, entries: forkedEntries, leafId });
+		const session = toSession(storage);
+		this.sessions.set(info.id, session);
+		return session;
+	}
+}

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

@@ -1,24 +1,44 @@
 import { v7 as uuidv7 } from "uuid";
-import type { SessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js";
+import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js";
 
-export class InMemorySessionTreeStorage implements SessionTreeStorage {
-	private readonly sessionInfo: SessionInfo;
+function updateLabelCache(labelsById: Map<string, string>, entry: SessionTreeEntry): void {
+	if (entry.type !== "label") return;
+	const label = entry.label?.trim();
+	if (label) {
+		labelsById.set(entry.targetId, label);
+	} else {
+		labelsById.delete(entry.targetId);
+	}
+}
+
+function buildLabelsById(entries: SessionTreeEntry[]): Map<string, string> {
+	const labelsById = new Map<string, string>();
+	for (const entry of entries) {
+		updateLabelCache(labelsById, entry);
+	}
+	return labelsById;
+}
+
+export class InMemorySessionStorage implements SessionStorage {
+	private readonly metadata: SessionMetadata;
 	private entries: SessionTreeEntry[];
 	private byId: Map<string, SessionTreeEntry>;
+	private labelsById: Map<string, string>;
 	private leafId: string | null;
 
-	constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; sessionInfo?: SessionInfo }) {
+	constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; metadata?: SessionMetadata }) {
 		this.entries = options?.entries ? [...options.entries] : [];
 		this.byId = new Map(this.entries.map((entry) => [entry.id, entry]));
+		this.labelsById = buildLabelsById(this.entries);
 		this.leafId = options?.leafId ?? this.entries[this.entries.length - 1]?.id ?? null;
 		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() };
+		this.metadata = options?.metadata ?? { id: uuidv7(), createdAt: new Date().toISOString() };
 	}
 
-	async getSessionInfo(): Promise<SessionInfo> {
-		return this.sessionInfo;
+	async getMetadata(): Promise<SessionMetadata> {
+		return this.metadata;
 	}
 
 	async getLeafId(): Promise<string | null> {
@@ -35,6 +55,7 @@ export class InMemorySessionTreeStorage implements SessionTreeStorage {
 	async appendEntry(entry: SessionTreeEntry): Promise<void> {
 		this.entries.push(entry);
 		this.byId.set(entry.id, entry);
+		updateLabelCache(this.labelsById, entry);
 		this.leafId = entry.id;
 	}
 
@@ -42,6 +63,10 @@ export class InMemorySessionTreeStorage implements SessionTreeStorage {
 		return this.byId.get(id);
 	}
 
+	async getLabel(id: string): Promise<string | undefined> {
+		return this.labelsById.get(id);
+	}
+
 	async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
 		if (leafId === null) return [];
 		const path: SessionTreeEntry[] = [];

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

@@ -1,41 +1,25 @@
-import { existsSync, mkdirSync, readdirSync, rmSync } from "node:fs";
-import { join, resolve } from "node:path";
 import { v7 as uuidv7 } from "uuid";
-import type {
-	JsonlSessionInfo,
-	JsonlSessionRepo,
-	Session,
-	SessionInfo,
-	SessionRepo,
-	SessionTreeEntry,
-	SessionTreeStorage,
-} from "../types.js";
-import { JsonlSessionTreeStorage } from "./jsonl-session-storage.js";
-import { InMemorySessionTreeStorage } from "./memory-session-storage.js";
-import { DefaultSessionTree } from "./session-tree.js";
+import type { Session, SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js";
+import { DefaultSession } from "./session-tree.js";
 
-function createSessionId(): string {
+export function createSessionId(): string {
 	return uuidv7();
 }
 
-function createTimestamp(): string {
+export function createTimestamp(): string {
 	return new Date().toISOString();
 }
 
-function toSession<TInfo extends SessionInfo>(
-	storage: SessionTreeStorage<TInfo>,
-	tree: DefaultSessionTree<TInfo>,
-): Session<TInfo> {
-	return { storage, tree };
+export function toSession<TMetadata extends SessionMetadata>(storage: SessionStorage<TMetadata>): Session<TMetadata> {
+	return new DefaultSession(storage);
 }
 
-function getPathEntriesToFork(
-	entries: SessionTreeEntry[],
+export async function getPathEntriesToFork(
+	storage: SessionStorage,
 	entryId: string,
 	position: "before" | "at",
-): SessionTreeEntry[] {
-	const byId = new Map<string, SessionTreeEntry>(entries.map((entry) => [entry.id, entry]));
-	const target = byId.get(entryId);
+): Promise<SessionTreeEntry[]> {
+	const target = await storage.getEntry(entryId);
 	if (!target) {
 		throw new Error(`Entry ${entryId} not found`);
 	}
@@ -48,172 +32,5 @@ function getPathEntriesToFork(
 		}
 		effectiveLeafId = target.parentId;
 	}
-	if (effectiveLeafId === null) {
-		return [];
-	}
-	const path: SessionTreeEntry[] = [];
-	let current = byId.get(effectiveLeafId);
-	while (current) {
-		path.unshift(current);
-		current = current.parentId ? byId.get(current.parentId) : undefined;
-	}
-	return path;
-}
-
-export class InMemorySessionRepo implements SessionRepo<string> {
-	private sessions = new Map<string, Session<SessionInfo>>();
-
-	async create(options?: { id?: string }): Promise<Session<SessionInfo>> {
-		const info: SessionInfo = {
-			id: options?.id ?? createSessionId(),
-			createdAt: createTimestamp(),
-		};
-		const storage = new InMemorySessionTreeStorage({ sessionInfo: info });
-		const session = toSession(storage, new DefaultSessionTree(storage));
-		this.sessions.set(info.id, session);
-		return session;
-	}
-
-	async open(ref: string): Promise<Session<SessionInfo>> {
-		const session = this.sessions.get(ref);
-		if (!session) {
-			throw new Error(`Session not found: ${ref}`);
-		}
-		return session;
-	}
-
-	async list(): Promise<Array<Session<SessionInfo>>> {
-		return [...this.sessions.values()];
-	}
-
-	async delete(ref: string): Promise<void> {
-		this.sessions.delete(ref);
-	}
-
-	async fork(
-		ref: string,
-		options: { entryId: string; position?: "before" | "at"; id?: string },
-	): Promise<Session<SessionInfo>> {
-		const source = await this.open(ref);
-		const entries = await source.tree.getEntries();
-		const forkedEntries = getPathEntriesToFork(entries, options.entryId, options.position ?? "before");
-		const info: SessionInfo = {
-			id: options.id ?? createSessionId(),
-			createdAt: createTimestamp(),
-		};
-		const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null;
-		const storage = new InMemorySessionTreeStorage({ sessionInfo: info, entries: forkedEntries, leafId });
-		const session = toSession(storage, new DefaultSessionTree(storage));
-		this.sessions.set(info.id, session);
-		return session;
-	}
-}
-
-export class JsonlSessionFileRepo implements JsonlSessionRepo<string> {
-	private sessionDir: string;
-	private cwd: string;
-
-	constructor(options: { sessionDir: string; cwd: string }) {
-		this.sessionDir = resolve(options.sessionDir);
-		this.cwd = options.cwd;
-		mkdirSync(this.sessionDir, { recursive: true });
-	}
-
-	private createSessionFilePath(sessionId: string, timestamp: string): string {
-		return join(this.sessionDir, `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`);
-	}
-
-	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 = await JsonlSessionTreeStorage.create(filePath, {
-			cwd: this.cwd,
-			sessionId: id,
-			parentSessionPath: options?.parentSessionPath,
-		});
-		return toSession(storage, new DefaultSessionTree(storage));
-	}
-
-	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 = await JsonlSessionTreeStorage.open(filePath);
-		return toSession(storage, new DefaultSessionTree(storage));
-	}
-
-	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<JsonlSessionInfo>> = [];
-		for (const filePath of files) {
-			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<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<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> {
-		const filePath = ref.includes("/") || ref.endsWith(".jsonl") ? resolve(ref) : join(this.sessionDir, ref);
-		if (existsSync(filePath)) {
-			rmSync(filePath, { force: true });
-		}
-	}
-
-	async fork(
-		ref: string,
-		options: { entryId: string; position?: "before" | "at"; id?: string },
-	): 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 sourceInfo = await source.storage.getSessionInfo();
-		const storage = await JsonlSessionTreeStorage.create(filePath, {
-			cwd: sourceInfo.cwd,
-			sessionId: id,
-			parentSessionPath: sourceInfo.path,
-		});
-		for (const entry of forkedEntries) {
-			await storage.appendEntry(entry);
-		}
-		if (forkedEntries.length === 0) {
-			await storage.getSessionInfo();
-		}
-		return toSession(storage, new DefaultSessionTree(storage));
-	}
+	return storage.getPathToRoot(effectiveLeafId);
 }

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

@@ -10,12 +10,12 @@ import type {
 	LabelEntry,
 	MessageEntry,
 	ModelChangeEntry,
+	Session,
 	SessionContext,
-	SessionInfo,
 	SessionInfoEntry,
-	SessionTree,
+	SessionMetadata,
+	SessionStorage,
 	SessionTreeEntry,
-	SessionTreeStorage,
 	ThinkingLevelChangeEntry,
 } from "../types.js";
 
@@ -27,12 +27,12 @@ function generateId(byId: { has(id: string): boolean }): string {
 	return randomUUID();
 }
 
-export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext {
+export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionContext {
 	let thinkingLevel = "off";
 	let model: { provider: string; modelId: string } | null = null;
 	let compaction: CompactionEntry | null = null;
 
-	for (const entry of entries) {
+	for (const entry of pathEntries) {
 		if (entry.type === "thinking_level_change") {
 			thinkingLevel = entry.thinkingLevel;
 		} else if (entry.type === "model_change") {
@@ -47,10 +47,16 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext
 	const messages: AgentMessage[] = [];
 	const appendMessage = (entry: SessionTreeEntry) => {
 		if (entry.type === "message") {
-			messages.push(entry.message);
+			messages.push(entry.message as AgentMessage);
 		} else if (entry.type === "custom_message") {
 			messages.push(
-				createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp),
+				createCustomMessage(
+					entry.customType,
+					entry.content as string | (TextContent | ImageContent)[],
+					entry.display,
+					entry.details,
+					entry.timestamp,
+				),
 			);
 		} else if (entry.type === "branch_summary" && entry.summary) {
 			messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp));
@@ -59,18 +65,18 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext
 
 	if (compaction) {
 		messages.push(createCompactionSummaryMessage(compaction.summary, compaction.tokensBefore, compaction.timestamp));
-		const compactionIdx = entries.findIndex((e) => e.type === "compaction" && e.id === compaction.id);
+		const compactionIdx = pathEntries.findIndex((e) => e.type === "compaction" && e.id === compaction.id);
 		let foundFirstKept = false;
 		for (let i = 0; i < compactionIdx; i++) {
-			const entry = entries[i]!;
+			const entry = pathEntries[i]!;
 			if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true;
 			if (foundFirstKept) appendMessage(entry);
 		}
-		for (let i = compactionIdx + 1; i < entries.length; i++) {
-			appendMessage(entries[i]!);
+		for (let i = compactionIdx + 1; i < pathEntries.length; i++) {
+			appendMessage(pathEntries[i]!);
 		}
 	} else {
-		for (const entry of entries) {
+		for (const entry of pathEntries) {
 			appendMessage(entry);
 		}
 	}
@@ -78,13 +84,21 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext
 	return { messages, thinkingLevel, model };
 }
 
-export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> implements SessionTree {
-	private storage: SessionTreeStorage<TInfo>;
+export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata> implements Session<TMetadata> {
+	private storage: SessionStorage<TMetadata>;
 
-	constructor(storage: SessionTreeStorage<TInfo>) {
+	constructor(storage: SessionStorage<TMetadata>) {
 		this.storage = storage;
 	}
 
+	getMetadata(): Promise<TMetadata> {
+		return this.storage.getMetadata();
+	}
+
+	getStorage(): SessionStorage<TMetadata> {
+		return this.storage;
+	}
+
 	getLeafId(): Promise<string | null> {
 		return this.storage.getLeafId();
 	}
@@ -106,19 +120,8 @@ export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> impleme
 		return buildSessionContext(await this.getBranch());
 	}
 
-	getSessionInfo(): Promise<TInfo> {
-		return this.storage.getSessionInfo();
-	}
-
-	async getLabel(id: string): Promise<string | undefined> {
-		const entries = await this.storage.getEntries();
-		for (let i = entries.length - 1; i >= 0; i--) {
-			const entry = entries[i]!;
-			if (entry.type === "label" && entry.targetId === id) {
-				return entry.label?.trim() || undefined;
-			}
-		}
-		return undefined;
+	getLabel(id: string): Promise<string | undefined> {
+		return this.storage.getLabel(id);
 	}
 
 	async getSessionName(): Promise<string | undefined> {
@@ -193,24 +196,6 @@ export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> impleme
 		} satisfies CompactionEntry<T>);
 	}
 
-	async appendBranchSummary<T = unknown>(
-		fromId: string,
-		summary: string,
-		details?: T,
-		fromHook?: boolean,
-	): Promise<string> {
-		return this.appendTypedEntry({
-			type: "branch_summary",
-			id: await this.makeEntryId(),
-			parentId: await this.storage.getLeafId(),
-			timestamp: new Date().toISOString(),
-			fromId,
-			summary,
-			details,
-			fromHook,
-		} satisfies BranchSummaryEntry<T>);
-	}
-
 	async appendCustomEntry(customType: string, data?: unknown): Promise<string> {
 		return this.appendTypedEntry({
 			type: "custom",
@@ -240,7 +225,10 @@ export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> impleme
 		} satisfies CustomMessageEntry<T>);
 	}
 
-	async appendLabelChange(targetId: string, label: string | undefined): Promise<string> {
+	async appendLabel(targetId: string, label: string | undefined): Promise<string> {
+		if (!(await this.storage.getEntry(targetId))) {
+			throw new Error(`Entry ${targetId} not found`);
+		}
 		return this.appendTypedEntry({
 			type: "label",
 			id: await this.makeEntryId(),
@@ -251,7 +239,7 @@ export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> impleme
 		} satisfies LabelEntry);
 	}
 
-	async appendSessionInfo(name: string): Promise<string> {
+	async appendSessionName(name: string): Promise<string> {
 		return this.appendTypedEntry({
 			type: "session_info",
 			id: await this.makeEntryId(),
@@ -261,10 +249,24 @@ export class DefaultSessionTree<TInfo extends SessionInfo = SessionInfo> impleme
 		} satisfies SessionInfoEntry);
 	}
 
-	async moveTo(entryId: string | null): Promise<void> {
+	async moveTo(
+		entryId: string | null,
+		summary?: { summary: string; details?: unknown; fromHook?: boolean },
+	): Promise<string | undefined> {
 		if (entryId !== null && !(await this.storage.getEntry(entryId))) {
 			throw new Error(`Entry ${entryId} not found`);
 		}
 		await this.storage.setLeafId(entryId);
+		if (!summary) return undefined;
+		return this.appendTypedEntry({
+			type: "branch_summary",
+			id: await this.makeEntryId(),
+			parentId: entryId,
+			timestamp: new Date().toISOString(),
+			fromId: entryId ?? "root",
+			summary: summary.summary,
+			details: summary.details,
+			fromHook: summary.fromHook,
+		} satisfies BranchSummaryEntry);
 	}
 }