import { mkdtemp, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import type { AgentMessage } from "@mariozechner/pi-agent-core"; import { AppDatabase } from "../../src/main/database"; import { SqliteSessionStore } from "../../src/main/agent-runtime/sqlite-session-store"; /** * Verifies the SqliteSessionStore round-trips pi `AgentMessage` values * correctly through the `session_entries` table. */ describe("SqliteSessionStore", () => { let tempDir: string; let db: AppDatabase; const TASK_SESSION_ID = "task-session-test-1"; beforeEach(async () => { tempDir = await mkdtemp(join(tmpdir(), "agentcompany-session-store-")); db = new AppDatabase(tempDir); db.init(); }); afterEach(async () => { db.close(); await rm(tempDir, { recursive: true, force: true }); }); it("appends and reloads pi user/assistant messages preserving order", () => { const store = new SqliteSessionStore(db, TASK_SESSION_ID); const userMessage: AgentMessage = { role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now(), } as AgentMessage; const assistantMessage: AgentMessage = { role: "assistant", content: [{ type: "text", text: "hi" }], stopReason: "stop", provider: "anthropic", model: { id: "claude-opus-4-5" }, usage: { input: 10, output: 5, cacheRead: 0, cacheWrite: 0, totalTokens: 15, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } }, timestamp: Date.now(), } as unknown as AgentMessage; store.appendMessage(userMessage); store.appendMessage(assistantMessage); const loaded = store.loadMessages(); expect(loaded).toHaveLength(2); expect((loaded[0] as { role: string }).role).toBe("user"); expect((loaded[1] as { role: string }).role).toBe("assistant"); }); it("writes schema versioned envelopes for new messages", () => { const store = new SqliteSessionStore(db, TASK_SESSION_ID); store.appendMessage({ role: "user", content: [{ type: "text", text: "hello" }], timestamp: 1, } as AgentMessage); const row = db.projection.sessionEntries.listActive(TASK_SESSION_ID)[0]; const payload = JSON.parse(row?.payloadJson ?? "{}") as { schemaVersion?: number; kind?: string; message?: { role?: string; timestamp?: number }; }; expect(payload.schemaVersion).toBe(1); expect(payload.kind).toBe("agent_message"); expect(payload.message).toMatchObject({ role: "user", timestamp: 1 }); }); it("reads legacy raw AgentMessage payloads", () => { db.projection.sessionEntries.append({ taskSessionId: TASK_SESSION_ID, kind: "message", role: "user", payloadJson: JSON.stringify({ role: "user", content: [{ type: "text", text: "legacy" }], }), }); const store = new SqliteSessionStore(db, TASK_SESSION_ID); const loaded = store.loadMessages(); expect(loaded).toHaveLength(1); expect((loaded[0] as { role?: string }).role).toBe("user"); expect((loaded[0] as { timestamp?: number }).timestamp).toEqual(expect.any(Number)); }); it("rejects invalid message shapes before writing session_entries", () => { const store = new SqliteSessionStore(db, TASK_SESSION_ID); expect(() => store.appendMessage({ role: "bad-role", content: [{ type: "text", text: "invalid" }], } as unknown as AgentMessage)).toThrow('invalid role "bad-role"'); expect(db.projection.sessionEntries.listActive(TASK_SESSION_ID)).toHaveLength(0); }); it("rejects new messages without a finite timestamp before writing session_entries", () => { const store = new SqliteSessionStore(db, TASK_SESSION_ID); expect(() => store.appendMessage({ role: "user", content: [{ type: "text", text: "missing timestamp" }], } as AgentMessage)).toThrow("timestamp must be a finite number"); expect(() => store.appendMessage({ role: "user", content: [{ type: "text", text: "bad timestamp" }], timestamp: Number.NaN, } as AgentMessage)).toThrow("timestamp must be a finite number"); expect(db.projection.sessionEntries.listActive(TASK_SESSION_ID)).toHaveLength(0); }); it("rejects new messages with invalid content before writing session_entries", () => { const store = new SqliteSessionStore(db, TASK_SESSION_ID); expect(() => store.appendMessage({ role: "user", content: null, timestamp: 1, } as unknown as AgentMessage)).toThrow("content must be a string or an array"); expect(db.projection.sessionEntries.listActive(TASK_SESSION_ID)).toHaveLength(0); }); it("fails reads with the corrupt row id when payload JSON is invalid", () => { db.projection.sessionEntries.append({ taskSessionId: TASK_SESSION_ID, kind: "message", role: "user", payloadJson: "{", }); const row = db.projection.sessionEntries.listActive(TASK_SESSION_ID)[0]; const store = new SqliteSessionStore(db, TASK_SESSION_ID); expect(() => store.loadMessages()).toThrow(`Invalid session entry ${row?.id} payload JSON`); }); it("fails reads with the corrupt row id when payload shape drifts", () => { db.projection.sessionEntries.append({ taskSessionId: TASK_SESSION_ID, kind: "message", role: null, payloadJson: JSON.stringify({ schemaVersion: 1, kind: "agent_message", message: { role: "assistant" }, persistedAt: new Date().toISOString() }), }); const row = db.projection.sessionEntries.listActive(TASK_SESSION_ID)[0]; const store = new SqliteSessionStore(db, TASK_SESSION_ID); expect(() => store.loadMessages()).toThrow(`Invalid session entry ${row?.id} AgentMessage`); }); it("durably compacts oversized active transcripts into a summary plus recent rows", () => { const store = new SqliteSessionStore(db, TASK_SESSION_ID); for (let i = 0; i < 46; i += 1) { const role = i % 2 === 0 ? "user" : "assistant"; store.appendMessage({ role, content: [{ type: "text", text: `${role} turn ${i}\n${"context ".repeat(6_000)}`, }], timestamp: i, } as AgentMessage); } const loaded = store.loadMessages(); const activeRows = db.projection.sessionEntries.listActive(TASK_SESSION_ID); const superseded = db.queryAll<{ count: number }>( "select count(*) as count from session_entries where task_session_id = ? and is_superseded = 1", TASK_SESSION_ID, )[0]?.count ?? 0; expect(loaded.length).toBeLessThan(46); expect(activeRows).toHaveLength(loaded.length); expect(activeRows[0]?.sequence).toBe(1); expect(activeRows[0]?.role).toBe("user"); expect(activeRows.slice(1).every((row, index, rows) => index === 0 || row.sequence > (rows[index - 1]?.sequence ?? 0))).toBe(true); expect(superseded).toBeGreaterThan(0); expect((loaded[0] as { role?: string }).role).toBe("user"); const firstContent = (loaded[0] as { content?: Array<{ text?: string }> }).content; expect(firstContent?.[0]?.text).toContain("Context summary"); store.appendMessage({ role: "user", content: [{ type: "text", text: "continue" }], timestamp: 10_000, } as AgentMessage); const appended = db.projection.sessionEntries.listActive(TASK_SESSION_ID).at(-1); const allMaxSequence = db.queryAll<{ maxSeq: number }>( "select max(sequence) as maxSeq from session_entries where task_session_id = ?", TASK_SESSION_ID, )[0]?.maxSeq; expect(appended?.sequence).toBe(47); expect(allMaxSequence).toBe(47); }); });