import { basename, join, resolve } 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 { nanoid } from "nanoid"; import { getAgentDir as getDefaultAgentDir } from "../config"; import { resizeImage } from "../utils/image-resize"; import { logger } from "./logger"; import { type BashExecutionMessage, type CustomMessage, createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage, type FileMentionMessage, type HookMessage, } 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 extends SessionEntryBase { type: "compaction"; summary: string; firstKeptEntryId: string; tokensBefore: number; /** Extension-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ details?: T; /** True if generated by an extension, undefined/false if pi-generated (backward compatible) */ fromExtension?: boolean; } export interface BranchSummaryEntry 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 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[]; } /** * 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 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; /** 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; /** 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(); 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[] { const entries: FileEntry[] = []; const lines = content.trim().split("\n"); for (const line of lines) { if (!line.trim()) continue; try { const entry = JSON.parse(line) as FileEntry; entries.push(entry); } catch { // Skip malformed lines } } return entries; } 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, ): SessionContext { // Build uuid index if not available if (!byId) { byId = new Map(); 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 = {}; let compaction: CompactionEntry | null = null; const injectedTtsrRulesSet = new Set(); 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)); // 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 = join(getDefaultAgentDir(), "sessions", safePath); storage.ensureDirSync(sessionDir); return sessionDir; } /** Exported for testing */ export function loadEntriesFromFile(filePath: string, storage: SessionStorage = new FileSessionStorage()): FileEntry[] { if (!storage.existsSync(filePath)) return []; const content = storage.readTextSync(filePath); const entries: FileEntry[] = []; const lines = content.trim().split("\n"); for (const line of lines) { if (!line.trim()) continue; try { const entry = JSON.parse(line) as FileEntry; entries.push(entry); } catch { // Skip malformed lines } } // Validate session header if (entries.length === 0) return entries; const header = entries[0]; if (header.type !== "session" || typeof (header as any).id !== "string") { return []; } return entries; } /** * Lightweight metadata for a session file, used in session picker UI. * Uses lazy getters to defer string formatting until actually displayed. */ class RecentSessionInfo { readonly path: string; readonly mtime: number; #fullName: string | undefined; #name: string | undefined; #timeAgo: string | undefined; constructor(path: string, mtime: number, header: Record) { this.path = path; this.mtime = mtime; // Extract title from session header, falling back to id if title is missing const trystr = (v: unknown) => (typeof v === "string" ? v : undefined); this.#fullName = trystr(header.title) ?? 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; } } /** * 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 512 bytes of each file * to extract the JSON header without loading entire session logs into memory. */ function getSortedSessions(sessionDir: string, storage: SessionStorage): RecentSessionInfo[] { try { const buf = Buffer.alloc(512); const files: string[] = storage.listFilesSync(sessionDir, "*.jsonl"); return files .map((path: string) => { try { const length = storage.readTextPrefixSync(path, buf); const content = buf.toString("utf-8", 0, length); const firstLine = content.split("\n")[0]; if (!firstLine || !firstLine.trim()) return null; const header = JSON.parse(firstLine) as Record; if (header.type !== "session" || typeof header.id !== "string") return null; const mtime = storage.statSync(path).mtimeMs; return new RecentSessionInfo(path, mtime, header); } catch { return null; } }) .filter((item): item is RecentSessionInfo => item !== null) .sort((a, b) => b.mtime - a.mtime); } catch { return []; } } /** Exported for testing */ export function findMostRecentSession( sessionDir: string, storage: SessionStorage = new FileSessionStorage(), ): string | null { const sessions = 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 { 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(obj: T, key?: string): Promise { 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 = {}; for (const [k, v] of Object.entries(obj as Record)) { // 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 { return truncateForPersistence(entry); } class NdjsonFileWriter { private writer: SessionStorageWriter; private closed = false; private closing = false; private error: Error | undefined; private pendingWrites: Promise = 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): Promise { 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 { 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 { 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 { 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 { 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 { 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 function getRecentSessions( sessionDir: string, limit = 3, storage: SessionStorage = new FileSessionStorage(), ): RecentSessionInfo[] { return getSortedSessions(sessionDir, storage).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; const usage = record.usage; if (!usage || typeof usage !== "object") return undefined; return usage as Usage; } export class SessionManager { private sessionId: string = ""; private sessionTitle: 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 = new Map(); private labelsById: Map = 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 = 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 { 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 { await this._closePersistWriter(); this.persistError = undefined; this.persistErrorReported = false; this.sessionFile = resolve(sessionFile); if (this.storage.existsSync(this.sessionFile)) { this.fileEntries = loadEntriesFromFile(this.sessionFile!, this.storage); const header = this.fileEntries.find((e) => e.type === "session") as SessionHeader | undefined; this.sessionId = header?.id ?? nanoid(); this.sessionTitle = header?.title; if (migrateToCurrentVersion(this.fileEntries)) { await this._rewriteFile(); } this._buildIndex(); this.flushed = true; } else { this._newSessionSync(); } } /** Start a new session. Closes any existing writer first. */ async newSession(options?: NewSessionOptions): Promise { await this._closePersistWriter(); return this._newSessionSync(options); } /** 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.leafId = null; this.flushed = false; this.usageStatistics = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0 }; // Only generate filename if persisting and not already set (e.g., via --session flag) if (this.persist && !this.sessionFile) { const fileTimestamp = timestamp.replace(/[:.]/g, "-"); this.sessionFile = 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, options?: { ignoreError?: boolean }): Promise { 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 { if (this.persistWriter) { await this.persistWriter.close(); this.persistWriter = undefined; } this.persistWriterPath = undefined; } private async _closePersistWriter(): Promise { await this._queuePersistTask( async () => { await this._closePersistWriterInternal(); }, { ignoreError: true }, ); } private async _writeEntriesAtomically(entries: FileEntry[]): Promise { if (!this.sessionFile) return; const dir = resolve(this.sessionFile, ".."); const tempPath = join(dir, `.${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 { 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 { 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; } getSessionTitle(): string | undefined { return this.sessionTitle; } async setSessionTitle(title: string): Promise { this.sessionTitle = title; // 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 = title; } // 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) 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 | 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 a compaction summary as child of current leaf, then advance leaf. Returns entry id. */ appendCompaction( summary: string, firstKeptEntryId: string, tokensBefore: number, details?: T, fromExtension?: boolean, ): string { const entry: CompactionEntry = { type: "compaction", id: generateId(this.byId), parentId: this.leafId, timestamp: new Date().toISOString(), summary, firstKeptEntryId, tokensBefore, details, fromExtension, }; 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; } /** * 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( customType: string, content: string | (TextContent | ImageContent)[], display: boolean, details?: T, ): string { const entry: CustomMessageEntry = { 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(); 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(); 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 path = this.getBranch(leafId); if (path.length === 0) { throw new Error(`Entry ${leafId} not found`); } // Filter out LabelEntry from path - we'll recreate them from the resolved map const pathWithoutLabels = path.filter((e) => e.type !== "label"); const newSessionId = nanoid(); const timestamp = new Date().toISOString(); const fileTimestamp = timestamp.replace(/[:.]/g, "-"); const newSessionFile = join(this.getSessionDir(), `${fileTimestamp}_${newSessionId}.jsonl`); const header: SessionHeader = { type: "session", version: CURRENT_SESSION_VERSION, id: newSessionId, timestamp, cwd: this.cwd, 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//). */ 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; } /** * 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( path: string, sessionDir?: string, storage: SessionStorage = new FileSessionStorage(), ): Promise { // Extract cwd from session header if possible, otherwise use process.cwd() const entries = loadEntriesFromFile(path, 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 ?? resolve(path, ".."); const manager = new SessionManager(cwd, dir, true, storage); await manager._initSessionFile(path); 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//). */ static async continueRecent( cwd: string, sessionDir?: string, storage: SessionStorage = new FileSessionStorage(), ): Promise { const dir = sessionDir ?? getDefaultSessionDir(cwd, storage); const mostRecent = 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//). */ static list(cwd: string, sessionDir?: string, storage: SessionStorage = new FileSessionStorage()): SessionInfo[] { const dir = sessionDir ?? getDefaultSessionDir(cwd, storage); const sessions: SessionInfo[] = []; try { const files = storage.listFilesSync(dir, "*.jsonl"); for (const file of files) { try { const content = storage.readTextSync(file); const lines = content.trim().split("\n"); if (lines.length === 0) continue; // Check first line for valid session header let header: { type: string; id: string; cwd?: string; title?: string; timestamp: string } | null = null; try { const first = JSON.parse(lines[0]); if (first.type === "session" && first.id) { header = first; } } catch { // Not valid JSON } if (!header) continue; const stats = storage.statSync(file); let messageCount = 0; let firstMessage = ""; const allMessages: string[] = []; for (let i = 1; i < lines.length; i++) { try { const entry = JSON.parse(lines[i]); if (entry.type === "message") { messageCount++; if (entry.message.role === "user" || entry.message.role === "assistant") { const textContent = entry.message.content .filter((c: any) => c.type === "text") .map((c: any) => c.text) .join(" "); if (textContent) { allMessages.push(textContent); if (!firstMessage && entry.message.role === "user") { firstMessage = textContent; } } } } } catch { // Skip malformed lines } } sessions.push({ path: file, id: header.id, cwd: typeof header.cwd === "string" ? header.cwd : "", title: header.title, 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()); } catch { // Return empty list on error } return sessions; } }