Session tree structure with id/parentId linking

- Add TreeNode base type with id, parentId, timestamp
- Add *Content types for clean input/output separation
- Entry types are now TreeNode & *Content intersections
- SessionManager assigns id/parentId on save, tracks leafId
- Add migrateSessionEntries() for v1 to v2 conversion
- Migration runs on load, rewrites file
- buildSessionContext() uses tree traversal from leaf
- Compaction returns CompactionResult (content only)
- Hooks return compaction content, not full entries
- Add firstKeptEntryId to before_compact hook event
- Update mom package for tree fields
- Better error messages for compaction failures

packages/coding-agent/docs/session-tree-plan.md

@@ -0,0 +1,86 @@
+# Session Tree Implementation Plan
+
+Reference: [session-tree.md](./session-tree.md)
+
+## Phase 1: SessionManager Core
+
+- [x] Update entry types with `id`, `parentId` fields (using TreeNode intersection)
+- [x] Add `version` field to `SessionHeader`
+- [x] Change `CompactionEntry.firstKeptEntryIndex` → `firstKeptEntryId`
+- [x] Add `BranchSummaryEntry` type
+- [x] Add `byId: Map<string, ConversationEntry>` index
+- [x] Add `leafId: string` tracking
+- [x] Implement `getPath(fromId?)` tree traversal
+- [x] Implement `getEntry(id)` lookup
+- [x] Implement `getLeafId()` helper
+- [x] Update `_buildIndex()` to populate `byId` map
+- [x] Update `saveMessage()` to include id/parentId (returns id)
+- [x] Update `saveCompaction()` signature and fields (returns id)
+- [x] Update `saveThinkingLevelChange()` to include id/parentId (returns id)
+- [x] Update `saveModelChange()` to include id/parentId (returns id)
+- [x] Update `buildSessionContext()` to use `getPath()` traversal
+
+### Type Hierarchy
+
+```typescript
+// Tree fields (added by SessionManager)
+interface TreeNode { id, parentId, timestamp }
+
+// Content types (for input)
+interface MessageContent { type: "message"; message: AppMessage }
+interface CompactionContent { type: "compaction"; summary; firstKeptEntryId; tokensBefore }
+// etc...
+
+// Full entry types (TreeNode & Content)
+type SessionMessageEntry = TreeNode & MessageContent;
+type CompactionEntry = TreeNode & CompactionContent;
+// etc...
+```
+
+## Phase 2: Migration
+
+- [x] Add `CURRENT_SESSION_VERSION = 2` constant
+- [x] Implement `_migrateToV2()` for v1→v2
+- [x] Update `setSessionFile()` to detect version and migrate
+- [x] Implement `_rewriteFile()` for post-migration persistence
+- [x] Handle `firstKeptEntryIndex` → `firstKeptEntryId` conversion in migration
+
+## Phase 3: Branching
+
+- [x] Implement `branchInPlace(id)` - switch leaf pointer
+- [x] Implement `branchWithSummary(id, summary)` - create summary entry
+- [x] Update `branchToNewFile()` to use IDs (no remapping)
+- [ ] Update `AgentSession.branch()` to use new API
+
+## Phase 4: Compaction Integration
+
+- [x] Update `compaction.ts` to work with IDs
+- [x] Update `prepareCompaction()` to return `firstKeptEntryId`
+- [x] Update `compact()` to return `CompactionResult` with `firstKeptEntryId`
+- [x] Update `AgentSession` compaction methods
+- [x] Add `firstKeptEntryId` to `before_compact` hook event
+
+## Phase 5: Testing
+
+- [ ] Add test fixtures from existing sessions
+- [ ] Test migration of v1 sessions
+- [ ] Test context building with tree structure
+- [ ] Test branching operations
+- [ ] Test compaction with IDs
+- [x] Update existing tests for new types
+
+## Phase 6: UI Integration
+
+- [ ] Update `/branch` command for new API
+- [ ] Add `/branch-here` command for in-place branching
+- [ ] Add `/branches` command to list branches (future)
+- [ ] Update session display to show tree info (future)
+
+## Notes
+
+- All save methods return the new entry's ID
+- Migration rewrites file on first load if version < CURRENT_VERSION
+- Existing sessions become linear chains after migration (parentId = previous entry)
+- Tree features available immediately after migration
+- SessionHeader does NOT have id/parentId (it's metadata, not part of tree)
+- Content types allow clean input/output separation

packages/coding-agent/examples/hooks/custom-compaction.ts

@@ -94,14 +94,12 @@ Format the summary as structured markdown with clear sections.`,
 				return;
 			}
 
-			// Return a compaction entry that discards ALL messages
+			// Return compaction content - SessionManager adds id/parentId
-			// firstKeptEntryIndex points past all current entries
+			// Use firstKeptEntryId from event to keep recent messages
 			return {
-				compactionEntry: {
+				compaction: {
-					type: "compaction" as const,
-					timestamp: new Date().toISOString(),
 					summary,
-					firstKeptEntryIndex: entries.length,
+					firstKeptEntryId: event.firstKeptEntryId,
 					tokensBefore,
 				},
 			};

packages/coding-agent/src/core/agent-session.ts

@@ -754,8 +754,13 @@ export class AgentSession {
 
 			const preparation = prepareCompaction(entries, settings);
 			if (!preparation) {
+				// Check why we can't compact
+				const lastEntry = entries[entries.length - 1];
+				if (lastEntry?.type === "compaction") {
 					throw new Error("Already compacted");
 				}
+				throw new Error("Nothing to compact (session too small or needs migration)");
+			}
 
 			// Find previous compaction summary if any
 			let previousSummary: string | undefined;
@@ -766,7 +771,7 @@ export class AgentSession {
 				}
 			}
 
-			let compactionEntry: CompactionEntry | undefined;
+			let hookCompaction: { summary: string; firstKeptEntryId: string; tokensBefore: number } | undefined;
 			let fromHook = false;
 
 			if (this._hookRunner?.hasHandlers("session")) {
@@ -777,6 +782,7 @@ export class AgentSession {
 					previousSessionFile: null,
 					reason: "before_compact",
 					cutPoint: preparation.cutPoint,
+					firstKeptEntryId: preparation.firstKeptEntryId,
 					previousSummary,
 					messagesToSummarize: [...preparation.messagesToSummarize],
 					messagesToKeep: [...preparation.messagesToKeep],
@@ -791,14 +797,24 @@ export class AgentSession {
 					throw new Error("Compaction cancelled");
 				}
 
-				if (result?.compactionEntry) {
+				if (result?.compaction) {
-					compactionEntry = result.compactionEntry;
+					hookCompaction = result.compaction;
 					fromHook = true;
 				}
 			}
 
-			if (!compactionEntry) {
+			let summary: string;
-				compactionEntry = await compact(
+			let firstKeptEntryId: string;
+			let tokensBefore: number;
+
+			if (hookCompaction) {
+				// Hook provided compaction content
+				summary = hookCompaction.summary;
+				firstKeptEntryId = hookCompaction.firstKeptEntryId;
+				tokensBefore = hookCompaction.tokensBefore;
+			} else {
+				// Generate compaction result
+				const result = await compact(
 					entries,
 					this.model,
 					settings,
@@ -806,33 +822,41 @@ export class AgentSession {
 					this._compactionAbortController.signal,
 					customInstructions,
 				);
+				summary = result.summary;
+				firstKeptEntryId = result.firstKeptEntryId;
+				tokensBefore = result.tokensBefore;
 			}
 
 			if (this._compactionAbortController.signal.aborted) {
 				throw new Error("Compaction cancelled");
 			}
 
-			this.sessionManager.saveCompaction(compactionEntry);
+			this.sessionManager.saveCompaction(summary, firstKeptEntryId, tokensBefore);
 			const newEntries = this.sessionManager.getEntries();
 			const sessionContext = this.sessionManager.buildSessionContext();
 			this.agent.replaceMessages(sessionContext.messages);
 
-			if (this._hookRunner) {
+			// Get the saved compaction entry for the hook
+			const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as
+				| CompactionEntry
+				| undefined;
+
+			if (this._hookRunner && savedCompactionEntry) {
 				await this._hookRunner.emit({
 					type: "session",
 					entries: newEntries,
 					sessionFile: this.sessionFile,
 					previousSessionFile: null,
 					reason: "compact",
-					compactionEntry,
+					compactionEntry: savedCompactionEntry,
-					tokensBefore: compactionEntry.tokensBefore,
+					tokensBefore,
 					fromHook,
 				});
 			}
 
 			return {
-				tokensBefore: compactionEntry.tokensBefore,
+				tokensBefore,
-				summary: compactionEntry.summary,
+				summary,
 			};
 		} finally {
 			this._compactionAbortController = null;
@@ -928,7 +952,7 @@ export class AgentSession {
 				}
 			}
 
-			let compactionEntry: CompactionEntry | undefined;
+			let hookCompaction: { summary: string; firstKeptEntryId: string; tokensBefore: number } | undefined;
 			let fromHook = false;
 
 			if (this._hookRunner?.hasHandlers("session")) {
@@ -939,6 +963,7 @@ export class AgentSession {
 					previousSessionFile: null,
 					reason: "before_compact",
 					cutPoint: preparation.cutPoint,
+					firstKeptEntryId: preparation.firstKeptEntryId,
 					previousSummary,
 					messagesToSummarize: [...preparation.messagesToSummarize],
 					messagesToKeep: [...preparation.messagesToKeep],
@@ -954,20 +979,33 @@ export class AgentSession {
 					return;
 				}
 
-				if (hookResult?.compactionEntry) {
+				if (hookResult?.compaction) {
-					compactionEntry = hookResult.compactionEntry;
+					hookCompaction = hookResult.compaction;
 					fromHook = true;
 				}
 			}
 
-			if (!compactionEntry) {
+			let summary: string;
-				compactionEntry = await compact(
+			let firstKeptEntryId: string;
+			let tokensBefore: number;
+
+			if (hookCompaction) {
+				// Hook provided compaction content
+				summary = hookCompaction.summary;
+				firstKeptEntryId = hookCompaction.firstKeptEntryId;
+				tokensBefore = hookCompaction.tokensBefore;
+			} else {
+				// Generate compaction result
+				const compactResult = await compact(
 					entries,
 					this.model,
 					settings,
 					apiKey,
 					this._autoCompactionAbortController.signal,
 				);
+				summary = compactResult.summary;
+				firstKeptEntryId = compactResult.firstKeptEntryId;
+				tokensBefore = compactResult.tokensBefore;
 			}
 
 			if (this._autoCompactionAbortController.signal.aborted) {
@@ -975,27 +1013,32 @@ export class AgentSession {
 				return;
 			}
 
-			this.sessionManager.saveCompaction(compactionEntry);
+			this.sessionManager.saveCompaction(summary, firstKeptEntryId, tokensBefore);
 			const newEntries = this.sessionManager.getEntries();
 			const sessionContext = this.sessionManager.buildSessionContext();
 			this.agent.replaceMessages(sessionContext.messages);
 
-			if (this._hookRunner) {
+			// Get the saved compaction entry for the hook
+			const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as
+				| CompactionEntry
+				| undefined;
+
+			if (this._hookRunner && savedCompactionEntry) {
 				await this._hookRunner.emit({
 					type: "session",
 					entries: newEntries,
 					sessionFile: this.sessionFile,
 					previousSessionFile: null,
 					reason: "compact",
-					compactionEntry,
+					compactionEntry: savedCompactionEntry,
-					tokensBefore: compactionEntry.tokensBefore,
+					tokensBefore,
 					fromHook,
 				});
 			}
 
 			const result: CompactionResult = {
-				tokensBefore: compactionEntry.tokensBefore,
+				tokensBefore,
-				summary: compactionEntry.summary,
+				summary,
 			};
 			this._emit({ type: "auto_compaction_end", result, aborted: false, willRetry });
 

packages/coding-agent/src/core/compaction.ts

@@ -9,7 +9,14 @@ import type { AppMessage } from "@mariozechner/pi-agent-core";
 import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai";
 import { complete } from "@mariozechner/pi-ai";
 import { messageTransformer } from "./messages.js";
-import type { CompactionEntry, SessionEntry } from "./session-manager.js";
+import type { CompactionEntry, ConversationEntry, SessionEntry } from "./session-manager.js";
+
+/** Result from compact() - SessionManager adds uuid/parentUuid when saving */
+export interface CompactionResult {
+	summary: string;
+	firstKeptEntryId: string;
+	tokensBefore: number;
+}
 
 // ============================================================================
 // Types
@@ -327,6 +334,8 @@ export async function generateSummary(
 
 export interface CompactionPreparation {
 	cutPoint: CutPointResult;
+	/** UUID of first entry to keep */
+	firstKeptEntryId: string;
 	/** Messages that will be summarized and discarded */
 	messagesToSummarize: AppMessage[];
 	/** Messages that will be kept after the summary (recent turns) */
@@ -355,6 +364,16 @@ export function prepareCompaction(entries: SessionEntry[], settings: CompactionS
 
 	const cutPoint = findCutPoint(entries, boundaryStart, boundaryEnd, settings.keepRecentTokens);
 
+	// Get UUID of first kept entry
+	const firstKeptEntry = entries[cutPoint.firstKeptEntryIndex];
+	if (firstKeptEntry.type === "session") {
+		return null; // Can't compact if first kept is header
+	}
+	const firstKeptEntryId = (firstKeptEntry as ConversationEntry).id;
+	if (!firstKeptEntryId) {
+		return null; // Session needs migration
+	}
+
 	const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex;
 
 	// Messages to summarize (will be discarded after summary)
@@ -375,7 +394,7 @@ export function prepareCompaction(entries: SessionEntry[], settings: CompactionS
 		}
 	}
 
-	return { cutPoint, messagesToSummarize, messagesToKeep, tokensBefore, boundaryStart };
+	return { cutPoint, firstKeptEntryId, messagesToSummarize, messagesToKeep, tokensBefore, boundaryStart };
 }
 
 // ============================================================================
@@ -394,9 +413,9 @@ Be concise. Focus on information needed to understand the retained recent work.`
 
 /**
  * Calculate compaction and generate summary.
- * Returns the CompactionEntry to append to the session file.
+ * Returns CompactionResult - SessionManager adds uuid/parentUuid when saving.
  *
- * @param entries - All session entries
+ * @param entries - All session entries (must have uuid fields for v2)
  * @param model - Model to use for summarization
  * @param settings - Compaction settings
  * @param apiKey - API key for LLM
@@ -410,7 +429,7 @@ export async function compact(
 	apiKey: string,
 	signal?: AbortSignal,
 	customInstructions?: string,
-): Promise<CompactionEntry> {
+): Promise<CompactionResult> {
 	// Don't compact if the last entry is already a compaction
 	if (entries.length > 0 && entries[entries.length - 1].type === "compaction") {
 		throw new Error("Already compacted");
@@ -490,11 +509,19 @@ export async function compact(
 		);
 	}
 
+	// Get UUID of first kept entry
+	const firstKeptEntry = entries[cutResult.firstKeptEntryIndex];
+	if (firstKeptEntry.type === "session") {
+		throw new Error("Cannot compact: first kept entry is session header");
+	}
+	const firstKeptEntryId = (firstKeptEntry as ConversationEntry).id;
+	if (!firstKeptEntryId) {
+		throw new Error("First kept entry has no UUID - session may need migration");
+	}
+
 	return {
-		type: "compaction",
-		timestamp: new Date().toISOString(),
 		summary,
-		firstKeptEntryIndex: cutResult.firstKeptEntryIndex,
+		firstKeptEntryId,
 		tokensBefore,
 	};
 }

packages/coding-agent/src/core/hooks/types.ts

@@ -130,6 +130,8 @@ export type SessionEvent =
 	| (SessionEventBase & {
 			reason: "before_compact";
 			cutPoint: CutPointResult;
+			/** ID of first entry to keep (for hooks that return CompactionEntry) */
+			firstKeptEntryId: string;
 			/** Summary from previous compaction, if any. Include this in your summary to preserve context. */
 			previousSummary?: string;
 			/** Messages that will be summarized and discarded */
@@ -351,8 +353,12 @@ export interface SessionEventResult {
 	cancel?: boolean;
 	/** If true (for before_branch only), skip restoring conversation to branch point while still creating the branched session file */
 	skipConversationRestore?: boolean;
-	/** Custom compaction entry (for before_compact event) */
+	/** Custom compaction result (for before_compact event) - SessionManager adds id/parentId */
-	compactionEntry?: CompactionEntry;
+	compaction?: {
+		summary: string;
+		firstKeptEntryId: string;
+		tokensBefore: number;
+	};
 }
 
 // ============================================================================

packages/coding-agent/src/core/session-manager.ts

@@ -1,9 +1,11 @@
 import type { AppMessage } from "@mariozechner/pi-agent-core";
 import { randomBytes } from "crypto";
-import { appendFileSync, existsSync, mkdirSync, readdirSync, readFileSync, statSync } from "fs";
+import { appendFileSync, existsSync, mkdirSync, readdirSync, readFileSync, statSync, writeFileSync } from "fs";
 import { join, resolve } from "path";
 import { getAgentDir as getDefaultAgentDir } from "../config.js";
 
+export const CURRENT_SESSION_VERSION = 2;
+
 function uuidv4(): string {
 	const bytes = randomBytes(16);
 	bytes[6] = (bytes[6] & 0x0f) | 0x40;
@@ -12,47 +14,89 @@ function uuidv4(): string {
 	return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20, 32)}`;
 }
 
+// ============================================================================
+// Session Header (metadata, not part of conversation tree)
+// ============================================================================
+
 export interface SessionHeader {
 	type: "session";
+	version?: number; // v1 sessions don't have this
 	id: string;
 	timestamp: string;
 	cwd: string;
 	branchedFrom?: string;
 }
 
-export interface SessionMessageEntry {
+// ============================================================================
-	type: "message";
+// 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;
 }
 
-export interface ThinkingLevelChangeEntry {
+export interface ThinkingLevelContent {
 	type: "thinking_level_change";
-	timestamp: string;
 	thinkingLevel: string;
 }
 
-export interface ModelChangeEntry {
+export interface ModelChangeContent {
 	type: "model_change";
-	timestamp: string;
 	provider: string;
 	modelId: string;
 }
 
-export interface CompactionEntry {
+export interface CompactionContent {
 	type: "compaction";
-	timestamp: string;
 	summary: string;
-	firstKeptEntryIndex: number;
+	firstKeptEntryId: string;
 	tokensBefore: number;
 }
 
-export type SessionEntry =
+export interface BranchSummaryContent {
-	| SessionHeader
+	type: "branch_summary";
+	summary: string;
+}
+
+/** Union of all content types (for input) */
+export type ConversationContent =
+	| MessageContent
+	| ThinkingLevelContent
+	| ModelChangeContent
+	| CompactionContent
+	| BranchSummaryContent;
+
+// ============================================================================
+// Full Entry Types (TreeNode + Content - returned from SessionManager)
+// ============================================================================
+
+export type SessionMessageEntry = TreeNode & MessageContent;
+export type ThinkingLevelChangeEntry = TreeNode & ThinkingLevelContent;
+export type ModelChangeEntry = TreeNode & ModelChangeContent;
+export type CompactionEntry = TreeNode & CompactionContent;
+export type BranchSummaryEntry = TreeNode & BranchSummaryContent;
+
+/** Conversation entry - has id/parentId for tree structure */
+export type ConversationEntry =
 	| SessionMessageEntry
 	| ThinkingLevelChangeEntry
 	| ModelChangeEntry
-	| CompactionEntry;
+	| CompactionEntry
+	| BranchSummaryEntry;
+
+/** Any session entry (header or conversation) */
+export type SessionEntry = SessionHeader | ConversationEntry;
 
 export interface SessionContext {
 	messages: AppMessage[];
@@ -87,6 +131,45 @@ export function createSummaryMessage(summary: string): AppMessage {
 	};
 }
 
+/**
+ * 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: SessionEntry[]): void {
+	// Check if already migrated
+	const firstConv = entries.find((e) => e.type !== "session");
+	if (firstConv && "id" in firstConv && firstConv.id) {
+		return; // Already migrated
+	}
+
+	let prevId: string | null = null;
+	for (const entry of entries) {
+		if (entry.type === "session") {
+			entry.version = CURRENT_SESSION_VERSION;
+			continue;
+		}
+
+		// Add id/parentId to conversation entries
+		const convEntry = entry as ConversationEntry;
+		convEntry.id = uuidv4();
+		convEntry.parentId = prevId;
+		prevId = convEntry.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;
+				}
+				delete comp.firstKeptEntryIndex;
+			}
+		}
+	}
+}
+
 /** Exported for compaction.test.ts */
 export function parseSessionEntries(content: string): SessionEntry[] {
 	const entries: SessionEntry[] = [];
@@ -115,58 +198,107 @@ export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEnt
 }
 
 /**
- * Build the session context from entries. This is what gets sent to the LLM.
+ * Build the session context from entries using tree traversal.
- *
+ * If leafId is provided, walks from that entry to root.
- * If there's a compaction entry, returns the summary message plus messages
+ * Handles compaction and branch summaries along the path.
- * from `firstKeptEntryIndex` onwards. Otherwise returns all messages.
- *
- * Also extracts the current thinking level and model from the entries.
  */
-export function buildSessionContext(entries: SessionEntry[]): SessionContext {
+export function buildSessionContext(entries: SessionEntry[], leafId?: string): SessionContext {
+	// Build uuid index for conversation entries
+	const byId = new Map<string, ConversationEntry>();
+	for (const entry of entries) {
+		if (entry.type !== "session") {
+			byId.set(entry.id, entry);
+		}
+	}
+
+	// Find leaf
+	let leaf: ConversationEntry | undefined;
+	if (leafId) {
+		leaf = byId.get(leafId);
+	} else {
+		// Find last conversation entry
+		for (let i = entries.length - 1; i >= 0; i--) {
+			if (entries[i].type !== "session") {
+				leaf = entries[i] as ConversationEntry;
+				break;
+			}
+		}
+	}
+
+	if (!leaf) {
+		return { messages: [], thinkingLevel: "off", model: null };
+	}
+
+	// Walk from leaf to root, collecting path
+	const path: ConversationEntry[] = [];
+	let current: ConversationEntry | undefined = leaf;
+	while (current) {
+		path.unshift(current);
+		current = current.parentId ? byId.get(current.parentId) : undefined;
+	}
+
+	// Extract settings and find compaction
 	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 path) {
 		if (entry.type === "thinking_level_change") {
 			thinkingLevel = entry.thinkingLevel;
 		} else if (entry.type === "model_change") {
 			model = { provider: entry.provider, modelId: entry.modelId };
 		} else if (entry.type === "message" && entry.message.role === "assistant") {
 			model = { provider: entry.message.provider, modelId: entry.message.model };
+		} else if (entry.type === "compaction") {
+			compaction = entry;
 		}
 	}
 
-	let latestCompactionIndex = -1;
+	// Build messages - handle compaction ordering correctly
-	for (let i = entries.length - 1; i >= 0; i--) {
+	// When there's a compaction, we need to:
-		if (entries[i].type === "compaction") {
+	// 1. Emit summary first
-			latestCompactionIndex = i;
+	// 2. Emit kept messages (from firstKeptEntryId up to compaction)
-			break;
+	// 3. Emit messages after compaction
-		}
-	}
-
-	if (latestCompactionIndex === -1) {
 	const messages: AppMessage[] = [];
-		for (const entry of entries) {
+
-			if (entry.type === "message") {
+	if (compaction) {
+		// Emit summary first
+		messages.push(createSummaryMessage(compaction.summary));
+
+		// Find compaction index in path
+		const compactionIdx = path.findIndex((e) => e.type === "compaction" && e.id === compaction.id);
+
+		// Emit kept messages (before compaction, starting from firstKeptEntryId)
+		let foundFirstKept = false;
+		for (let i = 0; i < compactionIdx; i++) {
+			const entry = path[i];
+			if (entry.id === compaction.firstKeptEntryId) {
+				foundFirstKept = true;
+			}
+			if (foundFirstKept && entry.type === "message") {
 				messages.push(entry.message);
 			}
 		}
-		return { messages, thinkingLevel, model };
-	}
 
-	const compactionEvent = entries[latestCompactionIndex] as CompactionEntry;
+		// Emit messages after compaction
-
+		for (let i = compactionIdx + 1; i < path.length; i++) {
-	const keptMessages: AppMessage[] = [];
+			const entry = path[i];
-	for (let i = compactionEvent.firstKeptEntryIndex; i < entries.length; i++) {
-		const entry = entries[i];
 			if (entry.type === "message") {
-			keptMessages.push(entry.message);
+				messages.push(entry.message);
+			} else if (entry.type === "branch_summary") {
+				messages.push(createSummaryMessage(entry.summary));
+			}
+		}
+	} else {
+		// No compaction - emit all messages, handle branch summaries
+		for (const entry of path) {
+			if (entry.type === "message") {
+				messages.push(entry.message);
+			} else if (entry.type === "branch_summary") {
+				messages.push(createSummaryMessage(entry.summary));
+			}
 		}
 	}
-
-	const messages: AppMessage[] = [];
-	messages.push(createSummaryMessage(compactionEvent.summary));
-	messages.push(...keptMessages);
 
 	return { messages, thinkingLevel, model };
 }
@@ -229,6 +361,10 @@ export class SessionManager {
 	private flushed: boolean = false;
 	private inMemoryEntries: SessionEntry[] = [];
 
+	// Tree structure (v2)
+	private byId: Map<string, ConversationEntry> = new Map();
+	private leafId: string = "";
+
 	private constructor(cwd: string, sessionDir: string, sessionFile: string | null, persist: boolean) {
 		this.cwd = cwd;
 		this.sessionDir = sessionDir;
@@ -240,10 +376,7 @@ export class SessionManager {
 		if (sessionFile) {
 			this.setSessionFile(sessionFile);
 		} else {
-			this.sessionId = uuidv4();
+			this._initNewSession();
-			const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
-			const sessionFile = join(this.getSessionDir(), `${timestamp}_${this.sessionId}.jsonl`);
-			this.setSessionFile(sessionFile);
 		}
 	}
 
@@ -252,21 +385,59 @@ export class SessionManager {
 		this.sessionFile = resolve(sessionFile);
 		if (existsSync(this.sessionFile)) {
 			this.inMemoryEntries = loadEntriesFromFile(this.sessionFile);
-			const header = this.inMemoryEntries.find((e) => e.type === "session");
+			const header = this.inMemoryEntries.find((e) => e.type === "session") as SessionHeader | undefined;
-			this.sessionId = header ? (header as SessionHeader).id : uuidv4();
+			this.sessionId = header?.id ?? uuidv4();
+
+			// Migrate v1 to v2 if needed
+			const version = header?.version ?? 1;
+			if (version < CURRENT_SESSION_VERSION) {
+				this._migrateToV2();
+				this._rewriteFile();
+			}
+
+			this._buildIndex();
 			this.flushed = true;
 		} else {
+			this._initNewSession();
+		}
+	}
+
+	private _initNewSession(): void {
 		this.sessionId = uuidv4();
-			this.inMemoryEntries = [];
+		const timestamp = new Date().toISOString();
-			this.flushed = false;
+		const header: SessionHeader = {
-			const entry: SessionHeader = {
 			type: "session",
+			version: CURRENT_SESSION_VERSION,
 			id: this.sessionId,
-				timestamp: new Date().toISOString(),
+			timestamp,
 			cwd: this.cwd,
 		};
-			this.inMemoryEntries.push(entry);
+		this.inMemoryEntries = [header];
+		this.byId.clear();
+		this.leafId = "";
+		this.flushed = false;
+		const fileTimestamp = timestamp.replace(/[:.]/g, "-");
+		this.sessionFile = join(this.getSessionDir(), `${fileTimestamp}_${this.sessionId}.jsonl`);
 	}
+
+	private _migrateToV2(): void {
+		migrateSessionEntries(this.inMemoryEntries);
+	}
+
+	private _buildIndex(): void {
+		this.byId.clear();
+		this.leafId = "";
+		for (const entry of this.inMemoryEntries) {
+			if (entry.type === "session") continue;
+			this.byId.set(entry.id, entry);
+			this.leafId = entry.id;
+		}
+	}
+
+	private _rewriteFile(): void {
+		if (!this.persist) return;
+		const content = `${this.inMemoryEntries.map((e) => JSON.stringify(e)).join("\n")}\n`;
+		writeFileSync(this.sessionFile, content);
 	}
 
 	isPersisted(): boolean {
@@ -290,18 +461,7 @@ export class SessionManager {
 	}
 
 	reset(): void {
-		this.sessionId = uuidv4();
+		this._initNewSession();
-		this.flushed = false;
-		const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
-		this.sessionFile = join(this.getSessionDir(), `${timestamp}_${this.sessionId}.jsonl`);
-		this.inMemoryEntries = [
-			{
-				type: "session",
-				id: this.sessionId,
-				timestamp: new Date().toISOString(),
-				cwd: this.cwd,
-			},
-		];
 	}
 
 	_persist(entry: SessionEntry): void {
@@ -320,49 +480,93 @@ export class SessionManager {
 		}
 	}
 
-	saveMessage(message: AppMessage): void {
+	private _appendEntry(entry: ConversationEntry): void {
+		this.inMemoryEntries.push(entry);
+		this.byId.set(entry.id, entry);
+		this.leafId = entry.id;
+		this._persist(entry);
+	}
+
+	saveMessage(message: AppMessage): string {
 		const entry: SessionMessageEntry = {
 			type: "message",
+			id: uuidv4(),
+			parentId: this.leafId || null,
 			timestamp: new Date().toISOString(),
 			message,
 		};
-		this.inMemoryEntries.push(entry);
+		this._appendEntry(entry);
-		this._persist(entry);
+		return entry.id;
 	}
 
-	saveThinkingLevelChange(thinkingLevel: string): void {
+	saveThinkingLevelChange(thinkingLevel: string): string {
 		const entry: ThinkingLevelChangeEntry = {
 			type: "thinking_level_change",
+			id: uuidv4(),
+			parentId: this.leafId || null,
 			timestamp: new Date().toISOString(),
 			thinkingLevel,
 		};
-		this.inMemoryEntries.push(entry);
+		this._appendEntry(entry);
-		this._persist(entry);
+		return entry.id;
 	}
 
-	saveModelChange(provider: string, modelId: string): void {
+	saveModelChange(provider: string, modelId: string): string {
 		const entry: ModelChangeEntry = {
 			type: "model_change",
+			id: uuidv4(),
+			parentId: this.leafId || null,
 			timestamp: new Date().toISOString(),
 			provider,
 			modelId,
 		};
-		this.inMemoryEntries.push(entry);
+		this._appendEntry(entry);
-		this._persist(entry);
+		return entry.id;
 	}
 
-	saveCompaction(entry: CompactionEntry): void {
+	saveCompaction(summary: string, firstKeptEntryId: string, tokensBefore: number): string {
-		this.inMemoryEntries.push(entry);
+		const entry: CompactionEntry = {
-		this._persist(entry);
+			type: "compaction",
+			id: uuidv4(),
+			parentId: this.leafId || null,
+			timestamp: new Date().toISOString(),
+			summary,
+			firstKeptEntryId,
+			tokensBefore,
+		};
+		this._appendEntry(entry);
+		return entry.id;
+	}
+
+	// =========================================================================
+	// Tree Traversal
+	// =========================================================================
+
+	getLeafUuid(): string {
+		return this.leafId;
+	}
+
+	getEntry(id: string): ConversationEntry | undefined {
+		return this.byId.get(id);
+	}
+
+	/** Walk from entry to root, returning path (conversation entries only) */
+	getPath(fromId?: string): ConversationEntry[] {
+		const path: ConversationEntry[] = [];
+		let current = this.byId.get(fromId ?? this.leafId);
+		while (current) {
+			path.unshift(current);
+			current = current.parentId ? this.byId.get(current.parentId) : undefined;
+		}
+		return path;
 	}
 
 	/**
 	 * Build the session context (what gets sent to the LLM).
-	 * If compacted, returns summary + kept messages. Otherwise all messages.
+	 * Uses tree traversal from current leaf.
-	 * Includes thinking level and model.
 	 */
 	buildSessionContext(): SessionContext {
-		return buildSessionContext(this.getEntries());
+		return buildSessionContext(this.getEntries(), this.leafId);
 	}
 
 	/**
@@ -373,6 +577,35 @@ export class SessionManager {
 		return [...this.inMemoryEntries];
 	}
 
+	// =========================================================================
+	// Branching
+	// =========================================================================
+
+	/** Branch in-place by changing the leaf pointer */
+	branchInPlace(branchFromId: string): void {
+		if (!this.byId.has(branchFromId)) {
+			throw new Error(`Entry ${branchFromId} not found`);
+		}
+		this.leafId = branchFromId;
+	}
+
+	/** Branch with a summary of the abandoned path */
+	branchWithSummary(branchFromId: string, summary: string): string {
+		if (!this.byId.has(branchFromId)) {
+			throw new Error(`Entry ${branchFromId} not found`);
+		}
+		this.leafId = branchFromId;
+		const entry: BranchSummaryEntry = {
+			type: "branch_summary",
+			id: uuidv4(),
+			parentId: branchFromId,
+			timestamp: new Date().toISOString(),
+			summary,
+		};
+		this._appendEntry(entry);
+		return entry.id;
+	}
+
 	createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string | null {
 		const newSessionId = uuidv4();
 		const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
@@ -385,6 +618,7 @@ export class SessionManager {
 			if (entry.type === "session") {
 				newEntries.push({
 					...entry,
+					version: CURRENT_SESSION_VERSION,
 					id: newSessionId,
 					timestamp: new Date().toISOString(),
 					branchedFrom: this.persist ? this.sessionFile : undefined,
@@ -402,6 +636,7 @@ export class SessionManager {
 		}
 		this.inMemoryEntries = newEntries;
 		this.sessionId = newSessionId;
+		this._buildIndex();
 		return null;
 	}
 

packages/coding-agent/src/index.ts

@@ -107,11 +107,20 @@ export {
 	readOnlyTools,
 } from "./core/sdk.js";
 export {
+	type BranchSummaryContent,
+	type BranchSummaryEntry,
 	buildSessionContext,
+	type CompactionContent,
 	type CompactionEntry,
+	type ConversationContent,
+	type ConversationEntry,
+	CURRENT_SESSION_VERSION,
 	createSummaryMessage,
 	getLatestCompactionEntry,
+	type MessageContent,
+	type ModelChangeContent,
 	type ModelChangeEntry,
+	migrateSessionEntries,
 	parseSessionEntries,
 	type SessionContext as LoadedSession,
 	type SessionEntry,
@@ -122,6 +131,9 @@ export {
 	SUMMARY_PREFIX,
 	SUMMARY_SUFFIX,
 	type ThinkingLevelChangeEntry,
+	type ThinkingLevelContent,
+	// Tree types (v2)
+	type TreeNode,
 } from "./core/session-manager.js";
 export {
 	type CompactionSettings,

packages/coding-agent/test/agent-session-compaction.test.ts

@@ -156,9 +156,9 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
 		expect(compaction.type).toBe("compaction");
 		if (compaction.type === "compaction") {
 			expect(compaction.summary.length).toBeGreaterThan(0);
-			// firstKeptEntryIndex can be 0 if all messages fit within keepRecentTokens
+			// firstKeptEntryId can be 0 if all messages fit within keepRecentTokens
 			// (which is the case for small conversations)
-			expect(compaction.firstKeptEntryIndex).toBeGreaterThanOrEqual(0);
+			expect(compaction.firstKeptEntryId).toBeGreaterThanOrEqual(0);
 			expect(compaction.tokensBefore).toBeGreaterThan(0);
 		}
 	}, 120000);

packages/coding-agent/test/compaction-hooks-example.test.ts

@@ -4,7 +4,6 @@
 
 import { describe, expect, it } from "vitest";
 import type { HookAPI } from "../src/core/hooks/index.js";
-import type { CompactionEntry } from "../src/core/session-manager.js";
 
 describe("Documentation example", () => {
 	it("custom compaction example should type-check correctly", () => {
@@ -20,29 +19,30 @@ describe("Documentation example", () => {
 				const tokensBefore = event.tokensBefore;
 				const model = event.model;
 				const resolveApiKey = event.resolveApiKey;
+				const firstKeptEntryId = event.firstKeptEntryId;
 
 				// Verify types
 				expect(Array.isArray(messages)).toBe(true);
 				expect(Array.isArray(messagesToKeep)).toBe(true);
-				expect(typeof cutPoint.firstKeptEntryIndex).toBe("number");
+				expect(typeof cutPoint.firstKeptEntryIndex).toBe("number"); // cutPoint still uses index
 				expect(typeof tokensBefore).toBe("number");
 				expect(model).toBeDefined();
 				expect(typeof resolveApiKey).toBe("function");
+				expect(typeof firstKeptEntryId).toBe("string");
 
 				const summary = messages
 					.filter((m) => m.role === "user")
 					.map((m) => `- ${typeof m.content === "string" ? m.content.slice(0, 100) : "[complex]"}`)
 					.join("\n");
 
-				const compactionEntry: CompactionEntry = {
+				// Hooks return compaction content - SessionManager adds id/parentId
-					type: "compaction",
+				return {
-					timestamp: new Date().toISOString(),
+					compaction: {
 						summary: `User requests:\n${summary}`,
-					firstKeptEntryIndex: event.cutPoint.firstKeptEntryIndex,
+						firstKeptEntryId,
-					tokensBefore: event.tokensBefore,
+						tokensBefore,
+					},
 				};
-
-				return { compactionEntry };
 			});
 		};
 

packages/coding-agent/test/compaction.test.ts

@@ -3,7 +3,7 @@ import type { AssistantMessage, Usage } from "@mariozechner/pi-ai";
 import { getModel } from "@mariozechner/pi-ai";
 import { readFileSync } from "fs";
 import { join } from "path";
-import { describe, expect, it } from "vitest";
+import { beforeEach, describe, expect, it } from "vitest";
 import {
 	type CompactionSettings,
 	calculateContextTokens,
@@ -17,9 +17,12 @@ import {
 	buildSessionContext,
 	type CompactionEntry,
 	createSummaryMessage,
+	type ModelChangeEntry,
+	migrateSessionEntries,
 	parseSessionEntries,
 	type SessionEntry,
 	type SessionMessageEntry,
+	type ThinkingLevelChangeEntry,
 } from "../src/core/session-manager.js";
 
 // ============================================================================
@@ -29,7 +32,9 @@ import {
 function loadLargeSessionEntries(): SessionEntry[] {
 	const sessionPath = join(__dirname, "fixtures/large-session.jsonl");
 	const content = readFileSync(sessionPath, "utf-8");
-	return parseSessionEntries(content);
+	const entries = parseSessionEntries(content);
+	migrateSessionEntries(entries); // Add id/parentId for v1 fixtures
+	return entries;
 }
 
 function createMockUsage(input: number, output: number, cacheRead = 0, cacheWrite = 0): Usage {
@@ -60,18 +65,82 @@ function createAssistantMessage(text: string, usage?: Usage): AssistantMessage {
 	};
 }
 
-function createMessageEntry(message: AppMessage): SessionMessageEntry {
+let entryCounter = 0;
-	return { type: "message", timestamp: new Date().toISOString(), message };
+let lastId: string | null = null;
+
+function resetEntryCounter() {
+	entryCounter = 0;
+	lastId = null;
 }
 
-function createCompactionEntry(summary: string, firstKeptEntryIndex: number): CompactionEntry {
+// Reset counter before each test to get predictable IDs
+beforeEach(() => {
+	resetEntryCounter();
+});
+
+function createSessionHeader() {
 	return {
+		type: "session" as const,
+		version: 2,
+		id: "test-session",
+		timestamp: "",
+		cwd: "",
+	};
+}
+
+function createMessageEntry(message: AppMessage): SessionMessageEntry {
+	const id = `test-id-${entryCounter++}`;
+	const entry: SessionMessageEntry = {
+		type: "message",
+		id,
+		parentId: lastId,
+		timestamp: new Date().toISOString(),
+		message,
+	};
+	lastId = id;
+	return entry;
+}
+
+function createCompactionEntry(summary: string, firstKeptEntryId: string): CompactionEntry {
+	const id = `test-id-${entryCounter++}`;
+	const entry: CompactionEntry = {
 		type: "compaction",
+		id,
+		parentId: lastId,
 		timestamp: new Date().toISOString(),
 		summary,
-		firstKeptEntryIndex,
+		firstKeptEntryId,
 		tokensBefore: 10000,
 	};
+	lastId = id;
+	return entry;
+}
+
+function createModelChangeEntry(provider: string, modelId: string): ModelChangeEntry {
+	const id = `test-id-${entryCounter++}`;
+	const entry: ModelChangeEntry = {
+		type: "model_change",
+		id,
+		parentId: lastId,
+		timestamp: new Date().toISOString(),
+		provider,
+		modelId,
+	};
+	lastId = id;
+	return entry;
+}
+
+function createThinkingLevelEntry(thinkingLevel: string): ThinkingLevelChangeEntry {
+	const id = `test-id-${entryCounter++}`;
+	const entry: ThinkingLevelChangeEntry = {
+		type: "thinking_level_change",
+		id,
+		parentId: lastId,
+		timestamp: new Date().toISOString(),
+		thinkingLevel,
+	};
+	lastId = id;
+	return entry;
 }
 
 // ============================================================================
@@ -248,78 +317,59 @@ describe("buildSessionContext", () => {
 	});
 
 	it("should handle single compaction", () => {
-		// indices: 0=session, 1=u1, 2=a1, 3=u2, 4=a2, 5=compaction, 6=u3, 7=a3
+		// IDs: u1=test-id-0, a1=test-id-1, u2=test-id-2, a2=test-id-3, compaction=test-id-4, u3=test-id-5, a3=test-id-6
-		const entries: SessionEntry[] = [
+		const u1 = createMessageEntry(createUserMessage("1"));
-			{
+		const a1 = createMessageEntry(createAssistantMessage("a"));
-				type: "session",
+		const u2 = createMessageEntry(createUserMessage("2"));
-				id: "1",
+		const a2 = createMessageEntry(createAssistantMessage("b"));
-				timestamp: "",
+		const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id); // keep from u2 onwards
-				cwd: "",
+		const u3 = createMessageEntry(createUserMessage("3"));
-			},
+		const a3 = createMessageEntry(createAssistantMessage("c"));
-			createMessageEntry(createUserMessage("1")),
+
-			createMessageEntry(createAssistantMessage("a")),
+		const entries: SessionEntry[] = [createSessionHeader(), u1, a1, u2, a2, compaction, u3, a3];
-			createMessageEntry(createUserMessage("2")),
-			createMessageEntry(createAssistantMessage("b")),
-			createCompactionEntry("Summary of 1,a,2,b", 3), // keep from index 3 (u2) onwards
-			createMessageEntry(createUserMessage("3")),
-			createMessageEntry(createAssistantMessage("c")),
-		];
 
 		const loaded = buildSessionContext(entries);
-		// summary + kept (u2,a2 from idx 3-4) + after (u3,a3 from idx 6-7) = 5
+		// summary + kept (u2, a2) + after (u3, a3) = 5
 		expect(loaded.messages.length).toBe(5);
 		expect(loaded.messages[0].role).toBe("user");
 		expect((loaded.messages[0] as any).content).toContain("Summary of 1,a,2,b");
 	});
 
 	it("should handle multiple compactions (only latest matters)", () => {
-		// indices: 0=session, 1=u1, 2=a1, 3=compact1, 4=u2, 5=b, 6=u3, 7=c, 8=compact2, 9=u4, 10=d
+		// First batch
-		const entries: SessionEntry[] = [
+		const u1 = createMessageEntry(createUserMessage("1"));
-			{
+		const a1 = createMessageEntry(createAssistantMessage("a"));
-				type: "session",
+		const compact1 = createCompactionEntry("First summary", u1.id);
-				id: "1",
+		// Second batch
-				timestamp: "",
+		const u2 = createMessageEntry(createUserMessage("2"));
-				cwd: "",
+		const b = createMessageEntry(createAssistantMessage("b"));
-			},
+		const u3 = createMessageEntry(createUserMessage("3"));
-			createMessageEntry(createUserMessage("1")),
+		const c = createMessageEntry(createAssistantMessage("c"));
-			createMessageEntry(createAssistantMessage("a")),
+		const compact2 = createCompactionEntry("Second summary", u3.id); // keep from u3 onwards
-			createCompactionEntry("First summary", 1), // keep from index 1
+		// After second compaction
-			createMessageEntry(createUserMessage("2")),
+		const u4 = createMessageEntry(createUserMessage("4"));
-			createMessageEntry(createAssistantMessage("b")),
+		const d = createMessageEntry(createAssistantMessage("d"));
-			createMessageEntry(createUserMessage("3")),
+
-			createMessageEntry(createAssistantMessage("c")),
+		const entries: SessionEntry[] = [createSessionHeader(), u1, a1, compact1, u2, b, u3, c, compact2, u4, d];
-			createCompactionEntry("Second summary", 6), // keep from index 6 (u3) onwards
-			createMessageEntry(createUserMessage("4")),
-			createMessageEntry(createAssistantMessage("d")),
-		];
 
 		const loaded = buildSessionContext(entries);
-		// summary + kept from idx 6 (u3,c) + after (u4,d) = 5
+		// summary + kept from u3 (u3, c) + after (u4, d) = 5
 		expect(loaded.messages.length).toBe(5);
 		expect((loaded.messages[0] as any).content).toContain("Second summary");
 	});
 
-	it("should clamp firstKeptEntryIndex to valid range", () => {
+	it("should keep all messages when firstKeptEntryId is first entry", () => {
-		// indices: 0=session, 1=u1, 2=a1, 3=compact1, 4=u2, 5=b, 6=compact2
+		const u1 = createMessageEntry(createUserMessage("1"));
-		const entries: SessionEntry[] = [
+		const a1 = createMessageEntry(createAssistantMessage("a"));
-			{
+		const compact1 = createCompactionEntry("First summary", u1.id); // keep from first entry
-				type: "session",
+		const u2 = createMessageEntry(createUserMessage("2"));
-				id: "1",
+		const b = createMessageEntry(createAssistantMessage("b"));
-				timestamp: "",
+
-				cwd: "",
+		const entries: SessionEntry[] = [createSessionHeader(), u1, a1, compact1, u2, b];
-			},
-			createMessageEntry(createUserMessage("1")),
-			createMessageEntry(createAssistantMessage("a")),
-			createCompactionEntry("First summary", 1),
-			createMessageEntry(createUserMessage("2")),
-			createMessageEntry(createAssistantMessage("b")),
-			createCompactionEntry("Second summary", 0), // index 0 is before compaction1, should still work
-		];
 
 		const loaded = buildSessionContext(entries);
-		// Keeps from index 0, but compaction entries are skipped, so u1,a1,u2,b = 4 + summary = 5
+		// summary + all messages (u1, a1, u2, b) = 5
-		// Actually index 0 is session header, so messages are u1,a1,u2,b
+		expect(loaded.messages.length).toBe(5);
-		expect(loaded.messages.length).toBe(5); // summary + 4 messages
 	});
 
 	it("should track model and thinking level changes", () => {
@@ -331,9 +381,9 @@ describe("buildSessionContext", () => {
 				cwd: "",
 			},
 			createMessageEntry(createUserMessage("1")),
-			{ type: "model_change", timestamp: "", provider: "openai", modelId: "gpt-4" },
+			createModelChangeEntry("openai", "gpt-4"),
 			createMessageEntry(createAssistantMessage("a")),
-			{ type: "thinking_level_change", timestamp: "", thinkingLevel: "high" },
+			createThinkingLevelEntry("high"),
 		];
 
 		const loaded = buildSessionContext(entries);
@@ -380,27 +430,26 @@ describe("Large session fixture", () => {
 // ============================================================================
 
 describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
-	it("should generate a compaction event for the large session", async () => {
+	it("should generate a compaction result for the large session", async () => {
 		const entries = loadLargeSessionEntries();
 		const model = getModel("anthropic", "claude-sonnet-4-5")!;
 
-		const compactionEvent = await compact(
+		const compactionResult = await compact(
 			entries,
 			model,
 			DEFAULT_COMPACTION_SETTINGS,
 			process.env.ANTHROPIC_OAUTH_TOKEN!,
 		);
 
-		expect(compactionEvent.type).toBe("compaction");
+		expect(compactionResult.summary.length).toBeGreaterThan(100);
-		expect(compactionEvent.summary.length).toBeGreaterThan(100);
+		expect(compactionResult.firstKeptEntryId).toBeTruthy();
-		expect(compactionEvent.firstKeptEntryIndex).toBeGreaterThan(0);
+		expect(compactionResult.tokensBefore).toBeGreaterThan(0);
-		expect(compactionEvent.tokensBefore).toBeGreaterThan(0);
 
-		console.log("Summary length:", compactionEvent.summary.length);
+		console.log("Summary length:", compactionResult.summary.length);
-		console.log("First kept entry index:", compactionEvent.firstKeptEntryIndex);
+		console.log("First kept entry ID:", compactionResult.firstKeptEntryId);
-		console.log("Tokens before:", compactionEvent.tokensBefore);
+		console.log("Tokens before:", compactionResult.tokensBefore);
 		console.log("\n--- SUMMARY ---\n");
-		console.log(compactionEvent.summary);
+		console.log(compactionResult.summary);
 	}, 60000);
 
 	it("should produce valid session after compaction", async () => {
@@ -408,21 +457,30 @@ describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
 		const loaded = buildSessionContext(entries);
 		const model = getModel("anthropic", "claude-sonnet-4-5")!;
 
-		const compactionEvent = await compact(
+		const compactionResult = await compact(
 			entries,
 			model,
 			DEFAULT_COMPACTION_SETTINGS,
 			process.env.ANTHROPIC_OAUTH_TOKEN!,
 		);
 
-		// Simulate appending compaction to entries
+		// Simulate appending compaction to entries by creating a proper entry
-		const newEntries = [...entries, compactionEvent];
+		const lastEntry = entries[entries.length - 1];
+		const parentId = lastEntry.type === "session" ? null : lastEntry.id;
+		const compactionEntry: CompactionEntry = {
+			type: "compaction",
+			id: "compaction-test-id",
+			parentId,
+			timestamp: new Date().toISOString(),
+			...compactionResult,
+		};
+		const newEntries = [...entries, compactionEntry];
 		const reloaded = buildSessionContext(newEntries);
 
 		// Should have summary + kept messages
 		expect(reloaded.messages.length).toBeLessThan(loaded.messages.length);
 		expect(reloaded.messages[0].role).toBe("user");
-		expect((reloaded.messages[0] as any).content).toContain(compactionEvent.summary);
+		expect((reloaded.messages[0] as any).content).toContain(compactionResult.summary);
 
 		console.log("Original messages:", loaded.messages.length);
 		console.log("After compaction:", reloaded.messages.length);

packages/mom/src/context.ts

@@ -15,10 +15,14 @@ import {
 	buildSessionContext,
 	type CompactionEntry,
 	type LoadedSession,
+	type MessageContent,
+	type ModelChangeContent,
 	type ModelChangeEntry,
 	type SessionEntry,
 	type SessionMessageEntry,
 	type ThinkingLevelChangeEntry,
+	type ThinkingLevelContent,
+	type TreeNode,
 } from "@mariozechner/pi-coding-agent";
 import { randomBytes } from "crypto";
 import { appendFileSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs";
@@ -49,6 +53,7 @@ export class MomSessionManager {
 	private channelDir: string;
 	private flushed: boolean = false;
 	private inMemoryEntries: SessionEntry[] = [];
+	private leafId: string | null = null;
 
 	constructor(channelDir: string) {
 		this.channelDir = channelDir;
@@ -64,12 +69,14 @@ export class MomSessionManager {
 		if (existsSync(this.contextFile)) {
 			this.inMemoryEntries = this.loadEntriesFromFile();
 			this.sessionId = this.extractSessionId() || uuidv4();
+			this._updateLeafId();
 			this.flushed = true;
 		} else {
 			this.sessionId = uuidv4();
 			this.inMemoryEntries = [
 				{
 					type: "session",
+					version: 2,
 					id: this.sessionId,
 					timestamp: new Date().toISOString(),
 					cwd: this.channelDir,
@@ -79,6 +86,28 @@ export class MomSessionManager {
 		// Note: syncFromLog() is called explicitly from agent.ts with excludeTimestamp
 	}
 
+	private _updateLeafId(): void {
+		for (let i = this.inMemoryEntries.length - 1; i >= 0; i--) {
+			const entry = this.inMemoryEntries[i];
+			if (entry.type !== "session") {
+				this.leafId = entry.id;
+				return;
+			}
+		}
+		this.leafId = null;
+	}
+
+	private _createTreeNode(): TreeNode {
+		const id = uuidv4();
+		const node: TreeNode = {
+			id,
+			parentId: this.leafId,
+			timestamp: new Date().toISOString(),
+		};
+		this.leafId = id;
+		return node;
+	}
+
 	private _persist(entry: SessionEntry): void {
 		const hasAssistant = this.inMemoryEntries.some((e) => e.type === "message" && e.message.role === "assistant");
 		if (!hasAssistant) return;
@@ -206,11 +235,15 @@ export class MomSessionManager {
 		newMessages.sort((a, b) => new Date(a.timestamp).getTime() - new Date(b.timestamp).getTime());
 
 		for (const { timestamp, message } of newMessages) {
+			const id = uuidv4();
 			const entry: SessionMessageEntry = {
 				type: "message",
+				id,
+				parentId: this.leafId,
 				timestamp, // Use log date as entry timestamp for consistent deduplication
 				message,
 			};
+			this.leafId = id;
 
 			this.inMemoryEntries.push(entry);
 			appendFileSync(this.contextFile, `${JSON.stringify(entry)}\n`);
@@ -247,32 +280,22 @@ export class MomSessionManager {
 	}
 
 	saveMessage(message: AppMessage): void {
-		const entry: SessionMessageEntry = {
+		const content: MessageContent = { type: "message", message };
-			type: "message",
+		const entry: SessionMessageEntry = { ...this._createTreeNode(), ...content };
-			timestamp: new Date().toISOString(),
-			message,
-		};
 		this.inMemoryEntries.push(entry);
 		this._persist(entry);
 	}
 
 	saveThinkingLevelChange(thinkingLevel: string): void {
-		const entry: ThinkingLevelChangeEntry = {
+		const content: ThinkingLevelContent = { type: "thinking_level_change", thinkingLevel };
-			type: "thinking_level_change",
+		const entry: ThinkingLevelChangeEntry = { ...this._createTreeNode(), ...content };
-			timestamp: new Date().toISOString(),
-			thinkingLevel,
-		};
 		this.inMemoryEntries.push(entry);
 		this._persist(entry);
 	}
 
 	saveModelChange(provider: string, modelId: string): void {
-		const entry: ModelChangeEntry = {
+		const content: ModelChangeContent = { type: "model_change", provider, modelId };
-			type: "model_change",
+		const entry: ModelChangeEntry = { ...this._createTreeNode(), ...content };
-			timestamp: new Date().toISOString(),
-			provider,
-			modelId,
-		};
 		this.inMemoryEntries.push(entry);
 		this._persist(entry);
 	}