diff --git a/.changeset/chat-custom-agent-end-and-continue.md b/.changeset/chat-custom-agent-end-and-continue.md new file mode 100644 index 00000000000..6d7deffeeeb --- /dev/null +++ b/.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. diff --git a/apps/webapp/test/helpers/testChatAgent.ts b/apps/webapp/test/helpers/testChatAgent.ts index 8aebb713a96..e766bd3ee04 100644 --- a/apps/webapp/test/helpers/testChatAgent.ts +++ b/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 diff --git a/apps/webapp/test/session-agent.e2e.test.ts b/apps/webapp/test/session-agent.e2e.test.ts index f6a9a2a338c..6bc9d8c2c42 100644 --- a/apps/webapp/test/session-agent.e2e.test.ts +++ b/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 | 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 | 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(); + } + }); }); diff --git a/docs/ai-chat/custom-agents.mdx b/docs/ai-chat/custom-agents.mdx index 197bff6b5e1..cf588e8f5c1 100644 --- a/docs/ai-chat/custom-agents.mdx +++ b/docs/ai-chat/custom-agents.mdx @@ -144,6 +144,31 @@ 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 + +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 +// Detach any chat.messages.on() subscriptions you created. +stop.cleanup(); +await persistMessages(conversation.uiMessages); +await chat.writeTurnComplete(); +await chat.endAndContinue(); +return; +``` + +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`. + +If input has been dispatched to the old run but should be processed by the continuation, detach the listeners and do not write another turn-complete boundary before handing off. `chat.writeTurnComplete()` acknowledges the latest input dispatched to the old run; writing it after that dispatch would make the continuation resume after the input. + + + `chat.endAndContinue()` 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. If the handoff fails, the promise rejects. + + ### 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 +242,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 between turns, 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 | diff --git a/docs/ai-chat/patterns/version-upgrades.mdx b/docs/ai-chat/patterns/version-upgrades.mdx index 830f673e9a3..80e8911b439 100644 --- a/docs/ai-chat/patterns/version-upgrades.mdx +++ b/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()` is the managed upgrade signal for `chat.agent()` and the `chat.createSession()` iterator. Fully hand-rolled custom agents use `chat.endAndContinue()` between turns to immediately hand the Session to a new run. ## How it works @@ -151,14 +151,34 @@ 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 +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(); +await persistMessages(conversation.uiMessages); +await chat.writeTurnComplete(); +await chat.endAndContinue(); +return; +``` + +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`. + +If input has been dispatched to the old run but should be processed by the continuation, detach the old listeners and skip the final `chat.writeTurnComplete()`. A turn-complete boundary acknowledges the latest input dispatched to the old run, so writing one after that dispatch would cause the continuation to resume past the input. + + + `chat.endAndContinue()` starts the successor but does not stop the calling run. Perform no more Session reads or writes after calling it, and return from the task. If the handoff fails, the promise rejects. + ## Interaction with recovery boot -`chat.requestUpgrade()` is a graceful exit — the old run returns cleanly, never writing a partial assistant. The new continuation run boots with an empty `session.out` tail and the upgrade-trigger message on `session.in`. The trigger message dispatches as turn 1 on the new version via the normal continuation-wait path. [`onRecoveryBoot`](/ai-chat/patterns/recovery-boot) does NOT fire on this path — the hook is reserved for mid-stream interruptions (cancel / crash / OOM) where a partial assistant exists on the tail. +When `chat.requestUpgrade()` is handled before a turn starts, the SDK immediately hands the Session to a new run, which processes the same input on the latest version. When it is requested during a turn, including through `chat.createSession()`, the current turn finishes and the old run exits; the next input starts the continuation run. + +Both are graceful exits. [`onRecoveryBoot`](/ai-chat/patterns/recovery-boot) does not fire — the hook is reserved for mid-stream interruptions (cancel, crash, or OOM) where a partial assistant exists on the tail. ## See also diff --git a/docs/ai-chat/reference.mdx b/docs/ai-chat/reference.mdx index da08f9a0473..68812b3c21b 100644 --- a/docs/ai-chat/reference.mdx +++ b/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. Call between turns after detaching input listeners, then return immediately. The promise rejects if the handoff fails. | | `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 | diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index 844d506079b..c46dc457769 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -2538,6 +2538,10 @@ const chatOnCompactedKey = locals.create<(event: CompactedEvent) => Promise | void>("chat.onCompacted"); /** @internal Full task `ctx` for the active `chat.agent` run (for hooks invoked from nested compaction). */ const chatAgentRunContextKey = locals.create("chat.agentRunContext"); +/** @internal Marks the root run created by `chat.customAgent()`. */ +const chatCustomAgentRunKey = locals.create("chat.customAgentRun"); +/** @internal Number of active `chat.createSession()` iterators in this run. */ +const chatActiveSessionIteratorsKey = locals.create("chat.createSession.activeIterators"); const chatPrepareMessagesKey = locals.create<(event: PrepareMessagesEvent) => ModelMessage[] | Promise>( "chat.prepareMessages" @@ -5362,6 +5366,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 +5461,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 +8711,67 @@ 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. This method rejects while a + * `chat.createSession()` iterator is active. Close the iterator before calling + * 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. + * + * The server starts the continuation run but does not stop this run, so return + * from the task immediately after awaiting this function. The promise rejects + * if the server cannot complete the handoff. + * + * 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 + * // Detach any chat.messages.on() subscriptions you created. + * await persistMessages(); + * await chat.writeTurnComplete(); + * await chat.endAndContinue(); + * return; + * ``` + */ +async function endAndContinue(): Promise { + if (locals.get(chatCustomAgentRunKey) !== true) { + throw new Error( + "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. Close the iterator, then call chat.endAndContinue()." + ); + } + + await performEndAndContinue(); +} + +/** @internal Shared server handoff used by managed and custom agent loops. */ +async function performEndAndContinue(): Promise { + 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 @@ -9513,6 +9580,98 @@ export type ChatTurn = { | undefined; }; +function trackActiveChatSessionIterator( + iterator: AsyncIterator +): AsyncIterator { + locals.set(chatActiveSessionIteratorsKey, (locals.get(chatActiveSessionIteratorsKey) ?? 0) + 1); + let active = true; + let closing = false; + let activeNextCalls = 0; + let closePromise: Promise> | undefined; + const nextSettledWaiters = new Set<() => void>(); + + function finish() { + if (!active) return; + active = false; + const remaining = Math.max((locals.get(chatActiveSessionIteratorsKey) ?? 1) - 1, 0); + 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 { + if (activeNextCalls === 0) return Promise.resolve(); + return new Promise((resolve) => nextSettledWaiters.add(resolve)); + } + + function closeIterator(): Promise> { + 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; + try { + result = await iterator.next(); + } catch (error) { + settleNextCall(); + try { + await closeIterator(); + } catch { + // Preserve the original iterator error after best-effort cleanup. + } + throw error; + } + + settleNextCall(); + if (result.done) { + try { + await closeIterator(); + } 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; + }, + return() { + return closeIterator(); + }, + }; +} + /** * Create a chat session that yields turns as an async iterator. * @@ -9584,7 +9743,7 @@ function createChatSession( // top of next() in case user code threw without complete()/done(). let activeMsgSub: { off: () => void } | undefined; - return { + const iterator: AsyncIterator = { async next(): Promise> { activeMsgSub?.off(); activeMsgSub = undefined; @@ -10030,6 +10189,8 @@ function createChatSession( return { done: true, value: undefined }; }, }; + + return trackActiveChatSessionIterator(iterator); }, }; } @@ -10697,6 +10858,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 +11054,12 @@ async function writeTurnCompleteChunk( * @internal */ async function writeUpgradeRequiredChunk(): Promise { - 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 +11075,6 @@ async function writeUpgradeRequiredChunk(): Promise { 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. diff --git a/packages/trigger-sdk/test/chat-end-and-continue.test.ts b/packages/trigger-sdk/test/chat-end-and-continue.test.ts new file mode 100644 index 00000000000..fcb43dbcc4e --- /dev/null +++ b/packages/trigger-sdk/test/chat-end-and-continue.test.ts @@ -0,0 +1,10 @@ +import { describe, expect, it } from "vitest"; +import { chat } from "../src/v3/ai.js"; + +describe("chat.endAndContinue", () => { + it("rejects calls outside a custom agent run", async () => { + await expect(chat.endAndContinue()).rejects.toThrow( + "chat.endAndContinue() can only be called from inside a chat.customAgent() run" + ); + }); +});