oh-my-pi/packages/coding-agent/src/session/session-manager.ts

import * as path from "node:path";
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
import type { ImageContent, Message, TextContent, Usage } from "@oh-my-pi/pi-ai";
import { isEnoent, logger } from "@oh-my-pi/pi-utils";
import { nanoid } from "nanoid";
import { getAgentDir as getDefaultAgentDir } from "../config";
import { resizeImage } from "../utils/image-resize";
import {
	type BashExecutionMessage,
	type CustomMessage,
	createBranchSummaryMessage,
	createCompactionSummaryMessage,
	createCustomMessage,
	type FileMentionMessage,
	type HookMessage,
	type PythonExecutionMessage,
} from "./messages";
import type { SessionStorage, SessionStorageWriter } from "./session-storage";
import { FileSessionStorage, MemorySessionStorage } from "./session-storage";

export const CURRENT_SESSION_VERSION = 3;

export interface SessionHeader {
	type: "session";
	version?: number; // v1 sessions don't have this
	id: string;
	title?: string; // Auto-generated title from first message
	timestamp: string;
	cwd: string;
	parentSession?: string;
}

export interface NewSessionOptions {
	parentSession?: string;
}

export interface SessionEntryBase {
	type: string;
	id: string;
	parentId: string | null;
	timestamp: string;
}

export interface SessionMessageEntry extends SessionEntryBase {
	type: "message";
	message: AgentMessage;
}

export interface ThinkingLevelChangeEntry extends SessionEntryBase {
	type: "thinking_level_change";
	thinkingLevel: string;
}

export interface ModelChangeEntry extends SessionEntryBase {
	type: "model_change";
	/** Model in "provider/modelId" format */
	model: string;
	/** Role: "default", "smol", "slow", etc. Undefined treated as "default" */
	role?: string;
}

export interface CompactionEntry<T = unknown> extends SessionEntryBase {
	type: "compaction";
	summary: string;
	shortSummary?: string;
	firstKeptEntryId: string;
	tokensBefore: number;
	/** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */
	details?: T;
	/** Hook-provided data to persist across compaction */
	preserveData?: Record<string, unknown>;
	/** True if generated by an extension, undefined/false if pi-generated (backward compatible) */
	fromExtension?: boolean;
}

export interface BranchSummaryEntry<T = unknown> extends SessionEntryBase {
	type: "branch_summary";
	fromId: string;
	summary: string;
	/** Extension-specific data (not sent to LLM) */
	details?: T;
	/** True if generated by an extension, false if pi-generated */
	fromExtension?: boolean;
}

/**
 * Custom entry for extensions to store extension-specific data in the session.
 * Use customType to identify your extension's entries.
 *
 * Purpose: Persist extension state across session reloads. On reload, extensions can
 * scan entries for their customType and reconstruct internal state.
 *
 * Does NOT participate in LLM context (ignored by buildSessionContext).
 * For injecting content into context, see CustomMessageEntry.
 */
export interface CustomEntry<T = unknown> extends SessionEntryBase {
	type: "custom";
	customType: string;
	data?: T;
}

/** Label entry for user-defined bookmarks/markers on entries. */
export interface LabelEntry extends SessionEntryBase {
	type: "label";
	targetId: string;
	label: string | undefined;
}

/** TTSR injection entry - tracks which time-traveling rules have been injected this session. */
export interface TtsrInjectionEntry extends SessionEntryBase {
	type: "ttsr_injection";
	/** Names of rules that were injected */
	injectedRules: string[];
}

/** Session init entry - captures initial context for subagent sessions (debugging/replay). */
export interface SessionInitEntry extends SessionEntryBase {
	type: "session_init";
	/** Full system prompt sent to the model */
	systemPrompt: string;
	/** Initial task/user message */
	task: string;
	/** Tools available to the agent */
	tools: string[];
	/** Output schema if structured output was requested */
	outputSchema?: unknown;
}

/**
 * Custom message entry for extensions to inject messages into LLM context.
 * Use customType to identify your extension's entries.
 *
 * Unlike CustomEntry, this DOES participate in LLM context.
 * The content is converted to a user message in buildSessionContext().
 * Use details for extension-specific metadata (not sent to LLM).
 *
 * display controls TUI rendering:
 * - false: hidden entirely
 * - true: rendered with distinct styling (different from user messages)
 */
export interface CustomMessageEntry<T = unknown> extends SessionEntryBase {
	type: "custom_message";
	customType: string;
	content: string | (TextContent | ImageContent)[];
	details?: T;
	display: boolean;
}

/** Session entry - has id/parentId for tree structure (returned by "read" methods in SessionManager) */
export type SessionEntry =
	| SessionMessageEntry
	| ThinkingLevelChangeEntry
	| ModelChangeEntry
	| CompactionEntry
	| BranchSummaryEntry
	| CustomEntry
	| CustomMessageEntry
	| LabelEntry
	| TtsrInjectionEntry
	| SessionInitEntry;

/** 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[];
	/** Resolved label for this entry, if any */
	label?: string;
}

export interface SessionContext {
	messages: AgentMessage[];
	thinkingLevel: string;
	/** Model roles: { default: "provider/modelId", small: "provider/modelId", ... } */
	models: Record<string, string>;
	/** Names of TTSR rules that have been injected this session */
	injectedTtsrRules: string[];
}

export interface SessionInfo {
	path: string;
	id: string;
	/** Working directory where the session was started. Empty string for old sessions. */
	cwd: string;
	title?: string;
	created: Date;
	modified: Date;
	messageCount: number;
	firstMessage: string;
	allMessagesText: string;
}

export type ReadonlySessionManager = Pick<
	SessionManager,
	| "getCwd"
	| "getSessionDir"
	| "getSessionId"
	| "getSessionFile"
	| "getLeafId"
	| "getLeafEntry"
	| "getEntry"
	| "getLabel"
	| "getBranch"
	| "getHeader"
	| "getEntries"
	| "getTree"
	| "getUsageStatistics"
>;

/** Generate a unique short ID (8 hex chars, collision-checked) */
function generateId(byId: { has(id: string): boolean }): string {
	for (let i = 0; i < 100; i++) {
		const id = nanoid(8);
		if (!byId.has(id)) return id;
	}
	return nanoid(); // fallback to full nanoid
}

/** Migrate v1 → v2: add id/parentId tree structure. Mutates in place. */
function migrateV1ToV2(entries: FileEntry[]): void {
	const ids = new Set<string>();
	let prevId: string | null = null;

	for (const entry of entries) {
		if (entry.type === "session") {
			entry.version = 2;
			continue;
		}

		entry.id = generateId(ids);
		entry.parentId = prevId;
		prevId = entry.id;

		// Convert firstKeptEntryIndex to firstKeptEntryId for compaction
		if (entry.type === "compaction") {
			const comp = entry as CompactionEntry & { firstKeptEntryIndex?: number };
			if (typeof comp.firstKeptEntryIndex === "number") {
				const targetEntry = entries[comp.firstKeptEntryIndex];
				if (targetEntry && targetEntry.type !== "session") {
					comp.firstKeptEntryId = targetEntry.id;
				}
				delete comp.firstKeptEntryIndex;
			}
		}
	}
}

/** Migrate v2 → v3: rename hookMessage role to custom. Mutates in place. */
function migrateV2ToV3(entries: FileEntry[]): void {
	for (const entry of entries) {
		if (entry.type === "session") {
			entry.version = 3;
			continue;
		}

		if (entry.type === "message") {
			const msg = entry.message as { role?: string };
			if (msg.role === "hookMessage") {
				(entry.message as { role: string }).role = "custom";
			}
		}
	}
}

/**
 * Run all necessary migrations to bring entries to current version.
 * Mutates entries in place. Returns true if any migration was applied.
 */
function migrateToCurrentVersion(entries: FileEntry[]): boolean {
	const header = entries.find(e => e.type === "session") as SessionHeader | undefined;
	const version = header?.version ?? 1;

	if (version >= CURRENT_SESSION_VERSION) return false;

	if (version < 2) migrateV1ToV2(entries);
	if (version < 3) migrateV2ToV3(entries);

	return true;
}

function parseJsonlEntries<T>(content: string): T[] {
	if (!content.trim()) return [];
	const entries: T[] = [];
	let buffer = content;
	while (buffer.length > 0) {
		const result = Bun.JSONL.parseChunk(buffer);
		if (result.values.length > 0) {
			entries.push(...(result.values as T[]));
		}
		if (result.error) {
			const nextNewline = buffer.indexOf("\n", result.read);
			if (nextNewline === -1) break;
			buffer = buffer.slice(nextNewline + 1);
			continue;
		}
		if (result.read === 0) break;
		buffer = buffer.slice(result.read);
		if (result.done) break;
	}
	return entries;
}

/** Exported for testing */
export function migrateSessionEntries(entries: FileEntry[]): void {
	migrateToCurrentVersion(entries);
}

/** Exported for compaction.test.ts */
export function parseSessionEntries(content: string): FileEntry[] {
	return parseJsonlEntries<FileEntry>(content);
}

export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEntry | null {
	for (let i = entries.length - 1; i >= 0; i--) {
		if (entries[i].type === "compaction") {
			return entries[i] as CompactionEntry;
		}
	}
	return null;
}

function toError(value: unknown): Error {
	return value instanceof Error ? value : new Error(String(value));
}

/**
 * Build the session context from entries using tree traversal.
 * If leafId is provided, walks from that entry to root.
 * Handles compaction and branch summaries along the path.
 */
export function buildSessionContext(
	entries: SessionEntry[],
	leafId?: string | null,
	byId?: Map<string, SessionEntry>,
): SessionContext {
	// Build uuid index if not available
	if (!byId) {
		byId = new Map<string, SessionEntry>();
		for (const entry of entries) {
			byId.set(entry.id, entry);
		}
	}

	// Find leaf
	let leaf: SessionEntry | undefined;
	if (leafId === null) {
		// Explicitly null - return no messages (navigated to before first entry)
		return { messages: [], thinkingLevel: "off", models: {}, injectedTtsrRules: [] };
	}
	if (leafId) {
		leaf = byId.get(leafId);
	}
	if (!leaf) {
		// Fallback to last entry (when leafId is undefined)
		leaf = entries[entries.length - 1];
	}

	if (!leaf) {
		return { messages: [], thinkingLevel: "off", models: {}, injectedTtsrRules: [] };
	}

	// Walk from leaf to root, collecting path
	const path: SessionEntry[] = [];
	let current: SessionEntry | undefined = leaf;
	while (current) {
		path.unshift(current);
		current = current.parentId ? byId.get(current.parentId) : undefined;
	}

	// Extract settings and find compaction
	let thinkingLevel = "off";
	const models: Record<string, string> = {};
	let compaction: CompactionEntry | null = null;
	const injectedTtsrRulesSet = new Set<string>();

	for (const entry of path) {
		if (entry.type === "thinking_level_change") {
			thinkingLevel = entry.thinkingLevel;
		} else if (entry.type === "model_change") {
			// New format: { model: "provider/id", role?: string }
			if (entry.model) {
				const role = entry.role ?? "default";
				models[role] = entry.model;
			}
		} else if (entry.type === "message" && entry.message.role === "assistant") {
			// Infer default model from assistant messages
			models.default = `${entry.message.provider}/${entry.message.model}`;
		} else if (entry.type === "compaction") {
			compaction = entry;
		} else if (entry.type === "ttsr_injection") {
			// Collect injected TTSR rule names
			for (const ruleName of entry.injectedRules) {
				injectedTtsrRulesSet.add(ruleName);
			}
		}
	}

	const injectedTtsrRules = Array.from(injectedTtsrRulesSet);

	// Build messages and collect corresponding entries
	// When there's a compaction, we need to:
	// 1. Emit summary first (entry = compaction)
	// 2. Emit kept messages (from firstKeptEntryId up to compaction)
	// 3. Emit messages after compaction
	const messages: AgentMessage[] = [];

	const appendMessage = (entry: SessionEntry) => {
		if (entry.type === "message") {
			messages.push(entry.message);
		} else if (entry.type === "custom_message") {
			messages.push(
				createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp),
			);
		} else if (entry.type === "branch_summary" && entry.summary) {
			messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp));
		}
	};

	if (compaction) {
		// Emit summary first
		messages.push(
			createCompactionSummaryMessage(
				compaction.summary,
				compaction.tokensBefore,
				compaction.timestamp,
				compaction.shortSummary,
			),
		);

		// 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) {
				appendMessage(entry);
			}
		}

		// Emit messages after compaction
		for (let i = compactionIdx + 1; i < path.length; i++) {
			const entry = path[i];
			appendMessage(entry);
		}
	} else {
		// No compaction - emit all messages, handle branch summaries and custom messages
		for (const entry of path) {
			appendMessage(entry);
		}
	}

	return { messages, thinkingLevel, models, injectedTtsrRules };
}

/**
 * Compute the default session directory for a cwd.
 * Encodes cwd into a safe directory name under ~/.omp/agent/sessions/.
 */
function getDefaultSessionDir(cwd: string, storage: SessionStorage): string {
	const safePath = `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`;
	const sessionDir = path.join(getDefaultAgentDir(), "sessions", safePath);
	storage.ensureDirSync(sessionDir);
	return sessionDir;
}

/** Exported for testing */
export async function loadEntriesFromFile(
	filePath: string,
	storage: SessionStorage = new FileSessionStorage(),
): Promise<FileEntry[]> {
	let content: string;
	try {
		content = await storage.readText(filePath);
	} catch (err) {
		if (isEnoent(err)) return [];
		throw err;
	}
	const entries = parseJsonlEntries<FileEntry>(content);

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

	return entries;
}

/**
 * Lightweight metadata for a session file, used in session picker UI.
 * Uses lazy getters to defer string formatting until actually displayed.
 */
function sanitizeSessionName(value: string | undefined): string | undefined {
	if (!value) return undefined;
	const firstLine = value.split(/\r?\n/)[0] ?? "";
	const stripped = firstLine.replace(/[\x00-\x1F\x7F]/g, "");
	const trimmed = stripped.trim();
	return trimmed.length > 0 ? trimmed : undefined;
}

class RecentSessionInfo {
	readonly path: string;
	readonly mtime: number;

	#fullName: string | undefined;
	#name: string | undefined;
	#timeAgo: string | undefined;

	constructor(path: string, mtime: number, header: Record<string, unknown>, firstPrompt?: string) {
		this.path = path;
		this.mtime = mtime;

		// Extract title from session header, falling back to first user prompt, then id
		const trystr = (v: unknown) => (typeof v === "string" ? v : undefined);
		this.#fullName =
			sanitizeSessionName(trystr(header.title)) ??
			sanitizeSessionName(firstPrompt) ??
			sanitizeSessionName(trystr(header.id));
	}

	/** Full session name from header, or filename without extension as fallback */
	get fullName(): string {
		if (this.#fullName) return this.#fullName;
		this.#fullName = this.path.split("/").pop()?.replace(".jsonl", "") ?? "Unknown";
		return this.#fullName;
	}

	/** Truncated name for display (max 40 chars) */
	get name(): string {
		if (this.#name) return this.#name;
		const fullName = this.fullName;
		this.#name = fullName.length <= 40 ? fullName : `${fullName.slice(0, 37)}...`;
		return this.#name;
	}

	/** Human-readable relative time (e.g., "2 hours ago") */
	get timeAgo(): string {
		if (this.#timeAgo) return this.#timeAgo;
		this.#timeAgo = formatTimeAgo(new Date(this.mtime));
		return this.#timeAgo;
	}
}

/**
 * Extracts the text content from a user message entry.
 * Returns undefined if the entry is not a user message or has no text.
 */
function extractFirstUserPrompt(entries: Array<Record<string, unknown>>): string | undefined {
	for (const entry of entries) {
		if (entry.type !== "message") continue;
		const message = entry.message as Record<string, unknown> | undefined;
		if (message?.role !== "user") continue;
		const content = message.content;
		if (typeof content === "string") return content;
		if (Array.isArray(content)) {
			for (const block of content) {
				if (typeof block === "object" && block !== null && "text" in block) {
					const text = (block as { text: unknown }).text;
					if (typeof text === "string") return text;
				}
			}
		}
	}
	return undefined;
}

/**
 * Reads all session files from the directory and returns them sorted by mtime (newest first).
 * Uses low-level file I/O to efficiently read only the first 4KB of each file
 * to extract the JSON header and first user message without loading entire session logs into memory.
 */
async function getSortedSessions(sessionDir: string, storage: SessionStorage): Promise<RecentSessionInfo[]> {
	try {
		const files: string[] = storage.listFilesSync(sessionDir, "*.jsonl");
		const results = await Promise.all(
			files.map(async (path: string) => {
				try {
					const content = await storage.readTextPrefix(path, 4096);
					const entries = parseJsonlEntries<Record<string, unknown>>(content);
					if (entries.length === 0) return null;
					const header = entries[0] as Record<string, unknown>;
					if (header.type !== "session" || typeof header.id !== "string") return null;
					const mtime = storage.statSync(path).mtimeMs;
					const firstPrompt = header.title ? undefined : extractFirstUserPrompt(entries);
					return new RecentSessionInfo(path, mtime, header, firstPrompt);
				} catch {
					return null;
				}
			}),
		);
		return results.filter((item): item is RecentSessionInfo => item !== null).sort((a, b) => b.mtime - a.mtime);
	} catch {
		return [];
	}
}

/** Exported for testing */
export async function findMostRecentSession(
	sessionDir: string,
	storage: SessionStorage = new FileSessionStorage(),
): Promise<string | null> {
	const sessions = await getSortedSessions(sessionDir, storage);
	return sessions[0]?.path || null;
}

/** Format a time difference as a human-readable string */
function formatTimeAgo(date: Date): string {
	const now = Date.now();
	const diffMs = now - date.getTime();
	const diffMins = Math.floor(diffMs / 60000);
	const diffHours = Math.floor(diffMs / 3600000);
	const diffDays = Math.floor(diffMs / 86400000);

	if (diffMins < 1) return "just now";
	if (diffMins < 60) return `${diffMins}m ago`;
	if (diffHours < 24) return `${diffHours}h ago`;
	if (diffDays < 7) return `${diffDays}d ago`;
	return date.toLocaleDateString();
}

const MAX_PERSIST_CHARS = 500_000;
const TRUNCATION_NOTICE = "\n\n[Session persistence truncated large content]";
const PLACEHOLDER_IMAGE_DATA =
	"/9j/4AAQSkZJRgABAQEASABIAAD/2wBDAP//////////////////////////////////////////////////////////////////////////////////////2wBDAf//////////////////////////////////////////////////////////////////////////////////////wAARCAAQABADASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAf/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAgP/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwCkA//Z";

const TEXT_CONTENT_KEY = "content";

/**
 * Recursively truncate large strings in an object for session persistence.
 * - Truncates any oversized string fields (key-agnostic)
 * - Replaces oversized image blocks with text notices
 * - Updates lineCount when content is truncated
 * - Returns original object if no changes needed (structural sharing)
 */
function truncateString(value: string, maxLength: number): string {
	if (value.length <= maxLength) return value;
	let truncated = value.slice(0, maxLength);
	if (truncated.length > 0) {
		const last = truncated.charCodeAt(truncated.length - 1);
		if (last >= 0xd800 && last <= 0xdbff) {
			truncated = truncated.slice(0, -1);
		}
	}
	return truncated;
}

function isImageBlock(value: unknown): value is { type: "image"; data: string; mimeType?: string } {
	return (
		typeof value === "object" &&
		value !== null &&
		"type" in value &&
		(value as { type?: string }).type === "image" &&
		"data" in value &&
		typeof (value as { data?: string }).data === "string"
	);
}

async function compressImageForPersistence(image: ImageContent): Promise<ImageContent> {
	try {
		const maxBytes = Math.floor((MAX_PERSIST_CHARS * 3) / 4);
		const resized = await resizeImage(image, {
			maxWidth: 512,
			maxHeight: 512,
			maxBytes,
			jpegQuality: 70,
		});
		if (resized.data.length > MAX_PERSIST_CHARS) {
			return { type: "image", data: PLACEHOLDER_IMAGE_DATA, mimeType: "image/jpeg" };
		}
		return { type: "image", data: resized.data, mimeType: resized.mimeType };
	} catch {
		return { type: "image", data: PLACEHOLDER_IMAGE_DATA, mimeType: "image/jpeg" };
	}
}

async function truncateForPersistence<T>(obj: T, key?: string): Promise<T> {
	if (obj === null || obj === undefined) return obj;

	if (typeof obj === "string") {
		if (obj.length > MAX_PERSIST_CHARS) {
			const limit = Math.max(0, MAX_PERSIST_CHARS - TRUNCATION_NOTICE.length);
			return `${truncateString(obj, limit)}${TRUNCATION_NOTICE}` as T;
		}
		return obj;
	}

	if (Array.isArray(obj)) {
		let changed = false;
		const result = await Promise.all(
			obj.map(async item => {
				// Special handling: compress oversized images while preserving shape
				if (key === TEXT_CONTENT_KEY && isImageBlock(item)) {
					if (item.data.length > MAX_PERSIST_CHARS) {
						changed = true;
						return compressImageForPersistence({
							type: "image",
							data: item.data,
							mimeType: item.mimeType ?? "image/jpeg",
						});
					}
				}
				const newItem = await truncateForPersistence(item, key);
				if (newItem !== item) changed = true;
				return newItem;
			}),
		);
		return changed ? (result as T) : obj;
	}

	if (typeof obj === "object") {
		let changed = false;
		const result: Record<string, unknown> = {};
		for (const [k, v] of Object.entries(obj as Record<string, unknown>)) {
			// Strip transient/redundant properties that shouldn't be persisted
			// - partialJson: streaming accumulator for tool call JSON parsing
			// - jsonlEvents: raw subprocess streaming events (already saved to artifact files)
			if (k === "partialJson" || k === "jsonlEvents") {
				changed = true;
				continue;
			}
			const newV = await truncateForPersistence(v, k);
			result[k] = newV;
			if (newV !== v) changed = true;
		}
		// Update lineCount if content was truncated (for FileMentionFile)
		if (changed && "lineCount" in result && "content" in result && typeof result.content === "string") {
			result.lineCount = result.content.split("\n").length;
		}
		return changed ? (result as T) : obj;
	}

	return obj;
}

async function prepareEntryForPersistence(entry: FileEntry): Promise<FileEntry> {
	return truncateForPersistence(entry);
}

class NdjsonFileWriter {
	private writer: SessionStorageWriter;
	private closed = false;
	private closing = false;
	private error: Error | undefined;
	private pendingWrites: Promise<void> = Promise.resolve();
	private onError: ((err: Error) => void) | undefined;

	constructor(storage: SessionStorage, path: string, options?: { flags?: "a" | "w"; onError?: (err: Error) => void }) {
		this.onError = options?.onError;
		this.writer = storage.openWriter(path, {
			flags: options?.flags ?? "a",
			onError: (err: Error) => this.recordError(err),
		});
	}

	private recordError(err: unknown): Error {
		const writeErr = toError(err);
		if (!this.error) this.error = writeErr;
		this.onError?.(writeErr);
		return writeErr;
	}

	private enqueue(task: () => Promise<void>): Promise<void> {
		const run = async () => {
			if (this.error) throw this.error;
			await task();
		};
		const next = this.pendingWrites.then(run);
		void next.catch((err: unknown) => {
			if (!this.error) this.error = toError(err);
		});
		this.pendingWrites = next;
		return next;
	}

	private async writeLine(line: string): Promise<void> {
		if (this.error) throw this.error;
		try {
			await this.writer.writeLine(line);
		} catch (err) {
			throw this.recordError(err);
		}
	}

	/** Queue a write. Returns a promise so callers can await if needed. */
	write(entry: FileEntry): Promise<void> {
		if (this.closed || this.closing) throw new Error("Writer closed");
		if (this.error) throw this.error;
		const line = `${JSON.stringify(entry)}\n`;
		return this.enqueue(() => this.writeLine(line));
	}

	/** Flush all buffered data to disk. Waits for all queued writes. */
	async flush(): Promise<void> {
		if (this.closed) return;
		if (this.error) throw this.error;

		await this.enqueue(async () => {});

		if (this.error) throw this.error;

		try {
			await this.writer.flush();
		} catch (err) {
			throw this.recordError(err);
		}
	}

	/** Sync data to persistent storage. */
	async fsync(): Promise<void> {
		if (this.closed) return;
		if (this.error) throw this.error;
		try {
			await this.writer.fsync();
		} catch (err) {
			throw this.recordError(err);
		}
	}

	/** Close the writer, flushing all data. */
	async close(): Promise<void> {
		if (this.closed || this.closing) return;
		this.closing = true;

		let closeError: Error | undefined;
		try {
			await this.flush();
		} catch (err) {
			closeError = toError(err);
		}

		try {
			await this.pendingWrites;
		} catch (err) {
			if (!closeError) closeError = toError(err);
		}

		try {
			await this.writer.close();
		} catch (err) {
			const endErr = this.recordError(err);
			if (!closeError) closeError = endErr;
		}

		this.closed = true;

		if (!closeError && this.error) closeError = this.error;
		if (closeError) throw closeError;
	}

	/** Check if there's a stored error. */
	getError(): Error | undefined {
		return this.error;
	}
}

/** Get recent sessions for display in welcome screen */
export async function getRecentSessions(
	sessionDir: string,
	limit = 3,
	storage: SessionStorage = new FileSessionStorage(),
): Promise<RecentSessionInfo[]> {
	const sessions = await getSortedSessions(sessionDir, storage);
	return sessions.slice(0, limit);
}

/**
 * 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 interface UsageStatistics {
	input: number;
	output: number;
	cacheRead: number;
	cacheWrite: number;
	cost: number;
}

function getTaskToolUsage(details: unknown): Usage | undefined {
	if (!details || typeof details !== "object") return undefined;
	const record = details as Record<string, unknown>;
	const usage = record.usage;
	if (!usage || typeof usage !== "object") return undefined;
	return usage as Usage;
}

function extractTextFromContent(content: Message["content"]): string {
	if (typeof content === "string") return content;
	return content
		.filter((block): block is TextContent => block.type === "text")
		.map(block => block.text)
		.join(" ");
}

async function collectSessionsFromFiles(files: string[], storage: SessionStorage): Promise<SessionInfo[]> {
	const sessions: SessionInfo[] = [];

	for (const file of files) {
		try {
			const content = await storage.readText(file);
			const entries = parseJsonlEntries<Record<string, unknown>>(content);
			if (entries.length === 0) continue;

			// Check first entry for valid session header
			type SessionHeaderShape = { type: string; id: string; cwd?: string; title?: string; timestamp: string };
			const header = entries[0] as SessionHeaderShape;
			if (header.type !== "session" || !header.id) continue;

			const stats = storage.statSync(file);
			let messageCount = 0;
			let firstMessage = "";
			const allMessages: string[] = [];
			let shortSummary: string | undefined;

			for (let i = 1; i < entries.length; i++) {
				const entry = entries[i] as { type?: string; message?: Message; shortSummary?: string };

				if (entry.type === "compaction" && typeof entry.shortSummary === "string") {
					shortSummary = entry.shortSummary;
				}

				if (entry.type === "message" && entry.message) {
					messageCount++;

					if (entry.message.role === "user" || entry.message.role === "assistant") {
						const textContent = extractTextFromContent(entry.message.content);

						if (textContent) {
							allMessages.push(textContent);

							if (!firstMessage && entry.message.role === "user") {
								firstMessage = textContent;
							}
						}
					}
				}
			}

			sessions.push({
				path: file,
				id: header.id,
				cwd: typeof header.cwd === "string" ? header.cwd : "",
				title: header.title ?? shortSummary,
				created: new Date(header.timestamp),
				modified: stats.mtime,
				messageCount,
				firstMessage: firstMessage || "(no messages)",
				allMessagesText: allMessages.join(" "),
			});
		} catch {
			// Skip files that can't be read
		}
	}

	sessions.sort((a, b) => b.modified.getTime() - a.modified.getTime());
	return sessions;
}

export class SessionManager {
	private sessionId: string = "";
	private sessionName: string | undefined;
	private sessionFile: string | undefined;
	private sessionDir: string;
	private cwd: string;
	private persist: boolean;
	private flushed: boolean = false;
	private fileEntries: FileEntry[] = [];
	private byId: Map<string, SessionEntry> = new Map();
	private labelsById: Map<string, string> = new Map();
	private leafId: string | null = null;
	private usageStatistics: UsageStatistics = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0 };
	private persistWriter: NdjsonFileWriter | undefined;
	private persistWriterPath: string | undefined;
	private persistChain: Promise<void> = Promise.resolve();
	private persistError: Error | undefined;
	private persistErrorReported = false;
	private storage: SessionStorage;

	private constructor(cwd: string, sessionDir: string, persist: boolean, storage: SessionStorage) {
		this.cwd = cwd;
		this.sessionDir = sessionDir;
		this.persist = persist;
		this.storage = storage;
		if (persist && sessionDir) {
			this.storage.ensureDirSync(sessionDir);
		}
		// Note: call _initSession() or _initSessionFile() after construction
	}

	/** Initialize with a specific session file (used by factory methods) */
	private async _initSessionFile(sessionFile: string): Promise<void> {
		await this.setSessionFile(sessionFile);
	}

	/** Initialize with a new session (used by factory methods) */
	private _initNewSession(): void {
		this._newSessionSync();
	}

	/** Switch to a different session file (used for resume and branching) */
	async setSessionFile(sessionFile: string): Promise<void> {
		await this._closePersistWriter();
		this.persistError = undefined;
		this.persistErrorReported = false;
		this.sessionFile = path.resolve(sessionFile);
		this.fileEntries = await loadEntriesFromFile(this.sessionFile, this.storage);
		if (this.fileEntries.length > 0) {
			const header = this.fileEntries.find(e => e.type === "session") as SessionHeader | undefined;
			this.sessionId = header?.id ?? nanoid();
			this.sessionName = header?.title;

			if (migrateToCurrentVersion(this.fileEntries)) {
				await this._rewriteFile();
			}

			this._buildIndex();
			this.flushed = true;
		} else {
			const explicitPath = this.sessionFile;
			this._newSessionSync();
			this.sessionFile = explicitPath; // preserve explicit path from --session flag
			await this._rewriteFile();
			this.flushed = true;
			return;
		}
	}

	/** Start a new session. Closes any existing writer first. */
	async newSession(options?: NewSessionOptions): Promise<string | undefined> {
		await this._closePersistWriter();
		return this._newSessionSync(options);
	}

	/**
	 * Fork the current session, creating a new session file with the same entries.
	 * Returns both the old and new session file paths for artifact copying.
	 * @returns { oldSessionFile, newSessionFile } or undefined if not persisting
	 */
	async fork(): Promise<{ oldSessionFile: string; newSessionFile: string } | undefined> {
		if (!this.persist || !this.sessionFile) {
			return undefined;
		}

		const oldSessionFile = this.sessionFile;
		const oldSessionId = this.sessionId;

		// Close the current writer
		await this._closePersistWriter();
		this.persistChain = Promise.resolve();
		this.persistError = undefined;
		this.persistErrorReported = false;

		// Create new session ID and header
		this.sessionId = nanoid();
		const timestamp = new Date().toISOString();
		const fileTimestamp = timestamp.replace(/[:.]/g, "-");
		this.sessionFile = path.join(this.getSessionDir(), `${fileTimestamp}_${this.sessionId}.jsonl`);

		// Update the header with new ID but keep all entries
		const oldHeader = this.fileEntries.find(e => e.type === "session") as SessionHeader | undefined;
		const newHeader: SessionHeader = {
			type: "session",
			version: CURRENT_SESSION_VERSION,
			id: this.sessionId,
			title: oldHeader?.title ?? this.sessionName,
			timestamp,
			cwd: this.cwd,
			parentSession: oldSessionId,
		};
		this.sessionName = newHeader.title;

		// Replace the header in fileEntries
		const entries = this.fileEntries.filter(e => e.type !== "session") as SessionEntry[];
		this.fileEntries = [newHeader, ...entries];

		// Write the new session file
		this.flushed = false;
		await this._rewriteFile();

		return { oldSessionFile, newSessionFile: this.sessionFile };
	}

	/** Sync version for initial creation (no existing writer to close) */
	private _newSessionSync(options?: NewSessionOptions): string | undefined {
		this.persistChain = Promise.resolve();
		this.persistError = undefined;
		this.persistErrorReported = false;
		this.sessionId = nanoid();
		const timestamp = new Date().toISOString();
		const header: SessionHeader = {
			type: "session",
			version: CURRENT_SESSION_VERSION,
			id: this.sessionId,
			timestamp,
			cwd: this.cwd,
			parentSession: options?.parentSession,
		};
		this.fileEntries = [header];
		this.byId.clear();
		this.labelsById.clear();
		this.leafId = null;
		this.flushed = false;
		this.usageStatistics = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0 };

		if (this.persist) {
			const fileTimestamp = timestamp.replace(/[:.]/g, "-");
			this.sessionFile = path.join(this.getSessionDir(), `${fileTimestamp}_${this.sessionId}.jsonl`);
		}
		return this.sessionFile;
	}

	private _buildIndex(): void {
		this.byId.clear();
		this.labelsById.clear();
		this.leafId = null;
		this.usageStatistics = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0 };
		for (const entry of this.fileEntries) {
			if (entry.type === "session") continue;
			this.byId.set(entry.id, entry);
			this.leafId = entry.id;
			if (entry.type === "label") {
				if (entry.label) {
					this.labelsById.set(entry.targetId, entry.label);
				} else {
					this.labelsById.delete(entry.targetId);
				}
			}
			if (entry.type === "message" && entry.message.role === "assistant") {
				const usage = entry.message.usage;
				this.usageStatistics.input += usage.input;
				this.usageStatistics.output += usage.output;
				this.usageStatistics.cacheRead += usage.cacheRead;
				this.usageStatistics.cacheWrite += usage.cacheWrite;
				this.usageStatistics.cost += usage.cost.total;
			}

			if (entry.type === "message" && entry.message.role === "toolResult" && entry.message.toolName === "task") {
				const usage = getTaskToolUsage(entry.message.details);
				if (usage) {
					this.usageStatistics.input += usage.input;
					this.usageStatistics.output += usage.output;
					this.usageStatistics.cacheRead += usage.cacheRead;
					this.usageStatistics.cacheWrite += usage.cacheWrite;
					this.usageStatistics.cost += usage.cost.total;
				}
			}
		}
	}

	private _recordPersistError(err: unknown): Error {
		const normalized = toError(err);
		if (!this.persistError) this.persistError = normalized;
		if (!this.persistErrorReported) {
			this.persistErrorReported = true;
			logger.error("Session persistence error.", {
				sessionFile: this.sessionFile,
				error: normalized.message,
				stack: normalized.stack,
			});
		}
		return normalized;
	}

	private _queuePersistTask(task: () => Promise<void>, options?: { ignoreError?: boolean }): Promise<void> {
		const next = this.persistChain.then(async () => {
			if (this.persistError && !options?.ignoreError) throw this.persistError;
			await task();
		});
		this.persistChain = next.catch(err => {
			this._recordPersistError(err);
		});
		return next;
	}

	private _ensurePersistWriter(): NdjsonFileWriter | undefined {
		if (!this.persist || !this.sessionFile) return undefined;
		if (this.persistError) throw this.persistError;
		if (this.persistWriter && this.persistWriterPath === this.sessionFile) return this.persistWriter;
		// Note: caller must await _closePersistWriter() before calling this if switching files
		this.persistWriter = new NdjsonFileWriter(this.storage, this.sessionFile, {
			onError: err => {
				this._recordPersistError(err);
			},
		});
		this.persistWriterPath = this.sessionFile;
		return this.persistWriter;
	}

	private async _closePersistWriterInternal(): Promise<void> {
		if (this.persistWriter) {
			await this.persistWriter.close();
			this.persistWriter = undefined;
		}
		this.persistWriterPath = undefined;
	}

	private async _closePersistWriter(): Promise<void> {
		await this._queuePersistTask(
			async () => {
				await this._closePersistWriterInternal();
			},
			{ ignoreError: true },
		);
	}

	private async _writeEntriesAtomically(entries: FileEntry[]): Promise<void> {
		if (!this.sessionFile) return;
		const dir = path.resolve(this.sessionFile, "..");
		const tempPath = path.join(dir, `.${path.basename(this.sessionFile)}.${nanoid(6)}.tmp`);
		const writer = new NdjsonFileWriter(this.storage, tempPath, { flags: "w" });
		try {
			for (const entry of entries) {
				await writer.write(entry);
			}
			await writer.flush();
			await writer.fsync();
			await writer.close();
			await this.storage.rename(tempPath, this.sessionFile);
			this.storage.fsyncDirSync(dir);
		} catch (err) {
			try {
				await writer.close();
			} catch {
				// Ignore cleanup errors
			}
			try {
				await this.storage.unlink(tempPath);
			} catch {
				// Ignore cleanup errors
			}
			throw toError(err);
		}
	}

	private async _rewriteFile(): Promise<void> {
		if (!this.persist || !this.sessionFile) return;
		await this._queuePersistTask(async () => {
			await this._closePersistWriterInternal();
			const entries = await Promise.all(this.fileEntries.map(entry => prepareEntryForPersistence(entry)));
			await this._writeEntriesAtomically(entries);
			this.flushed = true;
		});
	}

	isPersisted(): boolean {
		return this.persist;
	}

	/** Flush pending writes to disk. Call before switching sessions or on shutdown. */
	async flush(): Promise<void> {
		if (!this.persistWriter) return;
		await this._queuePersistTask(async () => {
			if (this.persistWriter) {
				await this.persistWriter.flush();
				await this.persistWriter.fsync();
			}
		});
		if (this.persistError) throw this.persistError;
	}

	getCwd(): string {
		return this.cwd;
	}

	/** Get usage statistics across all assistant messages in the session. */
	getUsageStatistics(): UsageStatistics {
		return this.usageStatistics;
	}

	getSessionDir(): string {
		return this.sessionDir;
	}

	getSessionId(): string {
		return this.sessionId;
	}

	getSessionFile(): string | undefined {
		return this.sessionFile;
	}

	getSessionName(): string | undefined {
		return this.sessionName;
	}

	async setSessionName(name: string): Promise<void> {
		this.sessionName = name;

		// Update the in-memory header (so first flush includes title)
		const header = this.fileEntries.find(e => e.type === "session") as SessionHeader | undefined;
		if (header) {
			header.title = name;
		}

		// Update the session file header with the title (if already flushed)
		const sessionFile = this.sessionFile;
		if (this.persist && sessionFile && this.storage.existsSync(sessionFile)) {
			await this._rewriteFile();
		}
	}

	_persist(entry: SessionEntry): void {
		if (!this.persist || !this.sessionFile) return;
		if (this.persistError) throw this.persistError;

		const hasAssistant = this.fileEntries.some(e => e.type === "message" && e.message.role === "assistant");
		if (!hasAssistant && !this.flushed) return;

		if (!this.flushed) {
			this.flushed = true;
			void this._queuePersistTask(async () => {
				const writer = this._ensurePersistWriter();
				if (!writer) return;
				const entries = await Promise.all(this.fileEntries.map(e => prepareEntryForPersistence(e)));
				for (const persistedEntry of entries) {
					await writer.write(persistedEntry);
				}
			});
		} else {
			void this._queuePersistTask(async () => {
				const writer = this._ensurePersistWriter();
				if (!writer) return;
				const persistedEntry = await prepareEntryForPersistence(entry);
				await writer.write(persistedEntry);
			});
		}
	}

	private _appendEntry(entry: SessionEntry): void {
		this.fileEntries.push(entry);
		this.byId.set(entry.id, entry);
		this.leafId = entry.id;
		this._persist(entry);
		if (entry.type === "message" && entry.message.role === "assistant") {
			const usage = entry.message.usage;
			this.usageStatistics.input += usage.input;
			this.usageStatistics.output += usage.output;
			this.usageStatistics.cacheRead += usage.cacheRead;
			this.usageStatistics.cacheWrite += usage.cacheWrite;
			this.usageStatistics.cost += usage.cost.total;
		}

		if (entry.type === "message" && entry.message.role === "toolResult" && entry.message.toolName === "task") {
			const usage = getTaskToolUsage(entry.message.details);
			if (usage) {
				this.usageStatistics.input += usage.input;
				this.usageStatistics.output += usage.output;
				this.usageStatistics.cacheRead += usage.cacheRead;
				this.usageStatistics.cacheWrite += usage.cacheWrite;
				this.usageStatistics.cost += usage.cost.total;
			}
		}
	}

	/** Append a message as child of current leaf, then advance leaf. Returns entry id.
	 * Does not allow writing CompactionSummaryMessage and BranchSummaryMessage directly.
	 * Reason: we want these to be top-level entries in the session, not message session entries,
	 * so it is easier to find them.
	 * These need to be appended via appendCompaction() and appendBranchSummary() methods.
	 */
	appendMessage(
		message:
			| Message
			| CustomMessage
			| HookMessage
			| BashExecutionMessage
			| PythonExecutionMessage
			| FileMentionMessage,
	): string {
		const entry: SessionMessageEntry = {
			type: "message",
			id: generateId(this.byId),
			parentId: this.leafId,
			timestamp: new Date().toISOString(),
			message,
		};
		this._appendEntry(entry);
		return entry.id;
	}

	/** 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),
			parentId: this.leafId,
			timestamp: new Date().toISOString(),
			thinkingLevel,
		};
		this._appendEntry(entry);
		return entry.id;
	}

	/**
	 * Append a model change as child of current leaf, then advance leaf. Returns entry id.
	 * @param model Model in "provider/modelId" format
	 * @param role Optional role (default: "default")
	 */
	appendModelChange(model: string, role?: string): string {
		const entry: ModelChangeEntry = {
			type: "model_change",
			id: generateId(this.byId),
			parentId: this.leafId,
			timestamp: new Date().toISOString(),
			model,
			role,
		};
		this._appendEntry(entry);
		return entry.id;
	}

	/** Append session init metadata (for subagent debugging/replay). Returns entry id. */
	appendSessionInit(init: { systemPrompt: string; task: string; tools: string[]; outputSchema?: unknown }): string {
		const entry: SessionInitEntry = {
			type: "session_init",
			id: generateId(this.byId),
			parentId: this.leafId,
			timestamp: new Date().toISOString(),
			...init,
		};
		this._appendEntry(entry);
		return entry.id;
	}

	/** Append a compaction summary as child of current leaf, then advance leaf. Returns entry id. */
	appendCompaction<T = unknown>(
		summary: string,
		shortSummary: string | undefined,
		firstKeptEntryId: string,
		tokensBefore: number,
		details?: T,
		fromExtension?: boolean,
		preserveData?: Record<string, unknown>,
	): string {
		const entry: CompactionEntry<T> = {
			type: "compaction",
			id: generateId(this.byId),
			parentId: this.leafId,
			timestamp: new Date().toISOString(),
			summary,
			shortSummary,
			firstKeptEntryId,
			tokensBefore,
			details,
			fromExtension,
			preserveData,
		};
		this._appendEntry(entry);
		return entry.id;
	}

	/** Append a custom entry (for extensions) 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,
			timestamp: new Date().toISOString(),
		};
		this._appendEntry(entry);
		return entry.id;
	}

	/**
	 * Rewrite the session file after in-place entry updates.
	 * Use sparingly (e.g., pruning old tool outputs).
	 */
	async rewriteEntries(): Promise<void> {
		if (!this.persist || !this.sessionFile) return;
		await this._rewriteFile();
	}

	/**
	 * Rewrite tool call arguments in the most recent assistant message containing the toolCallId.
	 * Returns true if a tool call was updated.
	 */
	async rewriteAssistantToolCallArgs(toolCallId: string, args: Record<string, unknown>): Promise<boolean> {
		let updated = false;
		for (let i = this.fileEntries.length - 1; i >= 0; i--) {
			const entry = this.fileEntries[i];
			if (entry.type !== "message" || entry.message.role !== "assistant") continue;
			const message = entry.message as { content?: unknown };
			if (!Array.isArray(message.content)) continue;
			for (const block of message.content) {
				if (typeof block !== "object" || block === null) continue;
				if (!("type" in block) || (block as { type?: string }).type !== "toolCall") continue;
				const toolCall = block as { id?: string; arguments?: Record<string, unknown> };
				if (toolCall.id === toolCallId) {
					toolCall.arguments = args;
					updated = true;
					break;
				}
			}
			if (updated) break;
		}

		if (updated && this.persist && this.sessionFile) {
			await this._rewriteFile();
		}
		return updated;
	}

	/**
	 * Append a custom message entry (for extensions) that participates in LLM context.
	 * @param customType Hook identifier for filtering on reload
	 * @param content Message content (string or TextContent/ImageContent array)
	 * @param display Whether to show in TUI (true = styled display, false = hidden)
	 * @param details Optional extension-specific metadata (not sent to LLM)
	 * @returns Entry id
	 */
	appendCustomMessageEntry<T = unknown>(
		customType: string,
		content: string | (TextContent | ImageContent)[],
		display: boolean,
		details?: T,
	): string {
		const entry: CustomMessageEntry<T> = {
			type: "custom_message",
			customType,
			content,
			display,
			details,
			id: generateId(this.byId),
			parentId: this.leafId,
			timestamp: new Date().toISOString(),
		};
		this._appendEntry(entry);
		return entry.id;
	}

	// =========================================================================
	// TTSR (Time Traveling Stream Rules)
	// =========================================================================

	/**
	 * Append a TTSR injection entry recording which rules were injected.
	 * @param ruleNames Names of rules that were injected
	 * @returns Entry id
	 */
	appendTtsrInjection(ruleNames: string[]): string {
		const entry: TtsrInjectionEntry = {
			type: "ttsr_injection",
			id: generateId(this.byId),
			parentId: this.leafId,
			timestamp: new Date().toISOString(),
			injectedRules: ruleNames,
		};
		this._appendEntry(entry);
		return entry.id;
	}

	/**
	 * Get all unique TTSR rule names that have been injected in the current branch.
	 * Scans from root to current leaf for ttsr_injection entries.
	 */
	getInjectedTtsrRules(): string[] {
		const path = this.getBranch();
		const ruleNames = new Set<string>();
		for (const entry of path) {
			if (entry.type === "ttsr_injection") {
				for (const name of entry.injectedRules) {
					ruleNames.add(name);
				}
			}
		}
		return Array.from(ruleNames);
	}

	// =========================================================================
	// Tree Traversal
	// =========================================================================

	getLeafId(): string | null {
		return this.leafId;
	}

	getLeafEntry(): SessionEntry | undefined {
		return this.leafId ? this.byId.get(this.leafId) : undefined;
	}

	/**
	 * Get the most recent model role from the current session path.
	 * Returns undefined if no model change has been recorded.
	 */
	getLastModelChangeRole(): string | undefined {
		let current = this.getLeafEntry();
		while (current) {
			if (current.type === "model_change") {
				return current.role ?? "default";
			}
			current = current.parentId ? this.byId.get(current.parentId) : undefined;
		}
		return undefined;
	}

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

	/**
	 * Get all direct children of an entry.
	 */
	getChildren(parentId: string): SessionEntry[] {
		const children: SessionEntry[] = [];
		for (const entry of this.byId.values()) {
			if (entry.parentId === parentId) {
				children.push(entry);
			}
		}
		return children;
	}

	/**
	 * Get the label for an entry, if any.
	 */
	getLabel(id: string): string | undefined {
		return this.labelsById.get(id);
	}

	/**
	 * Set or clear a label on an entry.
	 * Labels are user-defined markers for bookmarking/navigation.
	 * Pass undefined or empty string to clear the label.
	 */
	appendLabelChange(targetId: string, label: string | undefined): string {
		if (!this.byId.has(targetId)) {
			throw new Error(`Entry ${targetId} not found`);
		}
		const entry: LabelEntry = {
			type: "label",
			id: generateId(this.byId),
			parentId: this.leafId,
			timestamp: new Date().toISOString(),
			targetId,
			label,
		};
		this._appendEntry(entry);
		if (label) {
			this.labelsById.set(targetId, label);
		} else {
			this.labelsById.delete(targetId);
		}
		return entry.id;
	}

	/**
	 * 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.
	 */
	getBranch(fromId?: string): SessionEntry[] {
		const path: SessionEntry[] = [];
		const startId = fromId ?? this.leafId;
		let current = startId ? this.byId.get(startId) : undefined;
		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).
	 * Uses tree traversal from current leaf.
	 */
	buildSessionContext(): SessionContext {
		return buildSessionContext(this.getEntries(), this.leafId, this.byId);
	}

	/**
	 * Get session header.
	 */
	getHeader(): SessionHeader | null {
		const h = this.fileEntries.find(e => e.type === "session");
		return h ? (h as SessionHeader) : null;
	}

	/**
	 * Get all session entries (excludes header). Returns a shallow copy.
	 * 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.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 with resolved labels
		for (const entry of entries) {
			const label = this.labelsById.get(entry.id);
			nodeMap.set(entry.id, { entry, children: [], label });
		}

		// Build tree
		for (const entry of entries) {
			const node = nodeMap.get(entry.id)!;
			if (entry.parentId === null || entry.parentId === entry.id) {
				roots.push(node);
			} else {
				const parent = nodeMap.get(entry.parentId);
				if (parent) {
					parent.children.push(node);
				} else {
					// Orphan - treat as root
					roots.push(node);
				}
			}
		}

		// Sort children by timestamp (oldest first, newest at bottom)
		// Use iterative approach to avoid stack overflow on deep trees
		const stack: SessionTreeNode[] = [...roots];
		while (stack.length > 0) {
			const node = stack.pop()!;
			node.children.sort((a, b) => new Date(a.entry.timestamp).getTime() - new Date(b.entry.timestamp).getTime());
			stack.push(...node.children);
		}

		return roots;
	}

	// =========================================================================
	// Branching
	// =========================================================================

	/**
	 * 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;
	}

	/**
	 * Reset the leaf pointer to null (before any entries).
	 * The next appendXXX() call will create a new root entry (parentId = null).
	 * Use this when navigating to re-edit the first user message.
	 */
	resetLeaf(): void {
		this.leafId = null;
	}

	/**
	 * 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 | null, summary: string, details?: unknown, fromExtension?: boolean): string {
		if (branchFromId !== null && !this.byId.has(branchFromId)) {
			throw new Error(`Entry ${branchFromId} not found`);
		}
		this.leafId = branchFromId;
		const entry: BranchSummaryEntry = {
			type: "branch_summary",
			id: generateId(this.byId),
			parentId: branchFromId,
			timestamp: new Date().toISOString(),
			fromId: branchFromId ?? "root",
			summary,
			details,
			fromExtension,
		};
		this._appendEntry(entry);
		return entry.id;
	}

	/**
	 * Create a new session file containing only the path from root to the specified leaf.
	 * Useful for extracting a single conversation path from a branched session.
	 * Returns the new session file path, or undefined if not persisting.
	 */
	createBranchedSession(leafId: string): string | undefined {
		const branchPath = this.getBranch(leafId);
		if (branchPath.length === 0) {
			throw new Error(`Entry ${leafId} not found`);
		}

		// Filter out LabelEntry from path - we'll recreate them from the resolved map
		const pathWithoutLabels = branchPath.filter(e => e.type !== "label");

		const newSessionId = nanoid();
		const timestamp = new Date().toISOString();
		const fileTimestamp = timestamp.replace(/[:.]/g, "-");
		const newSessionFile = path.join(this.getSessionDir(), `${fileTimestamp}_${newSessionId}.jsonl`);

		const header: SessionHeader = {
			type: "session",
			version: CURRENT_SESSION_VERSION,
			id: newSessionId,
			timestamp,
			cwd: this.cwd,
			parentSession: this.persist ? this.sessionFile : undefined,
		};

		// Collect labels for entries in the path
		const pathEntryIds = new Set(pathWithoutLabels.map(e => e.id));
		const labelsToWrite: Array<{ targetId: string; label: string }> = [];
		for (const [targetId, label] of this.labelsById) {
			if (pathEntryIds.has(targetId)) {
				labelsToWrite.push({ targetId, label });
			}
		}

		if (this.persist) {
			const lines: string[] = [];
			lines.push(JSON.stringify(header));
			for (const entry of pathWithoutLabels) {
				lines.push(JSON.stringify(entry));
			}
			// Write fresh label entries at the end
			const lastEntryId = pathWithoutLabels[pathWithoutLabels.length - 1]?.id || null;
			let parentId = lastEntryId;
			const labelEntries: LabelEntry[] = [];
			for (const { targetId, label } of labelsToWrite) {
				const labelEntry: LabelEntry = {
					type: "label",
					id: generateId(new Set(pathEntryIds)),
					parentId,
					timestamp: new Date().toISOString(),
					targetId,
					label,
				};
				lines.push(JSON.stringify(labelEntry));
				pathEntryIds.add(labelEntry.id);
				labelEntries.push(labelEntry);
				parentId = labelEntry.id;
			}
			this.storage.writeTextSync(newSessionFile, `${lines.join("\n")}\n`);
			this.fileEntries = [header, ...pathWithoutLabels, ...labelEntries];
			this.sessionId = newSessionId;
			this._buildIndex();
			return newSessionFile;
		}

		// In-memory mode: replace current session with the path + labels
		const labelEntries: LabelEntry[] = [];
		let parentId = pathWithoutLabels[pathWithoutLabels.length - 1]?.id || null;
		for (const { targetId, label } of labelsToWrite) {
			const labelEntry: LabelEntry = {
				type: "label",
				id: generateId(new Set([...pathEntryIds, ...labelEntries.map(e => e.id)])),
				parentId,
				timestamp: new Date().toISOString(),
				targetId,
				label,
			};
			labelEntries.push(labelEntry);
			parentId = labelEntry.id;
		}
		this.fileEntries = [header, ...pathWithoutLabels, ...labelEntries];
		this.sessionId = newSessionId;
		this._buildIndex();
		return undefined;
	}

	/**
	 * Create a new session.
	 * @param cwd Working directory (stored in session header)
	 * @param sessionDir Optional session directory. If omitted, uses default (~/.omp/agent/sessions/<encoded-cwd>/).
	 */
	static create(cwd: string, sessionDir?: string, storage: SessionStorage = new FileSessionStorage()): SessionManager {
		const dir = sessionDir ?? getDefaultSessionDir(cwd, storage);
		const manager = new SessionManager(cwd, dir, true, storage);
		manager._initNewSession();
		return manager;
	}

	/**
	 * Fork a session into the current project directory.
	 * Copies history from another session file while creating a new session file in the current sessionDir.
	 */
	static async forkFrom(
		sourcePath: string,
		cwd: string,
		sessionDir?: string,
		storage: SessionStorage = new FileSessionStorage(),
	): Promise<SessionManager> {
		const dir = sessionDir ?? getDefaultSessionDir(cwd, storage);
		const manager = new SessionManager(cwd, dir, true, storage);
		const forkEntries = structuredClone(await loadEntriesFromFile(sourcePath, storage)) as FileEntry[];
		migrateToCurrentVersion(forkEntries);
		const sourceHeader = forkEntries.find(e => e.type === "session") as SessionHeader | undefined;
		const historyEntries = forkEntries.filter(entry => entry.type !== "session") as SessionEntry[];
		manager._newSessionSync({ parentSession: sourceHeader?.id });
		const newHeader = manager.fileEntries[0] as SessionHeader;
		newHeader.title = sourceHeader?.title;
		manager.fileEntries = [newHeader, ...historyEntries];
		manager.sessionName = newHeader.title;
		manager._buildIndex();
		await manager._rewriteFile();
		return manager;
	}

	/**
	 * Open a specific session file.
	 * @param path Path to session file
	 * @param sessionDir Optional session directory for /new or /branch. If omitted, derives from file's parent.
	 */
	static async open(
		filePath: string,
		sessionDir?: string,
		storage: SessionStorage = new FileSessionStorage(),
	): Promise<SessionManager> {
		// Extract cwd from session header if possible, otherwise use process.cwd()
		const entries = await loadEntriesFromFile(filePath, storage);
		const header = entries.find(e => e.type === "session") as SessionHeader | undefined;
		const cwd = header?.cwd ?? process.cwd();
		// If no sessionDir provided, derive from file's parent directory
		const dir = sessionDir ?? path.resolve(filePath, "..");
		const manager = new SessionManager(cwd, dir, true, storage);
		await manager._initSessionFile(filePath);
		return manager;
	}

	/**
	 * Continue the most recent session, or create new if none.
	 * @param cwd Working directory
	 * @param sessionDir Optional session directory. If omitted, uses default (~/.omp/agent/sessions/<encoded-cwd>/).
	 */
	static async continueRecent(
		cwd: string,
		sessionDir?: string,
		storage: SessionStorage = new FileSessionStorage(),
	): Promise<SessionManager> {
		const dir = sessionDir ?? getDefaultSessionDir(cwd, storage);
		const mostRecent = await findMostRecentSession(dir, storage);
		const manager = new SessionManager(cwd, dir, true, storage);
		if (mostRecent) {
			await manager._initSessionFile(mostRecent);
		} else {
			manager._initNewSession();
		}
		return manager;
	}

	/** Create an in-memory session (no file persistence) */
	static inMemory(cwd: string = process.cwd(), storage: SessionStorage = new MemorySessionStorage()): SessionManager {
		const manager = new SessionManager(cwd, "", false, storage);
		manager._initNewSession();
		return manager;
	}

	/**
	 * List all sessions.
	 * @param cwd Working directory (used to compute default session directory)
	 * @param sessionDir Optional session directory. If omitted, uses default (~/.omp/agent/sessions/<encoded-cwd>/).
	 */
	static async list(
		cwd: string,
		sessionDir?: string,
		storage: SessionStorage = new FileSessionStorage(),
	): Promise<SessionInfo[]> {
		const dir = sessionDir ?? getDefaultSessionDir(cwd, storage);
		try {
			const files = storage.listFilesSync(dir, "*.jsonl");
			return await collectSessionsFromFiles(files, storage);
		} catch {
			return [];
		}
	}

	/**
	 * List all sessions across all project directories.
	 */
	static async listAll(storage: SessionStorage = new FileSessionStorage()): Promise<SessionInfo[]> {
		const sessionsRoot = path.join(getDefaultAgentDir(), "sessions");
		try {
			const files = Array.from(new Bun.Glob("**/*.jsonl").scanSync(sessionsRoot)).map(name =>
				path.join(sessionsRoot, name),
			);
			return await collectSessionsFromFiles(files, storage);
		} catch {
			return [];
		}
	}
}