docs(chat): clarify endAndContinue lifecycle

Author: gtremper

.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.

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.
+
 <Warning>
-  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.
 </Warning>
 
 ### 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                                            |

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.
+
 <Warning>
-  `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.
 </Warning>
 
 ## Interaction with recovery boot

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                                                                                             |

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<void> {

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<string, unknown>,
+  options: { ctx: unknown; signal: AbortSignal }
+) => Promise<unknown>;
+
+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<string, unknown>,
-      options: { ctx: unknown; signal: AbortSignal }
-    ) => Promise<unknown>;
+    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({