refactor(agent): harden harness session semantics

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

@@ -1,4 +1,10 @@
-import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js";
+import {
+	type LeafEntry,
+	SessionError,
+	type SessionMetadata,
+	type SessionStorage,
+	type SessionTreeEntry,
+} from "../types.js";
 import { uuidv7 } from "./uuid.js";
 
 function updateLabelCache(labelsById: Map<string, string>, entry: SessionTreeEntry): void {
@@ -27,36 +33,55 @@ function generateEntryId(byId: { has(id: string): boolean }): string {
 	return uuidv7();
 }
 
-export class InMemorySessionStorage implements SessionStorage {
-	private readonly metadata: SessionMetadata;
+function leafIdAfterEntry(entry: SessionTreeEntry): string | null {
+	return entry.type === "leaf" ? entry.targetId : entry.id;
+}
+
+export class InMemorySessionStorage<TMetadata extends SessionMetadata = SessionMetadata>
+	implements SessionStorage<TMetadata>
+{
+	private readonly metadata: TMetadata;
 	private entries: SessionTreeEntry[];
 	private byId: Map<string, SessionTreeEntry>;
 	private labelsById: Map<string, string>;
 	private leafId: string | null;
 
-	constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; metadata?: SessionMetadata }) {
+	constructor(options?: { entries?: SessionTreeEntry[]; metadata?: TMetadata }) {
 		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;
+		this.leafId = null;
+		for (const entry of this.entries) this.leafId = leafIdAfterEntry(entry);
 		if (this.leafId !== null && !this.byId.has(this.leafId)) {
-			throw new Error(`Entry ${this.leafId} not found`);
+			throw new SessionError("invalid_session", `Entry ${this.leafId} not found`);
 		}
-		this.metadata = options?.metadata ?? { id: uuidv7(), createdAt: new Date().toISOString() };
+		this.metadata = options?.metadata ?? ({ id: uuidv7(), createdAt: new Date().toISOString() } as TMetadata);
 	}
 
-	async getMetadata(): Promise<SessionMetadata> {
+	async getMetadata(): Promise<TMetadata> {
 		return this.metadata;
 	}
 
 	async getLeafId(): Promise<string | null> {
+		if (this.leafId !== null && !this.byId.has(this.leafId)) {
+			throw new SessionError("invalid_session", `Entry ${this.leafId} not found`);
+		}
 		return this.leafId;
 	}
 
 	async setLeafId(leafId: string | null): Promise<void> {
 		if (leafId !== null && !this.byId.has(leafId)) {
-			throw new Error(`Entry ${leafId} not found`);
+			throw new SessionError("not_found", `Entry ${leafId} not found`);
 		}
+		const entry: LeafEntry = {
+			type: "leaf",
+			id: generateEntryId(this.byId),
+			parentId: this.leafId,
+			timestamp: new Date().toISOString(),
+			targetId: leafId,
+		};
+		this.entries.push(entry);
+		this.byId.set(entry.id, entry);
 		this.leafId = leafId;
 	}
 
@@ -68,7 +93,7 @@ export class InMemorySessionStorage implements SessionStorage {
 		this.entries.push(entry);
 		this.byId.set(entry.id, entry);
 		updateLabelCache(this.labelsById, entry);
-		this.leafId = entry.id;
+		this.leafId = leafIdAfterEntry(entry);
 	}
 
 	async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
@@ -89,9 +114,13 @@ export class InMemorySessionStorage implements SessionStorage {
 		if (leafId === null) return [];
 		const path: SessionTreeEntry[] = [];
 		let current = this.byId.get(leafId);
+		if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`);
 		while (current) {
 			path.unshift(current);
-			current = current.parentId ? this.byId.get(current.parentId) : undefined;
+			if (!current.parentId) break;
+			const parent = this.byId.get(current.parentId);
+			if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`);
+			current = parent;
 		}
 		return path;
 	}