feat(chat): expose endAndContinue to custom agents (#4647)

## Summary

Raw `chat.customAgent()` loops can now call `chat.endAndContinue()` to
move the Session to a fresh run. The managed loop already used the same
server operation through `chat.requestUpgrade()`, but raw loops could
not call it directly.

Call the method between turns after detaching input listeners from the
old run. Await it and return immediately. Unconsumed `.in` records stay
on the Session for the continuation run.

I put this on the `chat` namespace next to the other raw chat
primitives. Happy to move it if maintainers prefer a different API
placement.

## Testing

- `pnpm exec vitest run` in `packages/trigger-sdk` (374 tests)
- Focused webapp Session E2E tests (3 tests)
- `pnpm run build` in `packages/trigger-sdk`
- Webapp typecheck
- `pnpm run format`
- `pnpm run lint`

## Checklist

- [x] I followed the contributing guide
- [x] The PR title follows the convention
- [x] I tested the change

## Changelog

Allow custom chat agents to rotate to a new task version without
dropping unconsumed Session input.

---------

Co-authored-by: Eric Allam <eallam@icloud.com>

Author: gtremper

.changeset/chat-custom-agent-end-and-continue.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Add `chat.endAndContinue()` so fully hand-rolled custom chat agents can hand a conversation off to a fresh run on the latest deployed task version while preserving unconsumed Session input.

apps/webapp/test/helpers/testChatAgent.ts

@@ -234,6 +234,119 @@ export const testUpgradeOnceChatAgent = chat.agent({
   },
 });
 
+/**
+ * Hands an unconsumed Session input record to a continuation run using the
+ * public custom-agent lifecycle primitive. The continuation echoes the input
+ * to `.out`, which lets the full-stack Session E2E assert durable delivery.
+ */
+export const testEndAndContinueCustomAgent = chat.customAgent({
+  id: "e2e-test-chat-custom-end-and-continue",
+  run: async (payload) => {
+    if (!payload.continuation) {
+      await chat.endAndContinue();
+      return;
+    }
+
+    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: "handoff-result" });
+        write({ type: "text-delta", id: "handoff-result", delta: `received:${text}` });
+        write({ type: "text-end", id: "handoff-result" });
+      },
+    });
+    await waitUntilComplete();
+    await chat.writeTurnComplete();
+  },
+});
+
+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 }) => {
+    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) {
+      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 return() to suppress the pending next() turn");
+    }
+    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,10 +29,13 @@ import {
 } from "./helpers/sessionStream";
 import { runChatAgentSession, runRealChatAgent } from "./helpers/agentHarness";
 import {
+  endAndContinueGuardEvents,
   suspendResumeEvents,
   testApprovalChatAgent,
   testChatAgent,
   testChatModelLocal,
+  testEndAndContinueCustomAgent,
+  testEndAndContinueIteratorGuardCustomAgent,
   testEndRunChatAgent,
   testHitlChatAgent,
   testHitlIdleChatAgent,
@@ -123,6 +126,34 @@ async function setupSession(agentId: string = testChatAgent.id) {
   return { addressingKey, token, apiKey, baseUrl: server.webapp.baseUrl };
 }
 
+async function setupStartedSession(agentId: string) {
+  const { environment, apiKey } = await seedTestEnvironment(server.prisma);
+  const addressingKey = `chat-${randomBytes(6).toString("hex")}`;
+  const createRes = await fetch(`${server.webapp.baseUrl}/api/v1/sessions`, {
+    method: "POST",
+    headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" },
+    body: JSON.stringify({
+      type: "chat.agent",
+      externalId: addressingKey,
+      taskIdentifier: agentId,
+      triggerConfig: { basePayload: {} },
+    }),
+  });
+
+  expect(createRes.ok).toBe(true);
+  const created = (await createRes.json()) as {
+    runId: string;
+    publicAccessToken: string;
+  };
+  return {
+    ...created,
+    addressingKey,
+    apiKey,
+    environment,
+    baseUrl: server.webapp.baseUrl,
+  };
+}
+
 function promptText(prompt: unknown): string {
   if (!Array.isArray(prompt)) return "";
   let out = "";
@@ -1533,4 +1564,211 @@ describe("session agent e2e (real chat.agent loop)", () => {
       await agent.close();
     }
   });
+
+  it("EA23: custom endAndContinue hands pending input to a fresh run", async () => {
+    const { addressingKey, publicAccessToken, runId, apiKey, environment, baseUrl } =
+      await setupStartedSession(testEndAndContinueCustomAgent.id);
+    const initialRun = await server.prisma.taskRun.findFirstOrThrow({
+      where: { friendlyId: runId },
+      select: { id: true },
+    });
+
+    const append = await appendInput({
+      baseUrl,
+      addressingKey,
+      token: publicAccessToken,
+      partId: "pending-handoff-input",
+      body: submitBody(
+        addressingKey,
+        userMessage("deliver after endAndContinue", "pending-handoff-input")
+      ),
+    });
+    expect(append.status).toBe(200);
+
+    const oldRun = runRealChatAgent({
+      agentId: testEndAndContinueCustomAgent.id,
+      baseUrl,
+      addressingKey,
+      secretKey: apiKey,
+      model: textModel("unused"),
+      modelLocal: testChatModelLocal,
+      runId,
+    });
+    let continuation: ReturnType<typeof runRealChatAgent> | undefined;
+
+    try {
+      await expect(oldRun.done).resolves.toBeUndefined();
+
+      const session = await server.prisma.session.findFirstOrThrow({
+        where: { runtimeEnvironmentId: environment.id, externalId: addressingKey },
+        select: { currentRunId: true, currentRunVersion: true },
+      });
+      expect(session.currentRunId).not.toBe(initialRun.id);
+      expect(session.currentRunVersion).toBeGreaterThan(1);
+
+      const successor = await server.prisma.taskRun.findFirstOrThrow({
+        where: { id: session.currentRunId! },
+        select: { friendlyId: true },
+      });
+      continuation = runRealChatAgent({
+        agentId: testEndAndContinueCustomAgent.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: (p) => p.some(isTurnComplete),
+        maxMs: 30_000,
+      });
+      expect(joinChunks(parts)).toContain("received:deliver after endAndContinue");
+      await expect(continuation.done).resolves.toBeUndefined();
+    } finally {
+      await continuation?.close();
+      await oldRun.close();
+    }
+  });
+
+  it("EA24: custom endAndContinue rejects when the server rejects the handoff", async () => {
+    const { addressingKey, apiKey, baseUrl } = await setupStartedSession(
+      testEndAndContinueCustomAgent.id
+    );
+    const agent = runRealChatAgent({
+      agentId: testEndAndContinueCustomAgent.id,
+      baseUrl,
+      addressingKey,
+      secretKey: apiKey,
+      model: textModel("unused"),
+      modelLocal: testChatModelLocal,
+      runId: "run_missing_end_and_continue",
+    });
+
+    try {
+      await expect(agent.done).rejects.toThrow("callingRunId not found in this environment");
+    } finally {
+      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;
+      }
+    );
+    let continuation: ReturnType<typeof runRealChatAgent> | undefined;
+
+    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);
+
+      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();
+    }
+  });
 });