diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index a0b85830212..f6005643b98 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -5128,6 +5128,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; @@ -5338,8 +5350,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`. */ @@ -5409,8 +5422,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[]; /** @@ -5983,9 +5997,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 @@ -6740,6 +6754,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; @@ -6890,6 +6916,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; /** @@ -7055,7 +7097,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. @@ -7065,18 +7107,22 @@ 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); + bootTranscriptState = loaded.state; + 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), @@ -7214,7 +7260,7 @@ function chatAgent< }); // ── Recovery boot + chain reconstruction ──────────────────────── - if (!hydrateMessages) { + { const settledMessages = mergeByIdReplaceWins( (bootSnapshot?.messages as TUIMessage[]) ?? [], replayedSettled @@ -7383,6 +7429,7 @@ function chatAgent< // and it's safe because the route handler isn't subject to the // `/in/append` 512 KiB cap. if ( + !hydrateMessages && accumulatedUIMessages.length === 0 && payload.trigger === "handover-prepare" && Array.isArray(payload.headStartMessages) && @@ -8069,11 +8116,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", @@ -8161,7 +8208,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) && @@ -8198,7 +8245,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 @@ -8221,7 +8268,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 6aca0564290..99d2bffd625 100644 --- a/packages/trigger-sdk/src/v3/transcriptStorage.ts +++ b/packages/trigger-sdk/src/v3/transcriptStorage.ts @@ -73,11 +73,31 @@ type TranscriptLoadResult = { nextCursor?: 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( @@ -85,6 +105,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..3e206b19c88 --- /dev/null +++ b/packages/trigger-sdk/test/transcript-gate-split.test.ts @@ -0,0 +1,183 @@ +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("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/); + }); +});