Refactor session manager: migration chain, validation, tests

- Add migrateV1ToV2/migrateToCurrentVersion for extensible migrations - createSummaryMessage now takes timestamp from entry - loadEntriesFromFile validates session header - findMostRecentSession only returns valid session files (reads first 512 bytes) - Remove ConversationEntry alias - Fix mom context.ts TreeNode type Tests: - migration.test.ts: v1 migration, idempotency - build-context.test.ts: 14 tests covering trivial, compaction, branches - file-operations.test.ts: loadEntriesFromFile, findMostRecentSession
2026-04-15 10:05:14 +00:00 · 2025-12-26 01:33:13 +01:00 · 2025-12-26 01:33:13 +01:00 · beb70f126d
commit beb70f126d
parent 95312e00bb
7 changed files with 606 additions and 102 deletions
--- a/packages/coding-agent/src/core/session-manager.ts
+++ b/packages/coding-agent/src/core/session-manager.ts
@ -1,15 +1,22 @@
 import type { AppMessage } from "@mariozechner/pi-agent-core";
 import { randomUUID } from "crypto";
-import { appendFileSync, existsSync, mkdirSync, readdirSync, readFileSync, statSync, writeFileSync } from "fs";
+import {
+	appendFileSync,
+	closeSync,
+	existsSync,
+	mkdirSync,
+	openSync,
+	readdirSync,
+	readFileSync,
+	readSync,
+	statSync,
+	writeFileSync,
+} from "fs";
 import { join, resolve } from "path";
 import { getAgentDir as getDefaultAgentDir } from "../config.js";

 export const CURRENT_SESSION_VERSION = 2;

-// ============================================================================
-// Session Header (metadata, not part of conversation tree)
-// ============================================================================
-
 export interface SessionHeader {
 	type: "session";
 	version?: number; // v1 sessions don't have this
@ -19,20 +26,6 @@ export interface SessionHeader {
 	branchedFrom?: string;
 }

-// ============================================================================
-// Tree Node (added by SessionManager to all conversation entries)
-// ============================================================================
-
-export interface TreeNode {
-	id: string;
-	parentId: string | null;
-	timestamp: string;
-}
-
-// ============================================================================
-// Content Types (what distinguishes entries - used for input)
-// ============================================================================
-
 export interface MessageContent {
 	type: "message";
 	message: AppMessage;
@ -61,17 +54,20 @@ export interface BranchSummaryContent {
 	summary: string;
 }

-/** Union of all content types (for input) */
-export type ConversationContent =
+/** Union of all content types (for "write" methods in SessionManager) */
+export type SessionContent =
 	| MessageContent
 	| ThinkingLevelContent
 	| ModelChangeContent
 	| CompactionContent
 	| BranchSummaryContent;

-// ============================================================================
-// Full Entry Types (TreeNode + Content - returned from SessionManager)
-// ============================================================================
+export interface TreeNode {
+	type: string;
+	id: string;
+	parentId: string | null;
+	timestamp: string;
+}

 export type SessionMessageEntry = TreeNode & MessageContent;
 export type ThinkingLevelChangeEntry = TreeNode & ThinkingLevelContent;
@ -79,7 +75,7 @@ export type ModelChangeEntry = TreeNode & ModelChangeContent;
 export type CompactionEntry = TreeNode & CompactionContent;
 export type BranchSummaryEntry = TreeNode & BranchSummaryContent;

-/** Session entry - has id/parentId for tree structure */
+/** Session entry - has id/parentId for tree structure (returned by "read" methods in SessionManager) */
 export type SessionEntry =
 	| SessionMessageEntry
 	| ThinkingLevelChangeEntry
@ -87,9 +83,6 @@ export type SessionEntry =
 	| CompactionEntry
 	| BranchSummaryEntry;

-/** @deprecated Use SessionEntry */
-export type ConversationEntry = SessionEntry;
-
 /** Raw file entry (includes header) */
 export type FileEntry = SessionHeader | SessionEntry;

@ -118,46 +111,46 @@ export const SUMMARY_SUFFIX = `
 </summary>`;

 /** Exported for compaction.test.ts */
-export function createSummaryMessage(summary: string): AppMessage {
+export function createSummaryMessage(summary: string, timestamp: string): AppMessage {
 	return {
 		role: "user",
 		content: SUMMARY_PREFIX + summary + SUMMARY_SUFFIX,
-		timestamp: Date.now(),
+		timestamp: new Date(timestamp).getTime(),
 	};
 }

-/**
- * Migrate v1 entries to v2 format by adding id/parentId fields.
- * Mutates entries in place. Safe to call on already-migrated entries.
- */
-export function migrateSessionEntries(entries: FileEntry[]): void {
-	// Check if already migrated
-	const firstConv = entries.find((e) => e.type !== "session");
-	if (firstConv && "id" in firstConv && firstConv.id) {
-		return; // Already migrated
+/** Generate a unique short ID (8 hex chars, collision-checked) */
+function generateId(byId: { has(id: string): boolean }): string {
+	for (let i = 0; i < 100; i++) {
+		const id = randomUUID().slice(0, 8);
+		if (!byId.has(id)) return id;
 	}
+	// Fallback to full UUID if somehow we have collisions
+	return randomUUID();
+}

+/** Migrate v1 → v2: add id/parentId tree structure. Mutates in place. */
+function migrateV1ToV2(entries: FileEntry[]): void {
+	const ids = new Set<string>();
 	let prevId: string | null = null;
+
 	for (const entry of entries) {
 		if (entry.type === "session") {
-			entry.version = CURRENT_SESSION_VERSION;
+			entry.version = 2;
 			continue;
 		}

-		// Add id/parentId to conversation entries
-		const convEntry = entry as ConversationEntry;
-		convEntry.id = randomUUID();
-		convEntry.parentId = prevId;
-		prevId = convEntry.id;
+		entry.id = generateId(ids);
+		entry.parentId = prevId;
+		prevId = entry.id;

 		// Convert firstKeptEntryIndex to firstKeptEntryId for compaction
 		if (entry.type === "compaction") {
 			const comp = entry as CompactionEntry & { firstKeptEntryIndex?: number };
 			if (typeof comp.firstKeptEntryIndex === "number") {
-				// Find the entry at that index and get its id
 				const targetEntry = entries[comp.firstKeptEntryIndex];
 				if (targetEntry && targetEntry.type !== "session") {
-					comp.firstKeptEntryId = (targetEntry as ConversationEntry).id;
+					comp.firstKeptEntryId = targetEntry.id;
 				}
 				delete comp.firstKeptEntryIndex;
 			}
@ -165,15 +158,39 @@ export function migrateSessionEntries(entries: FileEntry[]): void {
 	}
 }

+// Add future migrations here:
+// function migrateV2ToV3(entries: FileEntry[]): void { ... }
+
+/**
+ * Run all necessary migrations to bring entries to current version.
+ * Mutates entries in place. Returns true if any migration was applied.
+ */
+function migrateToCurrentVersion(entries: FileEntry[]): boolean {
+	const header = entries.find((e) => e.type === "session") as SessionHeader | undefined;
+	const version = header?.version ?? 1;
+
+	if (version >= CURRENT_SESSION_VERSION) return false;
+
+	if (version < 2) migrateV1ToV2(entries);
+	// if (version < 3) migrateV2ToV3(entries);
+
+	return true;
+}
+
+/** Exported for testing */
+export function migrateSessionEntries(entries: FileEntry[]): void {
+	migrateToCurrentVersion(entries);
+}
+
 /** Exported for compaction.test.ts */
 export function parseSessionEntries(content: string): FileEntry[] {
-	const entries: SessionEntry[] = [];
+	const entries: FileEntry[] = [];
 	const lines = content.trim().split("\n");

 	for (const line of lines) {
 		if (!line.trim()) continue;
 		try {
-			const entry = JSON.parse(line) as SessionEntry;
+			const entry = JSON.parse(line) as FileEntry;
 			entries.push(entry);
 		} catch {
 			// Skip malformed lines
@ -197,18 +214,26 @@ export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEnt
 * If leafId is provided, walks from that entry to root.
 * Handles compaction and branch summaries along the path.
 */
-export function buildSessionContext(entries: SessionEntry[], leafId?: string): SessionContext {
-	// Build uuid index
-	const byId = new Map<string, SessionEntry>();
-	for (const entry of entries) {
-		byId.set(entry.id, entry);
+export function buildSessionContext(
+	entries: SessionEntry[],
+	leafId?: string,
+	byId?: Map<string, SessionEntry>,
+): SessionContext {
+	// Build uuid index if not available
+	if (!byId) {
+		byId = new Map<string, SessionEntry>();
+		for (const entry of entries) {
+			byId.set(entry.id, entry);
+		}
 	}

 	// Find leaf
 	let leaf: SessionEntry | undefined;
 	if (leafId) {
 		leaf = byId.get(leafId);
-	} else {
+	}
+	if (!leaf) {
+		// Fallback to last entry
 		leaf = entries[entries.length - 1];
 	}

@ -250,7 +275,7 @@ export function buildSessionContext(entries: SessionEntry[], leafId?: string): S

 	if (compaction) {
 		// Emit summary first
-		messages.push(createSummaryMessage(compaction.summary));
+		messages.push(createSummaryMessage(compaction.summary, compaction.timestamp));

 		// Find compaction index in path
 		const compactionIdx = path.findIndex((e) => e.type === "compaction" && e.id === compaction.id);
@ -273,7 +298,7 @@ export function buildSessionContext(entries: SessionEntry[], leafId?: string): S
 			if (entry.type === "message") {
 				messages.push(entry.message);
 			} else if (entry.type === "branch_summary") {
-				messages.push(createSummaryMessage(entry.summary));
+				messages.push(createSummaryMessage(entry.summary, entry.timestamp));
 			}
 		}
 	} else {
@ -282,7 +307,7 @@ export function buildSessionContext(entries: SessionEntry[], leafId?: string): S
 			if (entry.type === "message") {
 				messages.push(entry.message);
 			} else if (entry.type === "branch_summary") {
-				messages.push(createSummaryMessage(entry.summary));
+				messages.push(createSummaryMessage(entry.summary, entry.timestamp));
 			}
 		}
 	}
@ -303,34 +328,57 @@ function getDefaultSessionDir(cwd: string): string {
 	return sessionDir;
 }

-function loadEntriesFromFile(filePath: string): FileEntry[] {
+/** Exported for testing */
+export function loadEntriesFromFile(filePath: string): FileEntry[] {
 	if (!existsSync(filePath)) return [];

 	const content = readFileSync(filePath, "utf8");
-	const entries: SessionEntry[] = [];
+	const entries: FileEntry[] = [];
 	const lines = content.trim().split("\n");

 	for (const line of lines) {
 		if (!line.trim()) continue;
 		try {
-			const entry = JSON.parse(line) as SessionEntry;
+			const entry = JSON.parse(line) as FileEntry;
 			entries.push(entry);
 		} catch {
 			// Skip malformed lines
 		}
 	}

+	// Validate session header
+	if (entries.length === 0) return entries;
+	const header = entries[0];
+	if (header.type !== "session" || typeof (header as any).id !== "string") {
+		return [];
+	}
+
 	return entries;
 }

-function findMostRecentSession(sessionDir: string): string | null {
+function isValidSessionFile(filePath: string): boolean {
+	try {
+		const fd = openSync(filePath, "r");
+		const buffer = Buffer.alloc(512);
+		const bytesRead = readSync(fd, buffer, 0, 512, 0);
+		closeSync(fd);
+		const firstLine = buffer.toString("utf8", 0, bytesRead).split("\n")[0];
+		if (!firstLine) return false;
+		const header = JSON.parse(firstLine);
+		return header.type === "session" && typeof header.id === "string";
+	} catch {
+		return false;
+	}
+}
+
+/** Exported for testing */
+export function findMostRecentSession(sessionDir: string): string | null {
 	try {
 		const files = readdirSync(sessionDir)
 			.filter((f) => f.endsWith(".jsonl"))
-			.map((f) => ({
-				path: join(sessionDir, f),
-				mtime: statSync(join(sessionDir, f)).mtime,
-			}))
+			.map((f) => join(sessionDir, f))
+			.filter(isValidSessionFile)
+			.map((path) => ({ path, mtime: statSync(path).mtime }))
 			.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());

 		return files[0]?.path || null;
@ -347,21 +395,9 @@ export class SessionManager {
 	private persist: boolean;
 	private flushed: boolean = false;
 	private inMemoryEntries: FileEntry[] = [];
-
-	// Tree structure (v2)
 	private byId: Map<string, SessionEntry> = new Map();
 	private leafId: string = "";

-	/** Generate a unique short ID (8 hex chars, collision-checked) */
-	private _generateId(): string {
-		for (let i = 0; i < 100; i++) {
-			const id = randomUUID().slice(0, 8);
-			if (!this.byId.has(id)) return id;
-		}
-		// Fallback to full UUID if somehow we have collisions
-		return randomUUID();
-	}
-
 	private constructor(cwd: string, sessionDir: string, sessionFile: string | null, persist: boolean) {
 		this.cwd = cwd;
 		this.sessionDir = sessionDir;
@ -385,10 +421,7 @@ export class SessionManager {
 			const header = this.inMemoryEntries.find((e) => e.type === "session") as SessionHeader | undefined;
 			this.sessionId = header?.id ?? randomUUID();

-			// Migrate v1 to v2 if needed
-			const version = header?.version ?? 1;
-			if (version < CURRENT_SESSION_VERSION) {
-				this._migrateToV2();
+			if (migrateToCurrentVersion(this.inMemoryEntries)) {
 				this._rewriteFile();
 			}

@ -420,10 +453,6 @@ export class SessionManager {
 		}
 	}

-	private _migrateToV2(): void {
-		migrateSessionEntries(this.inMemoryEntries);
-	}
-
 	private _buildIndex(): void {
 		this.byId.clear();
 		this.leafId = "";
@ -480,7 +509,7 @@ export class SessionManager {
 		}
 	}

-	private _appendEntry(entry: ConversationEntry): void {
+	private _appendEntry(entry: SessionEntry): void {
 		this.inMemoryEntries.push(entry);
 		this.byId.set(entry.id, entry);
 		this.leafId = entry.id;
@ -490,7 +519,7 @@ export class SessionManager {
 	saveMessage(message: AppMessage): string {
 		const entry: SessionMessageEntry = {
 			type: "message",
-			id: this._generateId(),
+			id: generateId(this.byId),
 			parentId: this.leafId || null,
 			timestamp: new Date().toISOString(),
 			message,
@ -502,7 +531,7 @@ export class SessionManager {
 	saveThinkingLevelChange(thinkingLevel: string): string {
 		const entry: ThinkingLevelChangeEntry = {
 			type: "thinking_level_change",
-			id: this._generateId(),
+			id: generateId(this.byId),
 			parentId: this.leafId || null,
 			timestamp: new Date().toISOString(),
 			thinkingLevel,
@ -514,7 +543,7 @@ export class SessionManager {
 	saveModelChange(provider: string, modelId: string): string {
 		const entry: ModelChangeEntry = {
 			type: "model_change",
-			id: this._generateId(),
+			id: generateId(this.byId),
 			parentId: this.leafId || null,
 			timestamp: new Date().toISOString(),
 			provider,
@ -527,7 +556,7 @@ export class SessionManager {
 	saveCompaction(summary: string, firstKeptEntryId: string, tokensBefore: number): string {
 		const entry: CompactionEntry = {
 			type: "compaction",
-			id: this._generateId(),
+			id: generateId(this.byId),
 			parentId: this.leafId || null,
 			timestamp: new Date().toISOString(),
 			summary,
@ -546,13 +575,13 @@ export class SessionManager {
 		return this.leafId;
 	}

-	getEntry(id: string): ConversationEntry | undefined {
+	getEntry(id: string): SessionEntry | undefined {
 		return this.byId.get(id);
 	}

 	/** Walk from entry to root, returning path (conversation entries only) */
-	getPath(fromId?: string): ConversationEntry[] {
-		const path: ConversationEntry[] = [];
+	getPath(fromId?: string): SessionEntry[] {
+		const path: SessionEntry[] = [];
 		let current = this.byId.get(fromId ?? this.leafId);
 		while (current) {
 			path.unshift(current);
@ -566,7 +595,7 @@ export class SessionManager {
 	 * Uses tree traversal from current leaf.
 	 */
 	buildSessionContext(): SessionContext {
-		return buildSessionContext(this.getEntries(), this.leafId);
+		return buildSessionContext(this.getEntries(), this.leafId, this.byId);
 	}

 	/**
@ -605,7 +634,7 @@ export class SessionManager {
 		this.leafId = branchFromId;
 		const entry: BranchSummaryEntry = {
 			type: "branch_summary",
-			id: this._generateId(),
+			id: generateId(this.byId),
 			parentId: branchFromId,
 			timestamp: new Date().toISOString(),
 			summary,
--- a/packages/coding-agent/src/index.ts
+++ b/packages/coding-agent/src/index.ts
@ -112,8 +112,6 @@ export {
 	buildSessionContext,
 	type CompactionContent,
 	type CompactionEntry,
-	type ConversationContent,
-	type ConversationEntry,
 	CURRENT_SESSION_VERSION,
 	createSummaryMessage,
 	type FileEntry,
@ -123,6 +121,7 @@ export {
 	type ModelChangeEntry,
 	migrateSessionEntries,
 	parseSessionEntries,
+	type SessionContent as ConversationContent,
 	type SessionContext as LoadedSession,
 	type SessionEntry,
 	type SessionHeader,
--- a/packages/coding-agent/test/compaction.test.ts
+++ b/packages/coding-agent/test/compaction.test.ts
@ -273,9 +273,11 @@ describe("findCutPoint", () => {
 });

 describe("createSummaryMessage", () => {
-	it("should create user message with prefix", () => {
-		const msg = createSummaryMessage("This is the summary");
+	it("should create user message with prefix and correct timestamp", () => {
+		const ts = "2025-01-01T12:00:00.000Z";
+		const msg = createSummaryMessage("This is the summary", ts);
 		expect(msg.role).toBe("user");
+		expect(msg.timestamp).toBe(new Date(ts).getTime());
 		if (msg.role === "user") {
 			expect(msg.content).toContain(
 				"The conversation history before this point was compacted into the following summary:",
--- a/packages/coding-agent/test/session-manager/build-context.test.ts
+++ b/packages/coding-agent/test/session-manager/build-context.test.ts
@ -0,0 +1,269 @@
+import { describe, expect, it } from "vitest";
+import {
+	type BranchSummaryEntry,
+	buildSessionContext,
+	type CompactionEntry,
+	type ModelChangeEntry,
+	type SessionEntry,
+	type SessionMessageEntry,
+	SUMMARY_PREFIX,
+	type ThinkingLevelChangeEntry,
+} from "../../src/core/session-manager.js";
+
+function msg(id: string, parentId: string | null, role: "user" | "assistant", text: string): SessionMessageEntry {
+	const base = { type: "message" as const, id, parentId, timestamp: "2025-01-01T00:00:00Z" };
+	if (role === "user") {
+		return { ...base, message: { role, content: text, timestamp: 1 } };
+	}
+	return {
+		...base,
+		message: {
+			role,
+			content: [{ type: "text", text }],
+			api: "anthropic-messages",
+			provider: "anthropic",
+			model: "claude-test",
+			usage: {
+				input: 1,
+				output: 1,
+				cacheRead: 0,
+				cacheWrite: 0,
+				totalTokens: 2,
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+			},
+			stopReason: "stop",
+			timestamp: 1,
+		},
+	};
+}
+
+function compaction(id: string, parentId: string | null, summary: string, firstKeptEntryId: string): CompactionEntry {
+	return {
+		type: "compaction",
+		id,
+		parentId,
+		timestamp: "2025-01-01T00:00:00Z",
+		summary,
+		firstKeptEntryId,
+		tokensBefore: 1000,
+	};
+}
+
+function branchSummary(id: string, parentId: string | null, summary: string): BranchSummaryEntry {
+	return { type: "branch_summary", id, parentId, timestamp: "2025-01-01T00:00:00Z", summary };
+}
+
+function thinkingLevel(id: string, parentId: string | null, level: string): ThinkingLevelChangeEntry {
+	return { type: "thinking_level_change", id, parentId, timestamp: "2025-01-01T00:00:00Z", thinkingLevel: level };
+}
+
+function modelChange(id: string, parentId: string | null, provider: string, modelId: string): ModelChangeEntry {
+	return { type: "model_change", id, parentId, timestamp: "2025-01-01T00:00:00Z", provider, modelId };
+}
+
+describe("buildSessionContext", () => {
+	describe("trivial cases", () => {
+		it("empty entries returns empty context", () => {
+			const ctx = buildSessionContext([]);
+			expect(ctx.messages).toEqual([]);
+			expect(ctx.thinkingLevel).toBe("off");
+			expect(ctx.model).toBeNull();
+		});
+
+		it("single user message", () => {
+			const entries: SessionEntry[] = [msg("1", null, "user", "hello")];
+			const ctx = buildSessionContext(entries);
+			expect(ctx.messages).toHaveLength(1);
+			expect(ctx.messages[0].role).toBe("user");
+		});
+
+		it("simple conversation", () => {
+			const entries: SessionEntry[] = [
+				msg("1", null, "user", "hello"),
+				msg("2", "1", "assistant", "hi there"),
+				msg("3", "2", "user", "how are you"),
+				msg("4", "3", "assistant", "great"),
+			];
+			const ctx = buildSessionContext(entries);
+			expect(ctx.messages).toHaveLength(4);
+			expect(ctx.messages.map((m) => m.role)).toEqual(["user", "assistant", "user", "assistant"]);
+		});
+
+		it("tracks thinking level changes", () => {
+			const entries: SessionEntry[] = [
+				msg("1", null, "user", "hello"),
+				thinkingLevel("2", "1", "high"),
+				msg("3", "2", "assistant", "thinking hard"),
+			];
+			const ctx = buildSessionContext(entries);
+			expect(ctx.thinkingLevel).toBe("high");
+			expect(ctx.messages).toHaveLength(2);
+		});
+
+		it("tracks model from assistant message", () => {
+			const entries: SessionEntry[] = [msg("1", null, "user", "hello"), msg("2", "1", "assistant", "hi")];
+			const ctx = buildSessionContext(entries);
+			expect(ctx.model).toEqual({ provider: "anthropic", modelId: "claude-test" });
+		});
+
+		it("tracks model from model change entry", () => {
+			const entries: SessionEntry[] = [
+				msg("1", null, "user", "hello"),
+				modelChange("2", "1", "openai", "gpt-4"),
+				msg("3", "2", "assistant", "hi"),
+			];
+			const ctx = buildSessionContext(entries);
+			// Assistant message overwrites model change
+			expect(ctx.model).toEqual({ provider: "anthropic", modelId: "claude-test" });
+		});
+	});
+
+	describe("with compaction", () => {
+		it("includes summary before kept messages", () => {
+			const entries: SessionEntry[] = [
+				msg("1", null, "user", "first"),
+				msg("2", "1", "assistant", "response1"),
+				msg("3", "2", "user", "second"),
+				msg("4", "3", "assistant", "response2"),
+				compaction("5", "4", "Summary of first two turns", "3"),
+				msg("6", "5", "user", "third"),
+				msg("7", "6", "assistant", "response3"),
+			];
+			const ctx = buildSessionContext(entries);
+
+			// Should have: summary + kept (3,4) + after (6,7) = 5 messages
+			expect(ctx.messages).toHaveLength(5);
+			expect((ctx.messages[0] as any).content).toContain("Summary of first two turns");
+			expect((ctx.messages[1] as any).content).toBe("second");
+			expect((ctx.messages[2] as any).content[0].text).toBe("response2");
+			expect((ctx.messages[3] as any).content).toBe("third");
+			expect((ctx.messages[4] as any).content[0].text).toBe("response3");
+		});
+
+		it("handles compaction keeping from first message", () => {
+			const entries: SessionEntry[] = [
+				msg("1", null, "user", "first"),
+				msg("2", "1", "assistant", "response"),
+				compaction("3", "2", "Empty summary", "1"),
+				msg("4", "3", "user", "second"),
+			];
+			const ctx = buildSessionContext(entries);
+
+			// Summary + all messages (1,2,4)
+			expect(ctx.messages).toHaveLength(4);
+			expect((ctx.messages[0] as any).content).toContain(SUMMARY_PREFIX);
+		});
+
+		it("multiple compactions uses latest", () => {
+			const entries: SessionEntry[] = [
+				msg("1", null, "user", "a"),
+				msg("2", "1", "assistant", "b"),
+				compaction("3", "2", "First summary", "1"),
+				msg("4", "3", "user", "c"),
+				msg("5", "4", "assistant", "d"),
+				compaction("6", "5", "Second summary", "4"),
+				msg("7", "6", "user", "e"),
+			];
+			const ctx = buildSessionContext(entries);
+
+			// Should use second summary, keep from 4
+			expect(ctx.messages).toHaveLength(4);
+			expect((ctx.messages[0] as any).content).toContain("Second summary");
+		});
+	});
+
+	describe("with branches", () => {
+		it("follows path to specified leaf", () => {
+			// Tree:
+			//   1 -> 2 -> 3 (branch A)
+			//         \-> 4 (branch B)
+			const entries: SessionEntry[] = [
+				msg("1", null, "user", "start"),
+				msg("2", "1", "assistant", "response"),
+				msg("3", "2", "user", "branch A"),
+				msg("4", "2", "user", "branch B"),
+			];
+
+			const ctxA = buildSessionContext(entries, "3");
+			expect(ctxA.messages).toHaveLength(3);
+			expect((ctxA.messages[2] as any).content).toBe("branch A");
+
+			const ctxB = buildSessionContext(entries, "4");
+			expect(ctxB.messages).toHaveLength(3);
+			expect((ctxB.messages[2] as any).content).toBe("branch B");
+		});
+
+		it("includes branch summary in path", () => {
+			const entries: SessionEntry[] = [
+				msg("1", null, "user", "start"),
+				msg("2", "1", "assistant", "response"),
+				msg("3", "2", "user", "abandoned path"),
+				branchSummary("4", "2", "Summary of abandoned work"),
+				msg("5", "4", "user", "new direction"),
+			];
+			const ctx = buildSessionContext(entries, "5");
+
+			expect(ctx.messages).toHaveLength(4);
+			expect((ctx.messages[2] as any).content).toContain("Summary of abandoned work");
+			expect((ctx.messages[3] as any).content).toBe("new direction");
+		});
+
+		it("complex tree with multiple branches and compaction", () => {
+			// Tree:
+			//   1 -> 2 -> 3 -> 4 -> compaction(5) -> 6 -> 7 (main path)
+			//              \-> 8 -> 9 (abandoned branch)
+			//                    \-> branchSummary(10) -> 11 (resumed from 3)
+			const entries: SessionEntry[] = [
+				msg("1", null, "user", "start"),
+				msg("2", "1", "assistant", "r1"),
+				msg("3", "2", "user", "q2"),
+				msg("4", "3", "assistant", "r2"),
+				compaction("5", "4", "Compacted history", "3"),
+				msg("6", "5", "user", "q3"),
+				msg("7", "6", "assistant", "r3"),
+				// Abandoned branch from 3
+				msg("8", "3", "user", "wrong path"),
+				msg("9", "8", "assistant", "wrong response"),
+				// Branch summary resuming from 3
+				branchSummary("10", "3", "Tried wrong approach"),
+				msg("11", "10", "user", "better approach"),
+			];
+
+			// Main path to 7: summary + kept(3,4) + after(6,7)
+			const ctxMain = buildSessionContext(entries, "7");
+			expect(ctxMain.messages).toHaveLength(5);
+			expect((ctxMain.messages[0] as any).content).toContain("Compacted history");
+			expect((ctxMain.messages[1] as any).content).toBe("q2");
+			expect((ctxMain.messages[2] as any).content[0].text).toBe("r2");
+			expect((ctxMain.messages[3] as any).content).toBe("q3");
+			expect((ctxMain.messages[4] as any).content[0].text).toBe("r3");
+
+			// Branch path to 11: 1,2,3 + branch_summary + 11
+			const ctxBranch = buildSessionContext(entries, "11");
+			expect(ctxBranch.messages).toHaveLength(5);
+			expect((ctxBranch.messages[0] as any).content).toBe("start");
+			expect((ctxBranch.messages[1] as any).content[0].text).toBe("r1");
+			expect((ctxBranch.messages[2] as any).content).toBe("q2");
+			expect((ctxBranch.messages[3] as any).content).toContain("Tried wrong approach");
+			expect((ctxBranch.messages[4] as any).content).toBe("better approach");
+		});
+	});
+
+	describe("edge cases", () => {
+		it("uses last entry when leafId not found", () => {
+			const entries: SessionEntry[] = [msg("1", null, "user", "hello"), msg("2", "1", "assistant", "hi")];
+			const ctx = buildSessionContext(entries, "nonexistent");
+			expect(ctx.messages).toHaveLength(2);
+		});
+
+		it("handles orphaned entries gracefully", () => {
+			const entries: SessionEntry[] = [
+				msg("1", null, "user", "hello"),
+				msg("2", "missing", "assistant", "orphan"), // parent doesn't exist
+			];
+			const ctx = buildSessionContext(entries, "2");
+			// Should only get the orphan since parent chain is broken
+			expect(ctx.messages).toHaveLength(1);
+		});
+	});
+});
--- a/packages/coding-agent/test/session-manager/file-operations.test.ts
+++ b/packages/coding-agent/test/session-manager/file-operations.test.ts
@ -0,0 +1,127 @@
+import { mkdirSync, rmSync, writeFileSync } from "fs";
+import { tmpdir } from "os";
+import { join } from "path";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { findMostRecentSession, loadEntriesFromFile } from "../../src/core/session-manager.js";
+
+describe("loadEntriesFromFile", () => {
+	let tempDir: string;
+
+	beforeEach(() => {
+		tempDir = join(tmpdir(), `session-test-${Date.now()}`);
+		mkdirSync(tempDir, { recursive: true });
+	});
+
+	afterEach(() => {
+		rmSync(tempDir, { recursive: true, force: true });
+	});
+
+	it("returns empty array for non-existent file", () => {
+		const entries = loadEntriesFromFile(join(tempDir, "nonexistent.jsonl"));
+		expect(entries).toEqual([]);
+	});
+
+	it("returns empty array for empty file", () => {
+		const file = join(tempDir, "empty.jsonl");
+		writeFileSync(file, "");
+		expect(loadEntriesFromFile(file)).toEqual([]);
+	});
+
+	it("returns empty array for file without valid session header", () => {
+		const file = join(tempDir, "no-header.jsonl");
+		writeFileSync(file, '{"type":"message","id":"1"}\n');
+		expect(loadEntriesFromFile(file)).toEqual([]);
+	});
+
+	it("returns empty array for malformed JSON", () => {
+		const file = join(tempDir, "malformed.jsonl");
+		writeFileSync(file, "not json\n");
+		expect(loadEntriesFromFile(file)).toEqual([]);
+	});
+
+	it("loads valid session file", () => {
+		const file = join(tempDir, "valid.jsonl");
+		writeFileSync(
+			file,
+			'{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n' +
+				'{"type":"message","id":"1","parentId":null,"timestamp":"2025-01-01T00:00:01Z","message":{"role":"user","content":"hi","timestamp":1}}\n',
+		);
+		const entries = loadEntriesFromFile(file);
+		expect(entries).toHaveLength(2);
+		expect(entries[0].type).toBe("session");
+		expect(entries[1].type).toBe("message");
+	});
+
+	it("skips malformed lines but keeps valid ones", () => {
+		const file = join(tempDir, "mixed.jsonl");
+		writeFileSync(
+			file,
+			'{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n' +
+				"not valid json\n" +
+				'{"type":"message","id":"1","parentId":null,"timestamp":"2025-01-01T00:00:01Z","message":{"role":"user","content":"hi","timestamp":1}}\n',
+		);
+		const entries = loadEntriesFromFile(file);
+		expect(entries).toHaveLength(2);
+	});
+});
+
+describe("findMostRecentSession", () => {
+	let tempDir: string;
+
+	beforeEach(() => {
+		tempDir = join(tmpdir(), `session-test-${Date.now()}`);
+		mkdirSync(tempDir, { recursive: true });
+	});
+
+	afterEach(() => {
+		rmSync(tempDir, { recursive: true, force: true });
+	});
+
+	it("returns null for empty directory", () => {
+		expect(findMostRecentSession(tempDir)).toBeNull();
+	});
+
+	it("returns null for non-existent directory", () => {
+		expect(findMostRecentSession(join(tempDir, "nonexistent"))).toBeNull();
+	});
+
+	it("ignores non-jsonl files", () => {
+		writeFileSync(join(tempDir, "file.txt"), "hello");
+		writeFileSync(join(tempDir, "file.json"), "{}");
+		expect(findMostRecentSession(tempDir)).toBeNull();
+	});
+
+	it("ignores jsonl files without valid session header", () => {
+		writeFileSync(join(tempDir, "invalid.jsonl"), '{"type":"message"}\n');
+		expect(findMostRecentSession(tempDir)).toBeNull();
+	});
+
+	it("returns single valid session file", () => {
+		const file = join(tempDir, "session.jsonl");
+		writeFileSync(file, '{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
+		expect(findMostRecentSession(tempDir)).toBe(file);
+	});
+
+	it("returns most recently modified session", async () => {
+		const file1 = join(tempDir, "older.jsonl");
+		const file2 = join(tempDir, "newer.jsonl");
+
+		writeFileSync(file1, '{"type":"session","id":"old","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
+		// Small delay to ensure different mtime
+		await new Promise((r) => setTimeout(r, 10));
+		writeFileSync(file2, '{"type":"session","id":"new","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
+
+		expect(findMostRecentSession(tempDir)).toBe(file2);
+	});
+
+	it("skips invalid files and returns valid one", async () => {
+		const invalid = join(tempDir, "invalid.jsonl");
+		const valid = join(tempDir, "valid.jsonl");
+
+		writeFileSync(invalid, '{"type":"not-session"}\n');
+		await new Promise((r) => setTimeout(r, 10));
+		writeFileSync(valid, '{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
+
+		expect(findMostRecentSession(tempDir)).toBe(valid);
+	});
+});
--- a/packages/coding-agent/test/session-manager/migration.test.ts
+++ b/packages/coding-agent/test/session-manager/migration.test.ts
@ -0,0 +1,78 @@
+import { describe, expect, it } from "vitest";
+import { type FileEntry, migrateSessionEntries } from "../../src/core/session-manager.js";
+
+describe("migrateSessionEntries", () => {
+	it("should add id/parentId to v1 entries", () => {
+		const entries: FileEntry[] = [
+			{ type: "session", id: "sess-1", timestamp: "2025-01-01T00:00:00Z", cwd: "/tmp" },
+			{ type: "message", timestamp: "2025-01-01T00:00:01Z", message: { role: "user", content: "hi", timestamp: 1 } },
+			{
+				type: "message",
+				timestamp: "2025-01-01T00:00:02Z",
+				message: {
+					role: "assistant",
+					content: [{ type: "text", text: "hello" }],
+					api: "test",
+					provider: "test",
+					model: "test",
+					usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 },
+					stopReason: "stop",
+					timestamp: 2,
+				},
+			},
+		] as FileEntry[];
+
+		migrateSessionEntries(entries);
+
+		// Header should have version set
+		expect((entries[0] as any).version).toBe(2);
+
+		// Entries should have id/parentId
+		const msg1 = entries[1] as any;
+		const msg2 = entries[2] as any;
+
+		expect(msg1.id).toBeDefined();
+		expect(msg1.id.length).toBe(8);
+		expect(msg1.parentId).toBeNull();
+
+		expect(msg2.id).toBeDefined();
+		expect(msg2.id.length).toBe(8);
+		expect(msg2.parentId).toBe(msg1.id);
+	});
+
+	it("should be idempotent (skip already migrated)", () => {
+		const entries: FileEntry[] = [
+			{ type: "session", id: "sess-1", version: 2, timestamp: "2025-01-01T00:00:00Z", cwd: "/tmp" },
+			{
+				type: "message",
+				id: "abc12345",
+				parentId: null,
+				timestamp: "2025-01-01T00:00:01Z",
+				message: { role: "user", content: "hi", timestamp: 1 },
+			},
+			{
+				type: "message",
+				id: "def67890",
+				parentId: "abc12345",
+				timestamp: "2025-01-01T00:00:02Z",
+				message: {
+					role: "assistant",
+					content: [{ type: "text", text: "hello" }],
+					api: "test",
+					provider: "test",
+					model: "test",
+					usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 },
+					stopReason: "stop",
+					timestamp: 2,
+				},
+			},
+		] as FileEntry[];
+
+		migrateSessionEntries(entries);
+
+		// IDs should be unchanged
+		expect((entries[1] as any).id).toBe("abc12345");
+		expect((entries[2] as any).id).toBe("def67890");
+		expect((entries[2] as any).parentId).toBe("abc12345");
+	});
+});
--- a/packages/mom/src/context.ts
+++ b/packages/mom/src/context.ts
@ -98,9 +98,9 @@ export class MomSessionManager {
 		this.leafId = null;
 	}

-	private _createTreeNode(): TreeNode {
+	private _createTreeNode(): Omit<TreeNode, "type"> {
 		const id = uuidv4();
-		const node: TreeNode = {
+		const node = {
 			id,
 			parentId: this.leafId,
 			timestamp: new Date().toISOString(),