fix: validate and compact session entries

This commit is contained in:
ZacharyZhang-NY committed 2026-04-23 21:21:59 -04:00
1 parent c850d70bfc
commit 84c00e0011
4 files changed
+329 -13

No files matched your search

@@ -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<string, unknown> {
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) : "<missing>"}`);
}
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;
}
+32 -9
View File
@@ -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),
});
}
}
+51 -4
View File
@@ -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 };
}