diff --git a/.changeset/chat-errored-turn-handler-leak.md b/.changeset/chat-errored-turn-handler-leak.md new file mode 100644 index 0000000000..7d9ac910be --- /dev/null +++ b/.changeset/chat-errored-turn-handler-leak.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Fix chat turns that throw (for example from an `onTurnStart` hook) leaking their message listener, which lost or duplicated messages sent during later turns. diff --git a/.changeset/chat-pending-message-drain.md b/.changeset/chat-pending-message-drain.md new file mode 100644 index 0000000000..8a847b4799 --- /dev/null +++ b/.changeset/chat-pending-message-drain.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Fix `chat.agent` and `chat.createSession` permanently dropping user messages when several arrived during a single turn: every buffered message is now dispatched as its own turn instead of only the first. diff --git a/.changeset/chat-session-in-resume-cursor.md b/.changeset/chat-session-in-resume-cursor.md new file mode 100644 index 0000000000..65b40aea29 --- /dev/null +++ b/.changeset/chat-session-in-resume-cursor.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Fix chat continuation runs replaying already-answered messages: turns delivered while the run was suspended now advance the session.in resume cursor, so a new run picks up exactly where the previous one left off. diff --git a/.changeset/chat-session-stop-window.md b/.changeset/chat-session-stop-window.md new file mode 100644 index 0000000000..4531c8055c --- /dev/null +++ b/.changeset/chat-session-stop-window.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Fix `chat.createSession` swallowing a message sent shortly after stopping a turn: the turn's message listener now detaches when the stream settles, so those messages run as the next turn. diff --git a/docs/ai-chat/pending-messages.mdx b/docs/ai-chat/pending-messages.mdx index 3e7c00b26f..5aa668099a 100644 --- a/docs/ai-chat/pending-messages.mdx +++ b/docs/ai-chat/pending-messages.mdx @@ -8,7 +8,9 @@ description: "Inject user messages mid-execution to steer agents between tool-ca When an AI agent is executing tool calls, users may want to send a message that **steers the agent mid-execution** — adding context, correcting course, or refining the request without waiting for the response to finish. -The `pendingMessages` option enables this by injecting user messages between tool-call steps via the AI SDK's `prepareStep`. Messages that arrive during streaming are queued and injected at the next step boundary. If there are no more step boundaries (single-step response or final text generation), the message becomes the next turn automatically. +By default (without `pendingMessages`), a message sent while the agent is responding never interrupts the in-flight response: it's buffered and processed as its own turn once the current turn completes, with multiple messages running sequentially in arrival order. + +The `pendingMessages` option enables steering instead, injecting user messages between tool-call steps via the AI SDK's `prepareStep`. Messages that arrive during streaming are queued and injected at the next step boundary. If there are no more step boundaries (single-step response or final text generation), the message becomes the next turn automatically. ## How it works diff --git a/packages/core/src/v3/test/test-session-stream-manager.ts b/packages/core/src/v3/test/test-session-stream-manager.ts index 9339f9be70..0e08441d4c 100644 --- a/packages/core/src/v3/test/test-session-stream-manager.ts +++ b/packages/core/src/v3/test/test-session-stream-manager.ts @@ -33,6 +33,7 @@ export class TestSessionStreamManager implements SessionStreamManager { private onceWaiters = new Map(); private buffer = new Map(); private seqNums = new Map(); + private dispatchedSeqNums = new Map(); on(sessionId: string, io: SessionChannelIO, handler: Handler): { off: () => void } { const key = keyFor(sessionId, io); @@ -150,15 +151,20 @@ export class TestSessionStreamManager implements SessionStreamManager { this.seqNums.set(keyFor(sessionId, io), seqNum); } - lastDispatchedSeqNum(_sessionId: string, _io: SessionChannelIO): number | undefined { - // The test harness drives records via `__sendFromTest` without seq - // numbers, so the committed-consume cursor stays undefined. Tests - // that need cursor behaviour exercise it via the real manager. - return undefined; + lastDispatchedSeqNum(sessionId: string, io: SessionChannelIO): number | undefined { + // `__sendFromTest` carries no seq numbers, so this only reflects + // explicit `setLastDispatchedSeqNum` calls (e.g. the waitpoint + // delivery path). Full cursor behaviour is exercised via the real + // manager. + return this.dispatchedSeqNums.get(keyFor(sessionId, io)); } - setLastDispatchedSeqNum(_sessionId: string, _io: SessionChannelIO, _seqNum: number): void { - // no-op — see comment on `lastDispatchedSeqNum`. + setLastDispatchedSeqNum(sessionId: string, io: SessionChannelIO, seqNum: number): void { + const key = keyFor(sessionId, io); + const current = this.dispatchedSeqNums.get(key); + if (current === undefined || seqNum > current) { + this.dispatchedSeqNums.set(key, seqNum); + } } setMinTimestamp( @@ -202,6 +208,7 @@ export class TestSessionStreamManager implements SessionStreamManager { this.handlers.clear(); this.buffer.clear(); this.seqNums.clear(); + this.dispatchedSeqNums.clear(); } disconnect(): void { diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index e13f7a74f3..ee70c79c75 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -5569,6 +5569,14 @@ function chatAgent< // `messagesInput.waitWithIdleTimeout` so recovered turns fire first. const bootInjectedQueue: ChatTaskWirePayload>[] = []; + // Messages consumed by a turn's `messagesInput.on` handler, dispatched + // one per turn by the end-of-turn pickup. Loop-level on purpose: + // consuming a record advances the committed `.in` cursor, so entries + // dropped with a turn-local buffer are lost permanently. + const pendingWireMessages: ChatTaskWirePayload< + TUIMessage, + inferSchemaIn + >[] = []; const couldHavePriorState = payload.continuation === true || ctx.attempt.number > 1; // `.in` resume cursor, computed at most once per boot. The boot @@ -6378,6 +6386,9 @@ function chatAgent< } for (let turn = 0; turn < maxTurns; turn++) { + // Declared here so the finally can detach it — a handler leaked past + // its turn duplicates every mid-stream message into the shared buffer. + let turnMsgSub: { off: () => void } | undefined; try { // Extract turn-level context before entering the span. Slim // wire: at most one delta message per record. `headStartMessages` @@ -6479,11 +6490,6 @@ function chatAgent< const cancelSignal = runSignal; const combinedSignal = AbortSignal.any([runSignal, stopController.signal]); - // Buffer messages that arrive during streaming - const pendingMessages: ChatTaskWirePayload< - TUIMessage, - inferSchemaIn - >[] = []; const pmConfig = locals.get(chatPendingMessagesKey); const msgSub = messagesInput.on(async (msg) => { // If pendingMessages is configured, route to the steering queue @@ -6532,10 +6538,11 @@ function chatAgent< } // No pendingMessages config — standard wire buffer for next turn - pendingMessages.push( + pendingWireMessages.push( msg as ChatTaskWirePayload> ); }); + turnMsgSub = msgSub; // Track new messages for this turn (user input + assistant response). const turnNewModelMessages: ModelMessage[] = []; @@ -7738,9 +7745,10 @@ function chatAgent< } // If messages arrived during streaming (without pendingMessages config), - // use the first one immediately as the next turn. - if (pendingMessages.length > 0) { - currentWirePayload = pendingMessages[0]!; + // dispatch the oldest as the next turn. The rest stay queued + // and drain one per turn. + if (pendingWireMessages.length > 0) { + currentWirePayload = pendingWireMessages.shift()!; return "continue"; } @@ -7847,6 +7855,9 @@ function chatAgent< // Turn error handler: write an error chunk + turn-complete to the stream // so the client sees the error, then wait for the next message instead // of killing the entire run. This keeps the conversation alive. + // Detach the turn's message handler first — left attached it would + // eat the very message the wait below is waiting for. + turnMsgSub?.off(); if ( turnError instanceof Error && turnError.name === "AbortError" && @@ -7982,6 +7993,12 @@ function chatAgent< continue; } + // Same for messages buffered during the errored turn — already consumed, idling strands them. + if (pendingWireMessages.length > 0) { + currentWirePayload = pendingWireMessages.shift()!; + continue; + } + // Wait for the next message — same as after a successful turn const effectiveIdleTimeout = (metadata.get(IDLE_TIMEOUT_METADATA_KEY) as number | undefined) ?? @@ -8004,6 +8021,8 @@ function chatAgent< inferSchemaIn >; // Continue to next iteration of the for loop + } finally { + turnMsgSub?.off(); } } } finally { @@ -9309,9 +9328,18 @@ function createChatSession( const accumulator = new ChatMessageAccumulator(); let previousTurnUsage: LanguageModelUsage | undefined; let cumulativeUsage: LanguageModelUsage = emptyUsage(); + // Messages consumed mid-turn, dispatched one per next(). Iterator-level + // for the same reason as the agent loop's `pendingWireMessages`: + // consumed records never replay, so a turn-local buffer loses them. + const sessionPendingWire: ChatTaskWirePayload[] = []; + // The current turn's message subscription — detached defensively at the + // top of next() in case user code threw without complete()/done(). + let activeMsgSub: { off: () => void } | undefined; return { async next(): Promise> { + activeMsgSub?.off(); + activeMsgSub = undefined; if (!booted) { booted = true; await seedSessionInResumeCursorForCustomLoop(currentPayload); @@ -9380,24 +9408,29 @@ function createChatSession( } } - // Subsequent turns: wait for the next message + // Subsequent turns: drain buffered mid-turn messages first (they + // were consumed and won't be re-delivered), then wait. if (turn > 0) { - // chat.requestUpgrade() / chat.endRun() — exit before waiting - if (locals.get(chatUpgradeRequestedKey) || locals.get(chatEndRunRequestedKey)) { - stop.cleanup(); - return { done: true, value: undefined }; - } + if (sessionPendingWire.length > 0) { + currentPayload = sessionPendingWire.shift()!; + } else { + // chat.requestUpgrade() / chat.endRun() — exit before waiting + if (locals.get(chatUpgradeRequestedKey) || locals.get(chatEndRunRequestedKey)) { + stop.cleanup(); + return { done: true, value: undefined }; + } - const next = await messagesInput.waitWithIdleTimeout({ - idleTimeoutInSeconds, - timeout, - spanName: "waiting for next message", - }); - if (!next.ok || runSignal.aborted) { - stop.cleanup(); - return { done: true, value: undefined }; + const next = await messagesInput.waitWithIdleTimeout({ + idleTimeoutInSeconds, + timeout, + spanName: "waiting for next message", + }); + if (!next.ok || runSignal.aborted) { + stop.cleanup(); + return { done: true, value: undefined }; + } + currentPayload = next.output; } - currentPayload = next.output; } // Check limits @@ -9426,11 +9459,10 @@ function createChatSession( }); // Listen for messages during streaming (steering + next-turn buffer) - const sessionPendingWire: ChatTaskWirePayload[] = []; const sessionMsgSub = messagesInput.on(async (msg) => { - sessionPendingWire.push(msg); - if (sessionPendingMessages) { + // Steering route — the frontend re-sends non-injected + // messages on turn complete, so don't also buffer the wire. // Slim wire: at most one delta message per record. Read // `msg.message` directly — no array slicing needed. const lastUIMessage = msg.message; @@ -9453,8 +9485,12 @@ function createChatSession( /* non-fatal */ } } + return; } + + sessionPendingWire.push(msg); }); + activeMsgSub = sessionMsgSub; // Accumulate messages. Slim wire: pass the single delta message as // a 0-or-1-length array. The accumulator's behavior is unchanged — @@ -9555,6 +9591,10 @@ function createChatSession( } else { throw error; } + } finally { + // Detach at stream end (like the agent loop): the steering queue + // can't inject anymore, so later arrivals must buffer for the next turn. + sessionMsgSub.off(); } if (response) { @@ -9726,6 +9766,8 @@ function createChatSession( }, async return() { + activeMsgSub?.off(); + activeMsgSub = undefined; // `stop` only exists once next() has booted the iterator. stop?.cleanup(); return { done: true, value: undefined }; diff --git a/packages/trigger-sdk/src/v3/sessions.ts b/packages/trigger-sdk/src/v3/sessions.ts index aad2900f27..8a01f8293c 100644 --- a/packages/trigger-sdk/src/v3/sessions.ts +++ b/packages/trigger-sdk/src/v3/sessions.ts @@ -753,11 +753,14 @@ export class SessionInputChannel { : undefined; if (waitResult.ok) { - // Advance the seq counter so the SSE tail doesn't replay the - // record that was consumed via the waitpoint. + // Advance both cursors past the record consumed via the + // waitpoint: the seq counter so the SSE tail doesn't replay + // it, and the consume cursor so turn-completes don't stamp a + // stale `session-in-event-id`. const prevSeq = sessionStreams.lastSeqNum(this.sessionId, "in"); const nextSeq = (prevSeq ?? -1) + 1; sessionStreams.setLastSeqNum(this.sessionId, "in", nextSeq); + sessionStreams.setLastDispatchedSeqNum(this.sessionId, "in", nextSeq); return { ok: true as const, output: data as T }; } else { diff --git a/packages/trigger-sdk/test/pending-message-drain.test.ts b/packages/trigger-sdk/test/pending-message-drain.test.ts new file mode 100644 index 0000000000..f5bd705751 --- /dev/null +++ b/packages/trigger-sdk/test/pending-message-drain.test.ts @@ -0,0 +1,273 @@ +// Import the test harness FIRST — this installs the resource catalog so +// `chat.agent()` calls below register their task functions correctly. +import { mockChatAgent } from "../src/v3/test/index.js"; + +import { describe, expect, it, vi } from "vitest"; +import { chat } from "../src/v3/ai.js"; +import { __setSessionOpenImplForTests, sessions } from "../src/v3/sessions.js"; +import { apiClientManager, sessionStreams } from "@trigger.dev/core/v3"; +import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; +import { simulateReadableStream, streamText } from "ai"; +import { MockLanguageModelV3 } from "ai/test"; +import type { LanguageModelV3StreamPart } from "@ai-sdk/provider"; + +// ── Helpers ──────────────────────────────────────────────────────────── + +function userMessage(text: string, id: string) { + return { + id, + role: "user" as const, + parts: [{ type: "text" as const, text }], + }; +} + +function textStreamChunks(text: string): LanguageModelV3StreamPart[] { + return [ + { 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 }, + }, + }, + ]; +} + +/** Model that answers `ANSWER()`, slowly enough that + * records sent right after the turn starts arrive mid-stream. */ +function echoModel() { + return new MockLanguageModelV3({ + doStream: async ({ prompt }) => { + const users = prompt.filter((m) => m.role === "user"); + const last = users[users.length - 1]; + const text = Array.isArray(last?.content) + ? last.content + .filter((p): p is { type: "text"; text: string } => p.type === "text") + .map((p) => p.text) + .join("") + : ""; + return { + stream: simulateReadableStream({ + chunks: textStreamChunks(`ANSWER(${text})`), + initialDelayInMs: 100, + chunkDelayInMs: 10, + }), + }; + }, + }); +} + +async function waitFor(check: () => boolean, timeoutMs = 10_000) { + const start = Date.now(); + while (Date.now() - start < timeoutMs) { + if (check()) return; + await new Promise((r) => setTimeout(r, 20)); + } + throw new Error("waitFor timed out"); +} + +function streamedText(harness: { allChunks: unknown[] }): string { + return (harness.allChunks as { type?: string; delta?: string }[]) + .filter((c) => c.type === "text-delta") + .map((c) => c.delta ?? "") + .join(""); +} + +function turnCompleteCount(harness: { allRawChunks: unknown[] }): number { + return (harness.allRawChunks as { type?: string }[]).filter( + (c) => c.type === "trigger:turn-complete" + ).length; +} + +// ── Tests ────────────────────────────────────────────────────────────── + +describe("chat.agent pending wire buffer", () => { + it("dispatches every message buffered during a turn, not just the first", async () => { + const agent = chat.agent({ + id: "pending-drain.agent", + run: async ({ messages, signal }) => { + return streamText({ model: echoModel(), messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "pending-drain-1" }); + try { + const first = harness.sendMessage(userMessage("m1", "u-1")); + // Once m1's turn is streaming, land two more records back-to-back — + // both are consumed into the turn's buffer before the turn ends. + await waitFor(() => streamedText(harness).includes("ANSWER(m1)")); + void harness.sendMessage(userMessage("m2", "u-2")); + void harness.sendMessage(userMessage("m3", "u-3")); + await first; + + await waitFor(() => turnCompleteCount(harness) >= 3); + + const text = streamedText(harness); + const m2At = text.indexOf("ANSWER(m2)"); + const m3At = text.indexOf("ANSWER(m3)"); + expect(m2At).toBeGreaterThan(-1); + expect(m3At).toBeGreaterThan(-1); + expect(m3At).toBeGreaterThan(m2At); + } finally { + await harness.close(); + } + }); +}); + +describe("chat.agent errored turn", () => { + it( + "does not duplicate messages buffered after a turn that threw", + { timeout: 20000 }, + async () => { + // Throw from a pre-stream hook: throws inside the streaming section are + // already covered by its finally, but a hook throw used to leak the + // turn's message handler into the loop-level buffer. + let turnStarts = 0; + const agent = chat.agent({ + id: "pending-drain.errored-turn", + onTurnStart: async () => { + turnStarts++; + if (turnStarts === 1) { + throw new Error("synthetic turn failure"); + } + }, + run: async ({ messages, signal }) => { + return streamText({ model: echoModel(), messages, abortSignal: signal }); + }, + }); + + const harness = mockChatAgent(agent, { chatId: "pending-drain-4" }); + try { + // Turn 1 throws — pre-fix its message handler leaked past the turn. + await harness.sendMessage(userMessage("boom", "u-1")); + const second = harness.sendMessage(userMessage("m2", "u-2")); + // m3 lands mid-turn; a leaked handler would push it twice. + await waitFor(() => streamedText(harness).includes("ANSWER(m2)")); + void harness.sendMessage(userMessage("m3", "u-3")); + await second; + + await waitFor(() => streamedText(harness).includes("ANSWER(m3)")); + await new Promise((r) => setTimeout(r, 500)); + const text = streamedText(harness); + expect(text.match(/ANSWER\(m3\)/g)).toHaveLength(1); + } finally { + await harness.close(); + } + } + ); +}); + +describe("chat.createSession pending wire buffer", () => { + it("dispatches messages buffered during a turn as subsequent turns", async () => { + const agent = chat.customAgent({ + id: "pending-drain.session", + run: async (payload) => { + const session = chat.createSession(payload, { + signal: new AbortController().signal, + idleTimeoutInSeconds: 2, + }); + for await (const turn of session) { + const result = streamText({ + model: echoModel(), + messages: turn.messages, + abortSignal: turn.signal, + }); + await turn.complete(result); + } + }, + }); + + const harness = mockChatAgent(agent, { chatId: "pending-drain-2" }); + try { + const first = harness.sendMessage(userMessage("m1", "u-1")); + await waitFor(() => streamedText(harness).includes("ANSWER(m1)")); + void harness.sendMessage(userMessage("m2", "u-2")); + void harness.sendMessage(userMessage("m3", "u-3")); + await first; + + await waitFor(() => turnCompleteCount(harness) >= 3); + + const text = streamedText(harness); + expect(text).toContain("ANSWER(m2)"); + expect(text).toContain("ANSWER(m3)"); + } finally { + await harness.close(); + } + }); +}); + +describe("chat.createSession stop + immediate send", () => { + it( + "dispatches a message that arrives right after a stopped turn", + { timeout: 20000 }, + async () => { + const agent = chat.customAgent({ + id: "pending-drain.session-stop", + run: async (payload) => { + const session = chat.createSession(payload, { + signal: new AbortController().signal, + idleTimeoutInSeconds: 2, + // Steering config active — the failure mode routed post-stream + // arrivals into the dead steering queue instead of the next turn. + pendingMessages: {}, + }); + for await (const turn of session) { + const result = streamText({ + model: echoModel(), + messages: turn.messages, + abortSignal: turn.signal, + }); + await turn.complete(result); + } + }, + }); + + const harness = mockChatAgent(agent, { chatId: "pending-drain-3" }); + try { + const first = harness.sendMessage(userMessage("write a long essay", "u-1")); + await waitFor(() => streamedText(harness).length > 0); + await harness.sendStop(); + // Land the next message inside the stopped turn's post-stream window + // (the ~2s totalUsage race), after the abort has settled — previously + // the still-attached handler steering-routed it into the dead queue. + await new Promise((r) => setTimeout(r, 150)); + void harness.sendMessage(userMessage("m2", "u-2")); + await first; + + await waitFor(() => turnCompleteCount(harness) >= 2); + await waitFor(() => streamedText(harness).includes("ANSWER(m2)")); + } finally { + await harness.close(); + } + } + ); +}); + +describe("session.in.wait() consume cursor", () => { + it("advances lastDispatchedSeqNum alongside lastSeqNum on waitpoint delivery", async () => { + __setSessionOpenImplForTests(undefined); + await runInMockTaskContext(async () => { + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + createSessionStreamWaitpoint: async () => ({ waitpointId: "wp_test_1" }), + waitForWaitpointToken: async () => ({ success: true }), + } as never); + + const sessionId = "cursor-sess"; + // Simulate records 0..4 already received via SSE before the suspend. + sessionStreams.setLastSeqNum(sessionId, "in", 4); + + const result = await sessions.open(sessionId).in.wait(); + + expect(result.ok).toBe(true); + expect(sessionStreams.lastSeqNum(sessionId, "in")).toBe(5); + // The waitpoint-delivered record was consumed by this caller, so the + // committed-consume cursor (what turn-completes persist as + // `session-in-event-id`) must advance with it. + expect(sessionStreams.lastDispatchedSeqNum(sessionId, "in")).toBe(5); + }); + }); +});