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

The model lane after a compaction cannot be rebuilt from the transcript, so every continuation used to re-read the whole conversation and summarise it again. The runtime now records the compacted lane in the storage's state slot, with the transcript id it covers and a fingerprint of that prefix, and rebuilds from it at boot when the prefix is unchanged. A rollback or edit that reconverts the lane clears the state in the same changeset as the truncate.

Conversational messages added with chat.inject are recorded the same way, anchored to the transcript message they followed, so they survive a continuation instead of living only in the worker that received them.

Adds an in-memory storage that logs the changesets it receives, and a test-only override for the storage the runtime persists through, so the exact changesets for a turn, a mid-turn steer, a compaction, a rollback and an injection are asserted.

Author: ericallam

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

@@ -76,10 +76,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,
@@ -2515,6 +2534,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.
@@ -4970,6 +4996,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;
@@ -6831,6 +6864,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
@@ -6870,8 +6920,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;
@@ -6921,6 +6981,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(
           {
@@ -6949,6 +7033,7 @@ function chatAgent<
           }
         );
         transcriptShadow = shadow;
+        persistedStateSet = runtimeState !== null;
       };
 
       /**
@@ -7033,6 +7118,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,
@@ -7377,7 +7464,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),
@@ -7906,6 +8007,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
@@ -8049,6 +8151,8 @@ function chatAgent<
                     );
                     accumulatedUIMessages = [...hydrated] as TUIMessage[];
                     accumulatedMessages = await toModelMessages(hydrated);
+                    laneCompacted = false;
+                    laneInjections = [];
                     locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                   }
 
@@ -8086,6 +8190,8 @@ function chatAgent<
                       locals.set(chatOverrideMessagesKey, undefined);
                       accumulatedUIMessages = [...actionOverride] as TUIMessage[];
                       accumulatedMessages = await toModelMessages(actionOverride);
+                      laneCompacted = false;
+                      laneInjections = [];
                       locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
 
                       actionChangedHistory = true;
@@ -8218,6 +8324,8 @@ function chatAgent<
 
                     accumulatedUIMessages = merged;
                     accumulatedMessages = await toModelMessages(merged);
+                    laneCompacted = false;
+                    laneInjections = [];
                     locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
 
                     // Track new messages for onTurnComplete.newUIMessages.
@@ -8267,6 +8375,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
@@ -8320,6 +8430,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);
@@ -8501,6 +8613,8 @@ function chatAgent<
                           locals.set(chatOverrideMessagesKey, undefined);
                           accumulatedUIMessages = [...turnStartOverride] as TUIMessage[];
                           accumulatedMessages = await toModelMessages(turnStartOverride);
+                          laneCompacted = false;
+                          laneInjections = [];
                           locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                         }
                       },
@@ -8566,7 +8680,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) {
@@ -8759,6 +8878,8 @@ function chatAgent<
                     accumulatedMessages = await toModelMessages(
                       runOverride.filter((m) => !pendingIds.has(m.id))
                     );
+                    laneCompacted = false;
+                    laneInjections = [];
                     locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                   }
 
@@ -8784,6 +8905,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) {
@@ -8802,9 +8925,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
@@ -8876,6 +9000,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);
@@ -8995,6 +9121,9 @@ function chatAgent<
                                     },
                                   ];
 
+                              laneCompacted = true;
+                              laneInjections = [];
+
                               // UI messages: callback or default (preserve all)
                               if (outerCompaction.compactUIMessages) {
                                 accumulatedUIMessages = (await outerCompaction.compactUIMessages(
@@ -9099,6 +9228,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;
@@ -9158,6 +9289,8 @@ function chatAgent<
                           locals.set(chatOverrideMessagesKey, undefined);
                           accumulatedUIMessages = [...turnCompleteOverride] as TUIMessage[];
                           accumulatedMessages = await toModelMessages(turnCompleteOverride);
+                          laneCompacted = false;
+                          laneInjections = [];
                           locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                         }
                       },
@@ -9484,6 +9617,7 @@ function chatAgent<
             let erroredNewModelMessages: ModelMessage[] = [];
 
             const reconciledSteer = reconcilePendingSteer();
+            const backgroundTailThisTurn = reconcilePendingBackground();
 
             if (!responseCommitted) {
               try {
@@ -9519,13 +9653,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;