feat: sqlite session storage (#6594)

This PR:

- Adds retainedTail to compaction entries in the new agent harness so we don't have to walk up the tree for the 2000 tokens before compaction,
- Changes getPathToRoot to getPathToRootOrCompaction to only load until last compaction, as unnecessary to access all nodes where it is called,
- Adds a SQLite storage backend, in a separate packages/session-backend-sqlite, with a migration system and schemas as per on-site discussions: sessions to match session header messages (except for metadata, which I couldn't understand what it's used for or where it gets written, so I omitted it), session_entries for shared entry types as columns plus payload as a json for what remains, session_sequences to represent the append-only, serialized nature of the jsonl files, branch_entries to attribute nodes to branches (relationship one-to-many), and session_materialized with the session info (see /session in TUI) to act as a "cache" or quick-access for costs, message count, token info, labels, session name, and model-thinking-level config (e.g. for fast resume).
- This is compatible with the new agent harness Session abstraction.

packages/agent/src/harness/agent-harness.ts

@@ -730,6 +730,7 @@ export class AgentHarness<
 				result.details,
 				provided !== undefined,
 				result.usage,
+				result.retainedTail,
 			);
 			const entry = await this.session.getEntry(entryId);
 			if (entry?.type === "compaction") {

packages/agent/src/harness/compaction/compaction.ts

@@ -97,12 +97,14 @@ function getMessageFromEntryForCompaction(entry: SessionTreeEntry): AgentMessage
 export interface CompactionResult<T = unknown> {
 	/** Summary text that replaces compacted history in future context. */
 	summary: string;
-	/** Entry id where retained history starts. */
-	firstKeptEntryId: string;
+	/** Entry id where retained history starts. Optional during Pi 2.0 transition. */
+	firstKeptEntryId?: string;
 	/** Estimated context tokens before compaction. */
 	tokensBefore: number;
 	/** Usage from the LLM call(s) that generated this summary, if available. */
 	usage?: Usage;
+	/** Retained recent messages stored directly on the compaction entry. Optional during Pi 2.0 transition. */
+	retainedTail?: AgentMessage[];
 	/** Optional implementation-specific details stored with the compaction entry. */
 	details?: T;
 }
@@ -583,6 +585,8 @@ export interface CompactionPreparation {
 	messagesToSummarize: AgentMessage[];
 	/** Prefix messages summarized separately when compaction splits a turn. */
 	turnPrefixMessages: AgentMessage[];
+	/** Recent messages retained after compaction and stored on the compaction entry. */
+	retainedTail: AgentMessage[];
 	/** Whether compaction splits a turn. */
 	isSplitTurn: boolean;
 	/** Estimated context tokens before compaction. */
@@ -617,7 +621,9 @@ export function prepareCompaction(
 	if (prevCompactionIndex >= 0) {
 		const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry;
 		previousSummary = prevCompaction.summary;
-		const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId);
+		const firstKeptEntryIndex = prevCompaction.firstKeptEntryId
+			? pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId)
+			: -1;
 		boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1;
 	}
 	const boundaryEnd = pathEntries.length;
@@ -644,6 +650,11 @@ export function prepareCompaction(
 			if (msg) turnPrefixMessages.push(msg);
 		}
 	}
+	const retainedTail: AgentMessage[] = [];
+	for (let i = cutPoint.firstKeptEntryIndex; i < boundaryEnd; i++) {
+		const msg = getMessageFromEntryForCompaction(pathEntries[i]);
+		if (msg) retainedTail.push(msg);
+	}
 	const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex);
 	if (cutPoint.isSplitTurn) {
 		for (const msg of turnPrefixMessages) {
@@ -655,6 +666,7 @@ export function prepareCompaction(
 		firstKeptEntryId,
 		messagesToSummarize,
 		turnPrefixMessages,
+		retainedTail,
 		isSplitTurn: cutPoint.isSplitTurn,
 		tokensBefore,
 		previousSummary,
@@ -693,6 +705,7 @@ export async function compact(
 		firstKeptEntryId,
 		messagesToSummarize,
 		turnPrefixMessages,
+		retainedTail,
 		isSplitTurn,
 		tokensBefore,
 		previousSummary,
@@ -762,6 +775,7 @@ export async function compact(
 		firstKeptEntryId,
 		tokensBefore,
 		usage: summaryUsage,
+		retainedTail,
 		details: { readFiles, modifiedFiles } as CompactionDetails,
 	});
 }

packages/agent/src/harness/session/jsonl-storage.ts

@@ -1,5 +1,12 @@
 import { uuidv7 } from "@earendil-works/pi-ai";
-import type { FileSystem, JsonlSessionMetadata, LeafEntry, SessionStorage, SessionTreeEntry } from "../types.ts";
+import type {
+	FileSystem,
+	JsonlSessionMetadata,
+	LeafEntry,
+	SessionEntryCursorOptions,
+	SessionStorage,
+	SessionTreeEntry,
+} from "../types.ts";
 import { SessionError, toError } from "../types.ts";
 import { getFileSystemResultOrThrow } from "./repo-utils.ts";
 
@@ -293,13 +300,66 @@ export class JsonlSessionStorage implements SessionStorage<JsonlSessionMetadata>
 		return this.labelsById.get(id);
 	}
 
-	async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
+	async getSessionName(): Promise<string | undefined> {
+		const entries = await this.findEntries("session_info");
+		return entries[entries.length - 1]?.name?.trim() || undefined;
+	}
+
+	async getSessionStats() {
+		let messageCount = 0;
+		let cachedTokens = 0;
+		let uncachedTokens = 0;
+		let totalTokens = 0;
+		let costTotal = 0;
+		for (const entry of this.entries) {
+			if (entry.type === "message") {
+				messageCount += 1;
+			}
+			const usage =
+				entry.type === "message"
+					? entry.message.role === "assistant"
+						? entry.message.usage
+						: undefined
+					: entry.type === "compaction" || entry.type === "branch_summary"
+						? entry.usage
+						: undefined;
+			if (
+				!usage ||
+				typeof usage.input !== "number" ||
+				typeof usage.output !== "number" ||
+				typeof usage.cacheRead !== "number" ||
+				typeof usage.cacheWrite !== "number" ||
+				typeof usage.cost?.total !== "number"
+			) {
+				continue;
+			}
+			cachedTokens += usage.cacheRead;
+			uncachedTokens += usage.input + usage.cacheWrite;
+			totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
+			costTotal += usage.cost.total;
+		}
+		return {
+			messageCount,
+			cachedTokens,
+			uncachedTokens,
+			totalTokens,
+			costTotal,
+		};
+	}
+
+	async getPathToRootOrCompaction(leafId: string | null): Promise<SessionTreeEntry[]> {
 		if (leafId === null) return [];
 		const path: SessionTreeEntry[] = [];
+		let stopAtEntryId: string | null = null;
 		let current = this.byId.get(leafId);
 		if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`);
 		while (current) {
 			path.unshift(current);
+			if (stopAtEntryId !== null && current.id === stopAtEntryId) break;
+			if (current.type === "compaction") {
+				if (current.retainedTail) break;
+				stopAtEntryId = current.firstKeptEntryId ?? null;
+			}
 			if (!current.parentId) break;
 			const parent = this.byId.get(current.parentId);
 			if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`);
@@ -308,7 +368,9 @@ export class JsonlSessionStorage implements SessionStorage<JsonlSessionMetadata>
 		return path;
 	}
 
-	async getEntries(): Promise<SessionTreeEntry[]> {
-		return [...this.entries];
+	async getEntries(options?: SessionEntryCursorOptions): Promise<SessionTreeEntry[]> {
+		const start = options?.afterEntrySeq ?? 0;
+		const end = options?.limit === undefined ? undefined : start + options.limit;
+		return this.entries.slice(start, end);
 	}
 }

packages/agent/src/harness/session/memory-storage.ts

@@ -1,6 +1,7 @@
 import { uuidv7 } from "@earendil-works/pi-ai";
 import {
 	type LeafEntry,
+	type SessionEntryCursorOptions,
 	SessionError,
 	type SessionMetadata,
 	type SessionStorage,
@@ -112,13 +113,66 @@ export class InMemorySessionStorage<TMetadata extends SessionMetadata = SessionM
 		return this.labelsById.get(id);
 	}
 
-	async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
+	async getSessionName(): Promise<string | undefined> {
+		const entries = await this.findEntries("session_info");
+		return entries[entries.length - 1]?.name?.trim() || undefined;
+	}
+
+	async getSessionStats() {
+		let messageCount = 0;
+		let cachedTokens = 0;
+		let uncachedTokens = 0;
+		let totalTokens = 0;
+		let costTotal = 0;
+		for (const entry of this.entries) {
+			if (entry.type === "message") {
+				messageCount += 1;
+			}
+			const usage =
+				entry.type === "message"
+					? entry.message.role === "assistant"
+						? entry.message.usage
+						: undefined
+					: entry.type === "compaction" || entry.type === "branch_summary"
+						? entry.usage
+						: undefined;
+			if (
+				!usage ||
+				typeof usage.input !== "number" ||
+				typeof usage.output !== "number" ||
+				typeof usage.cacheRead !== "number" ||
+				typeof usage.cacheWrite !== "number" ||
+				typeof usage.cost?.total !== "number"
+			) {
+				continue;
+			}
+			cachedTokens += usage.cacheRead;
+			uncachedTokens += usage.input + usage.cacheWrite;
+			totalTokens += usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
+			costTotal += usage.cost.total;
+		}
+		return {
+			messageCount,
+			cachedTokens,
+			uncachedTokens,
+			totalTokens,
+			costTotal,
+		};
+	}
+
+	async getPathToRootOrCompaction(leafId: string | null): Promise<SessionTreeEntry[]> {
 		if (leafId === null) return [];
 		const path: SessionTreeEntry[] = [];
+		let stopAtEntryId: string | null = null;
 		let current = this.byId.get(leafId);
 		if (!current) throw new SessionError("not_found", `Entry ${leafId} not found`);
 		while (current) {
 			path.unshift(current);
+			if (stopAtEntryId !== null && current.id === stopAtEntryId) break;
+			if (current.type === "compaction") {
+				if (current.retainedTail) break;
+				stopAtEntryId = current.firstKeptEntryId ?? null;
+			}
 			if (!current.parentId) break;
 			const parent = this.byId.get(current.parentId);
 			if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`);
@@ -127,7 +181,9 @@ export class InMemorySessionStorage<TMetadata extends SessionMetadata = SessionM
 		return path;
 	}
 
-	async getEntries(): Promise<SessionTreeEntry[]> {
-		return [...this.entries];
+	async getEntries(options?: SessionEntryCursorOptions): Promise<SessionTreeEntry[]> {
+		const start = options?.afterEntrySeq ?? 0;
+		const end = options?.limit === undefined ? undefined : start + options.limit;
+		return this.entries.slice(start, end);
 	}
 }

packages/agent/src/harness/session/repo-utils.ts

@@ -47,5 +47,5 @@ export async function getEntriesToFork(
 		}
 		effectiveLeafId = target.parentId;
 	}
-	return storage.getPathToRoot(effectiveLeafId);
+	return storage.getPathToRootOrCompaction(effectiveLeafId);
 }

packages/agent/src/harness/session/session.ts

@@ -11,8 +11,10 @@ import type {
 	MessageEntry,
 	ModelChangeEntry,
 	SessionContext,
+	SessionEntryCursorOptions,
 	SessionInfoEntry,
 	SessionMetadata,
+	SessionStats,
 	SessionStorage,
 	SessionTreeEntry,
 	ThinkingLevelChangeEntry,
@@ -67,11 +69,19 @@ export function defaultContextEntryTransform(pathEntries: readonly SessionTreeEn
 
 	const entries: SessionTreeEntry[] = [compaction];
 	const compactionIdx = pathEntries.findIndex((entry) => entry.type === "compaction" && entry.id === compaction.id);
-	let foundFirstKept = false;
-	for (let i = 0; i < compactionIdx; i++) {
-		const entry = pathEntries[i]!;
-		if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true;
-		if (foundFirstKept) entries.push(entry);
+	if (compaction.retainedTail) {
+		for (let i = compactionIdx + 1; i < pathEntries.length; i++) {
+			entries.push(pathEntries[i]!);
+		}
+		return entries;
+	}
+	if (compaction.firstKeptEntryId) {
+		let foundFirstKept = false;
+		for (let i = 0; i < compactionIdx; i++) {
+			const entry = pathEntries[i]!;
+			if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true;
+			if (foundFirstKept) entries.push(entry);
+		}
 	}
 	for (let i = compactionIdx + 1; i < pathEntries.length; i++) {
 		entries.push(pathEntries[i]!);
@@ -111,7 +121,10 @@ export function sessionEntryToContextMessages(
 		];
 	}
 	if (entry.type === "compaction") {
-		return [createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp)];
+		return [
+			createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp),
+			...(entry.retainedTail ?? []),
+		];
 	}
 	if (entry.type === "branch_summary" && entry.summary) {
 		return [createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)];
@@ -159,13 +172,13 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
 		return this.storage.getEntry(id);
 	}
 
-	getEntries(): Promise<SessionTreeEntry[]> {
-		return this.storage.getEntries();
+	getEntries(options?: SessionEntryCursorOptions): Promise<SessionTreeEntry[]> {
+		return this.storage.getEntries(options);
 	}
 
 	async getBranch(fromId?: string): Promise<SessionTreeEntry[]> {
 		const leafId = fromId ?? (await this.storage.getLeafId());
-		return this.storage.getPathToRoot(leafId);
+		return this.storage.getPathToRootOrCompaction(leafId);
 	}
 
 	async buildContextEntries(options: SessionContextBuildOptions = {}): Promise<SessionTreeEntry[]> {
@@ -190,9 +203,12 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
 		return this.storage.getLabel(id);
 	}
 
+	getSessionStats(): Promise<SessionStats> {
+		return this.storage.getSessionStats();
+	}
+
 	async getSessionName(): Promise<string | undefined> {
-		const entries = await this.storage.findEntries("session_info");
-		return entries[entries.length - 1]?.name?.trim() || undefined;
+		return this.storage.getSessionName();
 	}
 
 	private async appendTypedEntry<TEntry extends SessionTreeEntry>(entry: TEntry): Promise<string> {
@@ -243,11 +259,12 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
 
 	async appendCompaction<T = unknown>(
 		summary: string,
-		firstKeptEntryId: string,
+		firstKeptEntryId: string | undefined,
 		tokensBefore: number,
 		details?: T,
 		fromHook?: boolean,
 		usage?: Usage,
+		retainedTail?: AgentMessage[],
 	): Promise<string> {
 		return this.appendTypedEntry({
 			type: "compaction",
@@ -257,6 +274,7 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
 			summary,
 			firstKeptEntryId,
 			tokensBefore,
+			retainedTail,
 			details,
 			usage,
 			fromHook,

packages/agent/src/harness/types.ts

@@ -370,8 +370,9 @@ export interface ActiveToolsChangeEntry extends SessionTreeEntryBase {
 export interface CompactionEntry<T = unknown> extends SessionTreeEntryBase {
 	type: "compaction";
 	summary: string;
-	firstKeptEntryId: string;
+	firstKeptEntryId?: string;
 	tokensBefore: number;
+	retainedTail?: AgentMessage[];
 	details?: T;
 	usage?: Usage;
 	fromHook?: boolean;
@@ -436,6 +437,14 @@ export interface SessionContext {
 	activeToolNames: string[] | null;
 }
 
+export interface SessionStats {
+	messageCount: number;
+	cachedTokens: number;
+	uncachedTokens: number;
+	totalTokens: number;
+	costTotal: number;
+}
+
 export interface SessionMetadata {
 	id: string;
 	createdAt: string;
@@ -448,6 +457,11 @@ export interface JsonlSessionMetadata extends SessionMetadata {
 	metadata?: Record<string, unknown>;
 }
 
+export interface SessionEntryCursorOptions {
+	afterEntrySeq?: number;
+	limit?: number;
+}
+
 export interface SessionStorage<TMetadata extends SessionMetadata = SessionMetadata> {
 	getMetadata(): Promise<TMetadata>;
 	getLeafId(): Promise<string | null>;
@@ -460,8 +474,10 @@ export interface SessionStorage<TMetadata extends SessionMetadata = SessionMetad
 		type: TType,
 	): Promise<Array<Extract<SessionTreeEntry, { type: TType }>>>;
 	getLabel(id: string): Promise<string | undefined>;
-	getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]>;
-	getEntries(): Promise<SessionTreeEntry[]>;
+	getSessionName(): Promise<string | undefined>;
+	getSessionStats(): Promise<SessionStats>;
+	getPathToRootOrCompaction(leafId: string | null): Promise<SessionTreeEntry[]>;
+	getEntries(options?: SessionEntryCursorOptions): Promise<SessionTreeEntry[]>;
 }
 
 export type { Session } from "./session/session.ts";
@@ -753,10 +769,11 @@ export interface AbortResult {
 
 export interface CompactResult {
 	summary: string;
-	firstKeptEntryId: string;
+	firstKeptEntryId?: string;
 	tokensBefore: number;
 	/** Usage from the LLM call(s) that generated this summary, if available. */
 	usage?: Usage;
+	retainedTail?: AgentMessage[];
 	details?: unknown;
 }
 
@@ -776,6 +793,7 @@ export interface CompactionPreparation {
 	firstKeptEntryId: string;
 	messagesToSummarize: AgentMessage[];
 	turnPrefixMessages: AgentMessage[];
+	retainedTail: AgentMessage[];
 	isSplitTurn: boolean;
 	tokensBefore: number;
 	previousSummary?: string;

packages/agent/src/index.ts

@@ -1,4 +1,5 @@
 // Core Agent
+export { uuidv7 } from "@earendil-works/pi-ai";
 export * from "./agent.ts";
 // Loop functions
 export * from "./agent-loop.ts";