fix(chat): guard handoff during managed sessions

Author: gtremper

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:
 

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

@@ -2540,6 +2540,8 @@ const chatOnCompactedKey =
 const chatAgentRunContextKey = locals.create<TaskRunContext>("chat.agentRunContext");
 /** @internal Marks the root run created by `chat.customAgent()`. */
 const chatCustomAgentRunKey = locals.create<boolean>("chat.customAgentRun");
+/** @internal Number of active `chat.createSession()` iterators in this run. */
+const chatActiveSessionIteratorsKey = locals.create<number>("chat.createSession.activeIterators");
 const chatPrepareMessagesKey =
   locals.create<(event: PrepareMessagesEvent<unknown>) => ModelMessage[] | Promise<ModelMessage[]>>(
     "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<void> {
       "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<ChatTurn>
+): AsyncIterator<ChatTurn> {
+  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<ChatTurn> = {
         async next(): Promise<IteratorResult<ChatTurn>> {
           activeMsgSub?.off();
           activeMsgSub = undefined;
@@ -10087,6 +10136,8 @@ function createChatSession(
           return { done: true, value: undefined };
         },
       };
+
+      return trackActiveChatSessionIterator(iterator);
     },
   };
 }

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();
+    }
+  });
 });