From 1eb988cfe88fb0ff740ff62583d2f16359f7b6b0 Mon Sep 17 00:00:00 2001 From: muyiyr <940955635@qq.com> Date: Thu, 6 Aug 2026 23:14:48 +0800 Subject: [PATCH] feat(agent): allow blocked tool calls to terminate (#7715) Closes #5998 --- packages/agent/README.md | 8 +- packages/agent/src/agent-loop.ts | 6 +- packages/agent/src/types.ts | 6 + packages/agent/test/agent-loop.test.ts | 118 ++++++++++++++++++ packages/coding-agent/docs/extensions.md | 5 +- .../coding-agent/src/core/extensions/types.ts | 5 + .../5998-blocked-tool-terminate.test.ts | 53 ++++++++ 7 files changed, 194 insertions(+), 7 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/5998-blocked-tool-terminate.test.ts diff --git a/packages/agent/README.md b/packages/agent/README.md index d26229efa..879a47429 100644 --- a/packages/agent/README.md +++ b/packages/agent/README.md @@ -119,9 +119,9 @@ In parallel mode, tool completion events follow tool completion order, but persi The mode can be set globally via `toolExecution` in the agent config, or per-tool via `executionMode` on `AgentTool`. If any tool call in a batch targets a tool with `executionMode: "sequential"`, the entire batch executes sequentially regardless of the global setting. -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. +The `beforeToolCall` hook runs after `tool_execution_start` and validated argument parsing. It can block execution and attach `terminate: true` to the blocked result. 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. +Tools, blocked `beforeToolCall` results, and `afterToolCall` overrides can 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. The `Agent` class accepts `shouldStopAfterTurn` in `AgentOptions`. Low-level loop callers can set the same hook in `AgentLoopConfig`: @@ -213,7 +213,7 @@ const agent = new Agent({ // Preflight each tool call after args are validated. Can block execution. beforeToolCall: async ({ toolCall, args, context }) => { if (toolCall.name === "bash") { - return { block: true, reason: "bash is disabled" }; + return { block: true, reason: "bash is disabled", terminate: true }; } }, @@ -453,7 +453,7 @@ 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. +Return `terminate: true` from `execute()`, a blocked `beforeToolCall`, 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 diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 3a12506cb..a251fede0 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -634,9 +634,13 @@ async function prepareToolCall( }; } if (beforeResult?.block) { + const result = createErrorToolResult(beforeResult.reason || "Tool execution was blocked"); + if (beforeResult.terminate === true) { + result.terminate = true; + } return { kind: "immediate", - result: createErrorToolResult(beforeResult.reason || "Tool execution was blocked"), + result, isError: true, }; } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 52f34891a..5b20b21f4 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -61,6 +61,11 @@ export type AgentToolCall = Extract Promise; diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 8a2a7388d..e5b2f1db8 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -1250,6 +1250,124 @@ describe("agentLoop with AgentMessage", () => { expect(events.filter((event) => event.type === "turn_end")).toHaveLength(1); }); + it("should stop after a blocked tool call when beforeToolCall sets terminate=true", async () => { + const toolSchema = Type.Object({ value: Type.String() }); + let executed = false; + const tool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo tool", + parameters: toolSchema, + async execute() { + executed = true; + return { + content: [{ type: "text", text: "should not execute" }], + details: { value: "unexpected" }, + }; + }, + }; + const context: AgentContext = { + systemPrompt: "", + messages: [], + tools: [tool], + }; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + beforeToolCall: async () => ({ block: true, reason: "Blocked by policy", terminate: true }), + }; + + let llmCalls = 0; + const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { + llmCalls++; + const mockStream = new MockAssistantStream(); + queueMicrotask(() => { + const message = + llmCalls === 1 + ? createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], + "toolUse", + ) + : createAssistantMessage([{ type: "text", text: "should not run" }]); + mockStream.push({ type: "done", reason: llmCalls === 1 ? "toolUse" : "stop", message }); + }); + return mockStream; + }); + + for await (const _event of stream) { + // consume + } + + const messages = await stream.result(); + const toolResult = messages.find((message) => message.role === "toolResult"); + expect(executed).toBe(false); + expect(llmCalls).toBe(1); + expect(toolResult?.role === "toolResult" ? toolResult.isError : false).toBe(true); + expect(toolResult?.role === "toolResult" ? toolResult.content : []).toContainEqual({ + type: "text", + text: "Blocked by policy", + }); + }); + + it("should continue after a mixed batch with one terminating blocked call", 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], + }; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + toolExecution: "parallel", + beforeToolCall: async ({ args }) => { + const { value } = args as { value: string }; + return value === "first" ? { block: true, reason: "Blocked first", terminate: true } : undefined; + }, + }; + + let llmCalls = 0; + const stream = agentLoop([createUserMessage("echo both")], context, config, undefined, () => { + llmCalls++; + const mockStream = new MockAssistantStream(); + queueMicrotask(() => { + const message = + llmCalls === 1 + ? createAssistantMessage( + [ + { type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "first" } }, + { type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "second" } }, + ], + "toolUse", + ) + : createAssistantMessage([{ type: "text", text: "done" }]); + mockStream.push({ type: "done", reason: llmCalls === 1 ? "toolUse" : "stop", message }); + }); + return mockStream; + }); + + for await (const _event of stream) { + // consume + } + + expect(executed).toEqual(["second"]); + expect(llmCalls).toBe(2); + }); + it("should continue after parallel tool calls when not all tool results terminate", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool = { diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 48adbba9c..00e92a1d6 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -762,7 +762,8 @@ Behavior guarantees: - Mutations to `event.input` affect the actual tool execution - Later `tool_call` handlers see mutations made by earlier handlers - No re-validation is performed after your mutation -- Return values from `tool_call` only control blocking via `{ block: true, reason?: string }` +- Return values from `tool_call` control blocking via `{ block: true, reason?: string, terminate?: boolean }` +- `terminate` only applies to a blocked call; the agent stops early only when every finalized result in the batch is terminating ```typescript import { isToolCallEventType } from "@earendil-works/pi-coding-agent"; @@ -778,7 +779,7 @@ pi.on("tool_call", async (event, ctx) => { event.input.command = `source ~/.profile\n${event.input.command}`; if (event.input.command.includes("rm -rf")) { - return { block: true, reason: "Dangerous command" }; + return { block: true, reason: "Dangerous command", terminate: true }; } } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index e62dfa10f..a87322ce9 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -1072,6 +1072,11 @@ export interface ToolCallEventResult { /** Block tool execution. To modify arguments, mutate `event.input` in place instead. */ block?: boolean; reason?: string; + /** + * Hint that the agent should stop after the current tool batch when this call is blocked. + * Early termination only happens when every finalized tool result in the batch sets this to true. + */ + terminate?: boolean; } /** Result from user_bash event handler */ diff --git a/packages/coding-agent/test/suite/regressions/5998-blocked-tool-terminate.test.ts b/packages/coding-agent/test/suite/regressions/5998-blocked-tool-terminate.test.ts new file mode 100644 index 000000000..93413be83 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/5998-blocked-tool-terminate.test.ts @@ -0,0 +1,53 @@ +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; +import { Type } from "typebox"; +import { afterEach, describe, expect, it } from "vitest"; +import { createHarness, getAssistantTexts, type Harness } from "../harness.ts"; + +describe("#5998 blocked tool termination", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("lets a tool_call handler terminate the run after blocking execution", async () => { + const echoTool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo text back", + parameters: Type.Object({ text: Type.String() }), + execute: async () => { + throw new Error("tool should have been blocked"); + }, + }; + const harness = await createHarness({ + tools: [echoTool], + extensionFactories: [ + (pi) => { + pi.on("tool_call", async () => ({ + block: true, + reason: "Blocked by terminating policy", + terminate: true, + })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }), + fauxAssistantMessage("should not run"), + ]); + + await harness.session.prompt("hi"); + + expect(harness.getPendingResponseCount()).toBe(1); + expect(getAssistantTexts(harness)).not.toContain("should not run"); + expect(harness.eventsOfType("tool_execution_end")[0]?.result).toHaveProperty("terminate", true); + expect( + harness.session.messages.find((message) => message.role === "toolResult" && message.isError), + ).toBeDefined(); + }); +});