From cf82e2e2ae4485c6690972d76a7e26b76ffc0306 Mon Sep 17 00:00:00 2001 From: Cristina Poncela Cubeiro <140309543+cristinaponcela@users.noreply.github.com> Date: Thu, 13 Aug 2026 11:58:39 +0200 Subject: [PATCH] feat(sqlite): more helpers --- .../sqlite-node/src/sqlite/session/entries.ts | 136 ++++++++++++++++++ .../src/sqlite/session/registers.ts | 65 ++++++++- .../src/sqlite/session/session-stats.ts | 45 ++++++ .../src/sqlite/session/usage-ledger.ts | 35 +++++ 4 files changed, 276 insertions(+), 5 deletions(-) create mode 100644 packages/session-backends/sqlite-node/src/sqlite/session/entries.ts create mode 100644 packages/session-backends/sqlite-node/src/sqlite/session/session-stats.ts create mode 100644 packages/session-backends/sqlite-node/src/sqlite/session/usage-ledger.ts diff --git a/packages/session-backends/sqlite-node/src/sqlite/session/entries.ts b/packages/session-backends/sqlite-node/src/sqlite/session/entries.ts new file mode 100644 index 000000000..7c3d3558e --- /dev/null +++ b/packages/session-backends/sqlite-node/src/sqlite/session/entries.ts @@ -0,0 +1,136 @@ +import type { + BranchSummaryEntry, + CompactionEntry, + CustomEntry, + Entry, + EntryScan, + EntryStructure, + MessageEntry, +} from "@earendil-works/pi-agent-core"; +import { joinSqlFragments, type SqlQuery, sql } from "../sql.ts"; +import type { SqliteDatabase } from "../types.ts"; + +export interface EntryRow { + id: string; + parent_id: string | null; + seq: number; + type: Entry["type"]; + custom_type: string | null; + timestamp: number; + payload: string; +} + +type StoredEntryPayload = Omit< + TEntry, + "id" | "parentId" | "seq" | "timestamp" | "type" | "customType" +>; + +function entryPayload(entry: Entry): StoredEntryPayload { + switch (entry.type) { + case "message": { + const payload: StoredEntryPayload = { + message: entry.message, + ...(entry.terminate === undefined ? {} : { terminate: entry.terminate }), + }; + return payload; + } + case "compaction": { + const payload: StoredEntryPayload = { + summary: entry.summary, + retainedTail: entry.retainedTail, + tokensBefore: entry.tokensBefore, + ...(entry.details === undefined ? {} : { details: entry.details }), + ...(entry.usage === undefined ? {} : { usage: entry.usage }), + fromHook: entry.fromHook, + }; + return payload; + } + case "branch_summary": { + const payload: StoredEntryPayload = { + fromId: entry.fromId, + summary: entry.summary, + ...(entry.details === undefined ? {} : { details: entry.details }), + ...(entry.usage === undefined ? {} : { usage: entry.usage }), + fromHook: entry.fromHook, + }; + return payload; + } + case "custom": { + const payload: StoredEntryPayload = entry.data === undefined ? {} : { data: entry.data }; + return payload; + } + } +} + +function parsePayload(row: EntryRow): StoredEntryPayload { + return JSON.parse(row.payload) as StoredEntryPayload; +} + +export function insertEntryRow(db: SqliteDatabase, entry: Entry): void { + sql`INSERT INTO entries (id, parent_id, seq, type, custom_type, timestamp, payload) + VALUES ( + ${entry.id}, + ${entry.parentId}, + ${entry.seq}, + ${entry.type}, + ${entry.type === "custom" ? entry.customType : null}, + ${entry.timestamp}, + ${JSON.stringify(entryPayload(entry))} + )`.run(db); +} + +export function decodeEntryRow(row: EntryRow): Entry { + const base = { + id: row.id, + parentId: row.parent_id, + seq: row.seq, + timestamp: row.timestamp, + }; + switch (row.type) { + case "message": + return { ...base, type: "message", ...parsePayload(row) }; + case "compaction": + return { ...base, type: "compaction", ...parsePayload(row) }; + case "branch_summary": + return { ...base, type: "branch_summary", ...parsePayload(row) }; + case "custom": + if (row.custom_type === null) throw new Error(`Custom entry ${row.id} is missing custom_type`); + return { ...base, type: "custom", customType: row.custom_type, ...parsePayload(row) }; + } +} + +export function entryStructureFromRow(row: EntryRow): EntryStructure { + return { + id: row.id, + parentId: row.parent_id, + seq: row.seq, + timestamp: row.timestamp, + type: row.type, + ...(row.custom_type === null ? {} : { customType: row.custom_type }), + }; +} + +export function readEntryRows(db: SqliteDatabase, ids: readonly string[]): EntryRow[] { + if (ids.length === 0) return []; + const placeholders = joinSqlFragments( + ids.map((id) => sql`${id}`), + ", ", + ); + return sql`SELECT id, parent_id, seq, type, custom_type, timestamp, payload + FROM entries + WHERE id IN (${placeholders})`.all(db); +} + +export function scanEntryRows(db: SqliteDatabase, query: EntryScan): EntryRow[] { + const filters: SqlQuery[] = []; + if (query.type !== undefined) filters.push(sql`type = ${query.type}`); + if (query.customType !== undefined) filters.push(sql`custom_type = ${query.customType}`); + if (query.fromSeq !== undefined) filters.push(sql`seq >= ${query.fromSeq}`); + if (query.toSeq !== undefined) filters.push(sql`seq <= ${query.toSeq}`); + + const where = filters.length === 0 ? sql`` : sql`WHERE ${joinSqlFragments(filters, " AND ")}`; + const order = query.order === "desc" ? sql`ORDER BY seq DESC` : sql`ORDER BY seq ASC`; + const limit = query.limit === undefined ? sql`` : sql`LIMIT ${Math.max(0, query.limit)}`; + return sql`SELECT id, parent_id, seq, type, custom_type, timestamp, payload + FROM entries ${where} ${order} ${limit}`.all(db); +} diff --git a/packages/session-backends/sqlite-node/src/sqlite/session/registers.ts b/packages/session-backends/sqlite-node/src/sqlite/session/registers.ts index b352b960f..c3defa61d 100644 --- a/packages/session-backends/sqlite-node/src/sqlite/session/registers.ts +++ b/packages/session-backends/sqlite-node/src/sqlite/session/registers.ts @@ -1,11 +1,66 @@ -import type { LaneState } from "@earendil-works/pi-agent-core"; +import type { LaneState, Register, RegisterNamespace, RegisterValues } from "@earendil-works/pi-agent-core"; import { sql } from "../sql.ts"; import type { SqliteDatabase } from "../types.ts"; +export interface RegisterRow { + namespace: RegisterNamespace; + key: string; + seq: number; + value: string; +} + export function insertInitialMainLaneRegisters(db: SqliteDatabase): void { const laneState = { currentOperationId: null, pendingNextRun: [] } satisfies LaneState; - sql`INSERT INTO registers (namespace, key, seq, value) - VALUES (${"lane.leaf"}, ${"main"}, ${1}, ${JSON.stringify(null)})`.run(db); - sql`INSERT INTO registers (namespace, key, seq, value) - VALUES (${"lane.state"}, ${"main"}, ${2}, ${JSON.stringify(laneState)})`.run(db); + setRegisterRow(db, "lane.leaf", "main", 1, null); + setRegisterRow(db, "lane.state", "main", 2, laneState); +} + +export function setRegisterRow( + db: SqliteDatabase, + namespace: TNamespace, + key: string, + seq: number, + value: RegisterValues[TNamespace], +): void { + sql`INSERT INTO registers (namespace, key, seq, value) + VALUES (${namespace}, ${key}, ${seq}, ${JSON.stringify(value)}) + ON CONFLICT(namespace, key) DO UPDATE SET seq = excluded.seq, value = excluded.value`.run(db); +} + +export function deleteRegisterRow(db: SqliteDatabase, namespace: RegisterNamespace, key: string): void { + sql`DELETE FROM registers WHERE namespace = ${namespace} AND key = ${key}`.run(db); +} + +export function decodeRegisterRow( + namespace: TNamespace, + row: RegisterRow, +): Register { + if (row.namespace !== namespace) throw new Error(`Expected register namespace ${namespace}, found ${row.namespace}`); + return { + namespace, + key: row.key, + seq: row.seq, + value: JSON.parse(row.value) as RegisterValues[TNamespace], + }; +} + +export function readRegisterRow( + db: SqliteDatabase, + namespace: TNamespace, + key: string, +): Register | undefined { + const row = sql`SELECT namespace, key, seq, value FROM registers + WHERE namespace = ${namespace} AND key = ${key}`.get(db); + return row === undefined ? undefined : decodeRegisterRow(namespace, row); +} + +export function listRegisterRows( + db: SqliteDatabase, + namespace: TNamespace, + keyPrefix = "", +): Register[] { + const rows = sql`SELECT namespace, key, seq, value FROM registers + WHERE namespace = ${namespace} AND key >= ${keyPrefix} AND key < ${`${keyPrefix}\uffff`} + ORDER BY key ASC`.all(db); + return rows.map((row) => decodeRegisterRow(namespace, row)); } diff --git a/packages/session-backends/sqlite-node/src/sqlite/session/session-stats.ts b/packages/session-backends/sqlite-node/src/sqlite/session/session-stats.ts new file mode 100644 index 000000000..6949d751a --- /dev/null +++ b/packages/session-backends/sqlite-node/src/sqlite/session/session-stats.ts @@ -0,0 +1,45 @@ +import type { SessionStats, UsageRow } from "@earendil-works/pi-agent-core"; +import type { Usage } from "@earendil-works/pi-ai"; +import { sql } from "../sql.ts"; +import type { SqliteDatabase } from "../types.ts"; +import { readSingleSessionRow } from "./session-row.ts"; + +function addUsage(left: Usage, right: Usage): Usage { + return { + input: left.input + right.input, + output: left.output + right.output, + cacheRead: left.cacheRead + right.cacheRead, + cacheWrite: left.cacheWrite + right.cacheWrite, + ...(left.cacheWrite1h === undefined && right.cacheWrite1h === undefined + ? {} + : { cacheWrite1h: (left.cacheWrite1h ?? 0) + (right.cacheWrite1h ?? 0) }), + ...(left.reasoning === undefined && right.reasoning === undefined + ? {} + : { reasoning: (left.reasoning ?? 0) + (right.reasoning ?? 0) }), + totalTokens: left.totalTokens + right.totalTokens, + cost: { + input: left.cost.input + right.cost.input, + output: left.cost.output + right.cost.output, + cacheRead: left.cost.cacheRead + right.cost.cacheRead, + cacheWrite: left.cost.cacheWrite + right.cost.cacheWrite, + total: left.cost.total + right.cost.total, + }, + }; +} + +export function readSessionStats(db: SqliteDatabase): SessionStats { + const row = readSingleSessionRow(db); + return { + messageCount: row.message_count, + usage: JSON.parse(row.usage_payload) as Usage, + }; +} + +export function incrementMessageCount(db: SqliteDatabase): void { + sql`UPDATE session SET message_count = message_count + 1`.run(db); +} + +export function addUsageToSessionStats(db: SqliteDatabase, usage: UsageRow["usage"]): void { + const current = readSessionStats(db).usage; + sql`UPDATE session SET usage_payload = ${JSON.stringify(addUsage(current, usage))}`.run(db); +} diff --git a/packages/session-backends/sqlite-node/src/sqlite/session/usage-ledger.ts b/packages/session-backends/sqlite-node/src/sqlite/session/usage-ledger.ts new file mode 100644 index 000000000..ecb0307d2 --- /dev/null +++ b/packages/session-backends/sqlite-node/src/sqlite/session/usage-ledger.ts @@ -0,0 +1,35 @@ +import type { UsageRow } from "@earendil-works/pi-agent-core"; +import { sql } from "../sql.ts"; +import type { SqliteDatabase } from "../types.ts"; + +export interface UsageLedgerRow { + id: string; + seq: number; + entry_id: string | null; + adjustment: number; + usage: string; + details: string | null; +} + +export function insertUsageLedgerRow(db: SqliteDatabase, row: UsageRow): void { + sql`INSERT INTO usage_ledger (id, seq, entry_id, adjustment, usage, details) + VALUES ( + ${row.id}, + ${row.seq}, + ${row.entryId ?? null}, + ${row.adjustment ? 1 : 0}, + ${JSON.stringify(row.usage)}, + ${row.details === undefined ? null : JSON.stringify(row.details)} + )`.run(db); +} + +export function decodeUsageLedgerRow(row: UsageLedgerRow): UsageRow { + return { + id: row.id, + seq: row.seq, + usage: JSON.parse(row.usage) as UsageRow["usage"], + ...(row.entry_id === null ? {} : { entryId: row.entry_id }), + adjustment: row.adjustment !== 0, + ...(row.details === null ? {} : { details: JSON.parse(row.details) as UsageRow["details"] }), + }; +}