fix(chat): hold iterator guard through pending reads

Author: gtremper

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

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

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:
 

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.

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<void> {
   }
   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<ChatTurn> {
   locals.set(chatActiveSessionIteratorsKey, (locals.get(chatActiveSessionIteratorsKey) ?? 0) + 1);
   let active = true;
+  let closing = false;
+  let activeNextCalls = 0;
+  let closePromise: Promise<IteratorResult<ChatTurn>> | 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<void> {
+    if (activeNextCalls === 0) return Promise.resolve();
+    return new Promise((resolve) => nextSettledWaiters.add(resolve));
+  }
+
+  function closeIterator(): Promise<IteratorResult<ChatTurn>> {
+    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<ChatTurn>;
       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();
     },
   };
 }

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