From fcefcf52d0681318d53b88115b2bbcf4c6733388 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Sun, 16 Aug 2026 20:26:40 -0700 Subject: [PATCH 1/8] feat(chat): expose endAndContinue to custom agents --- .../chat-custom-agent-end-and-continue.md | 5 + docs/ai-chat/custom-agents.mdx | 20 +++ docs/ai-chat/patterns/version-upgrades.mdx | 21 ++- docs/ai-chat/reference.mdx | 1 + packages/trigger-sdk/src/v3/ai.ts | 76 ++++++--- .../test/chat-end-and-continue.test.ts | 144 ++++++++++++++++++ 6 files changed, 243 insertions(+), 24 deletions(-) create mode 100644 .changeset/chat-custom-agent-end-and-continue.md create mode 100644 packages/trigger-sdk/test/chat-end-and-continue.test.ts diff --git a/.changeset/chat-custom-agent-end-and-continue.md b/.changeset/chat-custom-agent-end-and-continue.md new file mode 100644 index 00000000000..b9b1cff29ef --- /dev/null +++ b/.changeset/chat-custom-agent-end-and-continue.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Allow custom chat agents to hand a conversation off to a fresh run on the latest deployed task version while preserving pending Session input. diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 197bff6b5e1..975e912ad15 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -144,6 +144,25 @@ for await (const turn of session) { Without this, a resumed chat silently loses its history: the model sees only the message that triggered the continuation. In a hand-rolled loop, seed by passing the stored history into the turn-0 `addIncoming` call — shown in the example below. +### Rotating to a new deployment + +`chat.createSession()` consumes `chat.requestUpgrade()` through its managed iterator. In a fully hand-rolled custom agent, hand the Session to a fresh run with `chat.endAndContinue()`. Finish the current turn and persist its state first, then make the handoff the last operation in `run()`: + +```ts +await chat.writeTurnComplete(); +await persistMessages(conversation.uiMessages); + +if (shouldRotateToLatestVersion()) { + return chat.endAndContinue(); +} +``` + +The server starts a continuation run using the Session's existing trigger configuration and atomically makes it the current run. The Session and its streams stay open, so input that the old run has not consumed remains on `.in` for the continuation run. The new run uses the latest deployed task version unless the Session's trigger configuration sets `lockToVersion`. + + + Call `chat.endAndContinue()` only at a completed turn boundary, after `chat.writeTurnComplete()` and after detaching the old run's input listeners. The operation starts the new run but does not stop the caller. Await it and return from `run()` immediately; continuing to read or write can race the new run on the same Session. + + ### turn.complete() vs manual control `turn.complete(result)` is the one-call path — it handles piping, capturing the response, accumulating messages, cleaning up aborted parts on a stop, and writing the turn-complete chunk. @@ -217,6 +236,7 @@ For full control, skip `createSession` and compose the primitives directly: | `chat.createStopSignal()` | Create a managed stop signal wired to the stop input stream | | `chat.pipeAndCapture(result)` | Pipe a stream and capture the response; returns `{ message, status, error }` | | `chat.writeTurnComplete()` | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors | +| `chat.endAndContinue()` | Hand off the Session to a continuation run; call at a completed turn boundary, then return | | `chat.MessageAccumulator` | Accumulates conversation messages across turns | | `chat.pipe(stream)` | Pipe a stream to the frontend (no response capture) | | `chat.cleanupAbortedParts(msg)` | Clean up incomplete parts from a stopped response | diff --git a/docs/ai-chat/patterns/version-upgrades.mdx b/docs/ai-chat/patterns/version-upgrades.mdx index 830f673e9a3..3bbe6c216f8 100644 --- a/docs/ai-chat/patterns/version-upgrades.mdx +++ b/docs/ai-chat/patterns/version-upgrades.mdx @@ -1,12 +1,12 @@ --- title: "Version upgrades" sidebarTitle: "Version upgrades" -description: "Gracefully migrate suspended chat agents to a new deployment using chat.requestUpgrade() and the continuation mechanism." +description: "Gracefully migrate chat agents to a new deployment using chat.requestUpgrade(), chat.endAndContinue(), and the continuation mechanism." --- Chat agent runs are pinned to the worker version they started on. When you deploy a new version, suspended runs resume on the **old** code. If your deploy includes breaking changes (new tools, changed schemas, updated API contracts), this can cause issues. -`chat.requestUpgrade()` lets the agent opt out of the current run so the transport triggers a new one on the latest version. +`chat.requestUpgrade()` lets `chat.agent()` and the `chat.createSession()` iterator opt out of the current run so the transport triggers a new one on the latest version. Fully hand-rolled custom agents use `chat.endAndContinue()` at a completed turn boundary for the same Session handoff. ## How it works @@ -151,10 +151,21 @@ export const myChat = chat This upgrades on **every** deploy, not just breaking changes. Good for fast-moving projects where you always want the latest code. -## Other agent types +## Custom agents -- **`chat.agent()`** and **`chat.createSession()`** — use `chat.requestUpgrade()` as shown above -- **`chat.customAgent()`** — you control the turn loop, so just `return` from `run()` when you want to exit +`chat.requestUpgrade()` is consumed by both `chat.agent()` and the `chat.createSession()` iterator. In a fully hand-rolled `chat.customAgent()` task, finish the turn, persist any application state, then call `chat.endAndContinue()` and return immediately: + +```ts +await chat.writeTurnComplete(); +await persistMessages(conversation.uiMessages); +return chat.endAndContinue(); +``` + +The continuation uses the same durable Session and receives `.in` records that the old run has not consumed. It starts on the latest deployed task version unless the Session's trigger configuration sets `lockToVersion`. + + + `chat.endAndContinue()` starts the successor but does not stop the calling run. Call it only after `chat.writeTurnComplete()` and after detaching the old run's input listeners, then perform no more Session reads or writes and return from the task. + ## Interaction with recovery boot diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index da08f9a0473..7f387996638 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -511,6 +511,7 @@ All methods available on the `chat` object from `@trigger.dev/sdk/ai`. | `chat.createStartSessionAction(taskId, options?)` | Returns a server action that creates a chat Session + triggers the first run + returns a session-scoped PAT. Idempotent on `(env, externalId)`. | | `chat.waitForHandover(options)` | Wait for a [`chat.headStart`](/ai-chat/fast-starts#handover-with-custom-agents) handover signal in a custom loop. Returns the signal or `null`. `chat.MessageAccumulator` wraps this as `consumeHandover()` / `applyHandover()` | | `chat.requestUpgrade()` | End the current run after this turn so the next message starts on the latest agent version. Server-orchestrated handoff. | +| `chat.endAndContinue()` | In a hand-rolled custom agent, hand off the Session to a fresh continuation run. Finish the turn, detach input listeners, call this method, then return immediately. | | `chat.setTurnTimeout(duration)` | Override turn timeout at runtime (e.g. `"2h"`) | | `chat.setTurnTimeoutInSeconds(seconds)` | Override turn timeout at runtime (in seconds) | | `chat.setIdleTimeoutInSeconds(seconds)` | Override idle timeout at runtime | diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 844d506079b..c638a4cbd04 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -2538,6 +2538,8 @@ const chatOnCompactedKey = locals.create<(event: CompactedEvent) => Promise | void>("chat.onCompacted"); /** @internal Full task `ctx` for the active `chat.agent` run (for hooks invoked from nested compaction). */ const chatAgentRunContextKey = locals.create("chat.agentRunContext"); +/** @internal Marks the root run created by `chat.customAgent()`. */ +const chatCustomAgentRunKey = locals.create("chat.customAgentRun"); const chatPrepareMessagesKey = locals.create<(event: PrepareMessagesEvent) => ModelMessage[] | Promise>( "chat.prepareMessages" @@ -5362,6 +5364,7 @@ function chatCustomAgent< locals.set(chatSessionHandleKey, sessions.open(payload.chatId)); locals.set(chatExternalIdKey, payload.chatId); locals.set(chatAgentRunContextKey, runOptions.ctx); + locals.set(chatCustomAgentRunKey, true); // Initialize the turn-complete trim slot so `chat.writeTurnComplete` // trims `session.out` back to the previous turn boundary. Without // this the slot is undefined and the trim never runs, so `.out` @@ -5456,6 +5459,7 @@ function chatAgent< { signal: runSignal, ctx } ) => { locals.set(chatAgentRunContextKey, ctx); + locals.set(chatCustomAgentRunKey, false); // On AI SDK 7, register the `@ai-sdk/otel` integration (once per process) // so `experimental_telemetry` spans flow into the run trace. Awaited here @@ -8705,6 +8709,54 @@ function requestUpgrade(): void { locals.set(chatUpgradeRequestedKey, true); } +/** + * Hand off the current custom agent Session to a fresh run. + * + * This is the low-level handoff for a fully hand-rolled + * `chat.customAgent()` loop. (`chat.createSession()` consumes + * {@link requestUpgrade} instead.) Call only after {@link chatWriteTurnComplete} + * and after detaching input listeners for the old run. The server starts the + * continuation run but does not stop this run, so return from the task + * immediately after awaiting this function. + * + * Pending Session input that the old run has not consumed remains on the + * durable `.in` stream and is delivered to the continuation run. + * + * @example + * ```ts + * await chat.writeTurnComplete(); + * + * if (shouldUpgrade) { + * return chat.endAndContinue(); + * } + * ``` + */ +async function endAndContinue(): Promise { + if (locals.get(chatCustomAgentRunKey) !== true) { + throw new Error( + "chat.endAndContinue() can only be called from inside a chat.customAgent() run" + ); + } + + await performEndAndContinue(); +} + +/** @internal Shared server handoff used by managed and custom agent loops. */ +async function performEndAndContinue(): Promise { + const chatId = locals.get(chatExternalIdKey); + const callingRunId = locals.get(chatAgentRunContextKey)?.run.id; + + if (!chatId || !callingRunId) { + throw new Error("Cannot end and continue without an active chat agent run"); + } + + const apiClient = apiClientManager.clientOrThrow(); + await apiClient.endAndContinueSession(chatId, { + callingRunId, + reason: "upgrade", + }); +} + /** * Exit the run after the current turn completes, without waiting for the * next message. Unlike {@link requestUpgrade}, no upgrade-required signal @@ -10697,6 +10749,8 @@ export const chat = { isStopped, /** Request that the run exits after the current turn so the next message starts on the latest version. See {@link requestUpgrade}. */ requestUpgrade, + /** Hand off a custom agent Session to a fresh run. See {@link endAndContinue}. */ + endAndContinue, /** Exit the run after the current turn completes, without any upgrade signal. See {@link endRun}. */ endRun, /** Clean up aborted parts from a UIMessage. See {@link cleanupAbortedParts}. */ @@ -10891,17 +10945,12 @@ async function writeTurnCompleteChunk( * @internal */ async function writeUpgradeRequiredChunk(): Promise { - const ctx = taskContext.ctx; - const chatId = ctx?.run.id ? getChatIdFromContext() : undefined; - const callingRunId = ctx?.run.id; + const chatId = locals.get(chatExternalIdKey); + const callingRunId = locals.get(chatAgentRunContextKey)?.run.id; if (chatId && callingRunId) { - const apiClient = apiClientManager.clientOrThrow(); try { - await apiClient.endAndContinueSession(chatId, { - callingRunId, - reason: "upgrade", - }); + await performEndAndContinue(); } catch (error) { // Non-fatal: the next `.in/append` re-triggers via the probe. // Swallow rather than throw so we still emit the chunk + exit. @@ -10917,17 +10966,6 @@ async function writeUpgradeRequiredChunk(): Promise { return session.out.writeControl(TRIGGER_CONTROL_SUBTYPE.UPGRADE_REQUIRED); } -/** - * Resolves the current chat's `chatId` (used as session externalId) from - * the bound session handle. Returns `undefined` if no agent is bound — - * shouldn't happen at the call sites that invoke - * `writeUpgradeRequiredChunk`, but defensive against misuse. - * @internal - */ -function getChatIdFromContext(): string | undefined { - return locals.get(chatSessionHandleKey)?.id; -} - /** * Extracts the text content of the last user message from a UIMessage array. * Returns undefined if no user message is found. diff --git a/packages/trigger-sdk/test/chat-end-and-continue.test.ts b/packages/trigger-sdk/test/chat-end-and-continue.test.ts new file mode 100644 index 00000000000..453b2871f8c --- /dev/null +++ b/packages/trigger-sdk/test/chat-end-and-continue.test.ts @@ -0,0 +1,144 @@ +// Import the test entry point first so chat.customAgent() registers its task. +import "../src/v3/test/index.js"; + +import { afterEach, describe, expect, it, vi } from "vitest"; +import { apiClientManager, resourceCatalog } from "@trigger.dev/core/v3"; +import { runInMockTaskContext, TestSessionStreamManager } from "@trigger.dev/core/v3/test"; +import { chat } from "../src/v3/ai.js"; + +const CHAT_ID = "chat-end-and-continue"; +const CALLING_RUN_ID = "run_before_handoff"; +const CONTINUATION_RUN_ID = "run_after_handoff"; + +class DurableTestSessionStreamManager extends TestSessionStreamManager { + override reset(): void { + // The Session stream outlives either task run. Drop run-local listeners, + // but preserve buffered input for the continuation run. + this.clearHandlers(); + } + + dispose(): void { + super.reset(); + } +} + +describe("chat.endAndContinue", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("ends cleanly and leaves pending input for the continuation run", async () => { + let continuationMessage: unknown; + + const agent = chat.customAgent({ + id: "end-and-continue-custom-agent", + run: async (payload) => { + if (!payload.continuation) { + return chat.endAndContinue(); + } + + const next = await chat.messages.waitWithIdleTimeout({ + idleTimeoutInSeconds: 1, + timeout: "1m", + }); + if (!next.ok) { + throw next.error; + } + + continuationMessage = next.output.message; + }, + }); + + const taskEntry = resourceCatalog.getTask(agent.id); + expect(taskEntry).toBeDefined(); + const runFn = taskEntry!.fns.run as ( + payload: Record, + options: { ctx: unknown; signal: AbortSignal } + ) => Promise; + + const readSessionStreamRecords = vi.fn(async () => ({ records: [] })); + const endAndContinueSession = vi.fn(async () => ({ + runId: CONTINUATION_RUN_ID, + swapped: true, + })); + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + readSessionStreamRecords, + endAndContinueSession, + } as never); + + const sessionStreams = new DurableTestSessionStreamManager(); + const pendingPayload = { + chatId: CHAT_ID, + trigger: "submit-message", + message: { + id: "pending-user-message", + role: "user", + parts: [{ type: "text", text: "deliver after handoff" }], + }, + metadata: {}, + }; + + try { + await runInMockTaskContext( + async (drivers) => { + // This record is durable Session input, not run-local input. It is + // written before the old run requests its handoff. + await drivers.sessions.in.send(CHAT_ID, { + kind: "message", + payload: pendingPayload, + }); + + await expect( + runFn( + { chatId: CHAT_ID, trigger: "preload", metadata: {} }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ) + ).resolves.toBeUndefined(); + }, + { + ctx: { run: { id: CALLING_RUN_ID } }, + sessionStreamManager: sessionStreams, + } + ); + + expect(endAndContinueSession).toHaveBeenCalledWith(CHAT_ID, { + callingRunId: CALLING_RUN_ID, + reason: "upgrade", + }); + + await runInMockTaskContext( + async (drivers) => { + await expect( + runFn( + { chatId: CHAT_ID, continuation: true, metadata: {} }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ) + ).resolves.toBeUndefined(); + }, + { + ctx: { run: { id: CONTINUATION_RUN_ID } }, + sessionStreamManager: sessionStreams, + } + ); + + expect(continuationMessage).toEqual(pendingPayload.message); + } finally { + sessionStreams.dispose(); + } + }); + + it("rejects calls outside a custom agent run", async () => { + const endAndContinueSession = vi.fn(); + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + endAndContinueSession, + } as never); + + await runInMockTaskContext(async () => { + await expect(chat.endAndContinue()).rejects.toThrow( + "chat.endAndContinue() can only be called from inside a chat.customAgent() run" + ); + }); + + expect(endAndContinueSession).not.toHaveBeenCalled(); + }); +}); From da596472d39e92151b21fbe4fd89c0a0430e0c5f Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Sun, 16 Aug 2026 21:22:01 -0700 Subject: [PATCH 2/8] docs(chat): clarify endAndContinue lifecycle --- .../chat-custom-agent-end-and-continue.md | 2 +- docs/ai-chat/custom-agents.mdx | 20 ++++--- docs/ai-chat/patterns/version-upgrades.mdx | 15 +++-- docs/ai-chat/reference.mdx | 2 +- packages/trigger-sdk/src/v3/ai.ts | 23 ++++--- .../test/chat-end-and-continue.test.ts | 60 ++++++++++++++++--- 6 files changed, 91 insertions(+), 31 deletions(-) diff --git a/.changeset/chat-custom-agent-end-and-continue.md b/.changeset/chat-custom-agent-end-and-continue.md index b9b1cff29ef..ef1bf090878 100644 --- a/.changeset/chat-custom-agent-end-and-continue.md +++ b/.changeset/chat-custom-agent-end-and-continue.md @@ -2,4 +2,4 @@ "@trigger.dev/sdk": patch --- -Allow custom chat agents to hand a conversation off to a fresh run on the latest deployed task version while preserving pending Session input. +Allow custom chat agents to hand a conversation off to a fresh run on the latest deployed task version while preserving unconsumed Session input. diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 975e912ad15..e373a7ed57e 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -146,21 +146,25 @@ Without this, a resumed chat silently loses its history: the model sees only the ### Rotating to a new deployment -`chat.createSession()` consumes `chat.requestUpgrade()` through its managed iterator. In a fully hand-rolled custom agent, hand the Session to a fresh run with `chat.endAndContinue()`. Finish the current turn and persist its state first, then make the handoff the last operation in `run()`: +With `chat.createSession()`, use `chat.requestUpgrade()` to leave the current run after the turn. In a fully hand-rolled custom agent, use `chat.endAndContinue()` to immediately hand the Session to a fresh run. + +Call it between turns, after detaching the old run's input listeners. If the old run completed its current turn, persist its state and write the turn-complete boundary before the handoff: ```ts -await chat.writeTurnComplete(); +messageSubscription.off(); +stop.cleanup(); await persistMessages(conversation.uiMessages); - -if (shouldRotateToLatestVersion()) { - return chat.endAndContinue(); -} +await chat.writeTurnComplete(); +await chat.endAndContinue(); +return; ``` The server starts a continuation run using the Session's existing trigger configuration and atomically makes it the current run. The Session and its streams stay open, so input that the old run has not consumed remains on `.in` for the continuation run. The new run uses the latest deployed task version unless the Session's trigger configuration sets `lockToVersion`. +If input arrives that the old run should leave for the continuation, detach the listeners and do not write another turn-complete boundary before handing off. `chat.writeTurnComplete()` acknowledges the latest input dispatched to the old run; writing it after receiving the deferred input would make the continuation resume after that input. + - Call `chat.endAndContinue()` only at a completed turn boundary, after `chat.writeTurnComplete()` and after detaching the old run's input listeners. The operation starts the new run but does not stop the caller. Await it and return from `run()` immediately; continuing to read or write can race the new run on the same Session. + `chat.endAndContinue()` starts the new run but does not stop the caller. Await it and return from `run()` immediately; continuing to read or write can race the new run on the same Session. If the handoff fails, the promise rejects. ### turn.complete() vs manual control @@ -236,7 +240,7 @@ For full control, skip `createSession` and compose the primitives directly: | `chat.createStopSignal()` | Create a managed stop signal wired to the stop input stream | | `chat.pipeAndCapture(result)` | Pipe a stream and capture the response; returns `{ message, status, error }` | | `chat.writeTurnComplete()` | Signal turn complete; returns `{ lastEventId, sessionInEventId }` resume cursors | -| `chat.endAndContinue()` | Hand off the Session to a continuation run; call at a completed turn boundary, then return | +| `chat.endAndContinue()` | Hand off the Session to a continuation run; call between turns, then return | | `chat.MessageAccumulator` | Accumulates conversation messages across turns | | `chat.pipe(stream)` | Pipe a stream to the frontend (no response capture) | | `chat.cleanupAbortedParts(msg)` | Clean up incomplete parts from a stopped response | diff --git a/docs/ai-chat/patterns/version-upgrades.mdx b/docs/ai-chat/patterns/version-upgrades.mdx index 3bbe6c216f8..eaf288b7b83 100644 --- a/docs/ai-chat/patterns/version-upgrades.mdx +++ b/docs/ai-chat/patterns/version-upgrades.mdx @@ -6,7 +6,7 @@ description: "Gracefully migrate chat agents to a new deployment using chat.requ Chat agent runs are pinned to the worker version they started on. When you deploy a new version, suspended runs resume on the **old** code. If your deploy includes breaking changes (new tools, changed schemas, updated API contracts), this can cause issues. -`chat.requestUpgrade()` lets `chat.agent()` and the `chat.createSession()` iterator opt out of the current run so the transport triggers a new one on the latest version. Fully hand-rolled custom agents use `chat.endAndContinue()` at a completed turn boundary for the same Session handoff. +`chat.requestUpgrade()` is the managed upgrade signal for `chat.agent()` and the `chat.createSession()` iterator. Fully hand-rolled custom agents use `chat.endAndContinue()` between turns to immediately hand the Session to a new run. ## How it works @@ -153,18 +153,23 @@ This upgrades on **every** deploy, not just breaking changes. Good for fast-movi ## Custom agents -`chat.requestUpgrade()` is consumed by both `chat.agent()` and the `chat.createSession()` iterator. In a fully hand-rolled `chat.customAgent()` task, finish the turn, persist any application state, then call `chat.endAndContinue()` and return immediately: +Use `chat.requestUpgrade()` with `chat.agent()` and the `chat.createSession()` iterator. In a fully hand-rolled `chat.customAgent()` task, detach input listeners, persist the completed turn, write its boundary, then call `chat.endAndContinue()` and return immediately: ```ts -await chat.writeTurnComplete(); +messageSubscription.off(); +stop.cleanup(); await persistMessages(conversation.uiMessages); -return chat.endAndContinue(); +await chat.writeTurnComplete(); +await chat.endAndContinue(); +return; ``` The continuation uses the same durable Session and receives `.in` records that the old run has not consumed. It starts on the latest deployed task version unless the Session's trigger configuration sets `lockToVersion`. +If input arrives that the continuation should process, detach the old listeners and skip the final `chat.writeTurnComplete()`. A turn-complete boundary acknowledges the latest input dispatched to the old run, so writing one after that input would cause the continuation to resume past it. + - `chat.endAndContinue()` starts the successor but does not stop the calling run. Call it only after `chat.writeTurnComplete()` and after detaching the old run's input listeners, then perform no more Session reads or writes and return from the task. + `chat.endAndContinue()` starts the successor but does not stop the calling run. Perform no more Session reads or writes after calling it, and return from the task. If the handoff fails, the promise rejects. ## Interaction with recovery boot diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index 7f387996638..68812b3c21b 100644 --- a/docs/ai-chat/reference.mdx +++ b/docs/ai-chat/reference.mdx @@ -511,7 +511,7 @@ All methods available on the `chat` object from `@trigger.dev/sdk/ai`. | `chat.createStartSessionAction(taskId, options?)` | Returns a server action that creates a chat Session + triggers the first run + returns a session-scoped PAT. Idempotent on `(env, externalId)`. | | `chat.waitForHandover(options)` | Wait for a [`chat.headStart`](/ai-chat/fast-starts#handover-with-custom-agents) handover signal in a custom loop. Returns the signal or `null`. `chat.MessageAccumulator` wraps this as `consumeHandover()` / `applyHandover()` | | `chat.requestUpgrade()` | End the current run after this turn so the next message starts on the latest agent version. Server-orchestrated handoff. | -| `chat.endAndContinue()` | In a hand-rolled custom agent, hand off the Session to a fresh continuation run. Finish the turn, detach input listeners, call this method, then return immediately. | +| `chat.endAndContinue()` | In a hand-rolled custom agent, hand off the Session to a fresh continuation run. Call between turns after detaching input listeners, then return immediately. The promise rejects if the handoff fails. | | `chat.setTurnTimeout(duration)` | Override turn timeout at runtime (e.g. `"2h"`) | | `chat.setTurnTimeoutInSeconds(seconds)` | Override turn timeout at runtime (in seconds) | | `chat.setIdleTimeoutInSeconds(seconds)` | Override idle timeout at runtime | diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index c638a4cbd04..18f89ab921f 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -8713,22 +8713,27 @@ function requestUpgrade(): void { * Hand off the current custom agent Session to a fresh run. * * This is the low-level handoff for a fully hand-rolled - * `chat.customAgent()` loop. (`chat.createSession()` consumes - * {@link requestUpgrade} instead.) Call only after {@link chatWriteTurnComplete} - * and after detaching input listeners for the old run. The server starts the - * continuation run but does not stop this run, so return from the task - * immediately after awaiting this function. + * `chat.customAgent()` loop. (Use {@link requestUpgrade} with + * `chat.createSession()` instead.) Call only between turns and after detaching + * input listeners for the old run. If the old run completed its current turn, + * persist its state and call {@link chatWriteTurnComplete} before handing off. + * Do not write a new turn boundary after receiving input that the continuation + * run should process: the boundary acknowledges the latest dispatched input. + * + * The server starts the continuation run but does not stop this run, so return + * from the task immediately after awaiting this function. The promise rejects + * if the server cannot complete the handoff. * * Pending Session input that the old run has not consumed remains on the * durable `.in` stream and is delivered to the continuation run. * * @example * ```ts + * messageSubscription.off(); + * await persistMessages(); * await chat.writeTurnComplete(); - * - * if (shouldUpgrade) { - * return chat.endAndContinue(); - * } + * await chat.endAndContinue(); + * return; * ``` */ async function endAndContinue(): Promise { diff --git a/packages/trigger-sdk/test/chat-end-and-continue.test.ts b/packages/trigger-sdk/test/chat-end-and-continue.test.ts index 453b2871f8c..4c706391b54 100644 --- a/packages/trigger-sdk/test/chat-end-and-continue.test.ts +++ b/packages/trigger-sdk/test/chat-end-and-continue.test.ts @@ -10,6 +10,20 @@ const CHAT_ID = "chat-end-and-continue"; const CALLING_RUN_ID = "run_before_handoff"; const CONTINUATION_RUN_ID = "run_after_handoff"; +type CustomAgentRun = ( + payload: Record, + options: { ctx: unknown; signal: AbortSignal } +) => Promise; + +function getCustomAgentRun(id: string): CustomAgentRun { + const taskEntry = resourceCatalog.getTask(id); + if (!taskEntry) { + throw new Error(`Task ${id} was not registered`); + } + + return taskEntry.fns.run as CustomAgentRun; +} + class DurableTestSessionStreamManager extends TestSessionStreamManager { override reset(): void { // The Session stream outlives either task run. Drop run-local listeners, @@ -27,7 +41,7 @@ describe("chat.endAndContinue", () => { vi.restoreAllMocks(); }); - it("ends cleanly and leaves pending input for the continuation run", async () => { + it("ends cleanly and leaves unconsumed input for the continuation run", async () => { let continuationMessage: unknown; const agent = chat.customAgent({ @@ -49,12 +63,7 @@ describe("chat.endAndContinue", () => { }, }); - const taskEntry = resourceCatalog.getTask(agent.id); - expect(taskEntry).toBeDefined(); - const runFn = taskEntry!.fns.run as ( - payload: Record, - options: { ctx: unknown; signal: AbortSignal } - ) => Promise; + const runFn = getCustomAgentRun(agent.id); const readSessionStreamRecords = vi.fn(async () => ({ records: [] })); const endAndContinueSession = vi.fn(async () => ({ @@ -127,6 +136,43 @@ describe("chat.endAndContinue", () => { } }); + it("rejects when the server handoff fails", async () => { + const agent = chat.customAgent({ + id: "end-and-continue-failure-agent", + run: async () => { + return chat.endAndContinue(); + }, + }); + + const runFn = getCustomAgentRun(agent.id); + + const readSessionStreamRecords = vi.fn(async () => ({ records: [] })); + const endAndContinueSession = vi.fn(async () => { + throw new Error("handoff failed"); + }); + vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + readSessionStreamRecords, + endAndContinueSession, + } as never); + + await runInMockTaskContext( + async (drivers) => { + await expect( + runFn( + { chatId: CHAT_ID, trigger: "preload", metadata: {} }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ) + ).rejects.toThrow("handoff failed"); + }, + { ctx: { run: { id: CALLING_RUN_ID } } } + ); + + expect(endAndContinueSession).toHaveBeenCalledWith(CHAT_ID, { + callingRunId: CALLING_RUN_ID, + reason: "upgrade", + }); + }); + it("rejects calls outside a custom agent run", async () => { const endAndContinueSession = vi.fn(); vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ From 03b84d1f9565d9d1b1168a79506a24bb46e95851 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 01:25:03 -0700 Subject: [PATCH 3/8] docs(chat): tighten endAndContinue guidance --- .changeset/chat-custom-agent-end-and-continue.md | 2 +- docs/ai-chat/custom-agents.mdx | 4 ++-- docs/ai-chat/patterns/version-upgrades.mdx | 8 +++++--- packages/trigger-sdk/src/v3/ai.ts | 6 +++--- 4 files changed, 11 insertions(+), 9 deletions(-) diff --git a/.changeset/chat-custom-agent-end-and-continue.md b/.changeset/chat-custom-agent-end-and-continue.md index ef1bf090878..6d7deffeeeb 100644 --- a/.changeset/chat-custom-agent-end-and-continue.md +++ b/.changeset/chat-custom-agent-end-and-continue.md @@ -2,4 +2,4 @@ "@trigger.dev/sdk": patch --- -Allow custom chat agents to hand a conversation off to a fresh run on the latest deployed task version while preserving unconsumed Session input. +Add `chat.endAndContinue()` so fully hand-rolled custom chat agents can hand a conversation off to a fresh run on the latest deployed task version while preserving unconsumed Session input. diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index e373a7ed57e..c7da25eda4b 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -151,7 +151,7 @@ With `chat.createSession()`, use `chat.requestUpgrade()` to leave the current ru Call it between turns, after detaching the old run's input listeners. If the old run completed its current turn, persist its state and write the turn-complete boundary before the handoff: ```ts -messageSubscription.off(); +// Detach any chat.messages.on() subscriptions you created. stop.cleanup(); await persistMessages(conversation.uiMessages); await chat.writeTurnComplete(); @@ -161,7 +161,7 @@ return; The server starts a continuation run using the Session's existing trigger configuration and atomically makes it the current run. The Session and its streams stay open, so input that the old run has not consumed remains on `.in` for the continuation run. The new run uses the latest deployed task version unless the Session's trigger configuration sets `lockToVersion`. -If input arrives that the old run should leave for the continuation, detach the listeners and do not write another turn-complete boundary before handing off. `chat.writeTurnComplete()` acknowledges the latest input dispatched to the old run; writing it after receiving the deferred input would make the continuation resume after that input. +If input has been dispatched to the old run but should be processed by the continuation, detach the listeners and do not write another turn-complete boundary before handing off. `chat.writeTurnComplete()` acknowledges the latest input dispatched to the old run; writing it after that dispatch would make the continuation resume after the input. `chat.endAndContinue()` starts the new run but does not stop the caller. Await it and return from `run()` immediately; continuing to read or write can race the new run on the same Session. If the handoff fails, the promise rejects. diff --git a/docs/ai-chat/patterns/version-upgrades.mdx b/docs/ai-chat/patterns/version-upgrades.mdx index eaf288b7b83..cd9c8571706 100644 --- a/docs/ai-chat/patterns/version-upgrades.mdx +++ b/docs/ai-chat/patterns/version-upgrades.mdx @@ -156,7 +156,7 @@ This upgrades on **every** deploy, not just breaking changes. Good for fast-movi Use `chat.requestUpgrade()` with `chat.agent()` and the `chat.createSession()` iterator. In a fully hand-rolled `chat.customAgent()` task, detach input listeners, persist the completed turn, write its boundary, then call `chat.endAndContinue()` and return immediately: ```ts -messageSubscription.off(); +// Detach any chat.messages.on() subscriptions you created. stop.cleanup(); await persistMessages(conversation.uiMessages); await chat.writeTurnComplete(); @@ -166,7 +166,7 @@ return; The continuation uses the same durable Session and receives `.in` records that the old run has not consumed. It starts on the latest deployed task version unless the Session's trigger configuration sets `lockToVersion`. -If input arrives that the continuation should process, detach the old listeners and skip the final `chat.writeTurnComplete()`. A turn-complete boundary acknowledges the latest input dispatched to the old run, so writing one after that input would cause the continuation to resume past it. +If input has been dispatched to the old run but should be processed by the continuation, detach the old listeners and skip the final `chat.writeTurnComplete()`. A turn-complete boundary acknowledges the latest input dispatched to the old run, so writing one after that dispatch would cause the continuation to resume past the input. `chat.endAndContinue()` starts the successor but does not stop the calling run. Perform no more Session reads or writes after calling it, and return from the task. If the handoff fails, the promise rejects. @@ -174,7 +174,9 @@ If input arrives that the continuation should process, detach the old listeners ## Interaction with recovery boot -`chat.requestUpgrade()` is a graceful exit — the old run returns cleanly, never writing a partial assistant. The new continuation run boots with an empty `session.out` tail and the upgrade-trigger message on `session.in`. The trigger message dispatches as turn 1 on the new version via the normal continuation-wait path. [`onRecoveryBoot`](/ai-chat/patterns/recovery-boot) does NOT fire on this path — the hook is reserved for mid-stream interruptions (cancel / crash / OOM) where a partial assistant exists on the tail. +When `chat.requestUpgrade()` is handled before a turn starts, the SDK immediately hands the Session to a new run, which processes the same input on the latest version. When it is requested during a turn, including through `chat.createSession()`, the current turn finishes and the old run exits; the next input starts the continuation run. + +Both are graceful exits. [`onRecoveryBoot`](/ai-chat/patterns/recovery-boot) does not fire — the hook is reserved for mid-stream interruptions (cancel, crash, or OOM) where a partial assistant exists on the tail. ## See also diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 18f89ab921f..c316af94ff6 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -8717,8 +8717,8 @@ function requestUpgrade(): void { * `chat.createSession()` instead.) Call only between turns and after detaching * input listeners for the old run. If the old run completed its current turn, * persist its state and call {@link chatWriteTurnComplete} before handing off. - * Do not write a new turn boundary after receiving input that the continuation - * run should process: the boundary acknowledges the latest dispatched input. + * Do not write a new turn boundary after input that the continuation run should + * process has been dispatched: the boundary acknowledges that input. * * The server starts the continuation run but does not stop this run, so return * from the task immediately after awaiting this function. The promise rejects @@ -8729,7 +8729,7 @@ function requestUpgrade(): void { * * @example * ```ts - * messageSubscription.off(); + * // Detach any chat.messages.on() subscriptions you created. * await persistMessages(); * await chat.writeTurnComplete(); * await chat.endAndContinue(); From e1a416eb9e209a187f58b7fe7d2765fdbe146987 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 10:42:54 -0700 Subject: [PATCH 4/8] test(chat): cover endAndContinue with real sessions --- apps/webapp/test/helpers/testChatAgent.ts | 35 ++++ apps/webapp/test/session-agent.e2e.test.ts | 122 ++++++++++++ .../test/chat-end-and-continue.test.ts | 188 +----------------- 3 files changed, 161 insertions(+), 184 deletions(-) diff --git a/apps/webapp/test/helpers/testChatAgent.ts b/apps/webapp/test/helpers/testChatAgent.ts index 8aebb713a96..29eab173472 100644 --- a/apps/webapp/test/helpers/testChatAgent.ts +++ b/apps/webapp/test/helpers/testChatAgent.ts @@ -234,6 +234,41 @@ export const testUpgradeOnceChatAgent = chat.agent({ }, }); +/** + * Hands an unconsumed Session input record to a continuation run using the + * public custom-agent lifecycle primitive. The continuation echoes the input + * to `.out`, which lets the full-stack Session E2E assert durable delivery. + */ +export const testEndAndContinueCustomAgent = chat.customAgent({ + id: "e2e-test-chat-custom-end-and-continue", + run: async (payload) => { + if (!payload.continuation) { + await chat.endAndContinue(); + return; + } + + const next = await chat.messages.waitWithIdleTimeout({ + idleTimeoutInSeconds: 2, + timeout: "1m", + }); + if (!next.ok) { + throw next.error; + } + + const message = next.output.message as UIMessage | undefined; + const text = message ? firstText(message) : ""; + const { waitUntilComplete } = chat.stream.writer({ + execute: ({ write }) => { + write({ type: "text-start", id: "handoff-result" }); + write({ type: "text-delta", id: "handoff-result", delta: `received:${text}` }); + write({ type: "text-end", id: "handoff-result" }); + }, + }); + await waitUntilComplete(); + await chat.writeTurnComplete(); + }, +}); + /** * A tool with a server-side `execute`: the agent runs it automatically and * feeds the result back to the model, so a single turn covers the whole diff --git a/apps/webapp/test/session-agent.e2e.test.ts b/apps/webapp/test/session-agent.e2e.test.ts index f6a9a2a338c..18a99862169 100644 --- a/apps/webapp/test/session-agent.e2e.test.ts +++ b/apps/webapp/test/session-agent.e2e.test.ts @@ -33,6 +33,7 @@ import { testApprovalChatAgent, testChatAgent, testChatModelLocal, + testEndAndContinueCustomAgent, testEndRunChatAgent, testHitlChatAgent, testHitlIdleChatAgent, @@ -123,6 +124,34 @@ async function setupSession(agentId: string = testChatAgent.id) { return { addressingKey, token, apiKey, baseUrl: server.webapp.baseUrl }; } +async function setupStartedSession(agentId: string) { + const { environment, apiKey } = await seedTestEnvironment(server.prisma); + const addressingKey = `chat-${randomBytes(6).toString("hex")}`; + const createRes = await fetch(`${server.webapp.baseUrl}/api/v1/sessions`, { + method: "POST", + headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" }, + body: JSON.stringify({ + type: "chat.agent", + externalId: addressingKey, + taskIdentifier: agentId, + triggerConfig: { basePayload: {} }, + }), + }); + + expect(createRes.ok).toBe(true); + const created = (await createRes.json()) as { + runId: string; + publicAccessToken: string; + }; + return { + ...created, + addressingKey, + apiKey, + environment, + baseUrl: server.webapp.baseUrl, + }; +} + function promptText(prompt: unknown): string { if (!Array.isArray(prompt)) return ""; let out = ""; @@ -1533,4 +1562,97 @@ describe("session agent e2e (real chat.agent loop)", () => { await agent.close(); } }); + + it("EA23: custom endAndContinue hands pending input to a fresh run", async () => { + const { addressingKey, publicAccessToken, runId, apiKey, environment, baseUrl } = + await setupStartedSession(testEndAndContinueCustomAgent.id); + const initialRun = await server.prisma.taskRun.findFirstOrThrow({ + where: { friendlyId: runId }, + select: { id: true }, + }); + + const append = await appendInput({ + baseUrl, + addressingKey, + token: publicAccessToken, + partId: "pending-handoff-input", + body: submitBody( + addressingKey, + userMessage("deliver after endAndContinue", "pending-handoff-input") + ), + }); + expect(append.status).toBe(200); + + const oldRun = runRealChatAgent({ + agentId: testEndAndContinueCustomAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("unused"), + modelLocal: testChatModelLocal, + runId, + }); + let continuation: ReturnType | undefined; + + try { + await expect(oldRun.done).resolves.toBeUndefined(); + + const session = await server.prisma.session.findFirstOrThrow({ + where: { runtimeEnvironmentId: environment.id, externalId: addressingKey }, + select: { currentRunId: true, currentRunVersion: true }, + }); + expect(session.currentRunId).not.toBe(initialRun.id); + expect(session.currentRunVersion).toBeGreaterThan(1); + + const successor = await server.prisma.taskRun.findUniqueOrThrow({ + where: { id: session.currentRunId! }, + select: { friendlyId: true }, + }); + continuation = runRealChatAgent({ + agentId: testEndAndContinueCustomAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("unused"), + modelLocal: testChatModelLocal, + runId: successor.friendlyId, + continuation: true, + previousRunId: runId, + }); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token: publicAccessToken, + until: (p) => p.some(isTurnComplete), + maxMs: 30_000, + }); + expect(joinChunks(parts)).toContain("received:deliver after endAndContinue"); + await expect(continuation.done).resolves.toBeUndefined(); + } finally { + await continuation?.close(); + await oldRun.close(); + } + }); + + it("EA24: custom endAndContinue rejects when the server rejects the handoff", async () => { + const { addressingKey, apiKey, baseUrl } = await setupStartedSession( + testEndAndContinueCustomAgent.id + ); + const agent = runRealChatAgent({ + agentId: testEndAndContinueCustomAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("unused"), + modelLocal: testChatModelLocal, + runId: "run_missing_end_and_continue", + }); + + try { + await expect(agent.done).rejects.toThrow("callingRunId not found in this environment"); + } finally { + await agent.close(); + } + }); }); diff --git a/packages/trigger-sdk/test/chat-end-and-continue.test.ts b/packages/trigger-sdk/test/chat-end-and-continue.test.ts index 4c706391b54..fcb43dbcc4e 100644 --- a/packages/trigger-sdk/test/chat-end-and-continue.test.ts +++ b/packages/trigger-sdk/test/chat-end-and-continue.test.ts @@ -1,190 +1,10 @@ -// Import the test entry point first so chat.customAgent() registers its task. -import "../src/v3/test/index.js"; - -import { afterEach, describe, expect, it, vi } from "vitest"; -import { apiClientManager, resourceCatalog } from "@trigger.dev/core/v3"; -import { runInMockTaskContext, TestSessionStreamManager } from "@trigger.dev/core/v3/test"; +import { describe, expect, it } from "vitest"; import { chat } from "../src/v3/ai.js"; -const CHAT_ID = "chat-end-and-continue"; -const CALLING_RUN_ID = "run_before_handoff"; -const CONTINUATION_RUN_ID = "run_after_handoff"; - -type CustomAgentRun = ( - payload: Record, - options: { ctx: unknown; signal: AbortSignal } -) => Promise; - -function getCustomAgentRun(id: string): CustomAgentRun { - const taskEntry = resourceCatalog.getTask(id); - if (!taskEntry) { - throw new Error(`Task ${id} was not registered`); - } - - return taskEntry.fns.run as CustomAgentRun; -} - -class DurableTestSessionStreamManager extends TestSessionStreamManager { - override reset(): void { - // The Session stream outlives either task run. Drop run-local listeners, - // but preserve buffered input for the continuation run. - this.clearHandlers(); - } - - dispose(): void { - super.reset(); - } -} - describe("chat.endAndContinue", () => { - afterEach(() => { - vi.restoreAllMocks(); - }); - - it("ends cleanly and leaves unconsumed input for the continuation run", async () => { - let continuationMessage: unknown; - - const agent = chat.customAgent({ - id: "end-and-continue-custom-agent", - run: async (payload) => { - if (!payload.continuation) { - return chat.endAndContinue(); - } - - const next = await chat.messages.waitWithIdleTimeout({ - idleTimeoutInSeconds: 1, - timeout: "1m", - }); - if (!next.ok) { - throw next.error; - } - - continuationMessage = next.output.message; - }, - }); - - const runFn = getCustomAgentRun(agent.id); - - const readSessionStreamRecords = vi.fn(async () => ({ records: [] })); - const endAndContinueSession = vi.fn(async () => ({ - runId: CONTINUATION_RUN_ID, - swapped: true, - })); - vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ - readSessionStreamRecords, - endAndContinueSession, - } as never); - - const sessionStreams = new DurableTestSessionStreamManager(); - const pendingPayload = { - chatId: CHAT_ID, - trigger: "submit-message", - message: { - id: "pending-user-message", - role: "user", - parts: [{ type: "text", text: "deliver after handoff" }], - }, - metadata: {}, - }; - - try { - await runInMockTaskContext( - async (drivers) => { - // This record is durable Session input, not run-local input. It is - // written before the old run requests its handoff. - await drivers.sessions.in.send(CHAT_ID, { - kind: "message", - payload: pendingPayload, - }); - - await expect( - runFn( - { chatId: CHAT_ID, trigger: "preload", metadata: {} }, - { ctx: drivers.ctx, signal: new AbortController().signal } - ) - ).resolves.toBeUndefined(); - }, - { - ctx: { run: { id: CALLING_RUN_ID } }, - sessionStreamManager: sessionStreams, - } - ); - - expect(endAndContinueSession).toHaveBeenCalledWith(CHAT_ID, { - callingRunId: CALLING_RUN_ID, - reason: "upgrade", - }); - - await runInMockTaskContext( - async (drivers) => { - await expect( - runFn( - { chatId: CHAT_ID, continuation: true, metadata: {} }, - { ctx: drivers.ctx, signal: new AbortController().signal } - ) - ).resolves.toBeUndefined(); - }, - { - ctx: { run: { id: CONTINUATION_RUN_ID } }, - sessionStreamManager: sessionStreams, - } - ); - - expect(continuationMessage).toEqual(pendingPayload.message); - } finally { - sessionStreams.dispose(); - } - }); - - it("rejects when the server handoff fails", async () => { - const agent = chat.customAgent({ - id: "end-and-continue-failure-agent", - run: async () => { - return chat.endAndContinue(); - }, - }); - - const runFn = getCustomAgentRun(agent.id); - - const readSessionStreamRecords = vi.fn(async () => ({ records: [] })); - const endAndContinueSession = vi.fn(async () => { - throw new Error("handoff failed"); - }); - vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ - readSessionStreamRecords, - endAndContinueSession, - } as never); - - await runInMockTaskContext( - async (drivers) => { - await expect( - runFn( - { chatId: CHAT_ID, trigger: "preload", metadata: {} }, - { ctx: drivers.ctx, signal: new AbortController().signal } - ) - ).rejects.toThrow("handoff failed"); - }, - { ctx: { run: { id: CALLING_RUN_ID } } } - ); - - expect(endAndContinueSession).toHaveBeenCalledWith(CHAT_ID, { - callingRunId: CALLING_RUN_ID, - reason: "upgrade", - }); - }); - it("rejects calls outside a custom agent run", async () => { - const endAndContinueSession = vi.fn(); - vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ - endAndContinueSession, - } as never); - - await runInMockTaskContext(async () => { - await expect(chat.endAndContinue()).rejects.toThrow( - "chat.endAndContinue() can only be called from inside a chat.customAgent() run" - ); - }); - - expect(endAndContinueSession).not.toHaveBeenCalled(); + await expect(chat.endAndContinue()).rejects.toThrow( + "chat.endAndContinue() can only be called from inside a chat.customAgent() run" + ); }); }); From adf37ab0ce6be9834e5eab89f6a4f2747fb8eeee Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 14:22:02 -0700 Subject: [PATCH 5/8] test(chat): use allowed run query --- apps/webapp/test/session-agent.e2e.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/webapp/test/session-agent.e2e.test.ts b/apps/webapp/test/session-agent.e2e.test.ts index 18a99862169..0690dd5b0ac 100644 --- a/apps/webapp/test/session-agent.e2e.test.ts +++ b/apps/webapp/test/session-agent.e2e.test.ts @@ -1604,7 +1604,7 @@ describe("session agent e2e (real chat.agent loop)", () => { expect(session.currentRunId).not.toBe(initialRun.id); expect(session.currentRunVersion).toBeGreaterThan(1); - const successor = await server.prisma.taskRun.findUniqueOrThrow({ + const successor = await server.prisma.taskRun.findFirstOrThrow({ where: { id: session.currentRunId! }, select: { friendlyId: true }, }); From b729865e5f8b01b586f475a714de0613455d58d3 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 14:53:50 -0700 Subject: [PATCH 6/8] fix(chat): guard handoff during managed sessions --- docs/ai-chat/custom-agents.mdx | 2 +- packages/trigger-sdk/src/v3/ai.ts | 59 +++++++++++++++++-- .../test/chat-end-and-continue.test.ts | 59 ++++++++++++++++++- 3 files changed, 114 insertions(+), 6 deletions(-) diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index c7da25eda4b..3454845e68f 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -146,7 +146,7 @@ Without this, a resumed chat silently loses its history: the model sees only the ### Rotating to a new deployment -With `chat.createSession()`, use `chat.requestUpgrade()` to leave the current run after the turn. In a fully hand-rolled custom agent, use `chat.endAndContinue()` to immediately hand the Session to a fresh run. +With `chat.createSession()`, use `chat.requestUpgrade()` to leave the current run after the turn. `chat.endAndContinue()` rejects while a `chat.createSession()` iterator is active. In a fully hand-rolled custom agent, use it to immediately hand the Session to a fresh run. Call it between turns, after detaching the old run's input listeners. If the old run completed its current turn, persist its state and write the turn-complete boundary before the handoff: diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index c316af94ff6..0d4aceb28ab 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -2540,6 +2540,8 @@ const chatOnCompactedKey = const chatAgentRunContextKey = locals.create("chat.agentRunContext"); /** @internal Marks the root run created by `chat.customAgent()`. */ const chatCustomAgentRunKey = locals.create("chat.customAgentRun"); +/** @internal Number of active `chat.createSession()` iterators in this run. */ +const chatActiveSessionIteratorsKey = locals.create("chat.createSession.activeIterators"); const chatPrepareMessagesKey = locals.create<(event: PrepareMessagesEvent) => ModelMessage[] | Promise>( "chat.prepareMessages" @@ -8714,9 +8716,10 @@ function requestUpgrade(): void { * * This is the low-level handoff for a fully hand-rolled * `chat.customAgent()` loop. (Use {@link requestUpgrade} with - * `chat.createSession()` instead.) Call only between turns and after detaching - * input listeners for the old run. If the old run completed its current turn, - * persist its state and call {@link chatWriteTurnComplete} before handing off. + * `chat.createSession()` instead; this method rejects while its iterator is + * active.) Call only between turns and after detaching input listeners for the + * old run. If the old run completed its current turn, persist its state and + * call {@link chatWriteTurnComplete} before handing off. * Do not write a new turn boundary after input that the continuation run should * process has been dispatched: the boundary acknowledges that input. * @@ -8742,6 +8745,11 @@ async function endAndContinue(): Promise { "chat.endAndContinue() can only be called from inside a chat.customAgent() run" ); } + if ((locals.get(chatActiveSessionIteratorsKey) ?? 0) > 0) { + throw new Error( + "chat.endAndContinue() cannot be called while a chat.createSession() iterator is active. Use chat.requestUpgrade() instead." + ); + } await performEndAndContinue(); } @@ -9570,6 +9578,47 @@ export type ChatTurn = { | undefined; }; +function trackActiveChatSessionIterator( + iterator: AsyncIterator +): AsyncIterator { + locals.set(chatActiveSessionIteratorsKey, (locals.get(chatActiveSessionIteratorsKey) ?? 0) + 1); + let active = true; + + function finish() { + if (!active) return; + active = false; + const remaining = Math.max((locals.get(chatActiveSessionIteratorsKey) ?? 1) - 1, 0); + locals.set(chatActiveSessionIteratorsKey, remaining); + } + + return { + async next() { + try { + const result = await iterator.next(); + if (result.done) finish(); + return result; + } catch (error) { + try { + await iterator.return?.(); + } catch { + // Preserve the original iterator error after best-effort cleanup. + } + finish(); + throw error; + } + }, + async return() { + try { + return iterator.return + ? await iterator.return() + : { done: true as const, value: undefined }; + } finally { + finish(); + } + }, + }; +} + /** * Create a chat session that yields turns as an async iterator. * @@ -9641,7 +9690,7 @@ function createChatSession( // top of next() in case user code threw without complete()/done(). let activeMsgSub: { off: () => void } | undefined; - return { + const iterator: AsyncIterator = { async next(): Promise> { activeMsgSub?.off(); activeMsgSub = undefined; @@ -10087,6 +10136,8 @@ function createChatSession( return { done: true, value: undefined }; }, }; + + return trackActiveChatSessionIterator(iterator); }, }; } diff --git a/packages/trigger-sdk/test/chat-end-and-continue.test.ts b/packages/trigger-sdk/test/chat-end-and-continue.test.ts index fcb43dbcc4e..b2567a10aab 100644 --- a/packages/trigger-sdk/test/chat-end-and-continue.test.ts +++ b/packages/trigger-sdk/test/chat-end-and-continue.test.ts @@ -1,4 +1,8 @@ -import { describe, expect, it } from "vitest"; +import "../src/v3/test/index.js"; + +import { apiClientManager, resourceCatalog } from "@trigger.dev/core/v3"; +import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; +import { describe, expect, it, vi } from "vitest"; import { chat } from "../src/v3/ai.js"; describe("chat.endAndContinue", () => { @@ -7,4 +11,57 @@ describe("chat.endAndContinue", () => { "chat.endAndContinue() can only be called from inside a chat.customAgent() run" ); }); + + it("rejects while a createSession iterator is active and allows handoff after return", async () => { + const chatId = "end-and-continue-active-session"; + const endAndContinueSession = vi.fn().mockResolvedValue({}); + const clientSpy = vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ + readSessionStreamRecords: async () => ({ records: [] }), + endAndContinueSession, + } as never); + + const agent = chat.customAgent({ + id: "end-and-continue-active-session-agent", + run: async (payload, { signal }) => { + const iterator = chat.createSession(payload, { signal })[Symbol.asyncIterator](); + const firstTurn = await iterator.next(); + expect(firstTurn.done).toBe(false); + + await expect(chat.endAndContinue()).rejects.toThrow( + "chat.endAndContinue() cannot be called while a chat.createSession() iterator is active. Use chat.requestUpgrade() instead." + ); + expect(endAndContinueSession).not.toHaveBeenCalled(); + + await iterator.return?.(); + await chat.endAndContinue(); + }, + }); + const run = resourceCatalog.getTask(agent.id)?.fns.run; + if (!run) throw new Error("custom agent was not registered"); + + try { + await runInMockTaskContext((drivers) => + run( + { + chatId, + trigger: "submit-message", + message: { + id: "user-1", + role: "user", + parts: [{ type: "text", text: "hello" }], + }, + }, + { ctx: drivers.ctx, signal: new AbortController().signal } + ) + ); + + expect(endAndContinueSession).toHaveBeenCalledOnce(); + expect(endAndContinueSession).toHaveBeenCalledWith(chatId, { + callingRunId: "run_test", + reason: "upgrade", + }); + } finally { + clientSpy.mockRestore(); + } + }); }); From 5839cc3a7badf116978f6e3be2593a148dfd352c Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 15:30:54 -0700 Subject: [PATCH 7/8] fix(chat): hold iterator guard through pending reads --- apps/webapp/test/helpers/testChatAgent.ts | 51 +++++++++++ apps/webapp/test/session-agent.e2e.test.ts | 88 +++++++++++++++++++ docs/ai-chat/custom-agents.mdx | 2 +- docs/ai-chat/patterns/version-upgrades.mdx | 2 +- packages/trigger-sdk/src/v3/ai.ts | 84 ++++++++++++++---- .../test/chat-end-and-continue.test.ts | 59 +------------ 6 files changed, 207 insertions(+), 79 deletions(-) diff --git a/apps/webapp/test/helpers/testChatAgent.ts b/apps/webapp/test/helpers/testChatAgent.ts index 29eab173472..9c3811515f0 100644 --- a/apps/webapp/test/helpers/testChatAgent.ts +++ b/apps/webapp/test/helpers/testChatAgent.ts @@ -269,6 +269,57 @@ export const testEndAndContinueCustomAgent = chat.customAgent({ }, }); +export const endAndContinueGuardEvents: Array<{ + chatId: string; + kind: "guard-held" | "return-settled"; +}> = []; + +const activeSessionIteratorError = + "chat.endAndContinue() cannot be called while a chat.createSession() iterator is active. Close the iterator, then call chat.endAndContinue()."; + +async function expectActiveSessionIteratorError() { + try { + await chat.endAndContinue(); + } catch (error) { + if (error instanceof Error && error.message === activeSessionIteratorError) return; + throw error; + } + throw new Error("Expected chat.endAndContinue() to reject while the iterator is active"); +} + +/** Exercises a return racing an already-started next() against real Session input. */ +export const testEndAndContinueIteratorGuardCustomAgent = chat.customAgent({ + id: "e2e-test-chat-custom-end-and-continue-iterator-guard", + run: async (payload, { signal }) => { + const iterator = chat.createSession(payload, { signal })[Symbol.asyncIterator](); + const firstTurn = await iterator.next(); + if (firstTurn.done) { + throw new Error("Expected an initial chat turn"); + } + await firstTurn.value.done(); + await expectActiveSessionIteratorError(); + + const pendingNext = iterator.next(); + if (!iterator.return) { + throw new Error("Expected the chat Session iterator to support return()"); + } + const pendingReturn = iterator.return(); + + // Let an immediately-resolving return() clear a broken guard before checking it. + await new Promise((resolve) => setTimeout(resolve, 0)); + await expectActiveSessionIteratorError(); + endAndContinueGuardEvents.push({ chatId: payload.chatId, kind: "guard-held" }); + + const [nextResult] = await Promise.all([pendingNext, pendingReturn]); + if (nextResult.done) { + throw new Error("Expected the pending next() call to receive the release input"); + } + endAndContinueGuardEvents.push({ chatId: payload.chatId, kind: "return-settled" }); + + await chat.endAndContinue(); + }, +}); + /** * A tool with a server-side `execute`: the agent runs it automatically and * feeds the result back to the model, so a single turn covers the whole diff --git a/apps/webapp/test/session-agent.e2e.test.ts b/apps/webapp/test/session-agent.e2e.test.ts index 0690dd5b0ac..a5ccd9478ed 100644 --- a/apps/webapp/test/session-agent.e2e.test.ts +++ b/apps/webapp/test/session-agent.e2e.test.ts @@ -29,11 +29,13 @@ import { } from "./helpers/sessionStream"; import { runChatAgentSession, runRealChatAgent } from "./helpers/agentHarness"; import { + endAndContinueGuardEvents, suspendResumeEvents, testApprovalChatAgent, testChatAgent, testChatModelLocal, testEndAndContinueCustomAgent, + testEndAndContinueIteratorGuardCustomAgent, testEndRunChatAgent, testHitlChatAgent, testHitlIdleChatAgent, @@ -1655,4 +1657,90 @@ describe("session agent e2e (real chat.agent loop)", () => { await agent.close(); } }); + + it("EA25: custom endAndContinue keeps the guard while iterator next is active", async () => { + const { addressingKey, publicAccessToken, runId, apiKey, environment, baseUrl } = + await setupStartedSession(testEndAndContinueIteratorGuardCustomAgent.id); + const initialRun = await server.prisma.taskRun.findFirstOrThrow({ + where: { friendlyId: runId }, + select: { id: true }, + }); + + const append = await appendInput({ + baseUrl, + addressingKey, + token: publicAccessToken, + partId: "iterator-guard-initial-input", + body: submitBody( + addressingKey, + userMessage("start iterator guard test", "iterator-guard-initial-input") + ), + }); + expect(append.status).toBe(200); + + const agent = runRealChatAgent({ + agentId: testEndAndContinueIteratorGuardCustomAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("unused"), + modelLocal: testChatModelLocal, + runId, + }); + let agentSettled = false; + let agentFailure: unknown; + void agent.done.then( + () => { + agentSettled = true; + }, + (error) => { + agentSettled = true; + agentFailure = error; + } + ); + + try { + await waitFor( + () => + agentSettled || + endAndContinueGuardEvents.some( + (event) => event.chatId === addressingKey && event.kind === "guard-held" + ), + 20_000 + ); + if (agentFailure) throw agentFailure; + expect(agentSettled).toBe(false); + expect( + endAndContinueGuardEvents.some( + (event) => event.chatId === addressingKey && event.kind === "guard-held" + ) + ).toBe(true); + + const release = await appendInput({ + baseUrl, + addressingKey, + token: publicAccessToken, + partId: "iterator-guard-release-input", + body: submitBody( + addressingKey, + userMessage("release pending next", "iterator-guard-release-input") + ), + }); + expect(release.status).toBe(200); + await expect(agent.done).resolves.toBeUndefined(); + + expect( + endAndContinueGuardEvents.some( + (event) => event.chatId === addressingKey && event.kind === "return-settled" + ) + ).toBe(true); + const session = await server.prisma.session.findFirstOrThrow({ + where: { runtimeEnvironmentId: environment.id, externalId: addressingKey }, + select: { currentRunId: true }, + }); + expect(session.currentRunId).not.toBe(initialRun.id); + } finally { + await agent.close(); + } + }); }); diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 3454845e68f..f52804bf988 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -146,7 +146,7 @@ Without this, a resumed chat silently loses its history: the model sees only the ### Rotating to a new deployment -With `chat.createSession()`, use `chat.requestUpgrade()` to leave the current run after the turn. `chat.endAndContinue()` rejects while a `chat.createSession()` iterator is active. In a fully hand-rolled custom agent, use it to immediately hand the Session to a fresh run. +With `chat.createSession()`, use `chat.requestUpgrade()` and let the iterator exit normally. For an immediate handoff, close the iterator before calling `chat.endAndContinue()`; the method rejects until the iterator and any active `next()` call have settled. In a fully hand-rolled custom agent, call it directly to hand the Session to a fresh run. Call it between turns, after detaching the old run's input listeners. If the old run completed its current turn, persist its state and write the turn-complete boundary before the handoff: diff --git a/docs/ai-chat/patterns/version-upgrades.mdx b/docs/ai-chat/patterns/version-upgrades.mdx index cd9c8571706..82f5c3160e1 100644 --- a/docs/ai-chat/patterns/version-upgrades.mdx +++ b/docs/ai-chat/patterns/version-upgrades.mdx @@ -153,7 +153,7 @@ This upgrades on **every** deploy, not just breaking changes. Good for fast-movi ## Custom agents -Use `chat.requestUpgrade()` with `chat.agent()` and the `chat.createSession()` iterator. In a fully hand-rolled `chat.customAgent()` task, detach input listeners, persist the completed turn, write its boundary, then call `chat.endAndContinue()` and return immediately: +Use `chat.requestUpgrade()` with `chat.agent()`. With `chat.createSession()`, call `chat.requestUpgrade()`, then advance the iterator once more so it can exit normally. For an immediate handoff, close the iterator before calling `chat.endAndContinue()`. In a fully hand-rolled `chat.customAgent()` task, detach input listeners, persist the completed turn, write its boundary, then call `chat.endAndContinue()` and return immediately: ```ts // Detach any chat.messages.on() subscriptions you created. diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 0d4aceb28ab..c0aecaeef03 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -8715,11 +8715,11 @@ function requestUpgrade(): void { * Hand off the current custom agent Session to a fresh run. * * This is the low-level handoff for a fully hand-rolled - * `chat.customAgent()` loop. (Use {@link requestUpgrade} with - * `chat.createSession()` instead; this method rejects while its iterator is - * active.) Call only between turns and after detaching input listeners for the - * old run. If the old run completed its current turn, persist its state and - * call {@link chatWriteTurnComplete} before handing off. + * `chat.customAgent()` loop. This method rejects while a + * `chat.createSession()` iterator is active. Close the iterator before calling + * it. Call only between turns and after detaching input listeners for the old + * run. If the old run completed its current turn, persist its state and call + * {@link chatWriteTurnComplete} before handing off. * Do not write a new turn boundary after input that the continuation run should * process has been dispatched: the boundary acknowledges that input. * @@ -8747,7 +8747,7 @@ async function endAndContinue(): Promise { } if ((locals.get(chatActiveSessionIteratorsKey) ?? 0) > 0) { throw new Error( - "chat.endAndContinue() cannot be called while a chat.createSession() iterator is active. Use chat.requestUpgrade() instead." + "chat.endAndContinue() cannot be called while a chat.createSession() iterator is active. Close the iterator, then call chat.endAndContinue()." ); } @@ -9583,6 +9583,10 @@ function trackActiveChatSessionIterator( ): AsyncIterator { locals.set(chatActiveSessionIteratorsKey, (locals.get(chatActiveSessionIteratorsKey) ?? 0) + 1); let active = true; + let closing = false; + let activeNextCalls = 0; + let closePromise: Promise> | undefined; + const nextSettledWaiters = new Set<() => void>(); function finish() { if (!active) return; @@ -9591,30 +9595,72 @@ function trackActiveChatSessionIterator( locals.set(chatActiveSessionIteratorsKey, remaining); } + function settleNextCall() { + activeNextCalls = Math.max(activeNextCalls - 1, 0); + if (activeNextCalls > 0) return; + + for (const resolve of nextSettledWaiters) { + resolve(); + } + nextSettledWaiters.clear(); + } + + function waitForNextCalls(): Promise { + if (activeNextCalls === 0) return Promise.resolve(); + return new Promise((resolve) => nextSettledWaiters.add(resolve)); + } + + function closeIterator(): Promise> { + closing = true; + if (!closePromise) { + closePromise = (async () => { + // A blocked next() can install a new message listener after it resumes. + // Let every started call settle, then make the inner cleanup final. + await waitForNextCalls(); + try { + return iterator.return + ? await iterator.return() + : { done: true as const, value: undefined }; + } finally { + finish(); + } + })(); + } + return closePromise; + } + return { async next() { + if (closing) { + return { done: true as const, value: undefined }; + } + + activeNextCalls++; + let result: IteratorResult; try { - const result = await iterator.next(); - if (result.done) finish(); - return result; + result = await iterator.next(); } catch (error) { + settleNextCall(); try { - await iterator.return?.(); + await closeIterator(); } catch { // Preserve the original iterator error after best-effort cleanup. } - finish(); throw error; } - }, - async return() { - try { - return iterator.return - ? await iterator.return() - : { done: true as const, value: undefined }; - } finally { - finish(); + + settleNextCall(); + if (result.done) { + try { + await closeIterator(); + } catch { + // The inner next() already ended cleanly; cleanup remains best-effort. + } } + return result; + }, + return() { + return closeIterator(); }, }; } diff --git a/packages/trigger-sdk/test/chat-end-and-continue.test.ts b/packages/trigger-sdk/test/chat-end-and-continue.test.ts index b2567a10aab..fcb43dbcc4e 100644 --- a/packages/trigger-sdk/test/chat-end-and-continue.test.ts +++ b/packages/trigger-sdk/test/chat-end-and-continue.test.ts @@ -1,8 +1,4 @@ -import "../src/v3/test/index.js"; - -import { apiClientManager, resourceCatalog } from "@trigger.dev/core/v3"; -import { runInMockTaskContext } from "@trigger.dev/core/v3/test"; -import { describe, expect, it, vi } from "vitest"; +import { describe, expect, it } from "vitest"; import { chat } from "../src/v3/ai.js"; describe("chat.endAndContinue", () => { @@ -11,57 +7,4 @@ describe("chat.endAndContinue", () => { "chat.endAndContinue() can only be called from inside a chat.customAgent() run" ); }); - - it("rejects while a createSession iterator is active and allows handoff after return", async () => { - const chatId = "end-and-continue-active-session"; - const endAndContinueSession = vi.fn().mockResolvedValue({}); - const clientSpy = vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({ - readSessionStreamRecords: async () => ({ records: [] }), - endAndContinueSession, - } as never); - - const agent = chat.customAgent({ - id: "end-and-continue-active-session-agent", - run: async (payload, { signal }) => { - const iterator = chat.createSession(payload, { signal })[Symbol.asyncIterator](); - const firstTurn = await iterator.next(); - expect(firstTurn.done).toBe(false); - - await expect(chat.endAndContinue()).rejects.toThrow( - "chat.endAndContinue() cannot be called while a chat.createSession() iterator is active. Use chat.requestUpgrade() instead." - ); - expect(endAndContinueSession).not.toHaveBeenCalled(); - - await iterator.return?.(); - await chat.endAndContinue(); - }, - }); - const run = resourceCatalog.getTask(agent.id)?.fns.run; - if (!run) throw new Error("custom agent was not registered"); - - try { - await runInMockTaskContext((drivers) => - run( - { - chatId, - trigger: "submit-message", - message: { - id: "user-1", - role: "user", - parts: [{ type: "text", text: "hello" }], - }, - }, - { ctx: drivers.ctx, signal: new AbortController().signal } - ) - ); - - expect(endAndContinueSession).toHaveBeenCalledOnce(); - expect(endAndContinueSession).toHaveBeenCalledWith(chatId, { - callingRunId: "run_test", - reason: "upgrade", - }); - } finally { - clientSpy.mockRestore(); - } - }); }); From 31810b738cb0122ee0031744a2aca3887ba3ed55 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Mon, 17 Aug 2026 16:13:51 -0700 Subject: [PATCH 8/8] fix(chat): suppress turns after iterator return --- apps/webapp/test/helpers/testChatAgent.ts | 31 ++++++++++++++++++++-- apps/webapp/test/session-agent.e2e.test.ts | 28 +++++++++++++++++++ docs/ai-chat/custom-agents.mdx | 2 ++ docs/ai-chat/patterns/version-upgrades.mdx | 2 ++ packages/trigger-sdk/src/v3/ai.ts | 13 ++++++--- 5 files changed, 71 insertions(+), 5 deletions(-) diff --git a/apps/webapp/test/helpers/testChatAgent.ts b/apps/webapp/test/helpers/testChatAgent.ts index 9c3811515f0..e766bd3ee04 100644 --- a/apps/webapp/test/helpers/testChatAgent.ts +++ b/apps/webapp/test/helpers/testChatAgent.ts @@ -291,6 +291,33 @@ async function expectActiveSessionIteratorError() { export const testEndAndContinueIteratorGuardCustomAgent = chat.customAgent({ id: "e2e-test-chat-custom-end-and-continue-iterator-guard", run: async (payload, { signal }) => { + if (payload.continuation) { + const next = await chat.messages.waitWithIdleTimeout({ + idleTimeoutInSeconds: 2, + timeout: "1m", + }); + if (!next.ok) { + throw next.error; + } + + const message = next.output.message as UIMessage | undefined; + const text = message ? firstText(message) : ""; + const { waitUntilComplete } = chat.stream.writer({ + execute: ({ write }) => { + write({ type: "text-start", id: "guard-continuation-result" }); + write({ + type: "text-delta", + id: "guard-continuation-result", + delta: `received:${text}`, + }); + write({ type: "text-end", id: "guard-continuation-result" }); + }, + }); + await waitUntilComplete(); + await chat.writeTurnComplete(); + return; + } + const iterator = chat.createSession(payload, { signal })[Symbol.asyncIterator](); const firstTurn = await iterator.next(); if (firstTurn.done) { @@ -311,8 +338,8 @@ export const testEndAndContinueIteratorGuardCustomAgent = chat.customAgent({ endAndContinueGuardEvents.push({ chatId: payload.chatId, kind: "guard-held" }); const [nextResult] = await Promise.all([pendingNext, pendingReturn]); - if (nextResult.done) { - throw new Error("Expected the pending next() call to receive the release input"); + if (!nextResult.done) { + throw new Error("Expected return() to suppress the pending next() turn"); } endAndContinueGuardEvents.push({ chatId: payload.chatId, kind: "return-settled" }); diff --git a/apps/webapp/test/session-agent.e2e.test.ts b/apps/webapp/test/session-agent.e2e.test.ts index a5ccd9478ed..6bc9d8c2c42 100644 --- a/apps/webapp/test/session-agent.e2e.test.ts +++ b/apps/webapp/test/session-agent.e2e.test.ts @@ -1698,6 +1698,7 @@ describe("session agent e2e (real chat.agent loop)", () => { agentFailure = error; } ); + let continuation: ReturnType | undefined; try { await waitFor( @@ -1739,7 +1740,34 @@ describe("session agent e2e (real chat.agent loop)", () => { select: { currentRunId: true }, }); expect(session.currentRunId).not.toBe(initialRun.id); + + const successor = await server.prisma.taskRun.findFirstOrThrow({ + where: { id: session.currentRunId! }, + select: { friendlyId: true }, + }); + continuation = runRealChatAgent({ + agentId: testEndAndContinueIteratorGuardCustomAgent.id, + baseUrl, + addressingKey, + secretKey: apiKey, + model: textModel("unused"), + modelLocal: testChatModelLocal, + runId: successor.friendlyId, + continuation: true, + previousRunId: runId, + }); + + const { parts } = await collectSessionOut({ + baseUrl, + addressingKey, + token: publicAccessToken, + until: (records) => records.filter(isTurnComplete).length >= 2, + maxMs: 30_000, + }); + expect(joinChunks(parts)).toContain("received:release pending next"); + await expect(continuation.done).resolves.toBeUndefined(); } finally { + await continuation?.close(); await agent.close(); } }); diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index f52804bf988..cf588e8f5c1 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -148,6 +148,8 @@ Without this, a resumed chat silently loses its history: the model sees only the With `chat.createSession()`, use `chat.requestUpgrade()` and let the iterator exit normally. For an immediate handoff, close the iterator before calling `chat.endAndContinue()`; the method rejects until the iterator and any active `next()` call have settled. In a fully hand-rolled custom agent, call it directly to hand the Session to a fresh run. +Close the iterator between reads. If `return()` races a `next()` that is already waiting for input, it waits for that read to settle before releasing the handoff guard. Input dispatched while the iterator is closing is not yielded as a turn and remains available to the continuation unless you write another turn-complete boundary. + Call it between turns, after detaching the old run's input listeners. If the old run completed its current turn, persist its state and write the turn-complete boundary before the handoff: ```ts diff --git a/docs/ai-chat/patterns/version-upgrades.mdx b/docs/ai-chat/patterns/version-upgrades.mdx index 82f5c3160e1..80e8911b439 100644 --- a/docs/ai-chat/patterns/version-upgrades.mdx +++ b/docs/ai-chat/patterns/version-upgrades.mdx @@ -155,6 +155,8 @@ This upgrades on **every** deploy, not just breaking changes. Good for fast-movi Use `chat.requestUpgrade()` with `chat.agent()`. With `chat.createSession()`, call `chat.requestUpgrade()`, then advance the iterator once more so it can exit normally. For an immediate handoff, close the iterator before calling `chat.endAndContinue()`. In a fully hand-rolled `chat.customAgent()` task, detach input listeners, persist the completed turn, write its boundary, then call `chat.endAndContinue()` and return immediately: +Close a `chat.createSession()` iterator between reads. If `return()` races a `next()` that is already waiting for input, it waits for that read to settle before the handoff can continue. Input dispatched while the iterator is closing is not yielded as a turn and remains available to the continuation unless you write another turn-complete boundary. + ```ts // Detach any chat.messages.on() subscriptions you created. stop.cleanup(); diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index c0aecaeef03..c46dc457769 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -8717,9 +8717,11 @@ function requestUpgrade(): void { * This is the low-level handoff for a fully hand-rolled * `chat.customAgent()` loop. This method rejects while a * `chat.createSession()` iterator is active. Close the iterator before calling - * it. Call only between turns and after detaching input listeners for the old - * run. If the old run completed its current turn, persist its state and call - * {@link chatWriteTurnComplete} before handing off. + * it. If `return()` races an active `next()`, it waits for that read to settle + * before releasing the handoff guard. Call only between turns and after + * detaching input listeners for the old run. If the old run completed its + * current turn, persist its state and call {@link chatWriteTurnComplete} before + * handing off. * Do not write a new turn boundary after input that the continuation run should * process has been dispatched: the boundary acknowledges that input. * @@ -9656,6 +9658,11 @@ function trackActiveChatSessionIterator( } catch { // The inner next() already ended cleanly; cleanup remains best-effort. } + } else if (closing) { + // return() won the race. Do not expose a turn the caller has already + // abandoned; without a new turn-complete boundary its input remains + // replayable by the continuation run. + return { done: true as const, value: undefined }; } return result; },