diff --git a/.changeset/chat-agent-recovery-message-loss.md b/.changeset/chat-agent-recovery-message-loss.md new file mode 100644 index 00000000000..c8d0ba3980d --- /dev/null +++ b/.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. diff --git a/packages/core/src/v3/sessionStreams/router.test.ts b/packages/core/src/v3/sessionStreams/router.test.ts index 7fd57465cae..f9a450a2543 100644 --- a/packages/core/src/v3/sessionStreams/router.test.ts +++ b/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" }); + }); +}); diff --git a/packages/core/src/v3/sessionStreams/router.ts b/packages/core/src/v3/sessionStreams/router.ts index 53669443a20..bb201aab711 100644 --- a/packages/core/src/v3/sessionStreams/router.ts +++ b/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(); + #owed = new Set(); #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): 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(); } } diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 4910b18cb09..049cb7ce06e 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/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 { 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); @@ -7270,6 +7267,19 @@ function chatAgent< // `messagesInput.waitWithIdleTimeout` so recovered turns fire first. const bootInjectedQueue: ChatTaskWirePayload>[] = []; + const recoveredSeqByPayload = new WeakMap< + ChatTaskWirePayload>, + number + >(); + const dispatchBootInjected = (): ChatTaskWirePayload< + TUIMessage, + inferSchemaIn + > => { + 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 @@ -7435,18 +7445,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, }); @@ -7536,7 +7539,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) { @@ -7545,11 +7548,20 @@ function chatAgent< seedChain = settledMessages; } if (hookRecoveredTurns !== undefined) { - recoveredTurns = hookRecoveredTurns; + const seqNumByRecoveredId = new Map(); + 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 @@ -7580,12 +7592,13 @@ function chatAgent< for (const entry of replayedInTail) { metadataById.set(entry.message.id, entry.metadata); } - for (const msg of recoveredTurns) { + const dispatchedRecoveredSeqs = new Set(); + 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, @@ -7594,7 +7607,17 @@ function chatAgent< messageId: msg.id, continuation: payload.continuation, previousRunId: payload.previousRunId, - } as ChatTaskWirePayload>); + } as ChatTaskWirePayload>; + 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; @@ -7778,7 +7801,7 @@ function chatAgent< */ let dispatchedRecoveredFirstTurn = false; if (preloaded && bootInjectedQueue.length > 0) { - currentWirePayload = bootInjectedQueue.shift()!; + currentWirePayload = dispatchBootInjected(); dispatchedRecoveredFirstTurn = true; } @@ -8029,7 +8052,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 = @@ -9610,7 +9633,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"; } @@ -9986,7 +10009,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; } diff --git a/packages/trigger-sdk/test/chat-agent-recovery-floor.test.ts b/packages/trigger-sdk/test/chat-agent-recovery-floor.test.ts new file mode 100644 index 00000000000..2af70357840 --- /dev/null +++ b/packages/trigger-sdk/test/chat-agent-recovery-floor.test.ts @@ -0,0 +1,115 @@ +import { mockChatAgent } from "../src/v3/test/index.js"; + +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream, streamText } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import { describe, expect, it } from "vitest"; +import { chat } from "../src/v3/ai.js"; + +function userMessage(text: string, id: string) { + return { id, role: "user" as const, parts: [{ type: "text" as const, text }] }; +} + +function textStream(text: string) { + const chunks: LanguageModelV3StreamPart[] = [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: text }, + { type: "text-end", id: "t1" }, + { + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined }, + outputTokens: { total: 10, text: 10, reasoning: undefined }, + }, + }, + ]; + return simulateReadableStream({ chunks }); +} + +function lastUserText(prompt: unknown): string { + const msgs = Array.isArray(prompt) ? prompt : []; + for (let i = msgs.length - 1; i >= 0; i--) { + const m = msgs[i] as { role?: string; content?: unknown }; + if (m?.role !== "user") continue; + if (typeof m.content === "string") return m.content; + if (Array.isArray(m.content)) + return (m.content as Array<{ text?: string }>).map((p) => p?.text ?? "").join(""); + } + return ""; +} + +describe("TRI-13752: chat.agent version handover duplicates messages and turns", () => { + it("effect #1: a continuation boot answers a handed-over session.in message exactly once", async () => { + const answered: string[] = []; + const model = new MockLanguageModelV3({ + doStream: async (options) => { + answered.push(lastUserText((options as { prompt?: unknown }).prompt)); + return { stream: textStream("ok") }; + }, + }); + const u1 = userMessage("the handed-over message", "u-1"); + const agent = chat.agent({ + id: "tri-13752.continuation-double-dispatch", + run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { + chatId: "tri-13752-cont", + continuation: true, + previousRunId: "run_prior", + }); + harness.seedSessionInTail([u1 as never]); + try { + await harness.deliverSessionInAtSeq(u1 as never, 1); + await new Promise((r) => setTimeout(r, 200)); + const u1Answers = answered.filter((t) => t.includes("the handed-over message")); + expect(u1Answers).toHaveLength(1); + } finally { + await harness.close(); + } + }); + + it("N>1: the floor published after the first recovered turn must not cover the un-dispatched second message", async () => { + const answered: string[] = []; + const model = new MockLanguageModelV3({ + doStream: async (options) => { + answered.push(lastUserText((options as { prompt?: unknown }).prompt)); + return { stream: textStream("ok") }; + }, + }); + const u1 = userMessage("first in-flight", "u-1"); + const u2 = userMessage("second in-flight", "u-2"); + const agent = chat.agent({ + id: "tri-13752.n-gt-1-recovery-floor", + run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }), + }); + const harness = mockChatAgent(agent, { + chatId: "tri-13752-nrec", + continuation: true, + previousRunId: "run_prior", + }); + harness.seedSessionInTail([u1 as never, u2 as never]); + try { + const deadline = Date.now() + 2000; + while ( + harness.allRawChunks.filter( + (c) => (c as { type?: string }).type === "trigger:turn-complete" + ).length < 2 && + Date.now() < deadline + ) { + await new Promise((r) => setTimeout(r, 20)); + } + + expect(answered.filter((t) => t.includes("first in-flight"))).toHaveLength(1); + expect(answered.filter((t) => t.includes("second in-flight"))).toHaveLength(1); + + const firstTurnComplete = harness.allRawChunks.find( + (c) => (c as { type?: string }).type === "trigger:turn-complete" + ) as { sessionInEventId?: string } | undefined; + const publishedFloor = Number(firstTurnComplete?.sessionInEventId); + expect(publishedFloor).toBeLessThan(2); + } finally { + await harness.close(); + } + }); +});