From 049e32057033d2fcc4e08a434cbcbaa75da580dc Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 22 Apr 2026 13:19:53 +0200 Subject: [PATCH] feat(agent): add terminating tool result hints closes #3525 --- packages/agent/CHANGELOG.md | 4 + packages/agent/README.md | 13 +- packages/agent/src/agent-loop.ts | 44 +++++-- packages/agent/src/types.ts | 12 ++ packages/agent/test/agent-loop.test.ts | 165 +++++++++++++++++++++++++ 5 files changed, 224 insertions(+), 14 deletions(-) diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 2de8bef30..24828d65b 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added `terminate: true` tool-result hints to skip the automatic follow-up LLM call when every finalized tool result in the current batch opts into early termination ([#3525](https://github.com/badlogic/pi-mono/issues/3525)) + ## [0.68.1] - 2026-04-22 ### Fixed diff --git a/packages/agent/README.md b/packages/agent/README.md index daf122d66..f3a5b9c11 100644 --- a/packages/agent/README.md +++ b/packages/agent/README.md @@ -110,6 +110,8 @@ The mode can be set globally via `toolExecution` in the agent config, or per-too The `beforeToolCall` hook runs after `tool_execution_start` and validated argument parsing. It can block execution. The `afterToolCall` hook runs after tool execution finishes and before `tool_execution_end` and final tool result message events are emitted. +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. + 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 @@ -137,8 +139,8 @@ The last message in context must be `user` or `toolResult` (not `assistant`). | `tool_execution_start` | Tool begins | | `tool_execution_update` | Tool streams progress | | `tool_execution_end` | Tool completes | -+ -+`Agent.subscribe()` listeners are awaited in registration order. `agent_end` means no more loop events will be emitted, but `await agent.waitForIdle()` and `await agent.prompt(...)` only settle after awaited `agent_end` listeners finish. + +`Agent.subscribe()` listeners are awaited in registration order. `agent_end` means no more loop events will be emitted, but `await agent.waitForIdle()` and `await agent.prompt(...)` only settle after awaited `agent_end` listeners finish. ## Agent Options @@ -186,6 +188,9 @@ const agent = new Agent({ // Postprocess each tool result before final tool events are emitted. afterToolCall: async ({ toolCall, result, isError, context }) => { + if (toolCall.name === "notify_done" && !isError) { + return { terminate: true }; + } if (!isError) { return { details: { ...result.details, audited: true } }; } @@ -382,6 +387,8 @@ const readFileTool: AgentTool = { // Optional: stream progress onUpdate?.({ content: [{ type: "text", text: "Reading..." }], details: {} }); + // Optional: add `terminate: true` here to skip the automatic follow-up LLM call + // when every finalized tool result in the batch does the same. return { content: [{ type: "text", text: content }], details: { path: params.path, size: content.length }, @@ -408,6 +415,8 @@ execute: async (toolCallId, params, signal, onUpdate) => { Thrown errors are caught by the agent and reported to the LLM as tool errors with `isError: true`. +Return `terminate: true` from `execute()` or `afterToolCall` to hint that the agent should stop after the current tool batch. This only takes effect when every finalized tool result in the batch is terminating. The hint is runtime-only; emitted `toolResult` transcript messages remain standard LLM tool results. + ## Proxy Usage For browser apps that proxy through a backend: diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 0c940a1ca..2248747a1 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -199,11 +199,13 @@ async function runLoop( // Check for tool calls const toolCalls = message.content.filter((c) => c.type === "toolCall"); - hasMoreToolCalls = toolCalls.length > 0; const toolResults: ToolResultMessage[] = []; - if (hasMoreToolCalls) { - toolResults.push(...(await executeToolCalls(currentContext, message, config, signal, emit))); + hasMoreToolCalls = false; + if (toolCalls.length > 0) { + const executedToolBatch = await executeToolCalls(currentContext, message, config, signal, emit); + toolResults.push(...executedToolBatch.messages); + hasMoreToolCalls = !executedToolBatch.terminate; for (const result of toolResults) { currentContext.messages.push(result); @@ -339,7 +341,7 @@ async function executeToolCalls( config: AgentLoopConfig, signal: AbortSignal | undefined, emit: AgentEventSink, -): Promise { +): Promise { const toolCalls = assistantMessage.content.filter((c) => c.type === "toolCall"); const hasSequentialToolCall = toolCalls.some( (tc) => currentContext.tools?.find((t) => t.name === tc.name)?.executionMode === "sequential", @@ -350,6 +352,11 @@ async function executeToolCalls( return executeToolCallsParallel(currentContext, assistantMessage, toolCalls, config, signal, emit); } +type ExecutedToolCallBatch = { + messages: ToolResultMessage[]; + terminate: boolean; +}; + async function executeToolCallsSequential( currentContext: AgentContext, assistantMessage: AssistantMessage, @@ -357,8 +364,9 @@ async function executeToolCallsSequential( config: AgentLoopConfig, signal: AbortSignal | undefined, emit: AgentEventSink, -): Promise { - const results: ToolResultMessage[] = []; +): Promise { + const finalizedCalls: FinalizedToolCallOutcome[] = []; + const messages: ToolResultMessage[] = []; for (const toolCall of toolCalls) { await emit({ @@ -391,10 +399,14 @@ async function executeToolCallsSequential( await emitToolExecutionEnd(finalized, emit); const toolResultMessage = createToolResultMessage(finalized); await emitToolResultMessage(toolResultMessage, emit); - results.push(toolResultMessage); + finalizedCalls.push(finalized); + messages.push(toolResultMessage); } - return results; + return { + messages, + terminate: shouldTerminateToolBatch(finalizedCalls), + }; } async function executeToolCallsParallel( @@ -404,7 +416,7 @@ async function executeToolCallsParallel( config: AgentLoopConfig, signal: AbortSignal | undefined, emit: AgentEventSink, -): Promise { +): Promise { const finalizedCalls: FinalizedToolCallEntry[] = []; for (const toolCall of toolCalls) { @@ -445,14 +457,17 @@ async function executeToolCallsParallel( const orderedFinalizedCalls = await Promise.all( finalizedCalls.map((entry) => (typeof entry === "function" ? entry() : Promise.resolve(entry))), ); - const results: ToolResultMessage[] = []; + const messages: ToolResultMessage[] = []; for (const finalized of orderedFinalizedCalls) { const toolResultMessage = createToolResultMessage(finalized); await emitToolResultMessage(toolResultMessage, emit); - results.push(toolResultMessage); + messages.push(toolResultMessage); } - return results; + return { + messages, + terminate: shouldTerminateToolBatch(orderedFinalizedCalls), + }; } type PreparedToolCall = { @@ -481,6 +496,10 @@ type FinalizedToolCallOutcome = { type FinalizedToolCallEntry = FinalizedToolCallOutcome | (() => Promise); +function shouldTerminateToolBatch(finalizedCalls: FinalizedToolCallOutcome[]): boolean { + return finalizedCalls.length > 0 && finalizedCalls.every((finalized) => finalized.result.terminate === true); +} + function prepareToolCallArguments(tool: AgentTool, toolCall: AgentToolCall): AgentToolCall { if (!tool.prepareArguments) { return toolCall; @@ -612,6 +631,7 @@ async function finalizeExecutedToolCall( result = { content: afterResult.content ?? result.content, details: afterResult.details ?? result.details, + terminate: afterResult.terminate ?? result.terminate, }; isError = afterResult.isError ?? isError; } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index c93134b33..193fc1f8e 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -56,6 +56,7 @@ export interface BeforeToolCallResult { * - `content`: if provided, replaces the tool result content array in full * - `details`: if provided, replaces the tool result details value in full * - `isError`: if provided, replaces the tool result error flag + * - `terminate`: if provided, replaces the early-termination hint * * Omitted fields keep the original executed tool result values. * There is no deep merge for `content` or `details`. @@ -64,6 +65,11 @@ export interface AfterToolCallResult { content?: (TextContent | ImageContent)[]; details?: unknown; isError?: boolean; + /** + * Hint that the agent should stop after the current tool batch. + * Early termination only happens when every finalized tool result in the batch sets this to true. + */ + terminate?: boolean; } /** Context passed to `beforeToolCall`. */ @@ -209,6 +215,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions { * - `content` replaces the full content array * - `details` replaces the full details payload * - `isError` replaces the error flag + * - `terminate` replaces the early-termination hint * * Any omitted fields keep their original values. No deep merge is performed. * The hook receives the agent abort signal and is responsible for honoring it. @@ -286,6 +293,11 @@ export interface AgentToolResult { content: (TextContent | ImageContent)[]; /** Arbitrary structured details for logs or UI rendering. */ details: T; + /** + * Hint that the agent should stop after the current tool batch. + * Early termination only happens when every finalized tool result in the batch sets this to true. + */ + terminate?: boolean; } /** Callback used by tools to stream partial execution updates. */ diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 8f5773dc5..33c1fb286 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -893,6 +893,171 @@ describe("agentLoop with AgentMessage", () => { // With executionMode=parallel, second tool should start before first finishes expect(parallelObserved).toBe(true); }); + + 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 = { + name: "echo", + label: "Echo", + description: "Echo tool", + parameters: toolSchema, + async execute(_toolCallId, params) { + return { + content: [{ type: "text", text: `echoed: ${params.value}` }], + details: { value: params.value }, + terminate: true, + }; + }, + }; + + const context: AgentContext = { + systemPrompt: "", + messages: [], + tools: [tool], + }; + + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + }; + + let llmCalls = 0; + const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { + llmCalls++; + const mockStream = new MockAssistantStream(); + queueMicrotask(() => { + const message = createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], + "toolUse", + ); + mockStream.push({ type: "done", reason: "toolUse", message }); + }); + return mockStream; + }); + + const events: AgentEvent[] = []; + for await (const event of stream) { + events.push(event); + } + + const messages = await stream.result(); + expect(llmCalls).toBe(1); + expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "toolResult"]); + expect(events.filter((event) => event.type === "turn_end")).toHaveLength(1); + }); + + it("should continue after parallel tool calls when not all tool results terminate", async () => { + const toolSchema = Type.Object({ value: Type.String() }); + const tool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo tool", + parameters: toolSchema, + async execute(_toolCallId, params) { + return { + content: [{ type: "text", text: `echoed: ${params.value}` }], + details: { value: params.value }, + terminate: params.value === "first", + }; + }, + }; + + const context: AgentContext = { + systemPrompt: "", + messages: [], + tools: [tool], + }; + + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + toolExecution: "parallel", + }; + + let callIndex = 0; + const stream = agentLoop([createUserMessage("echo both")], context, config, undefined, () => { + const mockStream = new MockAssistantStream(); + queueMicrotask(() => { + if (callIndex === 0) { + const message = createAssistantMessage( + [ + { type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "first" } }, + { type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "second" } }, + ], + "toolUse", + ); + mockStream.push({ type: "done", reason: "toolUse", message }); + } else { + const message = createAssistantMessage([{ type: "text", text: "done" }]); + mockStream.push({ type: "done", reason: "stop", message }); + } + callIndex++; + }); + return mockStream; + }); + + for await (const _event of stream) { + // consume + } + + const messages = await stream.result(); + expect(callIndex).toBe(2); + expect(messages.map((message) => message.role)).toEqual([ + "user", + "assistant", + "toolResult", + "toolResult", + "assistant", + ]); + }); + + it("should allow afterToolCall to mark a tool batch as terminating", async () => { + const toolSchema = Type.Object({ value: Type.String() }); + const tool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo tool", + parameters: toolSchema, + async execute(_toolCallId, params) { + return { + content: [{ type: "text", text: `echoed: ${params.value}` }], + details: { value: params.value }, + }; + }, + }; + + const context: AgentContext = { + systemPrompt: "", + messages: [], + tools: [tool], + }; + + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + afterToolCall: async () => ({ terminate: true }), + }; + + let llmCalls = 0; + const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { + llmCalls++; + const mockStream = new MockAssistantStream(); + queueMicrotask(() => { + const message = createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], + "toolUse", + ); + mockStream.push({ type: "done", reason: "toolUse", message }); + }); + return mockStream; + }); + + for await (const _event of stream) { + // consume + } + + expect(llmCalls).toBe(1); + }); }); describe("agentLoopContinue with AgentMessage", () => {