feat(chat): add chat.endRun()

Exit the loop after the current turn completes without the
upgrade-required signal that chat.requestUpgrade() sends. Use when an
agent finishes its work on its own terms — one-shot responses, goal
achieved, budget exhausted — instead of waiting idle for the next
user message. Callable from run(), chat.defer(),
onBeforeTurnComplete, or onTurnComplete.

Resolves TRI-8391.
This commit is contained in:
Eric Allam
2026-04-18 16:02:04 +01:00
parent fded2cd87e
commit 0a401d8bb7
3 changed files with 91 additions and 5 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"@trigger.dev/sdk": patch
---
Add `chat.endRun()` — exits the run after the current turn completes, without the upgrade-required signal that `chat.requestUpgrade()` sends. Use when an agent finishes its work on its own terms (one-shot responses, goal achieved, budget exhausted) instead of waiting idle for the next user message. Call from `run()`, `chat.defer()`, `onBeforeTurnComplete`, or `onTurnComplete`.
+54 -5
View File
@@ -872,6 +872,13 @@ const chatPrepareMessagesKey =
/** @internal Flag set by `chat.requestUpgrade()` to exit the loop after the current turn. */
const chatUpgradeRequestedKey = locals.create<boolean>("chat.upgradeRequested");
/**
* @internal Flag set by `chat.endRun()` to exit the loop after the current
* turn completes, without any upgrade semantics. Checked at the same
* post-turn / pre-wait sites as `chatUpgradeRequestedKey`.
*/
const chatEndRunRequestedKey = locals.create<boolean>("chat.endRunRequested");
/**
* Event passed to `summarize` callbacks.
*/
@@ -4160,7 +4167,11 @@ function chatAgent<
// chat.requestUpgrade() was called — exit the loop so the
// transport triggers a new run on the latest version.
if (locals.get(chatUpgradeRequestedKey)) {
// chat.endRun() — same exit, no upgrade semantics.
if (
locals.get(chatUpgradeRequestedKey) ||
locals.get(chatEndRunRequestedKey)
) {
return "exit";
}
@@ -4277,8 +4288,11 @@ function chatAgent<
// Best-effort — if stream write fails, let the run continue anyway
}
// chat.requestUpgrade() — exit after error turn too
if (locals.get(chatUpgradeRequestedKey)) {
// chat.requestUpgrade() / chat.endRun() — exit after error turn too
if (
locals.get(chatUpgradeRequestedKey) ||
locals.get(chatEndRunRequestedKey)
) {
return;
}
@@ -4861,6 +4875,36 @@ function requestUpgrade(): void {
locals.set(chatUpgradeRequestedKey, true);
}
/**
* Exit the run after the current turn completes, without waiting for the
* next message. Unlike {@link requestUpgrade}, no upgrade-required signal
* is sent to the client — the turn finishes normally, `onTurnComplete`
* fires, and the loop exits instead of going idle.
*
* Call from `run()`, `chat.defer()`, `onBeforeTurnComplete`, or
* `onTurnComplete` to end the run on your own terms (budget exhausted,
* task complete, goal achieved, etc.).
*
* The next user message on the same `chatId` starts a fresh run via the
* normal continuation mechanism.
*
* @example
* ```ts
* chat.agent({
* id: "one-shot-agent",
* run: async ({ messages, signal }) => {
* const result = streamText({ model: openai("gpt-4o"), messages, abortSignal: signal });
* // Single-response agent — exit after this turn.
* chat.endRun();
* return result;
* },
* });
* ```
*/
function endRun(): void {
locals.set(chatEndRunRequestedKey, true);
}
// ---------------------------------------------------------------------------
// Per-turn deferred work
// ---------------------------------------------------------------------------
@@ -5474,8 +5518,11 @@ function createChatSession(
// Subsequent turns: wait for the next message
if (turn > 0) {
// chat.requestUpgrade() — exit before waiting
if (locals.get(chatUpgradeRequestedKey)) {
// chat.requestUpgrade() / chat.endRun() — exit before waiting
if (
locals.get(chatUpgradeRequestedKey) ||
locals.get(chatEndRunRequestedKey)
) {
stop.cleanup();
return { done: true, value: undefined };
}
@@ -6108,6 +6155,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,
/** 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}. */
cleanupAbortedParts,
/** Register background work that runs in parallel with streaming. See {@link chatDefer}. */
@@ -236,6 +236,38 @@ describe("mockChatAgent", () => {
}
});
it("chat.endRun() exits the loop after the current turn", async () => {
const model = new MockLanguageModelV3({
doStream: async () => ({ stream: textStream("bye") }),
});
let turnCount = 0;
const agent = chat.agent({
id: "mockChatAgent.end-run",
run: async ({ messages, signal }) => {
turnCount++;
chat.endRun();
return streamText({ model, messages, abortSignal: signal });
},
});
const harness = mockChatAgent(agent, { chatId: "test-end-run" });
try {
await harness.sendMessage(userMessage("hello"));
// Give the loop a tick to exit after the turn-complete chunk
await new Promise((r) => setTimeout(r, 50));
expect(turnCount).toBe(1);
// Subsequent sends after endRun should not produce another run — the
// loop has exited. We can't easily assert this via sendMessage (it
// would block waiting for turn-complete), but we can verify the task
// has finished.
} finally {
// close() is a no-op here since the task already exited, but call
// for symmetry with other tests.
await harness.close();
}
});
it("exposes finishReason on the onTurnComplete event", async () => {
const model = new MockLanguageModelV3({
doStream: async () => ({ stream: textStream("hi") }),