fix(core,sdk): stop chat.agent losing messages during recovery (#4907)

## Summary

When a chat.agent run boots to continue a session (a version handover,
or a retry after a crash), it replays the unacknowledged user messages
off `session.in` and dispatches them itself. A previous change stopped
the live tail from re-answering those same messages by folding them into
the resume cursor in one step. That cursor is what the next boot reads
to know where to resume, and folding in every recovered message at once
let it advance past a message the run had not answered yet. So if the
run answered the first recovered message, wrote its turn boundary, then
crashed before dispatching the rest, the next boot resumed past those
messages and they were never answered.

## Fix

A recovered message is now claimed on the session-stream router instead
of folded into the cursor. A claim does two independent things:

- it drops the message however late the live tail re-delivers it, so a
recovered message is never answered twice;
- it holds the resume cursor behind that message until the boot has
dispatched it, so a turn boundary never publishes a cursor past a
message still waiting for a turn.

The boot settles each claim as it dispatches the message, or right away
for a message it folds into the seed chain or deliberately skips, so the
cursor only advances over messages that have actually been handled. A
claimed record whose route re-read never arrives over the tail degrades
to being answered twice on the next boot, never to being dropped.

Covered by router-level unit tests for the claim/settle floor and a
chat.agent boot test asserting the cursor published after the first
recovered turn stays behind the still-unanswered ones.

Author: ericallam

.changeset/chat-agent-recovery-message-loss.md

@@ -0,0 +1,6 @@
+---
+"@trigger.dev/sdk": patch
+"@trigger.dev/core": patch
+---
+
+`chat.agent`: a run that recovers a session with more than one in-flight user message no longer drops the unanswered ones if it restarts mid-recovery. Recovered messages now hold the resume cursor until each has been answered, so a restart re-answers the rest instead of resuming past them. Previously the cursor could advance past messages that were only held in memory, so a crash before they were dispatched lost them.

packages/core/src/v3/sessionStreams/router.test.ts

@@ -573,3 +573,76 @@ describe("SessionChannelRouter: untake", () => {
     expect(r.pendingCount("messages")).toBe(1);
   });
 });
+
+describe("SessionChannelRouter: recovered claim/settle floor", () => {
+  it("drops a claimed record on ingest instead of queueing it", () => {
+    const r = router();
+    r.restore({ resumeFrom: 0 });
+    r.markRecovered([1, 2]);
+    expect(r.ingest(rec(1, "message"))).toEqual({ action: "drop", reason: "recovered" });
+    expect(r.hasPending("messages")).toBe(false);
+  });
+
+  it("holds the resume floor below the earliest owed record until it is settled", () => {
+    const r = router();
+    r.restore({ resumeFrom: 0 });
+    r.markRecovered([1, 2]);
+    expect(r.resumeFloor()).toBe(0);
+    r.settleRecovered(1);
+    expect(r.resumeFloor()).toBe(1);
+    r.settleRecovered(2);
+    expect(r.resumeFloor()).toBe(2);
+  });
+
+  it("advances the floor after settling even when the tail never re-delivers", () => {
+    const r = router();
+    r.restore({ resumeFrom: 0 });
+    r.markRecovered([1, 2]);
+    r.settleRecovered(1);
+    r.settleRecovered(2);
+    expect(r.resumeFloor()).toBe(2);
+    expect(r.appliedThrough()).toBe(2);
+  });
+
+  it("keeps dropping a claimed record after it is settled, so a late tail re-read is never answered", () => {
+    const r = router();
+    r.restore({ resumeFrom: 0 });
+    r.markRecovered([1]);
+    r.settleRecovered(1);
+    expect(r.ingest(rec(1, "message"))).toEqual({ action: "drop", reason: "recovered" });
+    expect(r.hasPending("messages")).toBe(false);
+  });
+
+  it("queues a live record whose sequence was never claimed", () => {
+    const r = router();
+    r.restore({ resumeFrom: 0 });
+    r.markRecovered([1, 2]);
+    expect(r.ingest(rec(3, "message"))).toEqual({ action: "queue", route: "messages" });
+  });
+
+  it("advances only over the contiguous claimed run, holding the floor below a gap", () => {
+    const r = router();
+    r.restore({ resumeFrom: 0 });
+    r.markRecovered([1, 3]);
+    r.settleRecovered(1);
+    r.settleRecovered(3);
+    expect(r.resumeFloor()).toBe(1);
+  });
+
+  it("holds the floor below an unclaimed gap while later claims are owed and the tail is silent", () => {
+    const r = router();
+    r.restore({ resumeFrom: 0 });
+    r.markRecovered([1, 2, 5, 6]);
+    r.settleRecovered(1);
+    r.settleRecovered(2);
+    expect(r.resumeFloor()).toBe(2);
+  });
+
+  it("clears claims and owed records on reset", () => {
+    const r = router();
+    r.restore({ resumeFrom: 0 });
+    r.markRecovered([1, 2]);
+    r.reset();
+    expect(r.ingest(rec(1, "message"))).toEqual({ action: "queue", route: "messages" });
+  });
+});

packages/core/src/v3/sessionStreams/router.ts

@@ -61,7 +61,14 @@ export type RouterDropReason =
    */
   | "replayed"
   /** An `at-arrival` record with no handler attached right now. */
-  | "no-handler";
+  | "no-handler"
+  /**
+   * A record the boot took responsibility for over HTTP (see
+   * {@link SessionChannelRouter.markRecovered}). It is dropped however late the
+   * live tail re-delivers it, so a recovered message is never also answered as
+   * a router turn.
+   */
+  | "recovered";
 
 export type RouterDecision =
   /** Handed to a consumer that was already waiting, or to a live handler. */
@@ -143,6 +150,8 @@ export class SessionChannelRouter {
   #highestSeq: number | undefined;
   #resumeFrom: number | undefined;
   #appliedThrough: number | undefined;
+  #claimed = new Set<number>();
+  #owed = new Set<number>();
   #onDrop?: (record: SessionStreamRecord, reason: RouterDropReason, route?: string) => void;
 
   constructor(
@@ -197,6 +206,57 @@ export class SessionChannelRouter {
     return this.#resumeFrom;
   }
 
+  /**
+   * Declare the sequences a continuation boot already read over HTTP and took
+   * responsibility for dispatching itself, so the live tail's re-read of the
+   * same records does not answer them a second time.
+   *
+   * Two effects, deliberately separate:
+   *
+   * - Every claimed sequence is dropped in {@link ingest} however late it
+   *   arrives (`"recovered"`), so a record the boot owns never also becomes a
+   *   router turn. This survives the boot settling it — the boot owns its
+   *   disposition for the whole run, and the tail may re-deliver at any time.
+   * - Each claimed sequence is also *owed*: it holds the resume floor exactly
+   *   like a queued record would, until the boot {@link settleRecovered}s it.
+   *   Suppressing the second delivery without this would let a turn boundary
+   *   publish a floor past a message the boot has not answered yet, turning a
+   *   duplicate into a dropped message.
+   *
+   * `#highestSeq` is advanced over the contiguous run of claimed sequences from
+   * the current high water, so once every claim is settled the floor can move
+   * past them even if the tail never re-delivers (a silent tail then degrades
+   * to a duplicate, never a loss). The advance stops at the first gap so a
+   * record the boot left for the router is never skipped.
+   */
+  markRecovered(seqNums: Iterable<number>): void {
+    const nums = [...seqNums].filter((n) => Number.isFinite(n)).sort((a, b) => a - b);
+    if (nums.length === 0) return;
+    for (const n of nums) {
+      this.#claimed.add(n);
+      this.#owed.add(n);
+    }
+    let base = this.#highestSeq ?? nums[0]! - 1;
+    while (this.#claimed.has(base + 1)) base++;
+    if (this.#highestSeq === undefined || base > this.#highestSeq) {
+      this.#highestSeq = base;
+    }
+    const maxClaimed = nums[nums.length - 1]!;
+    if (this.#appliedThrough === undefined || maxClaimed > this.#appliedThrough) {
+      this.#appliedThrough = maxClaimed;
+    }
+  }
+
+  /**
+   * Release a claimed sequence's hold on the resume floor once the boot has
+   * decided its disposition (dispatched it as a turn, folded it into the seed
+   * chain, or deliberately dropped it). It stays claimed, so a late tail
+   * re-delivery is still dropped rather than answered again.
+   */
+  settleRecovered(seqNum: number): void {
+    this.#owed.delete(seqNum);
+  }
+
   /**
    * Classify one record and act on it. The record's destination is decided
    * here, once, and never by whichever consumer happens to be waiting.
@@ -213,6 +273,10 @@ export class SessionChannelRouter {
       }
     }
 
+    if (this.#claimed.has(record.seqNum)) {
+      return this.#drop(record, "recovered");
+    }
+
     const kind = this.#kindOf(record.data);
     if (kind === undefined) {
       return this.#drop(record, "malformed");
@@ -308,6 +372,9 @@ export class SessionChannelRouter {
       const pending = state.earliestUnrecovered();
       if (pending !== undefined) earliestPending = Math.min(earliestPending, pending);
     }
+    for (const owed of this.#owed) {
+      if (owed < earliestPending) earliestPending = owed;
+    }
 
     if (earliestPending === Infinity) return this.#highestSeq;
 
@@ -518,5 +585,7 @@ export class SessionChannelRouter {
     this.#highestSeq = undefined;
     this.#resumeFrom = undefined;
     this.#appliedThrough = undefined;
+    this.#claimed.clear();
+    this.#owed.clear();
   }
 }

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

@@ -2322,7 +2322,11 @@ async function findSessionInReplayWindowEnd(
  */
 async function installChatInputRouter(
   chatId: string,
-  options?: { fallbackResumeFrom?: number; recoveredThrough?: number; resuming?: boolean }
+  options?: {
+    fallbackResumeFrom?: number;
+    recoveredSeqNums?: readonly number[];
+    resuming?: boolean;
+  }
 ): Promise<SessionChannelRouter> {
   const entry = chatInputRouterEntry(chatId);
   if (entry.attached) return entry.router;
@@ -2353,20 +2357,13 @@ async function installChatInputRouter(
     }
   }
 
-  // A boot that replayed `.in` itself has already answered everything up to
-  // `recoveredThrough`, so the floor has to cover it before the tail opens.
-  if (options?.recoveredThrough !== undefined) {
-    const recovered = options.recoveredThrough;
-    checkpoint.resumeFrom = Math.max(checkpoint.resumeFrom ?? recovered, recovered);
-    checkpoint.appliedThrough = Math.max(
-      checkpoint.appliedThrough ?? checkpoint.resumeFrom,
-      checkpoint.resumeFrom
-    );
-  }
-
   const router = entry.router;
   router.restore(checkpoint);
 
+  if (options?.recoveredSeqNums && options.recoveredSeqNums.length > 0) {
+    router.markRecovered(options.recoveredSeqNums);
+  }
+
   const floor = router.resumeFrom();
   if (floor !== undefined) {
     sessionStreams.setLastSeqNum(chatId, "in", floor);
@@ -7272,6 +7269,20 @@ function chatAgent<
       // `messagesInput.waitWithIdleTimeout` so recovered turns fire first.
       const bootInjectedQueue: ChatTaskWirePayload<TUIMessage, inferSchemaIn<TClientDataSchema>>[] =
         [];
+      const recoveredSeqByPayload = new WeakMap<
+        ChatTaskWirePayload<TUIMessage, inferSchemaIn<TClientDataSchema>>,
+        number
+      >();
+      const dispatchBootInjected = (): ChatTaskWirePayload<
+        TUIMessage,
+        inferSchemaIn<TClientDataSchema>
+      > => bootInjectedQueue.shift()!;
+      const settleRecoveredTurn = (
+        wirePayload: ChatTaskWirePayload<TUIMessage, inferSchemaIn<TClientDataSchema>>
+      ) => {
+        const settledSeq = recoveredSeqByPayload.get(wirePayload);
+        if (settledSeq !== undefined) chatInputRouter().settleRecovered(settledSeq);
+      };
       const couldHavePriorState = payload.continuation === true || ctx.attempt.number > 1;
 
       // `.in` resume cursor, computed at most once per boot. The boot
@@ -7437,18 +7448,11 @@ function chatAgent<
 
       // ── session.in router ──────────────────────────────────────────
       //
-      // Reads the turn boundary and subscribes in one call. `bootInCursor` is
-      // only a fallback: the boot block above may already have resolved a
-      // cursor from the snapshot, which is used when the boundary itself
-      // carries none. Everything the boot replayed off `.in` is dispatched from
-      // `bootInjectedQueue` below, so it goes into the floor here — folded in
-      // after the subscription opens, the live tail re-delivers it as a turn.
-      const lastRecoveredInSeq =
-        replayedInTail.length > 0 ? replayedInTail[replayedInTail.length - 1]!.seqNum : undefined;
+      const recoveredSeqNums = replayedInTail.map((r) => r.seqNum);
 
       await installChatInputRouter(payload.chatId, {
         fallbackResumeFrom: bootInCursorResolved ? bootInCursor : undefined,
-        recoveredThrough: lastRecoveredInSeq,
+        recoveredSeqNums,
         resuming: Boolean(payload.continuation) || ctx.attempt.number > 1,
       });
 
@@ -7538,7 +7542,7 @@ function chatAgent<
         // branches: at n=1 the orphan partial is dropped and the interrupted
         // user is re-dispatched as a fresh turn instead.
         let seedChain: TUIMessage[];
-        let recoveredTurns: TUIMessage[];
+        let recoveredEntries: { message: TUIMessage; seqNum: number | undefined }[];
         if (hookChain !== undefined) {
           seedChain = hookChain;
         } else if (partialAssistant !== undefined && inFlightUsers.length > 1) {
@@ -7547,11 +7551,22 @@ function chatAgent<
           seedChain = settledMessages;
         }
         if (hookRecoveredTurns !== undefined) {
-          recoveredTurns = hookRecoveredTurns;
+          const seqNumsByRecoveredId = new Map<string, number[]>();
+          for (const entry of replayedInTail) {
+            const existing = seqNumsByRecoveredId.get(entry.message.id);
+            if (existing) existing.push(entry.seqNum);
+            else seqNumsByRecoveredId.set(entry.message.id, [entry.seqNum]);
+          }
+          recoveredEntries = hookRecoveredTurns.map((message) => ({
+            message,
+            seqNum: seqNumsByRecoveredId.get(message.id)?.shift(),
+          }));
         } else if (partialAssistant !== undefined && inFlightUsers.length > 1) {
-          recoveredTurns = inFlightUsers.slice(1);
+          recoveredEntries = replayedInTail
+            .slice(1)
+            .map((r) => ({ message: r.message, seqNum: r.seqNum }));
         } else {
-          recoveredTurns = inFlightUsers;
+          recoveredEntries = replayedInTail.map((r) => ({ message: r.message, seqNum: r.seqNum }));
         }
         // `beforeBoot` errors bubble — the customer opted into blocking
         // persistence and a failure there should fail the run rather than
@@ -7582,12 +7597,13 @@ function chatAgent<
         for (const entry of replayedInTail) {
           metadataById.set(entry.message.id, entry.metadata);
         }
-        for (const msg of recoveredTurns) {
+        const dispatchedRecoveredSeqs = new Set<number>();
+        for (const { message: msg, seqNum } of recoveredEntries) {
           if (wireMessageId && msg.id === wireMessageId) continue;
           const recoveredMetadata = metadataById.has(msg.id)
             ? metadataById.get(msg.id)
             : payload.metadata;
-          bootInjectedQueue.push({
+          const injectedPayload = {
             chatId: payload.chatId,
             sessionId: payload.sessionId,
             metadata: recoveredMetadata,
@@ -7596,7 +7612,17 @@ function chatAgent<
             messageId: msg.id,
             continuation: payload.continuation,
             previousRunId: payload.previousRunId,
-          } as ChatTaskWirePayload<TUIMessage, inferSchemaIn<TClientDataSchema>>);
+          } as ChatTaskWirePayload<TUIMessage, inferSchemaIn<TClientDataSchema>>;
+          bootInjectedQueue.push(injectedPayload);
+          if (seqNum !== undefined) {
+            recoveredSeqByPayload.set(injectedPayload, seqNum);
+            dispatchedRecoveredSeqs.add(seqNum);
+          }
+        }
+        for (const entry of replayedInTail) {
+          if (!dispatchedRecoveredSeqs.has(entry.seqNum)) {
+            chatInputRouter().settleRecovered(entry.seqNum);
+          }
         }
 
         accumulatedUIMessages = seedChain;
@@ -7780,7 +7806,7 @@ function chatAgent<
          */
         let dispatchedRecoveredFirstTurn = false;
         if (preloaded && bootInjectedQueue.length > 0) {
-          currentWirePayload = bootInjectedQueue.shift()!;
+          currentWirePayload = dispatchBootInjected();
           dispatchedRecoveredFirstTurn = true;
         }
 
@@ -8031,7 +8057,7 @@ function chatAgent<
           // waiting on the live session.in. Subsequent recovered turns
           // get drained by the end-of-turn picker below.
           if (bootInjectedQueue.length > 0) {
-            currentWirePayload = bootInjectedQueue.shift()!;
+            currentWirePayload = dispatchBootInjected();
           } else {
             const effectiveIdleTimeout = idleTimeoutInSeconds ?? payload.idleTimeoutInSeconds;
             const effectiveTurnTimeout =
@@ -8686,6 +8712,7 @@ function chatAgent<
                     chatId: currentWirePayload.chatId,
                     messageId: currentWirePayload.messageId,
                   });
+                  settleRecoveredTurn(currentWirePayload);
                   await writeTurnCompleteChunk(currentWirePayload.chatId);
                   // Not a turn — don't consume an iteration.
                   turn--;
@@ -9504,6 +9531,8 @@ function chatAgent<
                     locals.set(chatResponsePartsKey, []);
                   }
 
+                  settleRecoveredTurn(currentWirePayload);
+
                   // Write turn-complete control chunk — closes the frontend stream.
                   const turnCompleteResult = await writeTurnCompleteChunk(
                     currentWirePayload.chatId,
@@ -9634,7 +9663,7 @@ function chatAgent<
                 // produced these from in-flight user messages on session.in
                 // that the dead predecessor never acknowledged.
                 if (bootInjectedQueue.length > 0) {
-                  currentWirePayload = bootInjectedQueue.shift()!;
+                  currentWirePayload = dispatchBootInjected();
                   return "continue";
                 }
 
@@ -10012,7 +10041,7 @@ function chatAgent<
             // recovered turn shouldn't strand the rest of the boot queue
             // until an unrelated live message arrives.
             if (bootInjectedQueue.length > 0) {
-              currentWirePayload = bootInjectedQueue.shift()!;
+              currentWirePayload = dispatchBootInjected();
               continue;
             }