/** * SessionManager — persistent conversation state with branching and compaction. * * Provides: * - Multiple named sessions (conversations) * - Tree-structured messages (parent_id for branching) * - History retrieval following a branch path * - Compaction (summarize old messages to save context) * - Compatible with AI SDK's UIMessage type * * Usage: * const sessions = new SessionManager(agent); * const session = sessions.create("my-chat"); * sessions.append(session.id, { id: "msg1", role: "user", parts: [...] }); * const history = sessions.getHistory(session.id); // UIMessage[] */ import type { UIMessage } from "ai"; import { SessionStorage } from "./storage"; import type { Session, Compaction } from "./storage"; export type { Session, Compaction } from "./storage"; // ── Truncation utilities ────────────────────────────────────────── const DEFAULT_MAX_CHARS = 30_000; const ELLIPSIS = "\n\n... [truncated] ...\n\n"; /** * Truncate from the head (keep the end of the content). */ export function truncateHead( text: string, maxChars: number = DEFAULT_MAX_CHARS ): string { if (text.length <= maxChars) return text; const keep = maxChars - ELLIPSIS.length; if (keep <= 0) return text.slice(-maxChars); return ELLIPSIS + text.slice(-keep); } /** * Truncate from the tail (keep the start of the content). */ export function truncateTail( text: string, maxChars: number = DEFAULT_MAX_CHARS ): string { if (text.length <= maxChars) return text; const keep = maxChars - ELLIPSIS.length; if (keep <= 0) return text.slice(0, maxChars); return text.slice(0, keep) + ELLIPSIS; } /** * Truncate by line count (keep the first N lines). */ export function truncateLines(text: string, maxLines: number = 200): string { const lines = text.split("\n"); if (lines.length <= maxLines) return text; const kept = lines.slice(0, maxLines).join("\n"); const omitted = lines.length - maxLines; return kept + `\n\n... [${omitted} more lines truncated] ...`; } /** * Truncate from both ends, keeping the start and end. */ export function truncateMiddle( text: string, maxChars: number = DEFAULT_MAX_CHARS ): string { if (text.length <= maxChars) return text; const halfKeep = Math.floor((maxChars - ELLIPSIS.length) / 2); if (halfKeep <= 0) return text.slice(0, maxChars); return text.slice(0, halfKeep) + ELLIPSIS + text.slice(-halfKeep); } /** * Smart truncation for tool output. */ export function truncateToolOutput( output: string, options: { maxChars?: number; maxLines?: number; strategy?: "head" | "tail" | "middle"; } = {} ): string { const { maxChars = DEFAULT_MAX_CHARS, maxLines = 500, strategy = "tail" } = options; let result = truncateLines(output, maxLines); if (result.length > maxChars) { switch (strategy) { case "head": result = truncateHead(result, maxChars); break; case "middle": result = truncateMiddle(result, maxChars); break; case "tail": default: result = truncateTail(result, maxChars); break; } } return result; } // Mirrors Agent.sql — kept structural to avoid importing the 4k-line Agent class. interface AgentLike { sql: ( strings: TemplateStringsArray, ...values: (string | number | boolean | null)[] ) => Array>; } export interface SessionManagerOptions { /** * Maximum number of messages on the current branch before * needsCompaction() returns true. Default: 100. */ maxContextMessages?: number; /** * Raw SQL exec function for batch operations (e.g. DELETE ... WHERE id IN (...)). * When provided, batch deletes use a single query instead of N individual ones. * * Typically: `(query, ...values) => { agent.ctx.storage.sql.exec(query, ...values); }` */ exec?: ( query: string, ...values: (string | number | boolean | null)[] ) => void; } export class SessionManager { private _storage: SessionStorage; private _options: SessionManagerOptions; constructor(agent: AgentLike, options: SessionManagerOptions = {}) { this._storage = new SessionStorage(agent.sql.bind(agent), options.exec); this._options = { maxContextMessages: 100, ...options }; } // ── Session lifecycle ────────────────────────────────────────── /** * Create a new session with a name. */ create(name: string): Session { return this._storage.createSession(crypto.randomUUID(), name); } /** * Get a session by ID. */ get(sessionId: string): Session | null { return this._storage.getSession(sessionId); } /** * List all sessions, most recently updated first. */ list(): Session[] { return this._storage.listSessions(); } /** * Delete a session and all its messages and compactions. */ delete(sessionId: string): void { this._storage.deleteSession(sessionId); } /** * Clear all messages and compactions for a session without * deleting the session itself. */ clearMessages(sessionId: string): void { this._storage.clearSessionMessages(sessionId); } /** * Rename a session. */ rename(sessionId: string, name: string): void { this._storage.renameSession(sessionId, name); } // ── Messages ─────────────────────────────────────────────────── /** * Append a message to a session. If parentId is not provided, * the message is appended after the latest leaf. * * Idempotent — appending the same message.id twice is a no-op. * * Returns the stored message ID. */ append(sessionId: string, message: UIMessage, parentId?: string): string { const resolvedParent = parentId ?? this._storage.getLatestLeaf(sessionId)?.id ?? null; const id = message.id || crypto.randomUUID(); this._storage.appendMessage(id, sessionId, resolvedParent, message); return id; } /** * Insert or update a message. First call inserts, subsequent calls * update the content. Enables incremental persistence. * * Idempotent on insert, content-updating on subsequent calls. */ upsert(sessionId: string, message: UIMessage, parentId?: string): string { const resolvedParent = parentId ?? this._storage.getLatestLeaf(sessionId)?.id ?? null; const id = message.id || crypto.randomUUID(); this._storage.upsertMessage(id, sessionId, resolvedParent, message); return id; } /** * Delete a single message by ID. * Children of the deleted message naturally become path roots * (their parent_id points to a missing row, truncating the CTE walk). */ deleteMessage(messageId: string): void { this._storage.deleteMessage(messageId); } /** * Delete multiple messages by ID. */ deleteMessages(messageIds: string[]): void { this._storage.deleteMessages(messageIds); } /** * Append multiple messages in sequence (each parented to the previous). * Returns the ID of the last appended message. */ appendAll( sessionId: string, messages: UIMessage[], parentId?: string ): string | null { let lastId = parentId ?? null; for (const msg of messages) { const resolvedParent = lastId ?? this._storage.getLatestLeaf(sessionId)?.id ?? null; const id = msg.id || crypto.randomUUID(); this._storage.appendMessage(id, sessionId, resolvedParent, msg); lastId = id; } return lastId; } /** * Get the conversation history for a session as UIMessage[]. * * If leafId is provided, returns the path from root to that leaf * (a specific branch). Otherwise returns the path to the most * recent leaf (the "current" branch). * * If compactions exist, older messages covered by a compaction * are replaced with a system message containing the summary. */ getHistory(sessionId: string, leafId?: string): UIMessage[] { const leaf = leafId ? this._storage.getMessage(leafId) : this._storage.getLatestLeaf(sessionId); if (!leaf) return []; const storedPath = this._storage.getMessagePath(leaf.id); const compactions = this._storage.getCompactions(sessionId); if (compactions.length === 0) { return storedPath.map((m) => this._storage.parseMessage(m)); } return this._applyCompactions(storedPath, compactions); } /** * Get the total message count for a session (across all branches). */ getMessageCount(sessionId: string): number { return this._storage.getMessageCount(sessionId); } /** * Check if the session's current branch needs compaction. * Uses a count-only query — does not load message content. */ needsCompaction(sessionId: string): boolean { const leaf = this._storage.getLatestLeaf(sessionId); if (!leaf) return false; const pathLen = this._storage.getPathLength(leaf.id); return pathLen > (this._options.maxContextMessages ?? 100); } // ── Branching ────────────────────────────────────────────────── /** * Get the children of a message (branches from that point). */ getBranches(messageId: string): UIMessage[] { const children = this._storage.getChildren(messageId); return children.map((m) => this._storage.parseMessage(m)); } /** * Fork a session at a specific message, creating a new session * with the history up to that point copied over. */ fork(atMessageId: string, newName: string): Session { const newSession = this.create(newName); const path = this._storage.getMessagePath(atMessageId); let parentId: string | null = null; for (const stored of path) { const msg = this._storage.parseMessage(stored); const newId = crypto.randomUUID(); this._storage.appendMessage(newId, newSession.id, parentId, msg); parentId = newId; } return newSession; } // ── Compaction ───────────────────────────────────────────────── /** * Add a compaction record. The summary replaces messages from * fromMessageId to toMessageId in context assembly. * * Typically called after using an LLM to summarize older messages. */ addCompaction( sessionId: string, summary: string, fromMessageId: string, toMessageId: string ): Compaction { return this._storage.addCompaction( crypto.randomUUID(), sessionId, summary, fromMessageId, toMessageId ); } /** * Get all compaction records for a session. */ getCompactions(sessionId: string): Compaction[] { return this._storage.getCompactions(sessionId); } // ── Internal ─────────────────────────────────────────────────── private _applyCompactions( path: Array<{ id: string; content: string }>, compactions: Compaction[] ): UIMessage[] { const pathIds = path.map((m) => m.id); const result: UIMessage[] = []; let i = 0; while (i < path.length) { // Check if any compaction starts at this message const compaction = compactions.find( (c) => c.from_message_id === pathIds[i] ); if (compaction) { // Only apply if the compaction's end is also on this path const endIdx = pathIds.indexOf(compaction.to_message_id); if (endIdx >= i) { result.push({ id: `compaction_${compaction.id}`, role: "system", parts: [ { type: "text", text: `[Previous conversation summary]\n${compaction.summary}` } ] }); i = endIdx + 1; } else { // Compaction doesn't span this path — skip it, emit message as-is result.push(JSON.parse(path[i].content) as UIMessage); i++; } } else { result.push(JSON.parse(path[i].content) as UIMessage); i++; } } return result; } }