From 078d22b73beb4059af1723bf52fd8d0eefa0bbd7 Mon Sep 17 00:00:00 2001 From: Mike Aizatsky Date: Sat, 28 Mar 2026 19:18:50 -0700 Subject: [PATCH 1/2] Clone messages on insert to decouple session state --- .../0002_drop_incompatible_sessions.sql | 6 + packages/core/src/agent-session-do.ts | 115 ++++++------------ packages/core/src/db/entry-types.ts | 16 ++- packages/core/src/messages.ts | 43 +++++-- packages/core/src/session/context.ts | 16 ++- specs/core.md | 41 +++---- 6 files changed, 124 insertions(+), 113 deletions(-) create mode 100644 packages/core/migrations/0002_drop_incompatible_sessions.sql diff --git a/packages/core/migrations/0002_drop_incompatible_sessions.sql b/packages/core/migrations/0002_drop_incompatible_sessions.sql new file mode 100644 index 0000000..d6188fb --- /dev/null +++ b/packages/core/migrations/0002_drop_incompatible_sessions.sql @@ -0,0 +1,6 @@ +-- Message persistence format changed to store AI SDK ModelMessage JSON directly. +-- Existing sessions may contain incompatible message payloads. +-- Reset persisted conversation state so all sessions rehydrate with the new format. + +DELETE FROM entries; +DELETE FROM sessions; diff --git a/packages/core/src/agent-session-do.ts b/packages/core/src/agent-session-do.ts index 39cbe6e..71c3531 100644 --- a/packages/core/src/agent-session-do.ts +++ b/packages/core/src/agent-session-do.ts @@ -62,9 +62,9 @@ import { Messages } from "./messages.ts"; import { ObservableImpl } from "./observable-impl.ts"; import { buildSessionContext, walkToRoot } from "./session/context.ts"; import { + appendEntry as persistEntry, commitSession, deleteSession, - flushPendingEntries, forkSession, } from "./session/persistence.ts"; import { buildBasePrompt } from "./system-prompt.ts"; @@ -188,7 +188,6 @@ export class AgentSessionDO extends DurableObject implements ISession { // ─── In-memory session state ────────────────────────────────────────────── #messages = new Messages([]); - #pendingEntries: AnyEntry[] = []; #branchEntries: AnyEntry[] = []; #messageToEntryId = new Map(); @@ -359,16 +358,7 @@ export class AgentSessionDO extends DurableObject implements ISession { timestamp: new Date().toISOString(), data: { name }, }; - this.#appendEntry(entry); - if (this.#createdAt !== 0) { - await flushPendingEntries( - this.#pendingEntries, - this.#sessionId, - this.#requireLeafId(), - this.env.SESSIONS_DB, - ); - this.#pendingEntries = []; - } + await this.#appendEntry(entry); } // ─── ISession: Conversation ─────────────────────────────────────────────── @@ -433,9 +423,7 @@ export class AgentSessionDO extends DurableObject implements ISession { timestamp: new Date().toISOString(), data: userMessage, }; - this.#pendingEntries.push(userEntry); - this.#leafId = userEntryId; - this.#messageToEntryId.set(userMessage, userEntryId); + await this.#appendEntry(userEntry); // emitBeforeStart const beforeStart = (await this.#extensionRunner.emit( @@ -459,7 +447,13 @@ export class AgentSessionDO extends DurableObject implements ISession { this.#messagesAtTurnStart = this.#messages.length; // Start the agent turn — events flow via #runStream, which calls #onTurnEvent. - this.#startTurn([userMessage]); + const persistedUserMessages = this.#startTurn([userMessage]); + const persistedUserMessage = persistedUserMessages[0]; + if (persistedUserMessage !== undefined) { + // Messages stores structured-cloned snapshots; map the in-session object + // reference (not the caller's original) to the entry ID used in D1. + this.#messageToEntryId.set(persistedUserMessage, userEntryId); + } } catch (e) { // If anything above throws, release the turn reservation so the caller // can retry rather than being permanently locked out. @@ -616,16 +610,7 @@ export class AgentSessionDO extends DurableObject implements ISession { timestamp: new Date().toISOString(), data: { modelId }, }; - this.#appendEntry(entry); - if (this.#createdAt !== 0) { - await flushPendingEntries( - this.#pendingEntries, - this.#sessionId, - this.#requireLeafId(), - this.env.SESSIONS_DB, - ); - this.#pendingEntries = []; - } + await this.#appendEntry(entry); } async listModels(): Promise { @@ -652,7 +637,7 @@ export class AgentSessionDO extends DurableObject implements ISession { timestamp: new Date().toISOString(), data: { customType, content, display }, }; - this.#appendEntry(entry); + await this.#appendEntry(entry); } async appendCustomEntry(customType: string, data?: unknown): Promise { @@ -664,7 +649,7 @@ export class AgentSessionDO extends DurableObject implements ISession { timestamp: new Date().toISOString(), data: { customType, payload: data }, }; - this.#appendEntry(entry); + await this.#appendEntry(entry); } async getEntries(customType?: string): Promise { @@ -688,15 +673,6 @@ export class AgentSessionDO extends DurableObject implements ISession { async compact(options?: CompactOptions): Promise { await this.#compact(options ?? {}); - if (this.#createdAt !== 0 && this.#leafId !== null) { - await flushPendingEntries( - this.#pendingEntries, - this.#sessionId, - this.#leafId, - this.env.SESSIONS_DB, - ); - this.#pendingEntries = []; - } } // ─── ISession: System prompt ────────────────────────────────────────────── @@ -750,7 +726,6 @@ export class AgentSessionDO extends DurableObject implements ISession { await deleteSession(this.#sessionId, this.env.SESSIONS_DB); } this.#agentAbortController?.abort(); - this.#pendingEntries = []; this.#messages = new Messages([]); this.#leafId = null; } @@ -802,13 +777,17 @@ export class AgentSessionDO extends DurableObject implements ISession { * * Spec ref: specs/core.md §AgentSessionDO §Agent Loop */ - #startTurn(initialMessages: ModelMessage[]): void { + #startTurn(initialMessages: ModelMessage[]): ModelMessage[] { + let addedMessages: ModelMessage[] = []; if (initialMessages.length > 0) { - this.#messages.pushAll(initialMessages); + // Messages.pushAll() structured-clones inputs to detach mutable external + // references from internal session state. + addedMessages = this.#messages.pushAll(initialMessages); } const ac = new AbortController(); this.#agentAbortController = ac; this.ctx.waitUntil(this.#runStream(ac.signal)); + return addedMessages; } /** @@ -1043,10 +1022,12 @@ export class AgentSessionDO extends DurableObject implements ISession { this.#notifyListeners({ type: "turn_flushed" }); } - #appendEntry(entry: AnyEntry): void { - this.#pendingEntries.push(entry); + async #appendEntry(entry: AnyEntry): Promise { this.#branchEntries.push(entry); this.#leafId = entry.id; + await this.#ensureSessionCommitted(); + await persistEntry(entry, this.env.SESSIONS_DB); + this.#updatedAt = Date.now(); } #compactTokens(): number { @@ -1066,8 +1047,7 @@ export class AgentSessionDO extends DurableObject implements ISession { * * Lets extensions cancel or supply a pre-built summary, otherwise calls * agentCompact() to summarise old messages. Appends a CompactionEntry to - * #pendingEntries and replaces #messages with the compacted list. - * Does NOT flush to D1 — the caller flushes after the turn ends. + * D1 and replaces #messages with the compacted list. * * Spec ref: specs/core.md §Compaction algorithm */ @@ -1124,8 +1104,7 @@ export class AgentSessionDO extends DurableObject implements ISession { tokensBefore: this.#lastInputTokens, }, }; - this.#pendingEntries.push(compactionEntry); - this.#leafId = compactionEntry.id; + await this.#appendEntry(compactionEntry); // 5. Replace message history with summary + kept messages. // keptMessages are a validated subset of the prior #messages; summaryMessage @@ -1137,13 +1116,6 @@ export class AgentSessionDO extends DurableObject implements ISession { this.#messages = new Messages([summaryMessage, ...keptMessages]); } - #requireLeafId(): string { - if (this.#leafId === null) { - throw new Error("Session leafId is not initialized"); - } - return this.#leafId; - } - /** Persist all new agent messages to D1 since the last flush point. */ async #persistNewMessages(): Promise { const newMessages = this.#messages.get().slice(this.#messagesAtTurnStart); @@ -1158,35 +1130,26 @@ export class AgentSessionDO extends DurableObject implements ISession { timestamp: new Date().toISOString(), data: msg, }; - this.#pendingEntries.push(entry); - this.#leafId = entryId; + await this.#appendEntry(entry); this.#messageToEntryId.set(msg, entryId); } // Advance the start index so the next call only picks up genuinely new messages this.#messagesAtTurnStart = this.#messages.length; + } - if (this.#createdAt === 0 && this.#sessionId !== "") { - const now = Date.now(); - this.#createdAt = now; - this.#updatedAt = now; - await commitSession( - this.#sessionId, - this.#userId, - { ...(this.#name !== undefined ? { name: this.#name } : {}), modelId: this.#modelId }, - this.env.SESSIONS_DB, - ); - } - - if (this.#pendingEntries.length > 0 && this.#leafId !== null) { - await flushPendingEntries( - this.#pendingEntries, - this.#sessionId, - this.#leafId, - this.env.SESSIONS_DB, - ); - this.#pendingEntries = []; - this.#updatedAt = Date.now(); + async #ensureSessionCommitted(): Promise { + if (this.#createdAt !== 0 || this.#sessionId === "") { + return; } + const now = Date.now(); + this.#createdAt = now; + this.#updatedAt = now; + await commitSession( + this.#sessionId, + this.#userId, + { ...(this.#name !== undefined ? { name: this.#name } : {}), modelId: this.#modelId }, + this.env.SESSIONS_DB, + ); } } diff --git a/packages/core/src/db/entry-types.ts b/packages/core/src/db/entry-types.ts index a063ff6..ccec030 100644 --- a/packages/core/src/db/entry-types.ts +++ b/packages/core/src/db/entry-types.ts @@ -44,10 +44,24 @@ export interface EntryBase { // ─── Concrete entry types ───────────────────────────────────────────────────── +/** A single LLM message (user | assistant | tool | system). */ +export interface LegacySystemMessageData { + type: "system"; + content: string; +} + +/** + * Serialized payload stored in `entries.data` for `type: "message"` rows. + * + * Primary format is the AI SDK `ModelMessage` serialized as-is. + * `LegacySystemMessageData` is kept for old non-ModelMessage system entries. + */ +export type MessageEntryData = ModelMessage | LegacySystemMessageData; + /** A single LLM message (user | assistant | tool | system). */ export interface MessageEntry extends EntryBase { type: "message"; - data: ModelMessage; + data: MessageEntryData; } /** Active model was switched. */ diff --git a/packages/core/src/messages.ts b/packages/core/src/messages.ts index 6ef67a2..9e6b909 100644 --- a/packages/core/src/messages.ts +++ b/packages/core/src/messages.ts @@ -1,14 +1,11 @@ -import type { ModelMessage } from "@piccolo/api"; +import type { ModelMessage } from "ai"; import { modelMessageSchema } from "ai"; export class Messages { #messages: ModelMessage[] = []; - constructor(messages: ModelMessage[]) { - for (const message of messages) { - Messages.#validate(message); - } - this.#messages = messages; + constructor(messages: ModelMessage[] = []) { + this.#messages = Messages.#cloneAndValidateAll(messages); } get length() { @@ -28,11 +25,39 @@ export class Messages { } } - pushAll(contextMessages: ModelMessage[]) { - for (const message of contextMessages) { + /** + * Clone + validate at the boundary before entering session state. + * + * We intentionally structuredClone every incoming message to prevent hidden + * object identity/state coupling with external producers (AI SDK callbacks, + * extensions, gateway handlers). Session state should only hold detached + * snapshots. + */ + static #cloneAndValidateAll(messages: ModelMessage[]): ModelMessage[] { + const cloned = structuredClone(messages) as ModelMessage[]; + for (const message of cloned) { Messages.#validate(message); } - this.#messages.push(...contextMessages); + return cloned; + } + + push(...messages: ModelMessage[]): ModelMessage[] { + return this.pushAll(messages); + } + + pushAll(contextMessages: ModelMessage[]): ModelMessage[] { + const cloned = Messages.#cloneAndValidateAll(contextMessages); + this.#messages.push(...cloned); + return cloned; + } + + slice(start?: number, end?: number): ModelMessage[] { + return this.#messages.slice(start, end); + } + + replace(messages: ModelMessage[]): ModelMessage[] { + this.#messages = Messages.#cloneAndValidateAll(messages); + return this.#messages; } get(): ModelMessage[] { diff --git a/packages/core/src/session/context.ts b/packages/core/src/session/context.ts index 37d648d..e936cc6 100644 --- a/packages/core/src/session/context.ts +++ b/packages/core/src/session/context.ts @@ -17,6 +17,7 @@ import type { BranchSummaryEntry, CompactionEntry, CustomMessageEntry, + LegacySystemMessageData, MessageEntry, ModelChangeEntry, } from "../db/entry-types.ts"; @@ -81,8 +82,19 @@ export function walkToRoot(entries: AnyEntry[], leafId: string | null): AnyEntry */ function entryToMessage(entry: AnyEntry): ModelMessage | undefined { switch (entry.type) { - case "message": - return (entry as MessageEntry).data; + case "message": { + const messageData = (entry as MessageEntry).data; + if ( + typeof messageData === "object" && + messageData !== null && + "type" in messageData && + messageData.type === "system" + ) { + const legacySystem = messageData as LegacySystemMessageData; + return { role: "system", content: legacySystem.content }; + } + return messageData; + } case "custom_message": { const cm = entry as CustomMessageEntry; if (cm.data.display) { diff --git a/specs/core.md b/specs/core.md index 87975eb..ea365a5 100644 --- a/specs/core.md +++ b/specs/core.md @@ -149,7 +149,7 @@ type EntryType = // "message" — a single LLM message (user | assistant | tool | system) interface MessageEntry extends EntryBase { type: "message"; - data: ModelMessage; // from `ai` package + data: ModelMessage | { type: "system"; content: string }; // ModelMessage JSON (legacy system payload supported) } // "model_change" — active model switched @@ -241,7 +241,7 @@ class PiccoloCore extends WorkerEntrypoint { // 2. Resolve DO stub: env.AGENT_SESSION.idFromName(sessionId) // 3. Call stub.newSession(userId, options) — the DO initialises itself using // its own name (sessionId) from ctx.id.name; returns the SessionImpl RpcTarget - // Note: D1 row is NOT written here — lazy creation on first assistant response + // Note: D1 row is NOT written here — lazy creation on first persisted entry } async getSession(sessionId: string): Promise { @@ -301,11 +301,10 @@ interface DOState { updatedAt: number; // In-memory message list; rebuilt from D1 on cold start + // Every append clones messages via structuredClone before storing, so + // session state never retains external object references. messages: ModelMessage[]; - // Pending entries not yet flushed to D1 - pendingEntries: AnyEntry[]; - // Inlined agent state (no separate Agent class) model: LanguageModel; error: string | undefined; @@ -381,7 +380,7 @@ async newSession( e. Reconstruct the `LanguageModel` via `createModel(env, modelId)`. 2. Return `getSession(userId)` — the DO's own `SessionImpl` RpcTarget. -The D1 `sessions` row is **not** written here — it is written lazily on the first `finish` (see §Lazy session creation). +The D1 `sessions` row is **not** written here — it is written lazily on the first persisted entry append (see §Lazy session creation). --- @@ -396,7 +395,7 @@ The D1 `sessions` row is **not** written here — it is written lazily on the fi │ If "transform": use result.text as the prompt text │ ├─ 2. Build UserMessage from (text, attachments) -│ Append to agent.messages + pendingEntries +│ Append to agent.messages and persist immediately │ ├─ 3. Emit BeforeAgentStartEvent to ExtensionRunner │ → BeforeAgentStartResult { systemPrompt?, contextMessages? } @@ -428,27 +427,19 @@ off when the RPC call frame completes. #### Per-message append -When `AgentEvent.type === "finish"` fires: +When a persistent entry is created (user/assistant/tool/system message, model change, custom entry, compaction): ``` -1. For each new ModelMessage in agent.state.messages since last flush: +1. For each new ModelMessage in agent.state.messages since last persistence point: a. Create MessageEntry { id, sessionId, parentId: leafId, type: "message", ... } - b. Append to DO storage (synchronous) - c. leafId = newEntry.id - d. Add to pendingEntries - -2. Flush pendingEntries to D1: - INSERT INTO entries (id, session_id, parent_id, type, timestamp, data) - VALUES ... (batch insert, all pending entries) - -3. UPDATE sessions SET updated_at = ?, leaf_id = ?, model_id = ? WHERE id = ? - -4. Clear pendingEntries + b. Ensure `sessions` row is committed (if first persistent append) + c. INSERT entry into D1 immediately + d. UPDATE sessions SET updated_at = ?, leaf_id = ? WHERE id = ? ``` #### Lazy session creation -The D1 `sessions` row is not written until the first assistant response is committed. Before that point, session state lives only in the DO. On first flush: +The D1 `sessions` row is not written until the first persistent entry is appended. Before that point, session state lives only in the DO. On first append: ``` INSERT OR IGNORE INTO sessions (id, user_id, created_at, updated_at, name, model_id, leaf_id) @@ -649,7 +640,7 @@ Compaction is triggered in two cases: ### Compaction algorithm -Compaction is implemented as `#compact(options)` directly on `AgentSessionDO`. It reads and mutates `#messages`, `#leafId`, and `#pendingEntries` in place. +Compaction is implemented as `#compact(options)` directly on `AgentSessionDO`. It reads and mutates `#messages` and `#leafId` in place. ``` 1. Emit before_compact to extensions → BeforeCompactResult @@ -661,12 +652,12 @@ Compaction is implemented as `#compact(options)` directly on `AgentSessionDO`. I 3. Lookup firstKeptEntryId from #messageToEntryId (falls back to "" if not found) -4. Build CompactionEntry, push to #pendingEntries, advance #leafId +4. Build CompactionEntry, persist immediately, advance #leafId 5. Replace #messages = [summaryMessage, ...keptMessages] ``` -Caller (`compact()` public method or `#runStream` threshold check) is responsible for flushing `#pendingEntries` to D1 after the turn ends. +No deferred flush is required — compaction entries are persisted immediately. --- @@ -676,7 +667,7 @@ Caller (`compact()` public method or `#runStream` threshold check) is responsibl `SessionTarget extends RpcTarget` is a thin JSRPC-serialisable proxy that delegates every `ISession` method to the owning `AgentSessionDO`. One `SessionTarget` is created per DO lifetime (in `#initialize()`) and stored as `#rpcCtx`. It is passed to extension workers and threaded into tool `execute()` calls — this is the JSRPC-correct approach since `DurableObject` instances cannot be passed directly over dispatch RPC. -All mutations (model changes, custom entries, etc.) write directly into `#pendingEntries` / `#branchEntries` on the DO. No D1 flush happens during the turn — flushing occurs in `#persistNewMessages()` triggered by `#onAgentEnd()`. +All persistent mutations (model changes, custom entries, messages, compaction) write directly to D1 as soon as they are created, and also append to `#branchEntries` in-memory. ### Context injection into tool execute() From 98e021fcdc03402daf367320e845ddd61a3a6523 Mon Sep 17 00:00:00 2001 From: Mike Aizatsky Date: Sat, 28 Mar 2026 19:50:16 -0700 Subject: [PATCH 2/2] Fix message data typing guards and pass pnpm check --- packages/core/src/agent-session-do.ts | 9 ++++++--- packages/core/src/db/entry-types.ts | 11 +++++++++++ packages/core/src/session/context.ts | 14 ++++---------- 3 files changed, 21 insertions(+), 13 deletions(-) diff --git a/packages/core/src/agent-session-do.ts b/packages/core/src/agent-session-do.ts index 71c3531..7c1bf7d 100644 --- a/packages/core/src/agent-session-do.ts +++ b/packages/core/src/agent-session-do.ts @@ -54,7 +54,7 @@ import type { ModelChangeEntry, SessionInfoEntry, } from "./db/entry-types.ts"; -import { generateEntryId, parseEntry } from "./db/entry-types.ts"; +import { generateEntryId, isLegacySystemMessageData, parseEntry } from "./db/entry-types.ts"; import { getEntries, getSession } from "./db/schema.ts"; import { ExtensionRunner } from "./extension-runner.ts"; import { createModel } from "./gateway.ts"; @@ -62,10 +62,10 @@ import { Messages } from "./messages.ts"; import { ObservableImpl } from "./observable-impl.ts"; import { buildSessionContext, walkToRoot } from "./session/context.ts"; import { - appendEntry as persistEntry, commitSession, deleteSession, forkSession, + appendEntry as persistEntry, } from "./session/persistence.ts"; import { buildBasePrompt } from "./system-prompt.ts"; import { SystemPromptAssembler } from "./system-prompt-assembler.ts"; @@ -291,7 +291,10 @@ export class AgentSessionDO extends DurableObject implements ISession { } for (const entry of allEntries) { if (entry.type === "message") { - this.#messageToEntryId.set((entry as MessageEntry).data, entry.id); + const messageData = (entry as MessageEntry).data; + if (!isLegacySystemMessageData(messageData)) { + this.#messageToEntryId.set(messageData, entry.id); + } } } this.#branchEntries = walkToRoot(allEntries, this.#leafId); diff --git a/packages/core/src/db/entry-types.ts b/packages/core/src/db/entry-types.ts index ccec030..b8e3234 100644 --- a/packages/core/src/db/entry-types.ts +++ b/packages/core/src/db/entry-types.ts @@ -58,6 +58,17 @@ export interface LegacySystemMessageData { */ export type MessageEntryData = ModelMessage | LegacySystemMessageData; +export function isLegacySystemMessageData(data: MessageEntryData): data is LegacySystemMessageData { + return ( + typeof data === "object" && + data !== null && + "type" in data && + data.type === "system" && + "content" in data && + typeof data.content === "string" + ); +} + /** A single LLM message (user | assistant | tool | system). */ export interface MessageEntry extends EntryBase { type: "message"; diff --git a/packages/core/src/session/context.ts b/packages/core/src/session/context.ts index e936cc6..4259228 100644 --- a/packages/core/src/session/context.ts +++ b/packages/core/src/session/context.ts @@ -17,10 +17,10 @@ import type { BranchSummaryEntry, CompactionEntry, CustomMessageEntry, - LegacySystemMessageData, MessageEntry, ModelChangeEntry, } from "../db/entry-types.ts"; +import { isLegacySystemMessageData } from "../db/entry-types.ts"; // ─── Message schema validation ──────────────────────────────────────────────── @@ -84,16 +84,10 @@ function entryToMessage(entry: AnyEntry): ModelMessage | undefined { switch (entry.type) { case "message": { const messageData = (entry as MessageEntry).data; - if ( - typeof messageData === "object" && - messageData !== null && - "type" in messageData && - messageData.type === "system" - ) { - const legacySystem = messageData as LegacySystemMessageData; - return { role: "system", content: legacySystem.content }; + if (isLegacySystemMessageData(messageData)) { + return { role: "system", content: messageData.content }; } - return messageData; + return messageData as ModelMessage; } case "custom_message": { const cm = entry as CustomMessageEntry;