From 73b7b2c56456441f356e1fba016d14fc28e3c0d2 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sat, 2 May 2026 00:49:22 +0200 Subject: [PATCH] feat(agent): add post-turn stop callback --- packages/agent/CHANGELOG.md | 4 ++ packages/agent/README.md | 14 ++++ packages/agent/src/agent-loop.ts | 12 ++++ packages/agent/src/types.ts | 26 ++++++- packages/agent/test/agent-loop.test.ts | 97 ++++++++++++++++++++++++++ 5 files changed, 152 insertions(+), 1 deletion(-) diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 80436907..91a9c433 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added `shouldStopAfterTurn` to the low-level agent loop config for gracefully exiting after a completed turn before polling queued messages or starting another LLM call. + ## [0.71.1] - 2026-05-01 ## [0.71.0] - 2026-04-30 diff --git a/packages/agent/README.md b/packages/agent/README.md index 56ccc69f..65434768 100644 --- a/packages/agent/README.md +++ b/packages/agent/README.md @@ -112,6 +112,20 @@ The `beforeToolCall` hook runs after `tool_execution_start` and validated argume Tools can also return `terminate: true` to hint that the automatic follow-up LLM call should be skipped. The loop only stops early when every finalized tool result in that batch sets `terminate: true`. Mixed batches continue normally. +Low-level loop callers can set `shouldStopAfterTurn` to stop gracefully after the current turn completes: + +```typescript +const stream = agentLoop(prompts, context, { + model, + convertToLlm, + shouldStopAfterTurn: async ({ message, toolResults, context, newMessages }) => { + return shouldCompactBeforeNextTurn(context.messages); + }, +}); +``` + +`shouldStopAfterTurn` runs after `turn_end` is emitted and after the assistant response and any tool executions have completed normally. If it returns `true`, the loop emits `agent_end` and exits before polling steering or follow-up queues, and before starting another LLM call. It does not abort the provider stream, does not cancel running tools, and does not alter the assistant message stop reason. + When you use the `Agent` class, assistant `message_end` processing is treated as a barrier before tool preflight begins. That means `beforeToolCall` sees agent state that already includes the assistant message that requested the tool call. ### continue() Event Sequence diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 2248747a..23af04df 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -215,6 +215,18 @@ async function runLoop( await emit({ type: "turn_end", message, toolResults }); + if ( + await config.shouldStopAfterTurn?.({ + message, + toolResults, + context: currentContext, + newMessages, + }) + ) { + await emit({ type: "agent_end", messages: newMessages }); + return; + } + pendingMessages = (await config.getSteeringMessages?.()) || []; } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index ea18c068..803b9566 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -100,6 +100,18 @@ export interface AfterToolCallContext { context: AgentContext; } +/** Context passed to `shouldStopAfterTurn`. */ +export interface ShouldStopAfterTurnContext { + /** The assistant message that completed the turn. */ + message: AssistantMessage; + /** Tool result messages passed to the preceding `turn_end` event. */ + toolResults: ToolResultMessage[]; + /** Current agent context after the turn's assistant message and tool results have been appended. */ + context: AgentContext; + /** Messages that this loop invocation will return if it exits at this point. Prompt runs include the initial prompt messages; continuation runs do not include pre-existing context messages. */ + newMessages: AgentMessage[]; +} + export interface AgentLoopConfig extends SimpleStreamOptions { model: Model; @@ -163,10 +175,22 @@ export interface AgentLoopConfig extends SimpleStreamOptions { */ getApiKey?: (provider: string) => Promise | string | undefined; + /** + * Called after each turn fully completes and `turn_end` has been emitted. + * + * If it returns true, the loop emits `agent_end` and exits before polling steering or follow-up queues, + * without starting another LLM call. The current assistant response and any tool executions finish normally. + * + * Use this to request a graceful stop after the current turn, e.g. before context gets too full. + * + * Contract: must not throw or reject. Throwing interrupts the low-level agent loop without producing a normal event sequence. + */ + shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise; + /** * Returns steering messages to inject into the conversation mid-run. * - * Called after the current assistant turn finishes executing its tool calls. + * Called after the current assistant turn finishes executing its tool calls, unless `shouldStopAfterTurn` exits first. * If messages are returned, they are added to the context before the next LLM call. * Tool calls from the current assistant message are not skipped. * diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 56291787..fd5ea73f 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -894,6 +894,103 @@ describe("agentLoop with AgentMessage", () => { expect(parallelObserved).toBe(true); }); + it("should stop after the current turn when shouldStopAfterTurn returns true", async () => { + const toolSchema = Type.Object({ value: Type.String() }); + const executed: string[] = []; + const tool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo tool", + parameters: toolSchema, + async execute(_toolCallId, params) { + executed.push(params.value); + return { + content: [{ type: "text", text: `echoed: ${params.value}` }], + details: { value: params.value }, + }; + }, + }; + + const context: AgentContext = { + systemPrompt: "", + messages: [], + tools: [tool], + }; + + let steeringPolls = 0; + let followUpPolls = 0; + let callbackToolResultIds: string[] = []; + let callbackContextRoles: string[] = []; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + getSteeringMessages: async () => { + steeringPolls++; + return []; + }, + getFollowUpMessages: async () => { + followUpPolls++; + return [createUserMessage("follow up should stay queued")]; + }, + shouldStopAfterTurn: async ({ message, toolResults, context }) => { + expect(message.role).toBe("assistant"); + callbackToolResultIds = toolResults.map((toolResult) => toolResult.toolCallId); + callbackContextRoles = context.messages.map((contextMessage) => contextMessage.role); + return true; + }, + }; + + let llmCalls = 0; + const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { + llmCalls++; + const mockStream = new MockAssistantStream(); + queueMicrotask(() => { + if (llmCalls === 1) { + const message = createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], + "toolUse", + ); + mockStream.push({ type: "done", reason: "toolUse", message }); + } else { + mockStream.push({ + type: "done", + reason: "stop", + message: createAssistantMessage([{ type: "text", text: "should not run" }]), + }); + } + }); + return mockStream; + }); + + const events: AgentEvent[] = []; + for await (const event of stream) { + events.push(event); + } + + const messages = await stream.result(); + expect(llmCalls).toBe(1); + expect(executed).toEqual(["hello"]); + expect(steeringPolls).toBe(1); + expect(followUpPolls).toBe(0); + expect(callbackToolResultIds).toEqual(["tool-1"]); + expect(callbackContextRoles).toEqual(["user", "assistant", "toolResult"]); + expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "toolResult"]); + expect(events.map((event) => event.type)).toEqual([ + "agent_start", + "turn_start", + "message_start", + "message_end", + "message_start", + "message_end", + "tool_execution_start", + "tool_execution_end", + "message_start", + "message_end", + "turn_end", + "agent_end", + ]); + }); + it("should stop after a tool batch when every tool result sets terminate=true", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool = {