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

import * as fs from "node:fs";
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, parseJsonlLenient, Snowflake } from "@oh-my-pi/pi-utils";
import { getBlobsDir, getAgentDir as getDefaultAgentDir, getProjectDir } from "@oh-my-pi/pi-utils/dirs";
import { type BlobPutResult, BlobStore, externalizeImageData, isBlobRef, resolveImageData } from "./blob-store";
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;
}

/** Mode change entry - tracks agent mode transitions (e.g. plan mode). */
export interface ModeChangeEntry extends SessionEntryBase {
	type: "mode_change";
	/** Current mode name, or "none" when exiting a mode */
	mode: string;
	/** Optional mode-specific data (e.g. plan file path) */
	data?: Record<string, 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
	| ModeChangeEntry;

/** 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[];
	/** Active mode (e.g. "plan") or "none" if no special mode is active */
	mode: string;
	/** Mode-specific data from the last mode_change entry */
	modeData?: Record<string, unknown>;
}

export interface SessionInfo {
	path: string;
	id: string;
	/** Working directory where the session was started. Empty string for old sessions. */
	cwd: string;
	title?: string;
	/** Path to the parent session (if this session was forked). */
	parentSessionPath?: 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"
	| "putBlob"
>;

/** 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 = crypto.randomUUID().slice(-8);
		if (!byId.has(id)) return id;
	}
	return Snowflake.next(); // fallback to full snowflake id
}

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

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

/** Exported for compaction.test.ts */
export function parseSessionEntries(content: string): FileEntry[] {
	return parseJsonlLenient<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: [], mode: "none" };
	}
	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: [], mode: "none" };
	}

	// 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>();
	let mode = "none";
	let modeData: Record<string, unknown> | undefined;

	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);
			}
		} else if (entry.type === "mode_change") {
			mode = entry.mode;
			modeData = entry.data;
		}
	}

	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, mode, modeData };
}

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

// =============================================================================
// Terminal breadcrumbs: maps terminal (TTY) -> last session file for --continue
// =============================================================================

const TERMINAL_SESSIONS_DIR = "terminal-sessions";

/**
 * Get a stable identifier for the current terminal.
 * Uses the TTY device path (e.g., /dev/pts/3), falling back to environment
 * variables for terminal multiplexers or terminal emulators.
 * Returns null if no terminal can be identified (e.g., piped input).
 */
function getTerminalId(): string | null {
	// TTY device path — most reliable, unique per terminal tab
	if (process.stdin.isTTY) {
		try {
			// On Linux/macOS, /proc/self/fd/0 -> /dev/pts/N
			const ttyPath = fs.readlinkSync("/proc/self/fd/0");
			if (ttyPath.startsWith("/dev/")) {
				return ttyPath.slice(5).replace(/\//g, "-"); // /dev/pts/3 -> pts-3
			}
		} catch {}
	}

	// Fallback to terminal-specific env vars
	const kittyId = process.env.KITTY_WINDOW_ID;
	if (kittyId) return `kitty-${kittyId}`;

	const tmuxPane = process.env.TMUX_PANE;
	if (tmuxPane) return `tmux-${tmuxPane}`;

	const terminalSessionId = process.env.TERM_SESSION_ID; // macOS Terminal.app
	if (terminalSessionId) return `apple-${terminalSessionId}`;

	const wtSession = process.env.WT_SESSION; // Windows Terminal
	if (wtSession) return `wt-${wtSession}`;

	return null;
}

/**
 * Write a breadcrumb linking the current terminal to a session file.
 * The breadcrumb contains the cwd and session path so --continue can
 * find "this terminal's last session" even when running concurrent instances.
 */
function writeTerminalBreadcrumb(cwd: string, sessionFile: string): void {
	const terminalId = getTerminalId();
	if (!terminalId) return;

	try {
		const breadcrumbDir = path.join(getDefaultAgentDir(), TERMINAL_SESSIONS_DIR);
		const breadcrumbFile = path.join(breadcrumbDir, terminalId);
		const content = `${cwd}\n${sessionFile}\n`;
		// Bun.write auto-creates parent dirs
		void Bun.write(breadcrumbFile, content);
	} catch {
		// Best-effort — don't break session creation if breadcrumb fails
	}
}

/**
 * Read the terminal breadcrumb for the current terminal, scoped to a cwd.
 * Returns the session file path if it exists and matches the cwd, null otherwise.
 */
async function readTerminalBreadcrumb(cwd: string): Promise<string | null> {
	const terminalId = getTerminalId();
	if (!terminalId) return null;

	try {
		const breadcrumbFile = path.join(getDefaultAgentDir(), TERMINAL_SESSIONS_DIR, terminalId);
		const content = await Bun.file(breadcrumbFile).text();
		const lines = content.trim().split("\n");
		if (lines.length < 2) return null;

		const breadcrumbCwd = lines[0];
		const sessionFile = lines[1];

		// Only return if cwd matches (user might have cd'd)
		if (path.resolve(breadcrumbCwd) !== path.resolve(cwd)) return null;

		// Verify the session file still exists
		const stat = fs.statSync(sessionFile, { throwIfNoEntry: false });
		if (stat?.isFile()) return sessionFile;
	} catch (err) {
		if (!isEnoent(err)) logger.debug("Terminal breadcrumb read failed", { err });
		// Breadcrumb doesn't exist or is corrupt — fall through
	}
	return null;
}

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

/**
 * Resolve blob references in loaded entries, replacing `blob:sha256:<hash>` data fields
 * with the actual base64 content from the blob store. Mutates entries in place.
 */
async function resolveBlobRefsInEntries(entries: FileEntry[], blobStore: BlobStore): Promise<void> {
	const promises: Promise<void>[] = [];

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

		// Resolve image blocks in message content arrays
		let contentArray: unknown[] | undefined;
		if (entry.type === "message") {
			const content = (entry.message as { content?: unknown }).content;
			if (Array.isArray(content)) contentArray = content;
		} else if (entry.type === "custom_message" && Array.isArray(entry.content)) {
			contentArray = entry.content;
		}

		if (!contentArray) continue;

		for (const block of contentArray) {
			if (isImageBlock(block) && isBlobRef(block.data)) {
				promises.push(
					resolveImageData(blobStore, block.data).then(resolved => {
						(block as { data: string }).data = resolved;
					}),
				);
			}
		}
	}

	await Promise.all(promises);
}

/**
 * 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 {
	#fullName: string | undefined;
	#name: string | undefined;
	#timeAgo: string | undefined;

	constructor(
		readonly path: string,
		readonly mtime: number,
		header: Record<string, unknown>,
		firstPrompt?: string,
	) {
		// 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, 39)}…`;
		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 sessions: RecentSessionInfo[] = [];
		await Promise.all(
			files.map(async (path: string) => {
				try {
					const content = await storage.readTextPrefix(path, 4096);
					const entries = parseJsonlLenient<Record<string, unknown>>(content);
					if (entries.length === 0) return;
					const header = entries[0] as Record<string, unknown>;
					if (header.type !== "session" || typeof header.id !== "string") return;
					const mtime = storage.statSync(path).mtimeMs;
					const firstPrompt = header.title ? undefined : extractFirstUserPrompt(entries);
					sessions.push(new RecentSessionInfo(path, mtime, header, firstPrompt));
				} catch {}
			}),
		);
		return sessions.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]";
/** Minimum base64 length to externalize to blob store (skip tiny inline images) */
const BLOB_EXTERNALIZE_THRESHOLD = 1024;
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 truncateForPersistence<T>(obj: T, blobStore: BlobStore, 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 (!isBlobRef(item.data) && item.data.length >= BLOB_EXTERNALIZE_THRESHOLD) {
						changed = true;
						const blobRef = await externalizeImageData(blobStore, item.data);
						return { ...item, data: blobRef };
					}
				}
				const newItem = await truncateForPersistence(item, blobStore, 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, blobStore, 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, blobStore: BlobStore): Promise<FileEntry> {
	return truncateForPersistence(entry, blobStore);
}

class NdjsonFileWriter {
	#writer: SessionStorageWriter;
	#closed = false;
	#closing = false;
	#error: Error | undefined;
	#pendingWrites: Promise<void> = Promise.resolve();
	#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),
		});
	}

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

	#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;
	}

	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[] = [];

	// Collect session info for all files in parallel
	await Promise.all(
		files.map(async file => {
			try {
				const content = await storage.readText(file);
				const entries = parseJsonlLenient<Record<string, unknown>>(content);
				if (entries.length === 0) return;

				// 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) return;

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

				if (messageCount) {
					const stats = storage.statSync(file);
					sessions.push({
						path: file,
						id: header.id,
						cwd: typeof header.cwd === "string" ? header.cwd : "",
						title: header.title ?? shortSummary,
						parentSessionPath: (header as SessionHeader).parentSession,
						created: new Date(header.timestamp),
						modified: stats.mtime,
						messageCount,
						firstMessage: firstMessage || "(no messages)",
						allMessagesText: allMessages.join(" "),
					});
				}
			} catch {}
		}),
	);

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

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

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

	/** Puts a binary blob into the blob store and returns the blob reference */
	async putBlob(data: Buffer): Promise<BlobPutResult> {
		return this.#blobStore.put(data);
	}

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

	/** Initialize with a new session (used by factory methods) */
	#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);
		writeTerminalBreadcrumb(this.cwd, this.#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 ?? Snowflake.next();
			this.#sessionName = header?.title;

			if (migrateToCurrentVersion(this.#fileEntries)) {
				await this.#rewriteFile();
			}

			await resolveBlobRefsInEntries(this.#fileEntries, this.#blobStore);

			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 = Snowflake.next();
		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 };
	}

	/**
	 * Move the session to a new working directory.
	 * Moves session files and artifacts on disk, updates all internal references,
	 * and rewrites the session header with the new cwd.
	 */
	async moveTo(newCwd: string): Promise<void> {
		const resolvedCwd = path.resolve(newCwd);
		if (resolvedCwd === this.cwd) return;

		const newSessionDir = getDefaultSessionDir(resolvedCwd, this.storage);

		if (this.persist && this.#sessionFile) {
			// Close the persist writer before moving files
			await this.#closePersistWriter();
			this.#persistChain = Promise.resolve();
			this.#persistError = undefined;
			this.#persistErrorReported = false;

			const oldSessionFile = this.#sessionFile;
			const newSessionFile = path.join(newSessionDir, path.basename(oldSessionFile));
			const oldArtifactDir = oldSessionFile.slice(0, -6); // strip .jsonl
			const newArtifactDir = newSessionFile.slice(0, -6);
			let movedSessionFile = false;
			let movedArtifactDir = false;

			try {
				await fs.promises.rename(oldSessionFile, newSessionFile);
				movedSessionFile = true;

				try {
					const stat = await fs.promises.stat(oldArtifactDir);
					if (stat.isDirectory()) {
						await fs.promises.rename(oldArtifactDir, newArtifactDir);
						movedArtifactDir = true;
					}
				} catch (err) {
					if (!isEnoent(err)) throw err;
				}
			} catch (err) {
				if (movedArtifactDir) {
					try {
						await fs.promises.rename(newArtifactDir, oldArtifactDir);
					} catch (rollbackErr) {
						throw new Error(
							`Failed to move artifacts and rollback: ${rollbackErr instanceof Error ? rollbackErr.message : String(rollbackErr)}`,
						);
					}
				}
				if (movedSessionFile) {
					try {
						await fs.promises.rename(newSessionFile, oldSessionFile);
					} catch (rollbackErr) {
						throw new Error(
							`Failed to move session file and rollback: ${rollbackErr instanceof Error ? rollbackErr.message : String(rollbackErr)}`,
						);
					}
				}
				throw err;
			}
			this.#sessionFile = newSessionFile;
		}

		// Update cwd and sessionDir (controlled mutation of readonly fields)
		(this as unknown as { cwd: string }).cwd = resolvedCwd;
		(this as unknown as { sessionDir: string }).sessionDir = newSessionDir;

		// Update the session header in fileEntries
		const header = this.#fileEntries.find(e => e.type === "session") as SessionHeader | undefined;
		if (header) {
			header.cwd = resolvedCwd;
		}

		// Rewrite the session file at its new location with updated header
		if (this.persist && this.#sessionFile) {
			await this.#rewriteFile();
		}

		// Update terminal breadcrumb
		if (this.#sessionFile) {
			writeTerminalBreadcrumb(resolvedCwd, this.#sessionFile);
		}
	}

	/** Sync version for initial creation (no existing writer to close) */
	#newSessionSync(options?: NewSessionOptions): string | undefined {
		this.#persistChain = Promise.resolve();
		this.#persistError = undefined;
		this.#persistErrorReported = false;
		this.#sessionId = Snowflake.next();
		this.#sessionName = undefined;
		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`);
			writeTerminalBreadcrumb(this.cwd, this.#sessionFile);
		}
		return this.#sessionFile;
	}

	#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;
				}
			}
		}
	}

	#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;
	}

	#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;
	}

	#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;
	}

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

	async #closePersistWriter(): Promise<void> {
		await this.#queuePersistTask(
			async () => {
				await this.#closePersistWriterInternal();
			},
			{ ignoreError: true },
		);
	}

	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)}.${Snowflake.next()}.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);
		} catch (err) {
			try {
				await writer.close();
			} catch {
				// Ignore cleanup errors
			}
			try {
				await this.storage.unlink(tempPath);
			} catch {
				// Ignore cleanup errors
			}
			throw toError(err);
		}
	}

	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, this.#blobStore)),
			);
			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) {
			// Mark as not flushed so when assistant arrives, all entries get written
			this.#flushed = false;
			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, this.#blobStore)),
				);
				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, this.#blobStore);
				await writer.write(persistedEntry);
			});
		}
	}

	#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 mode change as child of current leaf, then advance leaf. Returns entry id. */
	appendModeChange(mode: string, data?: Record<string, unknown>): string {
		const entry: ModeChangeEntry = {
			type: "mode_change",
			id: generateId(this.#byId),
			parentId: this.#leafId,
			timestamp: new Date().toISOString(),
			mode,
			data,
		};
		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 previousSessionFile = this.#sessionFile;
		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 = Snowflake.next();
		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 ? previousSessionFile : 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.#sessionFile = newSessionFile;
			this.#flushed = true;
			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);
		await resolveBlobRefsInEntries(forkEntries, manager.#blobStore);
		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 getProjectDir()
		const entries = await loadEntriesFromFile(filePath, storage);
		const header = entries.find(e => e.type === "session") as SessionHeader | undefined;
		const cwd = header?.cwd ?? getProjectDir();
		// 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);
		// Prefer terminal-scoped breadcrumb (handles concurrent sessions correctly)
		const terminalSession = await readTerminalBreadcrumb(cwd);
		const mostRecent = terminalSession ?? (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 = getProjectDir(),
		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 [];
		}
	}
}