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

.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

@@ -2319,7 +2319,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;
@@ -2350,20 +2354,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);
@@ -7271,6 +7268,19 @@ 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>
+      > => {
+        const injected = bootInjectedQueue.shift()!;
+        const settledSeq = recoveredSeqByPayload.get(injected);
+        if (settledSeq !== undefined) chatInputRouter().settleRecovered(settledSeq);
+        return injected;
+      };
       const couldHavePriorState = payload.continuation === true || ctx.attempt.number > 1;
 
       // `.in` resume cursor, computed at most once per boot. The boot
@@ -7436,18 +7446,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,
       });
 
@@ -7537,7 +7540,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) {
@@ -7546,11 +7549,20 @@ function chatAgent<
           seedChain = settledMessages;
         }
         if (hookRecoveredTurns !== undefined) {
-          recoveredTurns = hookRecoveredTurns;
+          const seqNumByRecoveredId = new Map<string, number>();
+          for (const entry of replayedInTail) {
+            seqNumByRecoveredId.set(entry.message.id, entry.seqNum);
+          }
+          recoveredEntries = hookRecoveredTurns.map((message) => ({
+            message,
+            seqNum: seqNumByRecoveredId.get(message.id),
+          }));
         } 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
@@ -7581,12 +7593,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,
@@ -7595,7 +7608,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;
@@ -7779,7 +7802,7 @@ function chatAgent<
          */
         let dispatchedRecoveredFirstTurn = false;
         if (preloaded && bootInjectedQueue.length > 0) {
-          currentWirePayload = bootInjectedQueue.shift()!;
+          currentWirePayload = dispatchBootInjected();
           dispatchedRecoveredFirstTurn = true;
         }
 
@@ -8030,7 +8053,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 =
@@ -9611,7 +9634,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";
                 }
 
@@ -9987,7 +10010,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;
             }