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:
@@ -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`.
|
||||
@@ -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") }),
|
||||
|
||||
Reference in New Issue
Block a user