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,
@@ -6870,8 +6889,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;
@@ -6919,6 +6948,24 @@ function chatAgent<
         const { changes, shadow } = diffTranscript(transcriptShadow, opts.messages, {
           nonFinalIds: opts.nonFinalIds,
         });
+        const lastId = opts.messages.at(-1)?.id;
+        const runtimeState: TranscriptRuntimeState | null =
+          laneCompacted && lastId !== undefined
+            ? {
+                v: 1,
+                compaction: {
+                  modelMessages: accumulatedMessages,
+                  throughId: lastId,
+                  fingerprint: prefixFingerprint(shadow, lastId),
+                },
+              }
+            : laneInjections.length > 0
+              ? { v: 1, injections: laneInjections }
+              : null;
+        if (runtimeState !== null || persistedStateSet) {
+          changes.push({ op: "state", value: runtimeState } satisfies TranscriptChange);
+        }
+        transcriptState = runtimeState;
         const inCursor = chatInputRouter().resumeFloor();
         await transcriptStorage.save(
           {
@@ -6947,6 +6994,7 @@ function chatAgent<
           }
         );
         transcriptShadow = shadow;
+        persistedStateSet = runtimeState !== null;
       };
 
       /**
@@ -7031,6 +7079,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,
@@ -7375,7 +7425,14 @@ function chatAgent<
             }
           }
           try {
-            accumulatedMessages = await toModelMessages(accumulatedUIMessages);
+            const restored = await restoreModelLane(
+              accumulatedUIMessages,
+              parseTranscriptRuntimeState(bootTranscriptState),
+              (messages) => toModelMessages(messages)
+            );
+            accumulatedMessages = restored.messages;
+            laneCompacted = restored.compacted;
+            laneInjections = restored.injections;
           } catch (error) {
             logger.warn("chat.agent: toModelMessages failed at boot; starting empty", {
               error: error instanceof Error ? error.message : String(error),
@@ -8047,6 +8104,8 @@ function chatAgent<
                     );
                     accumulatedUIMessages = [...hydrated] as TUIMessage[];
                     accumulatedMessages = await toModelMessages(hydrated);
+                    laneCompacted = false;
+                    laneInjections = [];
                     locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                   }
 
@@ -8084,6 +8143,8 @@ function chatAgent<
                       locals.set(chatOverrideMessagesKey, undefined);
                       accumulatedUIMessages = [...actionOverride] as TUIMessage[];
                       accumulatedMessages = await toModelMessages(actionOverride);
+                      laneCompacted = false;
+                      laneInjections = [];
                       locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
 
                       actionChangedHistory = true;
@@ -8216,6 +8277,8 @@ function chatAgent<
 
                     accumulatedUIMessages = merged;
                     accumulatedMessages = await toModelMessages(merged);
+                    laneCompacted = false;
+                    laneInjections = [];
                     locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
 
                     // Track new messages for onTurnComplete.newUIMessages.
@@ -8265,6 +8328,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
@@ -8318,6 +8383,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);
@@ -8499,6 +8566,8 @@ function chatAgent<
                           locals.set(chatOverrideMessagesKey, undefined);
                           accumulatedUIMessages = [...turnStartOverride] as TUIMessage[];
                           accumulatedMessages = await toModelMessages(turnStartOverride);
+                          laneCompacted = false;
+                          laneInjections = [];
                           locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                         }
                       },
@@ -8564,7 +8633,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) {
@@ -8757,6 +8831,8 @@ function chatAgent<
                     accumulatedMessages = await toModelMessages(
                       runOverride.filter((m) => !pendingIds.has(m.id))
                     );
+                    laneCompacted = false;
+                    laneInjections = [];
                     locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                   }
 
@@ -8782,6 +8858,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) {
@@ -8874,6 +8952,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);
@@ -8993,6 +9073,9 @@ function chatAgent<
                                     },
                                   ];
 
+                              laneCompacted = true;
+                              laneInjections = [];
+
                               // UI messages: callback or default (preserve all)
                               if (outerCompaction.compactUIMessages) {
                                 accumulatedUIMessages = (await outerCompaction.compactUIMessages(
@@ -9097,6 +9180,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;
@@ -9156,6 +9241,8 @@ function chatAgent<
                           locals.set(chatOverrideMessagesKey, undefined);
                           accumulatedUIMessages = [...turnCompleteOverride] as TUIMessage[];
                           accumulatedMessages = await toModelMessages(turnCompleteOverride);
+                          laneCompacted = false;
+                          laneInjections = [];
                           locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                         }
                       },
@@ -9524,6 +9611,8 @@ function chatAgent<
                         "chat.agent: replaced partial not found at the model lane tail; reconverting the lane"
                       );
                       accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial);
+                      laneCompacted = false;
+                      laneInjections = [];
                     }
                   }
                   accumulatedUIMessages = erroredUIMessagesWithPartial;