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

A continuation boot replays unacknowledged user messages off session.in and dispatches them itself. Suppressing the tail's re-answer by folding every recovered message into the resume cursor let the cursor advance past a message the run had not answered yet, so a crash mid-recovery dropped the rest.

Recovered messages are now claimed on the session-stream router: a claim drops a late tail re-delivery (no double answer) and holds the resume cursor behind the message until the boot settles it (no drop). The boot settles each as it dispatches it, or immediately for ones folded into the seed chain or skipped.

Author: ericallam

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

@@ -573,3 +573,67 @@ 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("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

@@ -2472,7 +2472,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;
@@ -2503,20 +2507,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);
@@ -7211,6 +7208,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
@@ -7355,18 +7365,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,
       });
 
@@ -7456,7 +7459,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) {
@@ -7465,11 +7468,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
@@ -7500,12 +7512,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,
@@ -7514,7 +7527,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;
@@ -7683,7 +7706,7 @@ function chatAgent<
          */
         let dispatchedRecoveredFirstTurn = false;
         if (preloaded && bootInjectedQueue.length > 0) {
-          currentWirePayload = bootInjectedQueue.shift()!;
+          currentWirePayload = dispatchBootInjected();
           dispatchedRecoveredFirstTurn = true;
         }
 
@@ -7934,7 +7957,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 =
@@ -9478,7 +9501,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";
                 }
 
@@ -9849,7 +9872,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;
             }