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:
committed by
GitHub
parent
54fad505b9
commit
9e7582aa03
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user