fix(chat): suppress turns after iterator return

Author: gtremper

apps/webapp/test/helpers/testChatAgent.ts

@@ -291,6 +291,33 @@ async function expectActiveSessionIteratorError() {
 export const testEndAndContinueIteratorGuardCustomAgent = chat.customAgent({
   id: "e2e-test-chat-custom-end-and-continue-iterator-guard",
   run: async (payload, { signal }) => {
+    if (payload.continuation) {
+      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: "guard-continuation-result" });
+          write({
+            type: "text-delta",
+            id: "guard-continuation-result",
+            delta: `received:${text}`,
+          });
+          write({ type: "text-end", id: "guard-continuation-result" });
+        },
+      });
+      await waitUntilComplete();
+      await chat.writeTurnComplete();
+      return;
+    }
+
     const iterator = chat.createSession(payload, { signal })[Symbol.asyncIterator]();
     const firstTurn = await iterator.next();
     if (firstTurn.done) {
@@ -311,8 +338,8 @@ export const testEndAndContinueIteratorGuardCustomAgent = chat.customAgent({
     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");
+    if (!nextResult.done) {
+      throw new Error("Expected return() to suppress the pending next() turn");
     }
     endAndContinueGuardEvents.push({ chatId: payload.chatId, kind: "return-settled" });
 

apps/webapp/test/session-agent.e2e.test.ts

@@ -1698,6 +1698,7 @@ describe("session agent e2e (real chat.agent loop)", () => {
         agentFailure = error;
       }
     );
+    let continuation: ReturnType<typeof runRealChatAgent> | undefined;
 
     try {
       await waitFor(
@@ -1739,7 +1740,34 @@ describe("session agent e2e (real chat.agent loop)", () => {
         select: { currentRunId: true },
       });
       expect(session.currentRunId).not.toBe(initialRun.id);
+
+      const successor = await server.prisma.taskRun.findFirstOrThrow({
+        where: { id: session.currentRunId! },
+        select: { friendlyId: true },
+      });
+      continuation = runRealChatAgent({
+        agentId: testEndAndContinueIteratorGuardCustomAgent.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: (records) => records.filter(isTurnComplete).length >= 2,
+        maxMs: 30_000,
+      });
+      expect(joinChunks(parts)).toContain("received:release pending next");
+      await expect(continuation.done).resolves.toBeUndefined();
     } finally {
+      await continuation?.close();
       await agent.close();
     }
   });

docs/ai-chat/custom-agents.mdx

@@ -148,6 +148,8 @@ Without this, a resumed chat silently loses its history: the model sees only the
 
 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.
 
+Close the iterator between reads. If `return()` races a `next()` that is already waiting for input, it waits for that read to settle before releasing the handoff guard. Input dispatched while the iterator is closing is not yielded as a turn and remains available to the continuation unless you write another turn-complete boundary.
+
 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

docs/ai-chat/patterns/version-upgrades.mdx

@@ -155,6 +155,8 @@ This upgrades on **every** deploy, not just breaking changes. Good for fast-movi
 
 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:
 
+Close a `chat.createSession()` iterator between reads. If `return()` races a `next()` that is already waiting for input, it waits for that read to settle before the handoff can continue. Input dispatched while the iterator is closing is not yielded as a turn and remains available to the continuation unless you write another turn-complete boundary.
+
 ```ts
 // Detach any chat.messages.on() subscriptions you created.
 stop.cleanup();

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

@@ -8717,9 +8717,11 @@ function requestUpgrade(): void {
  * This is the low-level handoff for a fully hand-rolled
  * `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.
+ * it. If `return()` races an active `next()`, it waits for that read to settle
+ * before releasing the handoff guard. 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.
  *
@@ -9656,6 +9658,11 @@ function trackActiveChatSessionIterator(
         } catch {
           // The inner next() already ended cleanly; cleanup remains best-effort.
         }
+      } else if (closing) {
+        // return() won the race. Do not expose a turn the caller has already
+        // abandoned; without a new turn-complete boundary its input remains
+        // replayable by the continuation run.
+        return { done: true as const, value: undefined };
       }
       return result;
     },