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.
This commit is contained in:
Cristina Poncela Cubeiro
2026-07-21 11:36:31 +02:00
committed by GitHub
parent 54fad505b9
commit 9e7582aa03
40 changed files with 2659 additions and 145 deletions
@@ -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);
}
}