sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fix(agent): remove obsolete SQLite session internals

Browse source
This commit is contained in:
Christian Klotz 2026-07-31 22:45:02 +03:00
parent fbab971da2
commit 9f4dfdddbf
No known key found for this signature in database
5 changed files with 14 additions and 205 deletions
Download patch file Download diff file Expand all files Collapse all files

8
packages/agent/test/harness/sqlite-migrations.test.ts
View file

@ -416,7 +416,7 @@ END;
await expect(repo.open(metadata)).rejects.toMatchObject({ code: "invalid_entry" });
});
it("restores in-memory state when appendEntry fails after mutating caches", async () => {
it("rolls back state when appendEntry fails after mutating caches", async () => {
const root = createTempDir();
const databasePath = join(root, "sessions.sqlite");
const sqlite = createNodeSqliteFactory();
@ -445,8 +445,10 @@ END;
message: createUserMessage("root"),
}),
).rejects.toMatchObject({ code: "storage" });
expect(await storage.getLeafId()).toBeNull();
expect(await storage.getEntry("root")).toBeUndefined();
const sessionRow = await db
.prepare("SELECT active_leaf_id FROM sessions WHERE id = ?")
.get<{ active_leaf_id: string | null }>("session-1");
expect(sessionRow?.active_leaf_id).toBeNull();
expect(await storage.getEntries()).toEqual([]);
await db.close();
});

58
packages/storage/sqlite-node/src/sqlite/storage/branch-entries.ts
View file

@ -1,58 +0,0 @@
import type { SessionTreeEntry } from "@earendil-works/pi-agent-core";
import { toError } from "@earendil-works/pi-agent-core";
import type { SqliteDatabase } from "../types.ts";
import { decodeEntry, type SessionEntryRow } from "./session-entries.ts";
import { invalidSession } from "./shared.ts";
export interface BranchEntryRow {
entry_id: string;
entry_seq: number;
}
export async function getMaterializedBranchPathOrCompaction(
db: SqliteDatabase,
sessionId: string,
branchId: string,
byId: Map<string, SessionTreeEntry>,
): Promise<SessionTreeEntry[]> {
const branchRows = await db
.prepare(
"SELECT entry_id, entry_seq FROM branch_entries WHERE session_id = ? AND branch_id = ? ORDER BY entry_seq",
)
.all<BranchEntryRow>(sessionId, branchId);
if (branchRows.length === 0) {
return [];
}
const entryIds = branchRows.map((row) => row.entry_id);
const placeholders = entryIds.map(() => "?").join(", ");
const entryRows = await db
.prepare(
`SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND id IN (${placeholders})`,
)
.all<SessionEntryRow>(sessionId, ...entryIds);
const entryRowsById = new Map(entryRows.map((row) => [row.id, row]));
const entries: SessionTreeEntry[] = [];
for (const branchRow of branchRows) {
// leaf entries are navigation markers used to mark which branch became active;
// they are not part of the model/context path reconstructed from branch_entries.
const cached = byId.get(branchRow.entry_id);
if (cached) {
if (cached.type !== "leaf") {
entries.push(cached);
}
continue;
}
const entryRow = entryRowsById.get(branchRow.entry_id);
if (!entryRow) throw invalidSession(`missing entry row for branch entry ${branchRow.entry_id}`);
try {
const entry = decodeEntry(entryRow);
byId.set(entry.id, entry);
if (entry.type !== "leaf") {
entries.push(entry);
}
} catch (error) {
throw invalidSession(`invalid entry row for branch entry ${branchRow.entry_id}`, toError(error));
}
}
return entries;
}

126
packages/storage/sqlite-node/src/sqlite/storage/index.ts
View file

@ -1,8 +1,7 @@
import type { LeafEntry, SessionEntryCursorOptions, SessionTreeEntry } from "@earendil-works/pi-agent-core";
import type { SessionTreeEntry } from "@earendil-works/pi-agent-core";
import { SessionError, toError } from "@earendil-works/pi-agent-core";
import { uuidv7 } from "@earendil-works/pi-ai";
import type { SqliteDatabase, SqliteSessionMetadata } from "../types.ts";
import { getMaterializedBranchPathOrCompaction } from "./branch-entries.ts";
import { decodeEntry, encodeEntry, type SessionEntryRow } from "./session-entries.ts";
import {
applyEntryToMaterializedState,
@ -14,28 +13,22 @@ import {
type SessionMaterializedRow,
type SessionMaterializedState,
serializeSummary,
sessionStatsFromMaterializedState,
} from "./session-materialized.ts";
import { advanceSequence, getNextSequence } from "./session-sequences.ts";
import { rowToMetadata, type SessionRow } from "./sessions.ts";
import { generateEntryId, invalidEntry, invalidSession, leafIdAfterEntry } from "./shared.ts";
import { invalidEntry, invalidSession, leafIdAfterEntry } from "./shared.ts";
async function decodeEntryRows(entryRows: SessionEntryRow[]): Promise<{
entries: SessionTreeEntry[];
leafId: string | null;
}> {
function decodeEntryRows(entryRows: SessionEntryRow[]): SessionTreeEntry[] {
const entries: SessionTreeEntry[] = [];
let leafId: string | null = null;
for (const entryRow of entryRows) {
try {
const entry = decodeEntry(entryRow);
entries.push(entry);
leafId = leafIdAfterEntry(entry);
} catch (error) {
throw invalidEntry(`failed to decode entry ${entryRow.id}`, toError(error));
}
}
return { entries, leafId };
return entries;
}
async function loadEntryRowsByIds(
@ -112,7 +105,6 @@ export class SqliteSessionConnection {
private readonly db: SqliteDatabase;
private readonly metadata: SqliteSessionMetadata;
private byId: Map<string, SessionTreeEntry>;
private labelsById: Map<string, string>;
private currentLeafId: string | null;
private activeBranchId: string | null;
private materializedState: SessionMaterializedState;
@ -180,16 +172,14 @@ export class SqliteSessionConnection {
private constructor(
db: SqliteDatabase,
metadata: SqliteSessionMetadata,
entries: SessionTreeEntry[] | null,
leafId: string | null,
activeBranchId: string | null,
materializedState: SessionMaterializedState,
) {
this.db = db;
this.metadata = metadata;
this.byId = new Map((entries ?? []).map((entry) => [entry.id, entry]));
this.byId = new Map<string, SessionTreeEntry>();
this.materializedState = materializedState;
this.labelsById = materializedState.labelsById;
this.currentLeafId = leafId;
this.activeBranchId = activeBranchId;
}
@ -199,7 +189,6 @@ export class SqliteSessionConnection {
return new SqliteSessionConnection(
db,
rowToMetadata(loaded.row, metadata.path),
null,
loaded.leafId,
loaded.activeBranchId,
loaded.materializedState,
@ -243,7 +232,6 @@ export class SqliteSessionConnection {
parentSessionId: options.parentSessionId,
metadata: options.metadata,
},
[],
null,
null,
createEmptyMaterializedState(),
@ -254,36 +242,6 @@ export class SqliteSessionConnection {
return this.metadata;
}
async getLeafId(): Promise<string | null> {
return this.currentLeafId;
}
async setLeafId(leafId: string | null): Promise<LeafEntry> {
if (leafId !== null && !(await this.getEntry(leafId))) {
throw new SessionError("not_found", `Entry ${leafId} not found`);
}
const entry: LeafEntry = {
type: "leaf",
id: await this.createEntryId(),
parentId: this.currentLeafId,
timestamp: new Date().toISOString(),
targetId: leafId,
};
await this.appendEntry(entry);
return entry;
}
async createEntryId(): Promise<string> {
for (let i = 0; i < 100; i++) {
const id = generateEntryId(this.byId);
const existing = await this.db
.prepare("SELECT 1 AS found FROM session_entries WHERE session_id = ? AND id = ? LIMIT 1")
.get<{ found: number }>(this.metadata.id, id);
if (!existing) return id;
}
return uuidv7();
}
async appendEntry(entry: SessionTreeEntry, options: { transaction?: boolean } = {}): Promise<void> {
const encoded = encodeEntry(entry);
const previousMaterializedState: SessionMaterializedState = {
@ -334,7 +292,6 @@ export class SqliteSessionConnection {
else await this.db.transaction(write);
} catch (error) {
this.materializedState = previousMaterializedState;
this.labelsById = previousMaterializedState.labelsById;
this.byId = previousById;
this.currentLeafId = previousLeafId;
this.activeBranchId = previousActiveBranchId;
@ -343,7 +300,7 @@ export class SqliteSessionConnection {
}
}
async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
private async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
const cached = this.byId.get(id);
if (cached) return cached;
const row = await this.db
@ -361,80 +318,13 @@ export class SqliteSessionConnection {
}
}
async findEntries<TType extends SessionTreeEntry["type"]>(
type: TType,
): Promise<Array<Extract<SessionTreeEntry, { type: TType }>>> {
const rows = await this.db
.prepare(
"SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND type = ? ORDER BY entry_seq",
)
.all<SessionEntryRow>(this.metadata.id, type);
const entries: Array<Extract<SessionTreeEntry, { type: TType }>> = [];
for (const row of rows) {
try {
const entry = decodeEntry(row) as Extract<SessionTreeEntry, { type: TType }>;
this.byId.set(entry.id, entry);
entries.push(entry);
} catch (error) {
throw invalidEntry(`failed to decode entry ${row.id}`, toError(error));
}
}
return entries;
}
async getLabel(id: string): Promise<string | undefined> {
return this.labelsById.get(id);
}
async getSessionName(): Promise<string | undefined> {
return this.materializedState.name;
}
async getSessionStats() {
return sessionStatsFromMaterializedState(this.materializedState);
}
async getPathToRootOrCompaction(leafId: string | null): Promise<SessionTreeEntry[]> {
if (leafId === null) return [];
if (leafId === this.currentLeafId) {
if (!this.activeBranchId) {
throw invalidSession(`missing active branch for session ${this.metadata.id} leaf ${leafId}`);
}
return getMaterializedBranchPathOrCompaction(this.db, this.metadata.id, this.activeBranchId, this.byId);
}
return this.getPathToRootOrCompactionEntries(leafId);
}
async getEntries(options?: SessionEntryCursorOptions): Promise<SessionTreeEntry[]> {
const limit = options?.limit;
if (limit !== undefined) {
const beforeOrAtEntrySeq =
options?.afterEntrySeq ??
(
await this.db
.prepare("SELECT entry_seq FROM session_entries WHERE session_id = ? ORDER BY entry_seq DESC LIMIT 1")
.get<{ entry_seq: number }>(this.metadata.id)
)?.entry_seq;
if (beforeOrAtEntrySeq === undefined) {
return [];
}
const rows = await this.db
.prepare(
"SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? AND entry_seq <= ? ORDER BY entry_seq DESC LIMIT ?",
)
.all<SessionEntryRow>(this.metadata.id, beforeOrAtEntrySeq, limit);
const entries = (await decodeEntryRows(rows)).entries;
for (const entry of entries) {
this.byId.set(entry.id, entry);
}
return entries.reverse();
}
async getEntries(): Promise<SessionTreeEntry[]> {
const rows = await this.db
.prepare(
"SELECT session_id, id, entry_seq, parent_id, type, timestamp, payload FROM session_entries WHERE session_id = ? ORDER BY entry_seq",
)
.all<SessionEntryRow>(this.metadata.id);
const entries = (await decodeEntryRows(rows)).entries;
const entries = decodeEntryRows(rows);
for (const entry of entries) {
this.byId.set(entry.id, entry);
}

15
packages/storage/sqlite-node/src/sqlite/storage/session-materialized.ts
View file

@ -1,4 +1,4 @@
import type { SessionStats, SessionTreeEntry, ThinkingLevel } from "@earendil-works/pi-agent-core";
import type { SessionTreeEntry, ThinkingLevel } from "@earendil-works/pi-agent-core";
import { invalidSession, isRecord } from "./shared.ts";
export interface SessionMaterializedRow {
@ -311,24 +311,11 @@ export function materializedStateFromRows(
} else {
state.labelsById.delete(payload.targetId);
}
continue;
}
if (row.type !== "label") {
}
}
return state;
}
export function sessionStatsFromMaterializedState(state: SessionMaterializedState): SessionStats {
return {
messageCount: state.messageCount,
cachedTokens: state.cachedTokens,
uncachedTokens: state.uncachedTokens,
totalTokens: state.totalTokens,
costTotal: state.costTotal,
};
}
export function materializedStateValues(
sessionId: string,
state: SessionMaterializedState,

12
packages/storage/sqlite-node/src/sqlite/storage/shared.ts
View file

@ -1,17 +1,5 @@
import type { SessionTreeEntry } from "@earendil-works/pi-agent-core";
import { SessionError } from "@earendil-works/pi-agent-core";
import { uuidv7 } from "@earendil-works/pi-ai";
export function generateEntryId(byId: { has(id: string): boolean }): string {
for (let i = 0; i < 100; i++) {
// The uuidv7 prefix is timestamp-derived and nearly constant between calls,
// so short ids must come from the random tail.
const id = uuidv7().slice(0, 8);
if (!byId.has(id)) return id;
}
return uuidv7();
}
export function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null;
}