feat(chat): expose endAndContinue to custom agents

Author: gtremper

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

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

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

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                                                                                             |

packages/trigger-sdk/src/v3/ai.ts

@@ -2538,6 +2538,8 @@ const chatOnCompactedKey =
   locals.create<(event: CompactedEvent) => Promise<void> | void>("chat.onCompacted");
 /** @internal Full task `ctx` for the active `chat.agent` run (for hooks invoked from nested compaction). */
 const chatAgentRunContextKey = locals.create<TaskRunContext>("chat.agentRunContext");
+/** @internal Marks the root run created by `chat.customAgent()`. */
+const chatCustomAgentRunKey = locals.create<boolean>("chat.customAgentRun");
 const chatPrepareMessagesKey =
   locals.create<(event: PrepareMessagesEvent<unknown>) => ModelMessage[] | Promise<ModelMessage[]>>(
     "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<void> {
+  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<void> {
+  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<StreamWriteResult> {
-  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<StreamWriteResult> {
   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.

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