From 8292d1015d41a0fd4a5f94c623963c4ba84f0299 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Sat, 5 Sep 2026 10:21:02 +0100 Subject: [PATCH] feat(sdk): run tail recovery for every chat.agent and let a transcript storage own the model's context One condition used to decide three things at boot: whether to read the persisted transcript, whether to replay the session's output tail, and whether to replay unacknowledged input. Registering hydrateMessages switched all three off, so an app that owned its own context also lost crash recovery, and no application can rebuild the tail its dead run had already emitted. The replays and onRecoveryBoot now run for every agent; only the transcript read is skipped for hydrateMessages. The storage can now declare loadContext, which the runtime calls on every turn and action in place of the accumulated transcript, the role hydrateMessages played, while save keeps receiving every change. hydrateMessages is deprecated with a one-time warning, and configuring it together with a storage that has loadContext is an error. --- packages/trigger-sdk/src/v3/ai.ts | 103 ++++++-- .../trigger-sdk/src/v3/transcriptStorage.ts | 24 ++ .../test/transcript-gate-split.test.ts | 240 ++++++++++++++++++ 3 files changed, 341 insertions(+), 26 deletions(-) create mode 100644 packages/trigger-sdk/test/transcript-gate-split.test.ts diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index d6af68ff24b..40e82dc5b84 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -5191,6 +5191,18 @@ function isUIMessageStreamable(value: unknown): value is UIMessageStreamable { ); } +const warnedHydrateMessagesDeprecated = new Set(); +function warnHydrateMessagesDeprecatedOnce(agentId: string) { + if (warnedHydrateMessagesDeprecated.has(agentId)) return; + warnedHydrateMessagesDeprecated.add(agentId); + console.warn( + `[chat.agent] \`hydrateMessages\` on "${agentId}" is deprecated. Give the agent a transcript ` + + "storage instead: `save` receives every change to the conversation and `loadContext` " + + "lets the application own the model's context, with crash recovery and durable " + + "compaction that `hydrateMessages` never had." + ); +} + let warnedMissingOnAction = false; function warnMissingOnActionOnce() { if (warnedMissingOnAction) return; @@ -5401,8 +5413,9 @@ export type RecoveryPendingToolCall = { * `chat.endRun()` with no buffered user messages, fresh chat, OOM retry * after a successful turn-complete with no in-flight tail). * - * Does NOT fire when `hydrateMessages` is registered (the customer owns - * persistence; recovery decisions live in their own DB query). + * Fires regardless of who owns the model's context. With `hydrateMessages` + * or a storage `loadContext`, the recovered tail reaches that hook in + * `previousMessages` on the next turn. */ export type RecoveryBootEvent = { /** Task run context — same as `task({ run })` second-argument `ctx`. */ @@ -5472,8 +5485,9 @@ export type RecoveryBootResult = { * context, mutate its tool parts to inject synthesized results, * collapse history, etc. * - * Ignored when `hydrateMessages` is registered (the hydrate hook - * runs per-turn and overwrites the chain). + * With `hydrateMessages` or a storage `loadContext`, this chain is what + * the hook receives as `previousMessages` on the next turn; the hook's + * return value is the chain the model sees. */ chain?: TUIM[]; /** @@ -6058,9 +6072,9 @@ export type ChatAgentOptions< * continuation after `chat.endRun()` with no buffered user, a fresh * chat, or an OOM retry on top of a complete snapshot. * - * Does NOT fire when `hydrateMessages` is registered — that hook owns - * the per-turn chain and overlapping recovery decisions belong in the - * customer's DB. + * Fires regardless of who owns the model's context; a `hydrateMessages` + * hook or a storage `loadContext` receives the recovered tail in + * `previousMessages` on the next turn. * * Defaults (returned when the hook is omitted or returns no field): * - With two or more in-flight users, the partial and the user it @@ -6832,6 +6846,18 @@ function chatAgent< ...restOptions } = options; + if (hydrateMessages) { + const storageAtDefinition = transcriptStorageOverride ?? defaultStorage; + if (typeof storageAtDefinition.loadContext === "function") { + throw new Error( + `chat.agent: "${options.id}" sets \`hydrateMessages\` and uses a transcript storage with ` + + "`loadContext`. Both would own the model's context; keep one. `hydrateMessages` is " + + "deprecated, so prefer `loadContext` on the storage." + ); + } + warnHydrateMessagesDeprecatedOnce(options.id); + } + const parseClientData = clientDataSchema ? getSchemaParseFn(clientDataSchema) : undefined; const parseAction = actionSchema ? getSchemaParseFn(actionSchema) : undefined; @@ -6999,6 +7025,22 @@ function chatAgent< // swallow errors internally; the agent stays available either way. const sessionIdForSnapshot = payload.sessionId ?? payload.chatId; const transcriptStorage = transcriptStorageOverride ?? defaultStorage; + const storageLoadContext = transcriptStorage.loadContext?.bind(transcriptStorage); + /** + * Who supplies the model's context each turn: the deprecated + * `hydrateMessages` hook, the storage's `loadContext`, or (undefined) + * the runtime's own transcript. + */ + const loadContextHook = hydrateMessages + ? (event: HydrateMessagesEvent, TUIMessage>) => + hydrateMessages(event) + : storageLoadContext + ? (event: HydrateMessagesEvent, TUIMessage>) => + storageLoadContext( + { chatId: event.chatId, clientData: event.clientData }, + event + ) + : undefined; let transcriptShadow: TranscriptShadow = createTranscriptShadow([]); let bootTranscriptState: unknown = null; /** @@ -7181,7 +7223,7 @@ function chatAgent< let bootInCursor: number | undefined; let bootInCursorResolved = false; - if (!hydrateMessages && couldHavePriorState) { + if (couldHavePriorState) { // Single parent span for the whole boot read phase — snapshot // read, session.out replay, session.in replay. Per-phase timing // + result counts are attributes on the span. @@ -7191,18 +7233,26 @@ function chatAgent< // snapshot read const snapStart = Date.now(); try { - const loaded = await transcriptStorage.load({ - chatId: payload.chatId, - clientData: bootClientData, - }); - transcriptShadow = createTranscriptShadow(loaded.messages); - bootTranscriptState = loaded.state; - persistedStateSet = loaded.state !== null && loaded.state !== undefined; - bootSnapshot = { - messages: loaded.messages, - lastOutEventId: loaded.cursors?.lastOutEventId, - lastInEventId: loaded.cursors?.lastInEventId, - }; + const loaded = hydrateMessages + ? undefined + : await transcriptStorage.load({ + chatId: payload.chatId, + clientData: bootClientData, + }); + if (loaded) { + transcriptShadow = createTranscriptShadow( + loaded.messages, + new Set(loaded.nonFinalIds ?? []) + ); + bootTranscriptState = loaded.state; + transcriptState = loaded.state ?? null; + persistedStateSet = loaded.state !== null && loaded.state !== undefined; + bootSnapshot = { + messages: loaded.messages, + lastOutEventId: loaded.cursors?.lastOutEventId, + lastInEventId: loaded.cursors?.lastInEventId, + }; + } } catch (error) { logger.warn("chat.agent: transcript load failed; continuing from the stream tail", { error: error instanceof Error ? error.message : String(error), @@ -7346,7 +7396,7 @@ function chatAgent< }); // ── Recovery boot + chain reconstruction ──────────────────────── - if (!hydrateMessages) { + { const settledMessages = mergeByIdReplaceWins( (bootSnapshot?.messages as TUIMessage[]) ?? [], replayedSettled @@ -7504,6 +7554,7 @@ function chatAgent< // and it's safe because the route handler isn't subject to the // `/in/append` 512 KiB cap. if ( + !loadContextHook && accumulatedUIMessages.length === 0 && payload.trigger === "handover-prepare" && Array.isArray(payload.headStartMessages) && @@ -8201,11 +8252,11 @@ function chatAgent< : currentWirePayload.action; // Hydrate messages from backend if configured - if (hydrateMessages) { + if (loadContextHook) { const hydrated = await tracer.startActiveSpan( "hydrateMessages()", async () => { - return hydrateMessages({ + return loadContextHook({ chatId: currentWirePayload.chatId, turn, trigger: "action", @@ -8293,7 +8344,7 @@ function chatAgent< // incoming messages instead (gated on the pending handover). if ( turn === 0 && - hydrateMessages && + loadContextHook && cleanedUIMessages.length === 0 && (locals.get(chatHandoverPartialKey)?.length ?? 0) > 0 && Array.isArray(payload.headStartMessages) && @@ -8330,7 +8381,7 @@ function chatAgent< )) as TUIMessage[]; } - if (hydrateMessages) { + if (loadContextHook) { // Snapshot the ids the accumulator knew BEFORE this // turn ran — used below to decide whether an // incoming wire message is genuinely new or just a @@ -8353,7 +8404,7 @@ function chatAgent< const hydrated = await tracer.startActiveSpan( "hydrateMessages()", async () => { - return hydrateMessages({ + return loadContextHook({ chatId: currentWirePayload.chatId, turn, trigger: currentWirePayload.trigger as diff --git a/packages/trigger-sdk/src/v3/transcriptStorage.ts b/packages/trigger-sdk/src/v3/transcriptStorage.ts index 87c44906866..d4e6939372f 100644 --- a/packages/trigger-sdk/src/v3/transcriptStorage.ts +++ b/packages/trigger-sdk/src/v3/transcriptStorage.ts @@ -84,11 +84,31 @@ type TranscriptLoadResult = { nonFinalIds?: string[]; }; +/** What `loadContext` receives on every turn and action. */ +type LoadContextEvent = { + chatId: string; + /** The turn number (0-indexed). */ + turn: number; + trigger: "submit-message" | "regenerate-message" | "action"; + /** The messages the frontend sent for this turn. Empty for actions. */ + incomingMessages: TUIMessage[]; + /** The runtime's transcript before this turn, including any tail it recovered. */ + previousMessages: TUIMessage[]; + clientData?: TClientData; + continuation: boolean; + previousRunId?: string; +}; + /** * A persistence adapter for a `chat.agent` transcript. The runtime calls * `load` once at a continuation boot and `save` after every change to the * conversation. Both are best-effort from the runtime's point of view: an * error is logged and the turn continues. + * + * `loadContext` is optional. Its presence declares that the application + * owns the model's context: the runtime calls it on every turn and action + * and uses what it returns as the conversation, instead of the transcript + * it accumulated. Tail recovery still runs and `save` is still called. */ export type TranscriptStorage = { load( @@ -96,6 +116,10 @@ export type TranscriptStorage = { opts?: TranscriptLoadOptions ): Promise>; save(ctx: TranscriptStorageContext, changeset: TranscriptChangeset): Promise; + loadContext?( + scope: TranscriptScope, + event: LoadContextEvent + ): Promise | TUIMessage[]; }; /** An in-memory transcript: ordered entries plus the opaque state record. */ diff --git a/packages/trigger-sdk/test/transcript-gate-split.test.ts b/packages/trigger-sdk/test/transcript-gate-split.test.ts new file mode 100644 index 00000000000..afc364e93bd --- /dev/null +++ b/packages/trigger-sdk/test/transcript-gate-split.test.ts @@ -0,0 +1,240 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import type { UIMessage } from "ai"; +import { simulateReadableStream, streamText } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { __setTranscriptStorageForTests, chat } from "../src/v3/ai.js"; +import { + memoryTranscriptStorage, + type MemoryTranscriptStorage, + type TranscriptStorage, +} from "../src/v3/transcriptStorage.js"; + +const usage = { + inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 10, reasoning: undefined }, +}; + +function userMessage(text: string, id: string): UIMessage { + return { id, role: "user", parts: [{ type: "text", text }] }; +} + +function textChunks(text: string): LanguageModelV3StreamPart[] { + return [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: text }, + { type: "text-end", id: "t1" }, + { type: "finish", finishReason: { unified: "stop", raw: "stop" }, usage }, + ]; +} + +function recordingModel(prompts: unknown[]) { + return new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + prompts.push(prompt); + return { stream: simulateReadableStream({ chunks: textChunks("ack") }) }; + }, + }); +} + +async function waitFor(check: () => boolean, label: string, timeoutMs = 8_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 10)); + } + throw new Error(`waitFor timed out: ${label}`); +} + +let storage: MemoryTranscriptStorage; + +beforeEach(() => { + storage = memoryTranscriptStorage(); + __setTranscriptStorageForTests(storage); +}); + +afterEach(() => { + __setTranscriptStorageForTests(undefined); + vi.restoreAllMocks(); +}); + +describe("the persistence gate split", () => { + it("fires onRecoveryBoot for a hydrateMessages agent when a partial assistant is in the tail", async () => { + vi.spyOn(console, "warn").mockImplementation(() => {}); + const recoveryEvents: { partialAssistant?: UIMessage }[] = []; + const onRecoveryBoot = async (event: { partialAssistant?: UIMessage }) => { + recoveryEvents.push(event); + return {}; + }; + const hydrated: UIMessage[] = [ + userMessage("from my database", "db-u1"), + { id: "db-a1", role: "assistant", parts: [{ type: "text", text: "stored answer" }] }, + ]; + const hydrateCalls: { previousMessages: UIMessage[] }[] = []; + const prompts: unknown[] = []; + const agent = chat.agent({ + id: "gate-split-hydrate-recovery", + onRecoveryBoot, + hydrateMessages: async ({ previousMessages, incomingMessages }) => { + hydrateCalls.push({ previousMessages }); + return [...hydrated, ...incomingMessages]; + }, + run: async ({ messages, signal }) => + streamText({ model: recordingModel(prompts), messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { + chatId: "gate-split-hydrate-recovery", + continuation: true, + previousRunId: "run_prior", + }); + harness.seedSessionOutPartial({ + id: "a-orphan", + role: "assistant", + parts: [{ type: "text", text: "half an ans" }], + }); + try { + await harness.sendMessage(userMessage("next", "u2")); + await waitFor(() => prompts.length === 1, "turn"); + + expect(recoveryEvents).toHaveLength(1); + expect(recoveryEvents[0]!.partialAssistant?.id).toBe("a-orphan"); + + expect(hydrateCalls).toHaveLength(1); + expect(JSON.stringify(prompts[0])).toContain("from my database"); + expect(storage.changesets).toHaveLength(0); + } finally { + await harness.close(); + } + }); + + it("uses the storage's loadContext for the model's context and still saves the transcript", async () => { + const contextCalls: { trigger: string; previousMessages: UIMessage[] }[] = []; + const loadContext = async ( + _scope: unknown, + event: { trigger: string; previousMessages: UIMessage[]; incomingMessages: UIMessage[] } + ) => { + contextCalls.push({ trigger: event.trigger, previousMessages: event.previousMessages }); + return [userMessage("only what the app chose", "ctx-u1"), ...event.incomingMessages]; + }; + const withContext: TranscriptStorage = { + load: storage.load.bind(storage), + save: storage.save.bind(storage), + loadContext: loadContext as TranscriptStorage["loadContext"], + }; + __setTranscriptStorageForTests(withContext); + + const prompts: unknown[] = []; + const agent = chat.agent({ + id: "gate-split-load-context", + run: async ({ messages, signal }) => + streamText({ model: recordingModel(prompts), messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { chatId: "gate-split-load-context" }); + try { + await harness.sendMessage(userMessage("first", "u1")); + await waitFor(() => storage.changesets.length === 1, "save"); + + expect(contextCalls).toHaveLength(1); + expect(contextCalls[0]!.trigger).toBe("submit-message"); + const prompt = JSON.stringify(prompts[0]); + expect(prompt).toContain("only what the app chose"); + expect(prompt).toContain('"first"'); + + const ids = storage.changesets[0]!.changeset.changes.flatMap((c) => + c.op === "put" ? [c.message.id] : [] + ); + expect(ids).toEqual(["ctx-u1", "u1", expect.any(String)]); + } finally { + await harness.close(); + } + }); + + it("hands a head-start first turn to loadContext as incoming messages, without seeding them twice", async () => { + const calls: { incoming: string[]; previous: string[] }[] = []; + const stored: UIMessage[] = []; + const loadContext = async ( + _scope: unknown, + event: { incomingMessages: UIMessage[]; previousMessages: UIMessage[] } + ) => { + calls.push({ + incoming: event.incomingMessages.map((m) => m.id), + previous: event.previousMessages.map((m) => m.id), + }); + for (const m of event.incomingMessages) { + if (!stored.some((s) => s.id === m.id)) stored.push(m); + } + return [...stored]; + }; + __setTranscriptStorageForTests({ + load: storage.load.bind(storage), + save: storage.save.bind(storage), + loadContext: loadContext as TranscriptStorage["loadContext"], + }); + + let roles: string[] | undefined; + const agent = chat.agent({ + id: "gate-split-head-start-load-context", + onTurnComplete: ({ uiMessages }) => { + roles = uiMessages.map((m) => m.role); + }, + run: async ({ messages, signal }) => + streamText({ model: recordingModel([]), messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { + chatId: "gate-split-head-start-load-context", + mode: "handover-prepare", + headStartMessages: [ + { id: "hs-user-1", role: "user", parts: [{ type: "text", text: "say hi" }] }, + ], + }); + try { + await harness.sendHandover({ + partialAssistantMessage: [ + { role: "assistant", content: [{ type: "text", text: "Hi there." }] }, + ], + messageId: "asst-1", + isFinal: true, + }); + await waitFor(() => roles !== undefined, "turn complete"); + + expect(calls).toHaveLength(1); + expect(calls[0]!.incoming).toEqual(["hs-user-1"]); + expect(calls[0]!.previous).toEqual([]); + expect(roles).toEqual(["user", "assistant"]); + } finally { + await harness.close(); + } + }); + + it("refuses an agent that sets both hydrateMessages and a storage with loadContext", () => { + vi.spyOn(console, "warn").mockImplementation(() => {}); + __setTranscriptStorageForTests({ + load: storage.load.bind(storage), + save: storage.save.bind(storage), + loadContext: async () => [], + }); + expect(() => + chat.agent({ + id: "gate-split-both", + hydrateMessages: async () => [], + run: async ({ messages, signal }) => + streamText({ model: recordingModel([]), messages, abortSignal: signal }), + }) + ).toThrow(/hydrateMessages/); + }); + + it("warns once that hydrateMessages is deprecated", () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + chat.agent({ + id: "gate-split-deprecated", + hydrateMessages: async () => [], + run: async ({ messages, signal }) => + streamText({ model: recordingModel([]), messages, abortSignal: signal }), + }); + const deprecations = warn.mock.calls.filter((c) => String(c[0]).includes("hydrateMessages")); + expect(deprecations).toHaveLength(1); + expect(String(deprecations[0]![0])).toMatch(/deprecated/); + }); +});