Session tree: simplify types, add branching API, comprehensive tests

Types:
- SessionEntryBase with type field, extended by all entry types
- CustomEntry for hooks (type: 'custom', customType, data)
- Remove XXXContent types and TreeNode (redundant)

API:
- Rename saveXXX to appendXXX with JSDoc explaining tree semantics
- Rename branchInPlace to branch() with better docs
- Add createBranchedSession(leafId) replacing index-based version
- Add getTree() returning SessionTreeNode[] for tree traversal
- Add appendCustomEntry(customType, data) for hooks

Tests:
- tree-traversal.test.ts: 28 tests covering append, getPath, getTree,
  branch, branchWithSummary, createBranchedSession
- save-entry.test.ts: custom entry integration

Docs:
- Class-level JSDoc explaining append-only tree model
- Method docs explaining leaf advancement and branching
- CHANGELOG.md entry for all changes

packages/coding-agent/CHANGELOG.md

@@ -2,13 +2,41 @@
 
 ## [Unreleased]
 
+### Breaking Changes
+
+- **Session tree structure (v2)**: Sessions now store entries as a tree with `id`/`parentId` fields, enabling in-place branching without creating new files. Existing v1 sessions are auto-migrated on load.
+- **SessionManager API**:
+  - `saveXXX()` renamed to `appendXXX()` (e.g., `appendMessage`, `appendCompaction`)
+  - `branchInPlace()` renamed to `branch()`
+  - `reset()` renamed to `newSession()`
+  - `createBranchedSessionFromEntries(entries, index)` replaced with `createBranchedSession(leafId)`
+  - `saveCompaction(entry)` replaced with `appendCompaction(summary, firstKeptEntryId, tokensBefore)`
+  - `getEntries()` now excludes the session header (use `getHeader()` separately)
+  - New methods: `getTree()`, `getPath()`, `getLeafUuid()`, `getLeafEntry()`, `getEntry()`, `branchWithSummary()`
+  - New `appendCustomEntry(customType, data)` for hooks to store custom data
+- **Compaction API**:
+  - `compact()` now returns `CompactionResult` (`{ summary, firstKeptEntryId, tokensBefore }`) instead of `CompactionEntry`
+  - `CompactionEntry.firstKeptEntryIndex` replaced with `firstKeptEntryId`
+  - `prepareCompaction()` now returns `firstKeptEntryId` in its result
+- **Hook types**:
+  - `SessionEventResult.compactionEntry` replaced with `SessionEventResult.compaction` (content only, SessionManager adds id/parentId)
+  - `before_compact` event now includes `firstKeptEntryId` field for hooks that return custom compaction
+
 ### Added
 
 - **`enabledModels` setting**: Configure whitelisted models in `settings.json` (same format as `--models` CLI flag). CLI `--models` takes precedence over the setting.
 
+### Changed
+
+- **Entry IDs**: Session entries now use short 8-character hex IDs instead of full UUIDs
+- **API key priority**: `ANTHROPIC_OAUTH_TOKEN` now takes precedence over `ANTHROPIC_API_KEY`
+- **New entry types**: `BranchSummaryEntry` for branch context, `CustomEntry` for hook data
+
 ### Fixed
 
 - **Edit tool fails on Windows due to CRLF line endings**: Files with CRLF line endings now match correctly when LLMs send LF-only text. Line endings are normalized before matching and restored to original style on write. ([#355](https://github.com/badlogic/pi-mono/issues/355))
+- **Session file validation**: `findMostRecentSession()` now validates session headers before returning, preventing non-session JSONL files from being loaded
+- **Compaction error handling**: `generateSummary()` and `generateTurnPrefixSummary()` now throw on LLM errors instead of returning empty strings
 
 ## [0.30.2] - 2025-12-26
 

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

@@ -211,7 +211,7 @@ export class AgentSession {
 
 		// Handle session persistence
 		if (event.type === "message_end") {
-			this.sessionManager.saveMessage(event.message);
+			this.sessionManager.appendMessage(event.message);
 
 			// Track assistant message for auto-compaction (checked on agent_end)
 			if (event.message.role === "assistant") {
@@ -535,7 +535,7 @@ export class AgentSession {
 		this._disconnectFromAgent();
 		await this.abort();
 		this.agent.reset();
-		this.sessionManager.reset();
+		this.sessionManager.newSession();
 		this._queuedMessages = [];
 		this._reconnectToAgent();
 
@@ -572,7 +572,7 @@ export class AgentSession {
 		}
 
 		this.agent.setModel(model);
-		this.sessionManager.saveModelChange(model.provider, model.id);
+		this.sessionManager.appendModelChange(model.provider, model.id);
 		this.settingsManager.setDefaultModelAndProvider(model.provider, model.id);
 
 		// Re-clamp thinking level for new model's capabilities
@@ -611,7 +611,7 @@ export class AgentSession {
 
 		// Apply model
 		this.agent.setModel(next.model);
-		this.sessionManager.saveModelChange(next.model.provider, next.model.id);
+		this.sessionManager.appendModelChange(next.model.provider, next.model.id);
 		this.settingsManager.setDefaultModelAndProvider(next.model.provider, next.model.id);
 
 		// Apply thinking level (setThinkingLevel clamps to model capabilities)
@@ -638,7 +638,7 @@ export class AgentSession {
 		}
 
 		this.agent.setModel(nextModel);
-		this.sessionManager.saveModelChange(nextModel.provider, nextModel.id);
+		this.sessionManager.appendModelChange(nextModel.provider, nextModel.id);
 		this.settingsManager.setDefaultModelAndProvider(nextModel.provider, nextModel.id);
 
 		// Re-clamp thinking level for new model's capabilities
@@ -671,7 +671,7 @@ export class AgentSession {
 			effectiveLevel = "high";
 		}
 		this.agent.setThinkingLevel(effectiveLevel);
-		this.sessionManager.saveThinkingLevelChange(effectiveLevel);
+		this.sessionManager.appendThinkingLevelChange(effectiveLevel);
 		this.settingsManager.setDefaultThinkingLevel(effectiveLevel);
 	}
 
@@ -831,7 +831,7 @@ export class AgentSession {
 				throw new Error("Compaction cancelled");
 			}
 
-			this.sessionManager.saveCompaction(summary, firstKeptEntryId, tokensBefore);
+			this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore);
 			const newEntries = this.sessionManager.getEntries();
 			const sessionContext = this.sessionManager.buildSessionContext();
 			this.agent.replaceMessages(sessionContext.messages);
@@ -1013,7 +1013,7 @@ export class AgentSession {
 				return;
 			}
 
-			this.sessionManager.saveCompaction(summary, firstKeptEntryId, tokensBefore);
+			this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore);
 			const newEntries = this.sessionManager.getEntries();
 			const sessionContext = this.sessionManager.buildSessionContext();
 			this.agent.replaceMessages(sessionContext.messages);
@@ -1271,7 +1271,7 @@ export class AgentSession {
 				this.agent.appendMessage(bashMessage);
 
 				// Save to session
-				this.sessionManager.saveMessage(bashMessage);
+				this.sessionManager.appendMessage(bashMessage);
 			}
 
 			return result;
@@ -1309,7 +1309,7 @@ export class AgentSession {
 			this.agent.appendMessage(bashMessage);
 
 			// Save to session
-			this.sessionManager.saveMessage(bashMessage);
+			this.sessionManager.appendMessage(bashMessage);
 		}
 
 		this._pendingBashMessages = [];
@@ -1431,8 +1431,12 @@ export class AgentSession {
 			skipConversationRestore = result?.skipConversationRestore ?? false;
 		}
 
-		// Create branched session (returns null in --no-session mode)
+		// Create branched session ending before the selected message (returns null in --no-session mode)
-		const newSessionFile = this.sessionManager.createBranchedSessionFromEntries(entries, entryIndex);
+		// User will re-enter/edit the selected message
+		if (!selectedEntry.parentId) {
+			throw new Error("Cannot branch from first message");
+		}
+		const newSessionFile = this.sessionManager.createBranchedSession(selectedEntry.parentId);
 
 		// Update session file if we have one (file-based mode)
 		if (newSessionFile !== null) {

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

@@ -26,54 +26,47 @@ export interface SessionHeader {
 	branchedFrom?: string;
 }
 
-export interface MessageContent {
+export interface SessionEntryBase {
-	type: "message";
-	message: AppMessage;
-}
-
-export interface ThinkingLevelContent {
-	type: "thinking_level_change";
-	thinkingLevel: string;
-}
-
-export interface ModelChangeContent {
-	type: "model_change";
-	provider: string;
-	modelId: string;
-}
-
-export interface CompactionContent {
-	type: "compaction";
-	summary: string;
-	firstKeptEntryId: string;
-	tokensBefore: number;
-}
-
-export interface BranchSummaryContent {
-	type: "branch_summary";
-	summary: string;
-}
-
-/** Union of all content types (for "write" methods in SessionManager) */
-export type SessionContent =
-	| MessageContent
-	| ThinkingLevelContent
-	| ModelChangeContent
-	| CompactionContent
-	| BranchSummaryContent;
-
-export interface TreeNode {
 	type: string;
 	id: string;
 	parentId: string | null;
 	timestamp: string;
 }
 
-export type SessionMessageEntry = TreeNode & MessageContent;
+export interface SessionMessageEntry extends SessionEntryBase {
-export type ThinkingLevelChangeEntry = TreeNode & ThinkingLevelContent;
+	type: "message";
-export type ModelChangeEntry = TreeNode & ModelChangeContent;
+	message: AppMessage;
-export type CompactionEntry = TreeNode & CompactionContent;
+}
-export type BranchSummaryEntry = TreeNode & BranchSummaryContent;
+
+export interface ThinkingLevelChangeEntry extends SessionEntryBase {
+	type: "thinking_level_change";
+	thinkingLevel: string;
+}
+
+export interface ModelChangeEntry extends SessionEntryBase {
+	type: "model_change";
+	provider: string;
+	modelId: string;
+}
+
+export interface CompactionEntry extends SessionEntryBase {
+	type: "compaction";
+	summary: string;
+	firstKeptEntryId: string;
+	tokensBefore: number;
+}
+
+export interface BranchSummaryEntry extends SessionEntryBase {
+	type: "branch_summary";
+	summary: string;
+}
+
+/** Custom entry for hooks. Use customType to identify your hook's entries. */
+export interface CustomEntry extends SessionEntryBase {
+	type: "custom";
+	customType: string;
+	data?: unknown;
+}
 
 /** Session entry - has id/parentId for tree structure (returned by "read" methods in SessionManager) */
 export type SessionEntry =
@@ -81,11 +74,18 @@ export type SessionEntry =
 	| ThinkingLevelChangeEntry
 	| ModelChangeEntry
 	| CompactionEntry
-	| BranchSummaryEntry;
+	| BranchSummaryEntry
+	| CustomEntry;
 
 /** Raw file entry (includes header) */
 export type FileEntry = SessionHeader | SessionEntry;
 
+/** Tree node for getTree() - defensive copy of session structure */
+export interface SessionTreeNode {
+	entry: SessionEntry;
+	children: SessionTreeNode[];
+}
+
 export interface SessionContext {
 	messages: AppMessage[];
 	thinkingLevel: string;
@@ -387,6 +387,17 @@ export function findMostRecentSession(sessionDir: string): string | null {
 	}
 }
 
+/**
+ * Manages conversation sessions as append-only trees stored in JSONL files.
+ *
+ * Each session entry has an id and parentId forming a tree structure. The "leaf"
+ * pointer tracks the current position. Appending creates a child of the current leaf.
+ * Branching moves the leaf to an earlier entry, allowing new branches without
+ * modifying history.
+ *
+ * Use buildSessionContext() to get the resolved message list for the LLM, which
+ * handles compaction summaries and follows the path from root to current leaf.
+ */
 export class SessionManager {
 	private sessionId: string = "";
 	private sessionFile: string = "";
@@ -394,7 +405,7 @@ export class SessionManager {
 	private cwd: string;
 	private persist: boolean;
 	private flushed: boolean = false;
-	private inMemoryEntries: FileEntry[] = [];
+	private fileEntries: FileEntry[] = [];
 	private byId: Map<string, SessionEntry> = new Map();
 	private leafId: string = "";
 
@@ -409,7 +420,7 @@ export class SessionManager {
 		if (sessionFile) {
 			this.setSessionFile(sessionFile);
 		} else {
-			this._initNewSession();
+			this.newSession();
 		}
 	}
 
@@ -417,22 +428,22 @@ export class SessionManager {
 	setSessionFile(sessionFile: string): void {
 		this.sessionFile = resolve(sessionFile);
 		if (existsSync(this.sessionFile)) {
-			this.inMemoryEntries = loadEntriesFromFile(this.sessionFile);
+			this.fileEntries = loadEntriesFromFile(this.sessionFile);
-			const header = this.inMemoryEntries.find((e) => e.type === "session") as SessionHeader | undefined;
+			const header = this.fileEntries.find((e) => e.type === "session") as SessionHeader | undefined;
 			this.sessionId = header?.id ?? randomUUID();
 
-			if (migrateToCurrentVersion(this.inMemoryEntries)) {
+			if (migrateToCurrentVersion(this.fileEntries)) {
 				this._rewriteFile();
 			}
 
 			this._buildIndex();
 			this.flushed = true;
 		} else {
-			this._initNewSession();
+			this.newSession();
 		}
 	}
 
-	private _initNewSession(): void {
+	newSession(): void {
 		this.sessionId = randomUUID();
 		const timestamp = new Date().toISOString();
 		const header: SessionHeader = {
@@ -442,7 +453,7 @@ export class SessionManager {
 			timestamp,
 			cwd: this.cwd,
 		};
-		this.inMemoryEntries = [header];
+		this.fileEntries = [header];
 		this.byId.clear();
 		this.leafId = "";
 		this.flushed = false;
@@ -456,7 +467,7 @@ export class SessionManager {
 	private _buildIndex(): void {
 		this.byId.clear();
 		this.leafId = "";
-		for (const entry of this.inMemoryEntries) {
+		for (const entry of this.fileEntries) {
 			if (entry.type === "session") continue;
 			this.byId.set(entry.id, entry);
 			this.leafId = entry.id;
@@ -465,7 +476,7 @@ export class SessionManager {
 
 	private _rewriteFile(): void {
 		if (!this.persist) return;
-		const content = `${this.inMemoryEntries.map((e) => JSON.stringify(e)).join("\n")}\n`;
+		const content = `${this.fileEntries.map((e) => JSON.stringify(e)).join("\n")}\n`;
 		writeFileSync(this.sessionFile, content);
 	}
 
@@ -489,18 +500,14 @@ export class SessionManager {
 		return this.sessionFile;
 	}
 
-	reset(): void {
-		this._initNewSession();
-	}
-
 	_persist(entry: SessionEntry): void {
 		if (!this.persist) return;
 
-		const hasAssistant = this.inMemoryEntries.some((e) => e.type === "message" && e.message.role === "assistant");
+		const hasAssistant = this.fileEntries.some((e) => e.type === "message" && e.message.role === "assistant");
 		if (!hasAssistant) return;
 
 		if (!this.flushed) {
-			for (const e of this.inMemoryEntries) {
+			for (const e of this.fileEntries) {
 				appendFileSync(this.sessionFile, `${JSON.stringify(e)}\n`);
 			}
 			this.flushed = true;
@@ -510,13 +517,14 @@ export class SessionManager {
 	}
 
 	private _appendEntry(entry: SessionEntry): void {
-		this.inMemoryEntries.push(entry);
+		this.fileEntries.push(entry);
 		this.byId.set(entry.id, entry);
 		this.leafId = entry.id;
 		this._persist(entry);
 	}
 
-	saveMessage(message: AppMessage): string {
+	/** Append a message as child of current leaf, then advance leaf. Returns entry id. */
+	appendMessage(message: AppMessage): string {
 		const entry: SessionMessageEntry = {
 			type: "message",
 			id: generateId(this.byId),
@@ -528,7 +536,8 @@ export class SessionManager {
 		return entry.id;
 	}
 
-	saveThinkingLevelChange(thinkingLevel: string): string {
+	/** Append a thinking level change as child of current leaf, then advance leaf. Returns entry id. */
+	appendThinkingLevelChange(thinkingLevel: string): string {
 		const entry: ThinkingLevelChangeEntry = {
 			type: "thinking_level_change",
 			id: generateId(this.byId),
@@ -540,7 +549,8 @@ export class SessionManager {
 		return entry.id;
 	}
 
-	saveModelChange(provider: string, modelId: string): string {
+	/** Append a model change as child of current leaf, then advance leaf. Returns entry id. */
+	appendModelChange(provider: string, modelId: string): string {
 		const entry: ModelChangeEntry = {
 			type: "model_change",
 			id: generateId(this.byId),
@@ -553,7 +563,8 @@ export class SessionManager {
 		return entry.id;
 	}
 
-	saveCompaction(summary: string, firstKeptEntryId: string, tokensBefore: number): string {
+	/** Append a compaction summary as child of current leaf, then advance leaf. Returns entry id. */
+	appendCompaction(summary: string, firstKeptEntryId: string, tokensBefore: number): string {
 		const entry: CompactionEntry = {
 			type: "compaction",
 			id: generateId(this.byId),
@@ -567,6 +578,20 @@ export class SessionManager {
 		return entry.id;
 	}
 
+	/** Append a custom entry (for hooks) as child of current leaf, then advance leaf. Returns entry id. */
+	appendCustomEntry(customType: string, data?: unknown): string {
+		const entry: CustomEntry = {
+			type: "custom",
+			customType,
+			data,
+			id: generateId(this.byId),
+			parentId: this.leafId || null,
+			timestamp: new Date().toISOString(),
+		};
+		this._appendEntry(entry);
+		return entry.id;
+	}
+
 	// =========================================================================
 	// Tree Traversal
 	// =========================================================================
@@ -575,11 +600,19 @@ export class SessionManager {
 		return this.leafId;
 	}
 
+	getLeafEntry(): SessionEntry | undefined {
+		return this.byId.get(this.leafId);
+	}
+
 	getEntry(id: string): SessionEntry | undefined {
 		return this.byId.get(id);
 	}
 
-	/** Walk from entry to root, returning path (conversation entries only) */
+	/**
+	 * Walk from entry to root, returning all entries in path order.
+	 * Includes all entry types (messages, compaction, model changes, etc.).
+	 * Use buildSessionContext() to get the resolved messages for the LLM.
+	 */
 	getPath(fromId?: string): SessionEntry[] {
 		const path: SessionEntry[] = [];
 		let current = this.byId.get(fromId ?? this.leafId);
@@ -602,31 +635,75 @@ export class SessionManager {
 	 * Get session header.
 	 */
 	getHeader(): SessionHeader | null {
-		const h = this.inMemoryEntries.find((e) => e.type === "session");
+		const h = this.fileEntries.find((e) => e.type === "session");
 		return h ? (h as SessionHeader) : null;
 	}
 
 	/**
-	 * Get all session entries (excludes header). Returns a defensive copy.
+	 * Get all session entries (excludes header). Returns a shallow copy.
-	 * Use buildSessionContext() if you need the messages for the LLM.
+	 * The session is append-only: use appendXXX() to add entries, branch() to
+	 * change the leaf pointer. Entries cannot be modified or deleted.
 	 */
 	getEntries(): SessionEntry[] {
-		return this.inMemoryEntries.filter((e): e is SessionEntry => e.type !== "session");
+		return this.fileEntries.filter((e): e is SessionEntry => e.type !== "session");
+	}
+
+	/**
+	 * Get the session as a tree structure. Returns a shallow defensive copy of all entries.
+	 * A well-formed session has exactly one root (first entry with parentId === null).
+	 * Orphaned entries (broken parent chain) are also returned as roots.
+	 */
+	getTree(): SessionTreeNode[] {
+		const entries = this.getEntries();
+		const nodeMap = new Map<string, SessionTreeNode>();
+		const roots: SessionTreeNode[] = [];
+
+		// Create nodes
+		for (const entry of entries) {
+			nodeMap.set(entry.id, { entry, children: [] });
+		}
+
+		// Build tree
+		for (const entry of entries) {
+			const node = nodeMap.get(entry.id)!;
+			if (entry.parentId === null) {
+				roots.push(node);
+			} else {
+				const parent = nodeMap.get(entry.parentId);
+				if (parent) {
+					parent.children.push(node);
+				} else {
+					// Orphan - treat as root
+					roots.push(node);
+				}
+			}
+		}
+
+		return roots;
 	}
 
 	// =========================================================================
 	// Branching
 	// =========================================================================
 
-	/** Branch in-place by changing the leaf pointer */
+	/**
-	branchInPlace(branchFromId: string): void {
+	 * Start a new branch from an earlier entry.
+	 * Moves the leaf pointer to the specified entry. The next appendXXX() call
+	 * will create a child of that entry, forming a new branch. Existing entries
+	 * are not modified or deleted.
+	 */
+	branch(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 */
+	/**
+	 * Start a new branch with a summary of the abandoned path.
+	 * Same as branch(), but also appends a branch_summary entry that captures
+	 * context from the abandoned conversation path.
+	 */
 	branchWithSummary(branchFromId: string, summary: string): string {
 		if (!this.byId.has(branchFromId)) {
 			throw new Error(`Entry ${branchFromId} not found`);
@@ -643,35 +720,41 @@ export class SessionManager {
 		return entry.id;
 	}
 
-	createBranchedSessionFromEntries(entries: FileEntry[], branchBeforeIndex: number): string | null {
+	/**
-		const newSessionId = randomUUID();
+	 * Create a new session file containing only the path from root to the specified leaf.
-		const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
+	 * Useful for extracting a single conversation path from a branched session.
-		const newSessionFile = join(this.getSessionDir(), `${timestamp}_${newSessionId}.jsonl`);
+	 * Returns the new session file path, or null if not persisting.
-
+	 */
-		const newEntries: FileEntry[] = [];
+	createBranchedSession(leafId: string): string | null {
-		for (let i = 0; i < branchBeforeIndex; i++) {
+		const path = this.getPath(leafId);
-			const entry = entries[i];
+		if (path.length === 0) {
-
+			throw new Error(`Entry ${leafId} not found`);
-			if (entry.type === "session") {
-				newEntries.push({
-					...entry,
-					version: CURRENT_SESSION_VERSION,
-					id: newSessionId,
-					timestamp: new Date().toISOString(),
-					branchedFrom: this.persist ? this.sessionFile : undefined,
-				});
-			} else {
-				newEntries.push(entry);
-			}
 		}
 
+		const newSessionId = randomUUID();
+		const timestamp = new Date().toISOString();
+		const fileTimestamp = timestamp.replace(/[:.]/g, "-");
+		const newSessionFile = join(this.getSessionDir(), `${fileTimestamp}_${newSessionId}.jsonl`);
+
+		const header: SessionHeader = {
+			type: "session",
+			version: CURRENT_SESSION_VERSION,
+			id: newSessionId,
+			timestamp,
+			cwd: this.cwd,
+			branchedFrom: this.persist ? this.sessionFile : undefined,
+		};
+
 		if (this.persist) {
-			for (const entry of newEntries) {
+			appendFileSync(newSessionFile, `${JSON.stringify(header)}\n`);
+			for (const entry of path) {
 				appendFileSync(newSessionFile, `${JSON.stringify(entry)}\n`);
 			}
 			return newSessionFile;
 		}
-		this.inMemoryEntries = newEntries;
+
+		// In-memory mode: replace current session with the path
+		this.fileEntries = [header, ...path];
 		this.sessionId = newSessionId;
 		this._buildIndex();
 		return null;

packages/coding-agent/src/index.ts

@@ -107,23 +107,19 @@ export {
 	readOnlyTools,
 } from "./core/sdk.js";
 export {
-	type BranchSummaryContent,
 	type BranchSummaryEntry,
 	buildSessionContext,
-	type CompactionContent,
 	type CompactionEntry,
 	CURRENT_SESSION_VERSION,
 	createSummaryMessage,
 	type FileEntry,
 	getLatestCompactionEntry,
-	type MessageContent,
-	type ModelChangeContent,
 	type ModelChangeEntry,
 	migrateSessionEntries,
 	parseSessionEntries,
-	type SessionContent as ConversationContent,
+	type SessionContext,
-	type SessionContext as LoadedSession,
 	type SessionEntry,
+	type SessionEntryBase,
 	type SessionHeader,
 	type SessionInfo,
 	SessionManager,
@@ -131,9 +127,6 @@ 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/session-manager/save-entry.test.ts

@@ -0,0 +1,55 @@
+import { describe, expect, it } from "vitest";
+import { type CustomEntry, SessionManager } from "../../src/core/session-manager.js";
+
+describe("SessionManager.saveCustomEntry", () => {
+	it("saves custom entries and includes them in tree traversal", () => {
+		const session = SessionManager.inMemory();
+
+		// Save a message
+		const msgId = session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
+
+		// Save a custom entry
+		const customId = session.appendCustomEntry("my_hook", { foo: "bar" });
+
+		// Save another message
+		const msg2Id = session.appendMessage({
+			role: "assistant",
+			content: [{ type: "text", text: "hi" }],
+			api: "anthropic-messages",
+			provider: "anthropic",
+			model: "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: 2,
+		});
+
+		// Custom entry should be in entries
+		const entries = session.getEntries();
+		expect(entries).toHaveLength(3);
+
+		const customEntry = entries.find((e) => e.type === "custom") as CustomEntry;
+		expect(customEntry).toBeDefined();
+		expect(customEntry.customType).toBe("my_hook");
+		expect(customEntry.data).toEqual({ foo: "bar" });
+		expect(customEntry.id).toBe(customId);
+		expect(customEntry.parentId).toBe(msgId);
+
+		// Tree structure should be correct
+		const path = session.getPath();
+		expect(path).toHaveLength(3);
+		expect(path[0].id).toBe(msgId);
+		expect(path[1].id).toBe(customId);
+		expect(path[2].id).toBe(msg2Id);
+
+		// buildSessionContext should work (custom entries skipped in messages)
+		const ctx = session.buildSessionContext();
+		expect(ctx.messages).toHaveLength(2); // only message entries
+	});
+});

packages/coding-agent/test/session-manager/tree-traversal.test.ts

@@ -0,0 +1,483 @@
+import { describe, expect, it } from "vitest";
+import { type CustomEntry, SessionManager } from "../../src/core/session-manager.js";
+
+function userMsg(text: string) {
+	return { role: "user" as const, content: text, timestamp: Date.now() };
+}
+
+function assistantMsg(text: string) {
+	return {
+		role: "assistant" as const,
+		content: [{ type: "text" as const, text }],
+		api: "anthropic-messages" as const,
+		provider: "anthropic",
+		model: "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" as const,
+		timestamp: Date.now(),
+	};
+}
+
+describe("SessionManager append and tree traversal", () => {
+	describe("append operations", () => {
+		it("appendMessage creates entry with correct parentId chain", () => {
+			const session = SessionManager.inMemory();
+
+			const id1 = session.appendMessage(userMsg("first"));
+			const id2 = session.appendMessage(assistantMsg("second"));
+			const id3 = session.appendMessage(userMsg("third"));
+
+			const entries = session.getEntries();
+			expect(entries).toHaveLength(3);
+
+			expect(entries[0].id).toBe(id1);
+			expect(entries[0].parentId).toBeNull();
+			expect(entries[0].type).toBe("message");
+
+			expect(entries[1].id).toBe(id2);
+			expect(entries[1].parentId).toBe(id1);
+
+			expect(entries[2].id).toBe(id3);
+			expect(entries[2].parentId).toBe(id2);
+		});
+
+		it("appendThinkingLevelChange integrates into tree", () => {
+			const session = SessionManager.inMemory();
+
+			const msgId = session.appendMessage(userMsg("hello"));
+			const thinkingId = session.appendThinkingLevelChange("high");
+			const _msg2Id = session.appendMessage(assistantMsg("response"));
+
+			const entries = session.getEntries();
+			expect(entries).toHaveLength(3);
+
+			const thinkingEntry = entries.find((e) => e.type === "thinking_level_change");
+			expect(thinkingEntry).toBeDefined();
+			expect(thinkingEntry!.id).toBe(thinkingId);
+			expect(thinkingEntry!.parentId).toBe(msgId);
+
+			expect(entries[2].parentId).toBe(thinkingId);
+		});
+
+		it("appendModelChange integrates into tree", () => {
+			const session = SessionManager.inMemory();
+
+			const msgId = session.appendMessage(userMsg("hello"));
+			const modelId = session.appendModelChange("openai", "gpt-4");
+			const _msg2Id = session.appendMessage(assistantMsg("response"));
+
+			const entries = session.getEntries();
+			const modelEntry = entries.find((e) => e.type === "model_change");
+			expect(modelEntry).toBeDefined();
+			expect(modelEntry?.id).toBe(modelId);
+			expect(modelEntry?.parentId).toBe(msgId);
+			if (modelEntry?.type === "model_change") {
+				expect(modelEntry.provider).toBe("openai");
+				expect(modelEntry.modelId).toBe("gpt-4");
+			}
+
+			expect(entries[2].parentId).toBe(modelId);
+		});
+
+		it("appendCompaction integrates into tree", () => {
+			const session = SessionManager.inMemory();
+
+			const id1 = session.appendMessage(userMsg("1"));
+			const id2 = session.appendMessage(assistantMsg("2"));
+			const compactionId = session.appendCompaction("summary", id1, 1000);
+			const _id3 = session.appendMessage(userMsg("3"));
+
+			const entries = session.getEntries();
+			const compactionEntry = entries.find((e) => e.type === "compaction");
+			expect(compactionEntry).toBeDefined();
+			expect(compactionEntry?.id).toBe(compactionId);
+			expect(compactionEntry?.parentId).toBe(id2);
+			if (compactionEntry?.type === "compaction") {
+				expect(compactionEntry.summary).toBe("summary");
+				expect(compactionEntry.firstKeptEntryId).toBe(id1);
+				expect(compactionEntry.tokensBefore).toBe(1000);
+			}
+
+			expect(entries[3].parentId).toBe(compactionId);
+		});
+
+		it("appendCustomEntry integrates into tree", () => {
+			const session = SessionManager.inMemory();
+
+			const msgId = session.appendMessage(userMsg("hello"));
+			const customId = session.appendCustomEntry("my_hook", { key: "value" });
+			const _msg2Id = session.appendMessage(assistantMsg("response"));
+
+			const entries = session.getEntries();
+			const customEntry = entries.find((e) => e.type === "custom") as CustomEntry;
+			expect(customEntry).toBeDefined();
+			expect(customEntry.id).toBe(customId);
+			expect(customEntry.parentId).toBe(msgId);
+			expect(customEntry.customType).toBe("my_hook");
+			expect(customEntry.data).toEqual({ key: "value" });
+
+			expect(entries[2].parentId).toBe(customId);
+		});
+
+		it("leaf pointer advances after each append", () => {
+			const session = SessionManager.inMemory();
+
+			expect(session.getLeafUuid()).toBe("");
+
+			const id1 = session.appendMessage(userMsg("1"));
+			expect(session.getLeafUuid()).toBe(id1);
+
+			const id2 = session.appendMessage(assistantMsg("2"));
+			expect(session.getLeafUuid()).toBe(id2);
+
+			const id3 = session.appendThinkingLevelChange("high");
+			expect(session.getLeafUuid()).toBe(id3);
+		});
+	});
+
+	describe("getPath", () => {
+		it("returns empty array for empty session", () => {
+			const session = SessionManager.inMemory();
+			expect(session.getPath()).toEqual([]);
+		});
+
+		it("returns single entry path", () => {
+			const session = SessionManager.inMemory();
+			const id = session.appendMessage(userMsg("hello"));
+
+			const path = session.getPath();
+			expect(path).toHaveLength(1);
+			expect(path[0].id).toBe(id);
+		});
+
+		it("returns full path from root to leaf", () => {
+			const session = SessionManager.inMemory();
+
+			const id1 = session.appendMessage(userMsg("1"));
+			const id2 = session.appendMessage(assistantMsg("2"));
+			const id3 = session.appendThinkingLevelChange("high");
+			const id4 = session.appendMessage(userMsg("3"));
+
+			const path = session.getPath();
+			expect(path).toHaveLength(4);
+			expect(path.map((e) => e.id)).toEqual([id1, id2, id3, id4]);
+		});
+
+		it("returns path from specified entry to root", () => {
+			const session = SessionManager.inMemory();
+
+			const id1 = session.appendMessage(userMsg("1"));
+			const id2 = session.appendMessage(assistantMsg("2"));
+			const _id3 = session.appendMessage(userMsg("3"));
+			const _id4 = session.appendMessage(assistantMsg("4"));
+
+			const path = session.getPath(id2);
+			expect(path).toHaveLength(2);
+			expect(path.map((e) => e.id)).toEqual([id1, id2]);
+		});
+	});
+
+	describe("getTree", () => {
+		it("returns empty array for empty session", () => {
+			const session = SessionManager.inMemory();
+			expect(session.getTree()).toEqual([]);
+		});
+
+		it("returns single root for linear session", () => {
+			const session = SessionManager.inMemory();
+
+			const id1 = session.appendMessage(userMsg("1"));
+			const id2 = session.appendMessage(assistantMsg("2"));
+			const id3 = session.appendMessage(userMsg("3"));
+
+			const tree = session.getTree();
+			expect(tree).toHaveLength(1);
+
+			const root = tree[0];
+			expect(root.entry.id).toBe(id1);
+			expect(root.children).toHaveLength(1);
+			expect(root.children[0].entry.id).toBe(id2);
+			expect(root.children[0].children).toHaveLength(1);
+			expect(root.children[0].children[0].entry.id).toBe(id3);
+			expect(root.children[0].children[0].children).toHaveLength(0);
+		});
+
+		it("returns tree with branches after branch", () => {
+			const session = SessionManager.inMemory();
+
+			// Build: 1 -> 2 -> 3
+			const id1 = session.appendMessage(userMsg("1"));
+			const id2 = session.appendMessage(assistantMsg("2"));
+			const id3 = session.appendMessage(userMsg("3"));
+
+			// Branch from id2, add new path: 2 -> 4
+			session.branch(id2);
+			const id4 = session.appendMessage(userMsg("4-branch"));
+
+			const tree = session.getTree();
+			expect(tree).toHaveLength(1);
+
+			const root = tree[0];
+			expect(root.entry.id).toBe(id1);
+			expect(root.children).toHaveLength(1);
+
+			const node2 = root.children[0];
+			expect(node2.entry.id).toBe(id2);
+			expect(node2.children).toHaveLength(2); // id3 and id4 are siblings
+
+			const childIds = node2.children.map((c) => c.entry.id).sort();
+			expect(childIds).toEqual([id3, id4].sort());
+		});
+
+		it("handles multiple branches at same point", () => {
+			const session = SessionManager.inMemory();
+
+			const _id1 = session.appendMessage(userMsg("root"));
+			const id2 = session.appendMessage(assistantMsg("response"));
+
+			// Branch A
+			session.branch(id2);
+			const idA = session.appendMessage(userMsg("branch-A"));
+
+			// Branch B
+			session.branch(id2);
+			const idB = session.appendMessage(userMsg("branch-B"));
+
+			// Branch C
+			session.branch(id2);
+			const idC = session.appendMessage(userMsg("branch-C"));
+
+			const tree = session.getTree();
+			const node2 = tree[0].children[0];
+			expect(node2.entry.id).toBe(id2);
+			expect(node2.children).toHaveLength(3);
+
+			const branchIds = node2.children.map((c) => c.entry.id).sort();
+			expect(branchIds).toEqual([idA, idB, idC].sort());
+		});
+
+		it("handles deep branching", () => {
+			const session = SessionManager.inMemory();
+
+			// Main path: 1 -> 2 -> 3 -> 4
+			const _id1 = session.appendMessage(userMsg("1"));
+			const id2 = session.appendMessage(assistantMsg("2"));
+			const id3 = session.appendMessage(userMsg("3"));
+			const _id4 = session.appendMessage(assistantMsg("4"));
+
+			// Branch from 2: 2 -> 5 -> 6
+			session.branch(id2);
+			const id5 = session.appendMessage(userMsg("5"));
+			const _id6 = session.appendMessage(assistantMsg("6"));
+
+			// Branch from 5: 5 -> 7
+			session.branch(id5);
+			const _id7 = session.appendMessage(userMsg("7"));
+
+			const tree = session.getTree();
+
+			// Verify structure
+			const node2 = tree[0].children[0];
+			expect(node2.children).toHaveLength(2); // id3 and id5
+
+			const node5 = node2.children.find((c) => c.entry.id === id5)!;
+			expect(node5.children).toHaveLength(2); // id6 and id7
+
+			const node3 = node2.children.find((c) => c.entry.id === id3)!;
+			expect(node3.children).toHaveLength(1); // id4
+		});
+	});
+
+	describe("branch", () => {
+		it("moves leaf pointer to specified entry", () => {
+			const session = SessionManager.inMemory();
+
+			const id1 = session.appendMessage(userMsg("1"));
+			const _id2 = session.appendMessage(assistantMsg("2"));
+			const id3 = session.appendMessage(userMsg("3"));
+
+			expect(session.getLeafUuid()).toBe(id3);
+
+			session.branch(id1);
+			expect(session.getLeafUuid()).toBe(id1);
+		});
+
+		it("throws for non-existent entry", () => {
+			const session = SessionManager.inMemory();
+			session.appendMessage(userMsg("hello"));
+
+			expect(() => session.branch("nonexistent")).toThrow("Entry nonexistent not found");
+		});
+
+		it("new appends become children of branch point", () => {
+			const session = SessionManager.inMemory();
+
+			const id1 = session.appendMessage(userMsg("1"));
+			const _id2 = session.appendMessage(assistantMsg("2"));
+
+			session.branch(id1);
+			const id3 = session.appendMessage(userMsg("branched"));
+
+			const entries = session.getEntries();
+			const branchedEntry = entries.find((e) => e.id === id3)!;
+			expect(branchedEntry.parentId).toBe(id1); // sibling of id2
+		});
+	});
+
+	describe("branchWithSummary", () => {
+		it("inserts branch summary and advances leaf", () => {
+			const session = SessionManager.inMemory();
+
+			const id1 = session.appendMessage(userMsg("1"));
+			const _id2 = session.appendMessage(assistantMsg("2"));
+			const _id3 = session.appendMessage(userMsg("3"));
+
+			const summaryId = session.branchWithSummary(id1, "Summary of abandoned work");
+
+			expect(session.getLeafUuid()).toBe(summaryId);
+
+			const entries = session.getEntries();
+			const summaryEntry = entries.find((e) => e.type === "branch_summary");
+			expect(summaryEntry).toBeDefined();
+			expect(summaryEntry?.parentId).toBe(id1);
+			if (summaryEntry?.type === "branch_summary") {
+				expect(summaryEntry.summary).toBe("Summary of abandoned work");
+			}
+		});
+
+		it("throws for non-existent entry", () => {
+			const session = SessionManager.inMemory();
+			session.appendMessage(userMsg("hello"));
+
+			expect(() => session.branchWithSummary("nonexistent", "summary")).toThrow("Entry nonexistent not found");
+		});
+	});
+
+	describe("getLeafEntry", () => {
+		it("returns undefined for empty session", () => {
+			const session = SessionManager.inMemory();
+			expect(session.getLeafEntry()).toBeUndefined();
+		});
+
+		it("returns current leaf entry", () => {
+			const session = SessionManager.inMemory();
+
+			session.appendMessage(userMsg("1"));
+			const id2 = session.appendMessage(assistantMsg("2"));
+
+			const leaf = session.getLeafEntry();
+			expect(leaf).toBeDefined();
+			expect(leaf!.id).toBe(id2);
+		});
+	});
+
+	describe("getEntry", () => {
+		it("returns undefined for non-existent id", () => {
+			const session = SessionManager.inMemory();
+			expect(session.getEntry("nonexistent")).toBeUndefined();
+		});
+
+		it("returns entry by id", () => {
+			const session = SessionManager.inMemory();
+
+			const id1 = session.appendMessage(userMsg("first"));
+			const id2 = session.appendMessage(assistantMsg("second"));
+
+			const entry1 = session.getEntry(id1);
+			expect(entry1).toBeDefined();
+			expect(entry1?.type).toBe("message");
+			if (entry1?.type === "message" && entry1.message.role === "user") {
+				expect(entry1.message.content).toBe("first");
+			}
+
+			const entry2 = session.getEntry(id2);
+			expect(entry2).toBeDefined();
+			if (entry2?.type === "message" && entry2.message.role === "assistant") {
+				expect((entry2.message.content as any)[0].text).toBe("second");
+			}
+		});
+	});
+
+	describe("buildSessionContext with branches", () => {
+		it("returns messages from current branch only", () => {
+			const session = SessionManager.inMemory();
+
+			// Main: 1 -> 2 -> 3
+			session.appendMessage(userMsg("msg1"));
+			const id2 = session.appendMessage(assistantMsg("msg2"));
+			session.appendMessage(userMsg("msg3"));
+
+			// Branch from 2: 2 -> 4
+			session.branch(id2);
+			session.appendMessage(assistantMsg("msg4-branch"));
+
+			const ctx = session.buildSessionContext();
+			expect(ctx.messages).toHaveLength(3); // msg1, msg2, msg4-branch (not msg3)
+
+			expect((ctx.messages[0] as any).content).toBe("msg1");
+			expect((ctx.messages[1] as any).content[0].text).toBe("msg2");
+			expect((ctx.messages[2] as any).content[0].text).toBe("msg4-branch");
+		});
+	});
+});
+
+describe("createBranchedSession", () => {
+	it("throws for non-existent entry", () => {
+		const session = SessionManager.inMemory();
+		session.appendMessage(userMsg("hello"));
+
+		expect(() => session.createBranchedSession("nonexistent")).toThrow("Entry nonexistent not found");
+	});
+
+	it("creates new session with path to specified leaf (in-memory)", () => {
+		const session = SessionManager.inMemory();
+
+		// Build: 1 -> 2 -> 3 -> 4
+		const id1 = session.appendMessage(userMsg("1"));
+		const id2 = session.appendMessage(assistantMsg("2"));
+		const id3 = session.appendMessage(userMsg("3"));
+		session.appendMessage(assistantMsg("4"));
+
+		// Branch from 3: 3 -> 5
+		session.branch(id3);
+		const _id5 = session.appendMessage(userMsg("5"));
+
+		// Create branched session from id2 (should only have 1 -> 2)
+		const result = session.createBranchedSession(id2);
+		expect(result).toBeNull(); // in-memory returns null
+
+		// Session should now only have entries 1 and 2
+		const entries = session.getEntries();
+		expect(entries).toHaveLength(2);
+		expect(entries[0].id).toBe(id1);
+		expect(entries[1].id).toBe(id2);
+	});
+
+	it("extracts correct path from branched tree", () => {
+		const session = SessionManager.inMemory();
+
+		// Build: 1 -> 2 -> 3
+		const id1 = session.appendMessage(userMsg("1"));
+		const id2 = session.appendMessage(assistantMsg("2"));
+		session.appendMessage(userMsg("3"));
+
+		// Branch from 2: 2 -> 4 -> 5
+		session.branch(id2);
+		const id4 = session.appendMessage(userMsg("4"));
+		const id5 = session.appendMessage(assistantMsg("5"));
+
+		// Create branched session from id5 (should have 1 -> 2 -> 4 -> 5)
+		session.createBranchedSession(id5);
+
+		const entries = session.getEntries();
+		expect(entries).toHaveLength(4);
+		expect(entries.map((e) => e.id)).toEqual([id1, id2, id4, id5]);
+	});
+});

packages/mom/src/agent.ts

@@ -441,7 +441,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
 	});
 
 	// Load existing messages
-	const loadedSession = sessionManager.loadSession();
+	const loadedSession = sessionManager.buildSessionContex();
 	if (loadedSession.messages.length > 0) {
 		agent.replaceMessages(loadedSession.messages);
 		log.logInfo(`[${channelId}] Loaded ${loadedSession.messages.length} messages from context.jsonl`);
@@ -628,7 +628,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
 
 			// Reload messages from context.jsonl
 			// This picks up any messages synced from log.jsonl before this run
-			const reloadedSession = sessionManager.loadSession();
+			const reloadedSession = sessionManager.buildSessionContex();
 			if (reloadedSession.messages.length > 0) {
 				agent.replaceMessages(reloadedSession.messages);
 				log.logInfo(`[${channelId}] Reloaded ${reloadedSession.messages.length} messages from context`);

packages/mom/src/context.ts

@@ -15,15 +15,12 @@ import {
 	buildSessionContext,
 	type CompactionEntry,
 	type FileEntry,
-	type LoadedSession,
-	type MessageContent,
-	type ModelChangeContent,
 	type ModelChangeEntry,
+	type SessionContext,
 	type SessionEntry,
+	type SessionEntryBase,
 	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";
@@ -98,15 +95,15 @@ export class MomSessionManager {
 		this.leafId = null;
 	}
 
-	private _createTreeNode(): Omit<TreeNode, "type"> {
+	private _createEntryBase(): Omit<SessionEntryBase, "type"> {
 		const id = uuidv4();
-		const node = {
+		const base = {
 			id,
 			parentId: this.leafId,
 			timestamp: new Date().toISOString(),
 		};
 		this.leafId = id;
-		return node;
+		return base;
 	}
 
 	private _persist(entry: SessionEntry): void {
@@ -281,22 +278,23 @@ export class MomSessionManager {
 	}
 
 	saveMessage(message: AppMessage): void {
-		const content: MessageContent = { type: "message", message };
+		const entry: SessionMessageEntry = { ...this._createEntryBase(), type: "message", message };
-		const entry: SessionMessageEntry = { ...this._createTreeNode(), ...content };
 		this.inMemoryEntries.push(entry);
 		this._persist(entry);
 	}
 
 	saveThinkingLevelChange(thinkingLevel: string): void {
-		const content: ThinkingLevelContent = { type: "thinking_level_change", thinkingLevel };
+		const entry: ThinkingLevelChangeEntry = {
-		const entry: ThinkingLevelChangeEntry = { ...this._createTreeNode(), ...content };
+			...this._createEntryBase(),
+			type: "thinking_level_change",
+			thinkingLevel,
+		};
 		this.inMemoryEntries.push(entry);
 		this._persist(entry);
 	}
 
 	saveModelChange(provider: string, modelId: string): void {
-		const content: ModelChangeContent = { type: "model_change", provider, modelId };
+		const entry: ModelChangeEntry = { ...this._createEntryBase(), type: "model_change", provider, modelId };
-		const entry: ModelChangeEntry = { ...this._createTreeNode(), ...content };
 		this.inMemoryEntries.push(entry);
 		this._persist(entry);
 	}
@@ -307,7 +305,7 @@ export class MomSessionManager {
 	}
 
 	/** Load session with compaction support */
-	loadSession(): LoadedSession {
+	buildSessionContex(): SessionContext {
 		const entries = this.loadEntries();
 		return buildSessionContext(entries);
 	}
@@ -354,15 +352,15 @@ export class MomSessionManager {
 	}
 
 	loadModel(): { provider: string; modelId: string } | null {
-		return this.loadSession().model;
+		return this.buildSessionContex().model;
 	}
 
 	loadThinkingLevel(): string {
-		return this.loadSession().thinkingLevel;
+		return this.buildSessionContex().thinkingLevel;
 	}
 
 	/** Not used by mom but required by AgentSession interface */
-	createBranchedSessionFromEntries(_entries: SessionEntry[], _branchBeforeIndex: number): string | null {
+	createBranchedSession(_leafId: string): string | null {
 		return null; // Mom doesn't support branching
 	}
 }