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

+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),
});
}
}