Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
19 kB · 564 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565import { appendFileSync, existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, symlinkSync, unlinkSync, writeFileSync,} from "node:fs";import { tmpdir } from "node:os";import { dirname, join } from "node:path";import type { Usage } from "@earendil-works/pi-ai";import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";import { openSessionCatalogStore, SESSION_CATALOG_DB_FILENAME } from "../src/core/session-catalog-db.js";import { readSessionEntryStructure } from "../src/core/session-entry-structure.js";import { CURRENT_SESSION_VERSION, closeSessionCatalogStores, forgetSessionScanStates, type SessionHeader, SessionManager,} from "../src/core/session-manager.js";import type * as fileLines from "../src/utils/file-lines.js";
interface ReadRange { start: number; end: number;}
const readRanges: ReadRange[] = [];
vi.mock("../src/utils/file-lines.js", async () => { const actual = await vi.importActual<typeof fileLines>("../src/utils/file-lines.js"); return { ...actual, readLinesAsBuffers: (filePath: string, range?: fileLines.ReadLinesRange) => { readRanges.push({ start: range?.start ?? 0, end: range?.end ?? Number.POSITIVE_INFINITY }); return actual.readLinesAsBuffers(filePath, range); }, };});
interface IndexedLine { id: string; parentId: string | null; type: string; role?: string | null; offset: number; length: number;}
/** Independent parse of the file, used to falsify the index rather than trust it. */function referenceIndex(filePath: string): IndexedLine[] { const lines = readFileSync(filePath, "utf8") .split("\n") .filter((line) => line.length > 0); const rows: IndexedLine[] = []; let offset = 0; for (const line of lines) { const length = Buffer.byteLength(line); const entry = JSON.parse(line) as { id?: string; parentId?: string | null; type: string; message?: { role?: string }; }; if (entry.type !== "session" && entry.id) { rows.push({ id: entry.id, parentId: entry.parentId ?? null, type: entry.type, role: entry.message?.role ?? null, offset, length, }); } offset += length + 1; } return rows;}
function indexedLines(filePath: string): IndexedLine[] { const store = openSessionCatalogStore(dirname(filePath)); if (!store) return []; const structure = store.readEntryStructure(filePath); store.close(); if (!structure?.blob) return []; return readSessionEntryStructure(structure.blob).map((record) => ({ id: record.id, parentId: record.parentId, type: record.type, role: record.role, offset: record.offset, length: record.length, }));}
const dirs: string[] = [];
function createDir(): string { const dir = mkdtempSync(join(tmpdir(), "pi-catalog-scan-")); dirs.push(dir); return dir;}
function usage(input: number, output: number): Usage { return { input, output, cacheRead: 0, cacheWrite: 0, totalTokens: input + output, cost: { input: 0.01, output: 0.02, cacheRead: 0, cacheWrite: 0, total: 0.03 }, };}
function writeSession(filePath: string, id: string, createdAt = Date.parse("2026-01-01T00:00:00.000Z")): void { const header: SessionHeader = { type: "session", version: CURRENT_SESSION_VERSION, id, timestamp: new Date(createdAt).toISOString(), cwd: "/tmp", }; writeFileSync(filePath, `${JSON.stringify(header)}\n`);}
function messageLine(id: string, role: "user" | "assistant", text: string, timestamp: number): string { const message = role === "user" ? { role, content: [{ type: "text", text }], timestamp } : { role, content: [{ type: "text", text }], api: "openai-completions", provider: "openai", model: "test-model", usage: usage(10, 5), stopReason: "stop", timestamp, }; return JSON.stringify({ type: "message", id, parentId: null, timestamp: new Date(timestamp).toISOString(), message, });}
function appendLines(filePath: string, lines: string[]): void { appendFileSync(filePath, `${lines.join("\n")}\n`);}
function createSession( filePath: string, id: string, messages: Array<{ role: "user" | "assistant"; text: string; ts: number }>,): void { writeSession(filePath, id); appendLines( filePath, messages.map((entry, index) => messageLine(`${id}-m${index}`, entry.role, entry.text, entry.ts)), );}
/** * Real transcripts number entries with `randomUUID().slice(0, 8)`, which the * packed structure stores as four bytes; fixtures that exercise the structure * must use the same id shape. */function entryId(index: number): string { return (index + 1).toString(16).padStart(8, "0");}
function createIndexedSession( filePath: string, sessionId: string, messages: Array<{ role: "user" | "assistant"; text: string; ts: number }>,): string[] { writeSession(filePath, sessionId); const ids = messages.map((_, index) => entryId(index)); appendLines( filePath, messages.map((entry, index) => messageLine(ids[index]!, entry.role, entry.text, entry.ts)), ); return ids;}
/** * A listing answered from the database persists no per-message search text: it * reports identity, the first message, and no message corpus. Live scans keep it. */function withoutSearchText<T extends { allMessagesText: string }>(infos: T[]): T[] { return infos.map((info) => ({ ...info, allMessagesText: "" }));}
/** Simulates a fresh process: no in-process scan state, no open store. */function restart(): void { forgetSessionScanStates(); closeSessionCatalogStores(); readRanges.length = 0;}
afterEach(() => { delete process.env.PRIME_AGENT_SESSION_CATALOG_DB; restart(); for (const dir of dirs.splice(0)) { rmSync(dir, { recursive: true, force: true }); }});
beforeEach(() => { readRanges.length = 0;});
describe("session catalog scan", () => { it("answers a repeated list from the persisted row without reading the file", async () => { const dir = createDir(); const file = join(dir, "a.jsonl"); createSession(file, "session-a", [ { role: "user", text: "hello world", ts: Date.parse("2026-01-02T00:00:00.000Z") }, { role: "assistant", text: "hi there", ts: Date.parse("2026-01-02T00:00:01.000Z") }, ]);
const first = await SessionManager.listAll(undefined, dir); expect(first).toHaveLength(1); expect(readRanges.length).toBeGreaterThan(0);
restart(); const second = await SessionManager.listAll(undefined, dir); expect(second).toEqual(withoutSearchText(first)); expect(first[0]!.allMessagesText).toBe("hello world hi there"); expect(second[0]!.allMessagesText).toBe(""); expect(readRanges).toEqual([]); expect(existsSync(join(dir, SESSION_CATALOG_DB_FILENAME))).toBe(true); });
it("resumes an appended session from the persisted offset", async () => { const dir = createDir(); const file = join(dir, "a.jsonl"); createSession(file, "session-a", [ { role: "user", text: "first question", ts: Date.parse("2026-01-02T00:00:00.000Z") }, { role: "assistant", text: "first answer", ts: Date.parse("2026-01-02T00:00:01.000Z") }, ]); const first = await SessionManager.listAll(undefined, dir); expect(first).toHaveLength(1); const store = openSessionCatalogStore(dir)!; const firstRowSize = store.read(file)!.size; store.close();
restart(); appendLines(file, [ messageLine("session-a-m2", "user", "second question", Date.parse("2026-01-03T00:00:00.000Z")), messageLine("session-a-m3", "assistant", "second answer", Date.parse("2026-01-03T00:00:01.000Z")), ]);
const second = await SessionManager.listAll(undefined, dir); expect(second).toHaveLength(1); expect(second[0]!.messageCount).toBe(4); expect(second[0]!.firstMessage).toBe("first question"); expect(second[0]!.modified.getTime()).toBe(Date.parse("2026-01-03T00:00:01.000Z")); expect(second[0]!.usage).toEqual({ inputTokens: 20, outputTokens: 10, cost: 0.06 }); // Only the appended bytes were read, never the consumed prefix. expect(readRanges).toHaveLength(1); expect(readRanges[0]!.start).toBe(firstRowSize); });
it("rescans a rewritten file instead of trusting the cached summary", async () => { const dir = createDir(); const file = join(dir, "a.jsonl"); createSession(file, "session-a", [ { role: "user", text: "original question", ts: Date.parse("2026-01-02T00:00:00.000Z") }, { role: "assistant", text: "original answer", ts: Date.parse("2026-01-02T00:00:01.000Z") }, ]); await SessionManager.listAll(undefined, dir);
restart(); createSession(file, "session-a", [ { role: "user", text: "rewritten question", ts: Date.parse("2026-01-04T00:00:00.000Z") }, ]);
const listed = await SessionManager.listAll(undefined, dir); expect(listed).toHaveLength(1); expect(listed[0]!.messageCount).toBe(1); expect(listed[0]!.firstMessage).toBe("rewritten question"); expect(listed[0]!.modified.getTime()).toBe(Date.parse("2026-01-04T00:00:00.000Z")); });
it("drops rows for deleted files", async () => { const dir = createDir(); const kept = join(dir, "kept.jsonl"); const gone = join(dir, "gone.jsonl"); createSession(kept, "session-kept", [{ role: "user", text: "kept", ts: Date.parse("2026-01-02T00:00:00.000Z") }]); createSession(gone, "session-gone", [{ role: "user", text: "gone", ts: Date.parse("2026-01-02T00:00:00.000Z") }]); await SessionManager.listAll(undefined, dir);
restart(); unlinkSync(gone); const listed = await SessionManager.listAll(undefined, dir); expect(listed.map((session) => session.path)).toEqual([kept]);
const store = openSessionCatalogStore(dir)!; expect(store.read(gone)).toBeUndefined(); expect(store.read(kept)).toBeDefined(); store.close(); });
it("keeps unscannable files unlisted without rescanning them", async () => { const dir = createDir(); const file = join(dir, "broken.jsonl"); appendLines(file, [messageLine("orphan", "user", "no header", Date.parse("2026-01-02T00:00:00.000Z"))]);
expect(await SessionManager.listAll(undefined, dir)).toEqual([]); restart(); expect(await SessionManager.listAll(undefined, dir)).toEqual([]); expect(readRanges).toEqual([]); });
it("matches a full scan field for field", async () => { const dir = createDir(); const first = join(dir, "a.jsonl"); const second = join(dir, "b.jsonl"); createSession(first, "session-a", [ { role: "user", text: "alpha question", ts: Date.parse("2026-01-02T00:00:00.000Z") }, { role: "assistant", text: "alpha answer", ts: Date.parse("2026-01-02T00:00:01.000Z") }, ]); createSession(second, "session-b", [ { role: "user", text: "beta question", ts: Date.parse("2026-01-05T00:00:00.000Z") }, ]); appendLines(first, [ JSON.stringify({ type: "session_info", id: "rename", parentId: null, timestamp: new Date().toISOString(), name: "renamed session", }), JSON.stringify({ type: "session_state", id: "archive", parentId: null, timestamp: new Date().toISOString(), state: { status: "archived" }, }), ]);
const cold = await SessionManager.listAll(undefined, dir); expect(cold).toHaveLength(2); expect(cold.find((session) => session.path === first)?.name).toBe("renamed session"); expect(cold.find((session) => session.path === first)?.state).toEqual({ status: "archived" });
restart(); const cached = await SessionManager.listAll(undefined, dir); expect(cached).toEqual(withoutSearchText(cold)); expect(readRanges).toEqual([]);
closeSessionCatalogStores(); rmSync(join(dir, SESSION_CATALOG_DB_FILENAME), { force: true }); const rescan = await SessionManager.listAll(undefined, dir); expect(rescan).toEqual(cold); });
it("picks up a session directory created after the first list", async () => { const dir = join(createDir(), "late"); expect(await SessionManager.listAll(undefined, dir)).toEqual([]);
mkdirSync(dir, { recursive: true }); const file = join(dir, "a.jsonl"); createSession(file, "session-a", [ { role: "user", text: "late question", ts: Date.parse("2026-01-02T00:00:00.000Z") }, ]); const listed = await SessionManager.listAll(undefined, dir);
expect(listed).toHaveLength(1); expect(listed[0]!.path).toBe(file); expect(existsSync(join(dir, SESSION_CATALOG_DB_FILENAME))).toBe(true); });
it("indexes one row per file when the directory is reached through a symlink", async () => { const parent = createDir(); const dir = join(parent, "real"); mkdirSync(dir, { recursive: true }); const alias = join(parent, "alias"); symlinkSync(dir, alias); const file = join(dir, "a.jsonl"); createSession(file, "session-a", [ { role: "user", text: "symlinked question", ts: Date.parse("2026-01-02T00:00:00.000Z") }, ]);
expect(await SessionManager.listAll(undefined, dir)).toHaveLength(1); restart(); const throughAlias = await SessionManager.listAll(undefined, alias); expect(throughAlias).toHaveLength(1); expect(readRanges).toEqual([]); });
it("indexes entry structure with byte ranges that resolve to the same entries", async () => { const dir = createDir(); const file = join(dir, "a.jsonl"); const ids = createIndexedSession(file, "session-a", [ { role: "user", text: "first question", ts: Date.parse("2026-01-02T00:00:00.000Z") }, { role: "assistant", text: "first answer", ts: Date.parse("2026-01-02T00:00:01.000Z") }, ]); appendLines(file, [ JSON.stringify({ type: "session_info", id: entryId(2), parentId: ids[1], timestamp: new Date().toISOString(), name: "renamed", }), JSON.stringify({ type: "session_state", id: entryId(3), parentId: entryId(2), timestamp: new Date().toISOString(), state: { status: "archived" }, }), ]); await SessionManager.listAll(undefined, dir);
const rows = indexedLines(file); expect(rows).toEqual(referenceIndex(file)); const bytes = readFileSync(file); for (const row of rows) { const entry = JSON.parse(bytes.subarray(row.offset, row.offset + row.length).toString("utf8")) as { id: string; }; expect(entry.id).toBe(row.id); } });
it("indexes only the appended bytes after a resume", async () => { const dir = createDir(); const file = join(dir, "a.jsonl"); createIndexedSession(file, "session-a", [ { role: "user", text: "first", ts: Date.parse("2026-01-02T00:00:00.000Z") }, ]); await SessionManager.listAll(undefined, dir); const before = indexedLines(file); const sizeBefore = readFileSync(file).length;
restart(); appendLines(file, [messageLine(entryId(1), "assistant", "second", Date.parse("2026-01-02T00:00:01.000Z"))]); await SessionManager.listAll(undefined, dir);
const after = indexedLines(file); expect(after.length).toBe(before.length + 1); expect(after.slice(0, before.length)).toEqual(before); expect(after[after.length - 1]!.offset).toBeGreaterThanOrEqual(sizeBefore); });
it("indexes an entry too large to parse without reading its payload", async () => { const dir = createDir(); const file = join(dir, "a.jsonl"); writeSession(file, "session-a"); const big = "x".repeat(1_100_000); appendLines(file, [ JSON.stringify({ type: "message", id: entryId(0), parentId: null, timestamp: new Date(Date.parse("2026-01-02T00:00:00.000Z")).toISOString(), message: { role: "toolResult", toolCallId: "call-1", toolName: "bash", content: [{ type: "text", text: big }], timestamp: Date.parse("2026-01-02T00:00:00.000Z"), }, }), ]); await SessionManager.listAll(undefined, dir);
const rows = indexedLines(file); expect(rows).toHaveLength(1); expect(rows[0]).toMatchObject({ id: entryId(0), parentId: null, type: "message", role: "toolResult", offset: 100, }); expect(rows[0]!.length).toBeGreaterThan(big.length); const bytes = readFileSync(file); const entry = JSON.parse(bytes.subarray(rows[0]!.offset, rows[0]!.offset + rows[0]!.length).toString("utf8")) as { id: string; message: { role: string; content: Array<{ text: string }> }; }; expect(entry.id).toBe(entryId(0)); expect(entry.message.role).toBe("toolResult"); expect(entry.message.content[0]!.text.length).toBe(big.length); });
it("replaces indexed entries when the file is rewritten shorter", async () => { const dir = createDir(); const file = join(dir, "a.jsonl"); createIndexedSession(file, "session-a", [ { role: "user", text: "one", ts: Date.parse("2026-01-02T00:00:00.000Z") }, { role: "assistant", text: "two", ts: Date.parse("2026-01-02T00:00:01.000Z") }, { role: "user", text: "three", ts: Date.parse("2026-01-02T00:00:02.000Z") }, ]); await SessionManager.listAll(undefined, dir); expect(indexedLines(file)).toHaveLength(3);
restart(); createIndexedSession(file, "session-a", [ { role: "user", text: "rewritten", ts: Date.parse("2026-01-03T00:00:00.000Z") }, ]); await SessionManager.listAll(undefined, dir);
expect(indexedLines(file)).toEqual(referenceIndex(file)); expect(indexedLines(file)).toHaveLength(1); });
it("drops indexed entries with the summary row when the file is gone", async () => { const dir = createDir(); const file = join(dir, "a.jsonl"); createIndexedSession(file, "session-a", [ { role: "user", text: "gone", ts: Date.parse("2026-01-02T00:00:00.000Z") }, ]); await SessionManager.listAll(undefined, dir); expect(indexedLines(file)).toHaveLength(1);
restart(); unlinkSync(file); await SessionManager.listAll(undefined, dir); expect(indexedLines(file)).toEqual([]); });
it("lists a session whose entry ids cannot be packed without an index", async () => { const dir = createDir(); const file = join(dir, "a.jsonl"); createSession(file, "session-a", [ { role: "user", text: "legacy id", ts: Date.parse("2026-01-02T00:00:00.000Z") }, { role: "assistant", text: "second", ts: Date.parse("2026-01-02T00:00:01.000Z") }, ]); const listed = await SessionManager.listAll(undefined, dir);
expect(listed.map((info) => info.id)).toEqual(["session-a"]); expect(listed[0]!.messageCount).toBe(2); const structure = indexedLines(file); expect(structure).toEqual([]); const store = openSessionCatalogStore(dir)!; expect(store.readEntryStructure(file)?.blob ?? null).toBeNull(); store.close(); });
it("does not create a database when the catalog is disabled", async () => { const dir = createDir(); const file = join(dir, "a.jsonl"); createSession(file, "session-a", [ { role: "user", text: "question", ts: Date.parse("2026-01-02T00:00:00.000Z") }, ]); process.env.PRIME_AGENT_SESSION_CATALOG_DB = "0";
const listed = await SessionManager.listAll(undefined, dir); expect(listed).toHaveLength(1); expect(existsSync(join(dir, SESSION_CATALOG_DB_FILENAME))).toBe(false); });});