feat(sdk): run tail recovery for every chat.agent and let a transcript storage own the model's context

packages/trigger-sdk/src/v3/ai.ts

@@ -5128,6 +5128,18 @@ function isUIMessageStreamable(value: unknown): value is UIMessageStreamable {
   );
 }
 
+const warnedHydrateMessagesDeprecated = new Set<string>();
+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<TUIM extends UIMessage = UIMessage> = {
   /** Task run context — same as `task({ run })` second-argument `ctx`. */
@@ -5409,8 +5422,9 @@ export type RecoveryBootResult<TUIM extends UIMessage = UIMessage> = {
    * 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<inferSchemaOut<TClientDataSchema>, TUIMessage>) =>
+            hydrateMessages(event)
+        : storageLoadContext
+          ? (event: HydrateMessagesEvent<inferSchemaOut<TClientDataSchema>, TUIMessage>) =>
+              storageLoadContext<TUIMessage>(
+                { 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<TUIMessage>({
-                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<TUIMessage>({
+                    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<TUIMessage>(
           (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

packages/trigger-sdk/src/v3/transcriptStorage.ts

@@ -73,18 +73,42 @@ type TranscriptLoadResult<TUIMessage extends UIMessage = UIMessage> = {
   nextCursor?: string;
 };
 
+/** What `loadContext` receives on every turn and action. */
+type LoadContextEvent<TClientData = unknown, TUIMessage extends UIMessage = UIMessage> = {
+  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<TClientData = unknown> = {
   load<TUIMessage extends UIMessage = UIMessage>(
     scope: TranscriptScope<TClientData>,
     opts?: TranscriptLoadOptions
   ): Promise<TranscriptLoadResult<TUIMessage>>;
   save(ctx: TranscriptStorageContext<TClientData>, changeset: TranscriptChangeset): Promise<void>;
+  loadContext?<TUIMessage extends UIMessage = UIMessage>(
+    scope: TranscriptScope<TClientData>,
+    event: LoadContextEvent<TClientData, TUIMessage>
+  ): Promise<TUIMessage[]> | TUIMessage[];
 };
 
 /** An in-memory transcript: ordered entries plus the opaque state record. */