From fcefcf52d0681318d53b88115b2bbcf4c6733388 Mon Sep 17 00:00:00 2001 From: Graham Tremper Date: Sun, 16 Aug 2026 20:26:40 -0700 Subject: [PATCH 1/5] 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 0000000000..b9b1cff29e --- /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 197bff6b5e..975e912ad1 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 830f673e9a..3bbe6c216f 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 da08f9a047..7f38799663 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 844d506079..c638a4cbd0 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 0000000000..453b2871f8 --- /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/5] 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 b9b1cff29e..ef1bf09087 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 975e912ad1..e373a7ed57 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 3bbe6c216f..eaf288b7b8 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 7f38799663..68812b3c21 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 c638a4cbd0..18f89ab921 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 453b2871f8..4c706391b5 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/5] 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 ef1bf09087..6d7deffeee 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 e373a7ed57..c7da25eda4 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 eaf288b7b8..cd9c857170 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 18f89ab921..c316af94ff 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/5] 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 8aebb713a9..29eab17347 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 f6a9a2a338..18a9986216 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 4c706391b5..fcb43dbcc4 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/5] 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 18a9986216..0690dd5b0a 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 }, });