From 84c00e0011cb04f23b0955e7236f20d8222ab150 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Thu, 23 Apr 2026 21:21:59 -0400 Subject: [PATCH] fix: validate and compact session entries --- .../agent-runtime/agent-message-schema.ts | 100 ++++++++++++ .../agent-runtime/sqlite-session-store.ts | 41 +++-- src/main/database/session-entries.ts | 55 ++++++- tests/unit/sqlite-session-store.test.ts | 146 ++++++++++++++++++ 4 files changed, 329 insertions(+), 13 deletions(-) create mode 100644 src/main/agent-runtime/agent-message-schema.ts diff --git a/src/main/agent-runtime/agent-message-schema.ts b/src/main/agent-runtime/agent-message-schema.ts new file mode 100644 index 00000000..0f7407ac --- /dev/null +++ b/src/main/agent-runtime/agent-message-schema.ts @@ -0,0 +1,100 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; + +const AGENT_MESSAGE_SCHEMA_VERSION = 1; +const ALLOWED_AGENT_MESSAGE_ROLES = new Set(["user", "assistant", "toolResult"]); + +interface PersistedAgentMessageEnvelope { + schemaVersion: 1; + kind: "agent_message"; + message: AgentMessage; + persistedAt: string; +} + +interface SessionEntryPayloadRow { + id: string; + payloadJson: string; +} + +function formatMessageError(prefix: string, message: string): Error { + return new Error(`${prefix}: ${message}`); +} + +function assertRecord(value: unknown, prefix: string): asserts value is Record { + if (!value || typeof value !== "object" || Array.isArray(value)) { + throw formatMessageError(prefix, "expected an object"); + } +} + +function normalizeAgentMessageForSession( + input: AgentMessage, + prefix = "Invalid AgentMessage", + options: { requireTimestamp: boolean } = { requireTimestamp: true }, +): AgentMessage { + assertRecord(input, prefix); + const role = input.role; + if (typeof role !== "string" || !ALLOWED_AGENT_MESSAGE_ROLES.has(role)) { + throw formatMessageError(prefix, `invalid role ${typeof role === "string" ? JSON.stringify(role) : ""}`); + } + const content = (input as { content?: unknown }).content; + if (typeof content !== "string" && !Array.isArray(content)) { + throw formatMessageError(prefix, "content must be a string or an array"); + } + const timestamp = (input as { timestamp?: unknown }).timestamp; + if (typeof timestamp !== "number" || !Number.isFinite(timestamp)) { + if (options.requireTimestamp) { + throw formatMessageError(prefix, "timestamp must be a finite number"); + } + return { + ...input, + timestamp: Date.now(), + } as AgentMessage; + } + return input; +} + +export function validateNewAgentMessageForSession(input: AgentMessage, prefix = "Invalid AgentMessage"): AgentMessage { + return normalizeAgentMessageForSession(input, prefix, { requireTimestamp: true }); +} + +function normalizeLegacyAgentMessageForSession(input: AgentMessage, prefix = "Invalid legacy AgentMessage"): AgentMessage { + return normalizeAgentMessageForSession(input, prefix, { requireTimestamp: false }); +} + +export function encodeAgentMessageForSession(message: AgentMessage, persistedAt = new Date().toISOString()): string { + const normalized = validateNewAgentMessageForSession(message); + const envelope: PersistedAgentMessageEnvelope = { + schemaVersion: AGENT_MESSAGE_SCHEMA_VERSION, + kind: "agent_message", + message: normalized, + persistedAt, + }; + return JSON.stringify(envelope); +} + +export function decodeAgentMessageFromSessionEntry(row: SessionEntryPayloadRow): AgentMessage { + let parsed: unknown; + try { + parsed = JSON.parse(row.payloadJson); + } catch (error) { + throw new Error( + `Invalid session entry ${row.id} payload JSON: ${error instanceof Error ? error.message : String(error)}`, + ); + } + assertRecord(parsed, `Invalid session entry ${row.id}`); + + if (parsed.schemaVersion === AGENT_MESSAGE_SCHEMA_VERSION || parsed.kind === "agent_message") { + if (parsed.schemaVersion !== AGENT_MESSAGE_SCHEMA_VERSION || parsed.kind !== "agent_message") { + throw new Error(`Invalid session entry ${row.id}: unsupported AgentMessage envelope`); + } + return validateNewAgentMessageForSession( + parsed.message as AgentMessage, + `Invalid session entry ${row.id} AgentMessage`, + ); + } + + return normalizeLegacyAgentMessageForSession(parsed as unknown as AgentMessage, `Invalid legacy session entry ${row.id}`); +} + +export function getAgentMessageRole(message: AgentMessage): string { + return validateNewAgentMessageForSession(message).role as string; +} diff --git a/src/main/agent-runtime/sqlite-session-store.ts b/src/main/agent-runtime/sqlite-session-store.ts index a97c443d..4caac801 100644 --- a/src/main/agent-runtime/sqlite-session-store.ts +++ b/src/main/agent-runtime/sqlite-session-store.ts @@ -1,15 +1,22 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { AppDatabase } from "@main/database"; +import { compactMessagesIfNeeded } from "./context-compactor"; +import { + decodeAgentMessageFromSessionEntry, + encodeAgentMessageForSession, + getAgentMessageRole, + validateNewAgentMessageForSession, +} from "./agent-message-schema"; /** * SQLite-backed pi `AgentMessage[]` store. * * One store instance per task_session. `appendMessage` runs exactly * once per pi `message_end` event; `loadMessages` reads the live - * transcript back at `prepareRun`. Compaction happens at LLM-call - * time via pi's `transformContext` hook and does not rewrite the - * persisted transcript — the session shadow is the durable truth, - * the compacted payload is a per-call view. + * transcript back at `prepareRun`. When the persisted transcript + * crosses the compaction threshold, `loadMessages` writes a durable + * summary row and supersedes the older rows inside one DB transaction, + * leaving the recent tail as the active transcript. */ export class SqliteSessionStore { constructor( @@ -19,21 +26,37 @@ export class SqliteSessionStore { loadMessages(): AgentMessage[] { const rows = this.db.projection.sessionEntries.listActive(this.taskSessionId); - const result: AgentMessage[] = []; + const messages: AgentMessage[] = []; for (const row of rows) { if (row.kind === "message") { - result.push(JSON.parse(row.payloadJson) as AgentMessage); + messages.push(decodeAgentMessageFromSessionEntry(row)); } } - return result; + const outcome = compactMessagesIfNeeded(messages); + if (!outcome.changed || !outcome.messages || !outcome.telemetry) { + return messages; + } + const summary = outcome.messages[0]; + if (!summary) { + return messages; + } + const olderRows = rows.slice(0, outcome.telemetry.compactedOlderMessageCount); + this.db.projection.sessionEntries.compactWithSummary({ + taskSessionId: this.taskSessionId, + supersedeEntryIds: olderRows.map((row) => row.id), + summaryRole: getAgentMessageRole(summary), + summaryPayloadJson: encodeAgentMessageForSession(summary), + }); + return outcome.messages; } appendMessage(message: AgentMessage): void { + const normalized = validateNewAgentMessageForSession(message); this.db.projection.sessionEntries.append({ taskSessionId: this.taskSessionId, kind: "message", - role: (message as { role?: string }).role ?? null, - payloadJson: JSON.stringify(message), + role: getAgentMessageRole(normalized), + payloadJson: encodeAgentMessageForSession(normalized), }); } } diff --git a/src/main/database/session-entries.ts b/src/main/database/session-entries.ts index 85728bbe..c3b64c0d 100644 --- a/src/main/database/session-entries.ts +++ b/src/main/database/session-entries.ts @@ -1,5 +1,5 @@ import { randomUUID } from "node:crypto"; -import type { DatabaseSync } from "node:sqlite"; +import type { DatabaseSync, SQLInputValue } from "node:sqlite"; import type { ReadAll, Transaction } from "./helpers"; import { nowIso } from "./helpers"; @@ -33,7 +33,7 @@ export interface AppendInput { export function buildSessionEntryMethods( getDb: () => DatabaseSync, readAll: ReadAll, - _transaction: Transaction, + transaction: Transaction, ) { function nextSequence(taskSessionId: string): number { const rows = readAll<{ maxSeq: number | null }>( @@ -82,7 +82,7 @@ export function buildSessionEntryMethods( payload_json: string; created_at: string; }>( - "select id, task_session_id, sequence, kind, role, payload_json, created_at from session_entries where task_session_id = ? and is_superseded = 0 order by created_at asc, sequence asc", + "select id, task_session_id, sequence, kind, role, payload_json, created_at from session_entries where task_session_id = ? and is_superseded = 0 order by sequence asc, created_at asc", taskSessionId, ); return rows.map((row) => ({ @@ -96,5 +96,52 @@ export function buildSessionEntryMethods( })); } - return { append, listActive }; + function compactWithSummary(input: { + taskSessionId: string; + supersedeEntryIds: string[]; + summaryRole: string | null; + summaryPayloadJson: string; + }): SessionEntryRow { + return transaction(() => { + const uniqueIds = [...new Set(input.supersedeEntryIds)]; + if (uniqueIds.length === 0) { + throw new Error("session_entries compaction requires at least one entry to supersede."); + } + const placeholders = uniqueIds.map(() => "?").join(", "); + const rows = readAll<{ id: string; sequence: number }>( + `select id, sequence from session_entries where task_session_id = ? and is_superseded = 0 and id in (${placeholders}) order by sequence asc`, + input.taskSessionId, + ...(uniqueIds as SQLInputValue[]), + ); + if (rows.length !== uniqueIds.length) { + throw new Error(`session_entries compaction expected ${uniqueIds.length} active row(s), found ${rows.length}.`); + } + const summarySequence = rows[0]?.sequence; + if (summarySequence === undefined) { + throw new Error("session_entries compaction could not resolve summary sequence."); + } + getDb() + .prepare(`update session_entries set is_superseded = 1 where task_session_id = ? and id in (${placeholders})`) + .run(input.taskSessionId, ...(uniqueIds as SQLInputValue[])); + + const id = randomUUID(); + const created = nowIso(); + getDb() + .prepare( + "insert into session_entries (id, task_session_id, sequence, kind, role, payload_json, is_superseded, created_at) values (?, ?, ?, 'message', ?, ?, 0, ?)", + ) + .run(id, input.taskSessionId, summarySequence, input.summaryRole, input.summaryPayloadJson, created); + return { + id, + taskSessionId: input.taskSessionId, + sequence: summarySequence, + kind: "message", + role: input.summaryRole, + payloadJson: input.summaryPayloadJson, + createdAt: created, + }; + }); + } + + return { append, listActive, compactWithSummary }; } diff --git a/tests/unit/sqlite-session-store.test.ts b/tests/unit/sqlite-session-store.test.ts index 4ad7568c..1c4febc8 100644 --- a/tests/unit/sqlite-session-store.test.ts +++ b/tests/unit/sqlite-session-store.test.ts @@ -49,4 +49,150 @@ describe("SqliteSessionStore", () => { 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); + }); });