feat(sdk): persist compaction and injected context through the transcript storage state (#4894)

## Summary

Makes a compaction summary and `chat.inject` context survive a
continuation run, for every storage including the default.

Until now the model lane after a compaction lived only in the running
worker. When the next run booted it rebuilt the lane from the
transcript, so every continuation re-read the whole conversation and
summarised it again. The same applied to conversational messages added
with `chat.inject`: they lived for the worker's life and vanished on a
continuation.

## Design

The runtime records what it cannot rebuild from the transcript in the
storage's `state` slot: after a compaction, the compacted model lane
together with the transcript id it covers and a fingerprint of that
prefix; for injections, the messages anchored to the transcript message
they followed. At boot the compacted lane is used when the covered
prefix is unchanged, otherwise the lane is converted from the transcript
as before, and injections are re-inserted after their anchors.

A rollback or edit that reconverts the lane clears the state in the same
changeset as the `truncateAfter`, so a storage never holds a summary for
a transcript it no longer matches. A mid-turn steering message reaches
the storage as a `put` in that turn's changeset.

An in-memory storage that logs the changesets it receives, and a
test-only override for the storage the runtime persists through, let the
tests assert the exact changesets for a turn, a steer, a compaction, a
rollback and an injection.

Author: ericallam

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

@@ -80,10 +80,29 @@ import {
   createTranscriptShadow,
   defaultStorage,
   diffTranscript,
+  parseTranscriptRuntimeState,
+  prefixFingerprint,
+  restoreModelLane,
+  type TranscriptChange,
   type TranscriptChangeReason,
+  type TranscriptRuntimeState,
   type TranscriptShadow,
+  type TranscriptStorage,
   type TranscriptStorageContext,
 } from "./transcriptStorage.js";
+
+let transcriptStorageOverride: TranscriptStorage<unknown> | undefined;
+
+/**
+ * Test-only override for the storage `chat.agent` persists through, so a
+ * test can capture the exact changesets the runtime produces.
+ * @internal
+ */
+export function __setTranscriptStorageForTests(
+  storage: TranscriptStorage<unknown> | undefined
+): void {
+  transcriptStorageOverride = storage;
+}
 import {
   type ChatInputChunk,
   type ChatTaskWirePayload,
@@ -2542,6 +2561,13 @@ function spliceHandoverPartial(
  * @internal
  */
 const chatBackgroundQueueKey = locals.create<ModelMessage[]>("chat.backgroundQueue");
+/**
+ * Background injections a step-boundary drain handed to the model this turn,
+ * with the transcript message they followed. Reconciled into the model lane
+ * and the persisted injections once the turn's response is in.
+ */
+const chatPendingBackgroundKey =
+  locals.create<{ afterId: string; messages: ModelMessage[] }[]>("chat.pendingBackground");
 
 /**
  * System-role context injected mid-conversation, held for the instructions lane.
@@ -5022,6 +5048,13 @@ function toStreamTextOptions(options?: ToStreamTextOptionsOptions): Record<strin
       if (bgQueue && bgQueue.length > 0) {
         const injected = bgQueue.splice(0); // drain
         resultMessages = [...(resultMessages ?? messages), ...injected];
+        const pendingBackground = locals.get(chatPendingBackgroundKey) ?? [];
+        pendingBackground.push({
+          afterId:
+            (locals.get(chatCurrentUIMessagesKey) as UIMessage[] | undefined)?.at(-1)?.id ?? "",
+          messages: injected,
+        });
+        locals.set(chatPendingBackgroundKey, pendingBackground);
       }
 
       return resultMessages ? { messages: resultMessages } : undefined;
@@ -6912,6 +6945,23 @@ function chatAgent<
       // durable snapshot + `session.out` replay (or `hydrateMessages` if
       // registered) — the wire is delta-only now, no longer a seed.
       let accumulatedMessages: ModelMessage[] = [];
+      /**
+       * Give the model accumulator the background injections a step-boundary
+       * drain handed to the model this turn, and record them for persistence.
+       * Returns how many model messages were appended.
+       */
+      const reconcilePendingBackground = (): number => {
+        const pending = locals.get(chatPendingBackgroundKey);
+        if (!pending || pending.length === 0) return 0;
+        locals.set(chatPendingBackgroundKey, []);
+        let appended = 0;
+        for (const entry of pending) {
+          accumulatedMessages.push(...entry.messages);
+          laneInjections.push(entry);
+          appended += entry.messages.length;
+        }
+        return appended;
+      };
       /**
        * Give the model accumulator the steering messages a drain consumed,
        * in the form the model actually received. Appended, never reconverted
@@ -6951,8 +7001,18 @@ function chatAgent<
       // collectively cost ~600ms on every first-message TTFC. Both reads
       // swallow errors internally; the agent stays available either way.
       const sessionIdForSnapshot = payload.sessionId ?? payload.chatId;
-      const transcriptStorage = defaultStorage;
+      const transcriptStorage = transcriptStorageOverride ?? defaultStorage;
       let transcriptShadow: TranscriptShadow = createTranscriptShadow([]);
+      let bootTranscriptState: unknown = null;
+      /**
+       * True while the model lane holds a compaction summary, so it cannot be
+       * rebuilt from the transcript and has to be persisted as state. Reset
+       * wherever the lane is reconverted from the UI lane.
+       */
+      let laneCompacted = false;
+      /** Conversational `chat.inject` messages in the lane, anchored to the transcript. */
+      let laneInjections: NonNullable<TranscriptRuntimeState["injections"]> = [];
+      let persistedStateSet = false;
       let bootSnapshot:
         | { messages: TUIMessage[]; lastOutEventId?: string; lastInEventId?: string }
         | undefined;
@@ -7002,6 +7062,30 @@ function chatAgent<
         const { changes, shadow } = diffTranscript(transcriptShadow, opts.messages, {
           nonFinalIds: opts.nonFinalIds,
         });
+        const throughId = opts.messages.at(-1)?.id ?? "";
+        const queued = locals.get(chatBackgroundQueueKey) ?? [];
+        const runtimeState: TranscriptRuntimeState | null =
+          laneCompacted || laneInjections.length > 0 || queued.length > 0
+            ? {
+                v: 1,
+                ...(laneCompacted
+                  ? {
+                      compaction: {
+                        modelMessages: accumulatedMessages,
+                        throughId,
+                        fingerprint: prefixFingerprint(shadow, throughId),
+                      },
+                    }
+                  : laneInjections.length > 0
+                    ? { injections: laneInjections }
+                    : {}),
+                ...(queued.length > 0 ? { queued: [...queued] } : {}),
+              }
+            : null;
+        if (runtimeState !== null || persistedStateSet) {
+          changes.push({ op: "state", value: runtimeState } satisfies TranscriptChange);
+        }
+        transcriptState = runtimeState;
         const inCursor = chatInputRouter().resumeFloor();
         await transcriptStorage.save(
           {
@@ -7030,6 +7114,7 @@ function chatAgent<
           }
         );
         transcriptShadow = shadow;
+        persistedStateSet = runtimeState !== null;
       };
 
       /**
@@ -7114,6 +7199,8 @@ function chatAgent<
                 clientData: bootClientData,
               });
               transcriptShadow = createTranscriptShadow(loaded.messages);
+              bootTranscriptState = loaded.state;
+              persistedStateSet = loaded.state !== null && loaded.state !== undefined;
               bootSnapshot = {
                 messages: loaded.messages,
                 lastOutEventId: loaded.cursors?.lastOutEventId,
@@ -7453,7 +7540,21 @@ function chatAgent<
             }
           }
           try {
-            accumulatedMessages = await toModelMessages(accumulatedUIMessages);
+            const bootRuntimeState = parseTranscriptRuntimeState(bootTranscriptState);
+            const restored = await restoreModelLane(
+              accumulatedUIMessages,
+              bootRuntimeState,
+              (messages) => toModelMessages(messages)
+            );
+            accumulatedMessages = restored.messages;
+            laneCompacted = restored.compacted;
+            laneInjections = restored.injections;
+            if (bootRuntimeState?.queued && bootRuntimeState.queued.length > 0) {
+              locals.set(chatBackgroundQueueKey, [
+                ...(locals.get(chatBackgroundQueueKey) ?? []),
+                ...bootRuntimeState.queued,
+              ]);
+            }
           } catch (error) {
             logger.warn("chat.agent: toModelMessages failed at boot; starting empty", {
               error: error instanceof Error ? error.message : String(error),
@@ -7982,6 +8083,7 @@ function chatAgent<
                 locals.set(chatDeferKey, new Set());
                 locals.set(chatCompactionStateKey, undefined);
                 locals.set(chatSteeringQueueKey, []);
+                locals.set(chatPendingBackgroundKey, []);
                 locals.set(chatResponsePartsKey, []);
                 // NOTE: chatBackgroundQueueKey is NOT reset here — messages injected
                 // by deferred work from the previous turn's onTurnComplete need to
@@ -8128,6 +8230,8 @@ function chatAgent<
                     );
                     accumulatedUIMessages = [...hydrated] as TUIMessage[];
                     accumulatedMessages = await toModelMessages(hydrated);
+                    laneCompacted = false;
+                    laneInjections = [];
                     locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                   }
 
@@ -8165,6 +8269,8 @@ function chatAgent<
                       locals.set(chatOverrideMessagesKey, undefined);
                       accumulatedUIMessages = [...actionOverride] as TUIMessage[];
                       accumulatedMessages = await toModelMessages(actionOverride);
+                      laneCompacted = false;
+                      laneInjections = [];
                       locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
 
                       actionChangedHistory = true;
@@ -8297,6 +8403,8 @@ function chatAgent<
 
                     accumulatedUIMessages = merged;
                     accumulatedMessages = await toModelMessages(merged);
+                    laneCompacted = false;
+                    laneInjections = [];
                     locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
 
                     // Track new messages for onTurnComplete.newUIMessages.
@@ -8346,6 +8454,8 @@ function chatAgent<
                         accumulatedUIMessages.pop();
                       }
                       accumulatedMessages = await toModelMessages(accumulatedUIMessages);
+                      laneCompacted = false;
+                      laneInjections = [];
                     } else if (cleanedUIMessages.length > 0) {
                       // Submit-message (and the special-cased
                       // handover-prepare → submit-message rewrite earlier in
@@ -8399,6 +8509,8 @@ function chatAgent<
                             "chat.agent: replaced message not found at the model lane tail; reconverting the lane"
                           );
                           accumulatedMessages = await toModelMessages(accumulatedUIMessages);
+                          laneCompacted = false;
+                          laneInjections = [];
                         }
                       } else {
                         const incomingModelMessages = await toModelMessages(cleanedUIMessages);
@@ -8608,6 +8720,8 @@ function chatAgent<
                           locals.set(chatOverrideMessagesKey, undefined);
                           accumulatedUIMessages = [...turnStartOverride] as TUIMessage[];
                           accumulatedMessages = await toModelMessages(turnStartOverride);
+                          laneCompacted = false;
+                          laneInjections = [];
                           locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                         }
                       },
@@ -8676,7 +8790,12 @@ function chatAgent<
                     const lastAccumulated = accumulatedMessages[accumulatedMessages.length - 1];
                     const bgQueue = locals.get(chatBackgroundQueueKey);
                     if (bgQueue && bgQueue.length > 0 && lastAccumulated?.role !== "tool") {
-                      accumulatedMessages.push(...bgQueue.splice(0));
+                      const injected = bgQueue.splice(0);
+                      accumulatedMessages.push(...injected);
+                      laneInjections.push({
+                        afterId: accumulatedUIMessages.at(-1)?.id ?? "",
+                        messages: injected,
+                      });
                     }
 
                     if (isHeadStartFinalTurn) {
@@ -8869,6 +8988,8 @@ function chatAgent<
                     accumulatedMessages = await toModelMessages(
                       runOverride.filter((m) => !pendingIds.has(m.id))
                     );
+                    laneCompacted = false;
+                    laneInjections = [];
                     locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                   }
 
@@ -8894,6 +9015,8 @@ function chatAgent<
                     accumulatedMessages = taskCompactionConfig?.compactModelMessages
                       ? await taskCompactionConfig.compactModelMessages(compactEvent)
                       : modelOnlyOverride;
+                    laneCompacted = true;
+                    laneInjections = [];
 
                     // Apply UI messages: callback or default (preserve all)
                     if (taskCompactionConfig?.compactUIMessages) {
@@ -8912,9 +9035,10 @@ function chatAgent<
                   // before the response is appended so the order stays
                   // steer-then-answer. Outside the `capturedResponseMessage`
                   // branches below, so a turn that captured no response is covered.
-                  const steerTailThisTurn = reconcilePendingSteer({
-                    turnNew: turnNewModelMessages,
-                  }).reduce((n, e) => n + e.model.length, 0);
+                  const steerTailThisTurn =
+                    reconcilePendingSteer({
+                      turnNew: turnNewModelMessages,
+                    }).reduce((n, e) => n + e.model.length, 0) + reconcilePendingBackground();
 
                   // Append the assistant's response (partial or complete) to the accumulator.
                   // The onFinish callback fires even on abort/stop, so partial responses
@@ -8986,6 +9110,8 @@ function chatAgent<
                             "chat.agent: replaced response not found at the model lane tail; reconverting the lane"
                           );
                           accumulatedMessages = await toModelMessages(accumulatedUIMessages);
+                          laneCompacted = false;
+                          laneInjections = [];
                         }
                       } else {
                         accumulatedMessages.push(...responseModelMessages);
@@ -9105,6 +9231,9 @@ function chatAgent<
                                     },
                                   ];
 
+                              laneCompacted = true;
+                              laneInjections = [];
+
                               // UI messages: callback or default (preserve all)
                               if (outerCompaction.compactUIMessages) {
                                 accumulatedUIMessages = (await outerCompaction.compactUIMessages(
@@ -9209,6 +9338,8 @@ function chatAgent<
                           locals.set(chatOverrideMessagesKey, undefined);
                           accumulatedUIMessages = [...override] as TUIMessage[];
                           accumulatedMessages = await toModelMessages(override);
+                          laneCompacted = false;
+                          laneInjections = [];
                           locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                           // Update event so onTurnComplete sees compacted messages
                           turnCompleteEvent.messages = accumulatedMessages;
@@ -9268,6 +9399,8 @@ function chatAgent<
                           locals.set(chatOverrideMessagesKey, undefined);
                           accumulatedUIMessages = [...turnCompleteOverride] as TUIMessage[];
                           accumulatedMessages = await toModelMessages(turnCompleteOverride);
+                          laneCompacted = false;
+                          laneInjections = [];
                           locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                         }
                       },
@@ -9602,6 +9735,7 @@ function chatAgent<
             let erroredNewModelMessages: ModelMessage[] = [];
 
             const reconciledSteer = reconcilePendingSteer();
+            const backgroundTailThisTurn = reconcilePendingBackground();
 
             if (!responseCommitted) {
               try {
@@ -9637,13 +9771,16 @@ function chatAgent<
                       accumulatedMessages,
                       erroredUIMessages[partialIdx]!,
                       partialResponse!,
-                      reconciledSteer.reduce((n, e) => n + e.model.length, 0)
+                      reconciledSteer.reduce((n, e) => n + e.model.length, 0) +
+                        backgroundTailThisTurn
                     );
                     if (!ok) {
                       logger.warn(
                         "chat.agent: replaced partial not found at the model lane tail; reconverting the lane"
                       );
                       accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial);
+                      laneCompacted = false;
+                      laneInjections = [];
                     }
                   }
                   accumulatedUIMessages = erroredUIMessagesWithPartial;