From c129a1a74ba247460c6c061776ceeb995c757ecf Mon Sep 17 00:00:00 2001 From: jbeno Date: Wed, 25 Feb 2026 22:01:13 -0800 Subject: [PATCH] fix(hooks): deduplicate after_tool_call hook in embedded runs --- ...adapter.after-tool-call.fires-once.test.ts | 227 ++++++++++++++++++ ...definition-adapter.after-tool-call.test.ts | 102 +++----- src/agents/pi-tool-definition-adapter.ts | 50 +--- 3 files changed, 264 insertions(+), 115 deletions(-) create mode 100644 src/agents/pi-tool-definition-adapter.after-tool-call.fires-once.test.ts diff --git a/src/agents/pi-tool-definition-adapter.after-tool-call.fires-once.test.ts b/src/agents/pi-tool-definition-adapter.after-tool-call.fires-once.test.ts new file mode 100644 index 00000000000..9f9af15427e --- /dev/null +++ b/src/agents/pi-tool-definition-adapter.after-tool-call.fires-once.test.ts @@ -0,0 +1,227 @@ +/** + * Integration test: after_tool_call fires exactly once when both the adapter + * (toToolDefinitions) and the subscription handler (handleToolExecutionEnd) + * are active — the production scenario for embedded runs. + * + * Regression guard for the double-fire bug fixed by removing the adapter-side + * after_tool_call invocation (see PR #15012 → dedup in this fix). + */ +import type { AgentTool } from "@mariozechner/pi-agent-core"; +import { Type } from "@sinclair/typebox"; +import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; + +const hookMocks = vi.hoisted(() => ({ + runner: { + hasHooks: vi.fn(() => true), + runAfterToolCall: vi.fn(async () => {}), + runBeforeToolCall: vi.fn(async () => {}), + }, +})); + +vi.mock("../plugins/hook-runner-global.js", () => ({ + getGlobalHookRunner: () => hookMocks.runner, +})); + +vi.mock("../infra/agent-events.js", () => ({ + emitAgentEvent: vi.fn(), +})); + +vi.mock("./pi-tools.before-tool-call.js", () => ({ + consumeAdjustedParamsForToolCall: vi.fn((_: string) => undefined), + isToolWrappedWithBeforeToolCallHook: vi.fn(() => false), + runBeforeToolCallHook: vi.fn(async ({ params }: { params: unknown }) => ({ + blocked: false, + params, + })), +})); + +function createTestTool(name: string) { + return { + name, + label: name, + description: `test tool: ${name}`, + parameters: Type.Object({}), + execute: vi.fn(async () => ({ + content: [{ type: "text" as const, text: "ok" }], + details: { ok: true }, + })), + } satisfies AgentTool; +} + +function createFailingTool(name: string) { + return { + name, + label: name, + description: `failing tool: ${name}`, + parameters: Type.Object({}), + execute: vi.fn(async () => { + throw new Error("tool failed"); + }), + } satisfies AgentTool; +} + +function createToolHandlerCtx() { + return { + params: { + runId: "integration-test", + session: { messages: [] }, + }, + hookRunner: hookMocks.runner, + state: { + toolMetaById: new Map(), + toolMetas: [] as Array<{ toolName?: string; meta?: string }>, + toolSummaryById: new Set(), + lastToolError: undefined, + pendingMessagingTexts: new Map(), + pendingMessagingTargets: new Map(), + pendingMessagingMediaUrls: new Map(), + messagingToolSentTexts: [] as string[], + messagingToolSentTextsNormalized: [] as string[], + messagingToolSentMediaUrls: [] as string[], + messagingToolSentTargets: [] as unknown[], + blockBuffer: "", + successfulCronAdds: 0, + }, + log: { debug: vi.fn(), warn: vi.fn() }, + flushBlockReplyBuffer: vi.fn(), + shouldEmitToolResult: () => false, + shouldEmitToolOutput: () => false, + emitToolSummary: vi.fn(), + emitToolOutput: vi.fn(), + trimMessagingToolSent: vi.fn(), + }; +} + +let toToolDefinitions: typeof import("./pi-tool-definition-adapter.js").toToolDefinitions; +let handleToolExecutionStart: typeof import("./pi-embedded-subscribe.handlers.tools.js").handleToolExecutionStart; +let handleToolExecutionEnd: typeof import("./pi-embedded-subscribe.handlers.tools.js").handleToolExecutionEnd; + +describe("after_tool_call fires exactly once in embedded runs", () => { + beforeAll(async () => { + ({ toToolDefinitions } = await import("./pi-tool-definition-adapter.js")); + ({ handleToolExecutionStart, handleToolExecutionEnd } = + await import("./pi-embedded-subscribe.handlers.tools.js")); + }); + + beforeEach(() => { + hookMocks.runner.hasHooks.mockClear(); + hookMocks.runner.hasHooks.mockReturnValue(true); + hookMocks.runner.runAfterToolCall.mockClear(); + hookMocks.runner.runAfterToolCall.mockResolvedValue(undefined); + hookMocks.runner.runBeforeToolCall.mockClear(); + hookMocks.runner.runBeforeToolCall.mockResolvedValue(undefined); + }); + + it("fires after_tool_call exactly once on success when both adapter and handler are active", async () => { + const tool = createTestTool("read"); + const defs = toToolDefinitions([tool]); + const def = defs[0]; + if (!def) { + throw new Error("missing tool definition"); + } + + const toolCallId = "integration-call-1"; + const args = { path: "/tmp/test.txt" }; + const ctx = createToolHandlerCtx(); + + // Step 1: Simulate tool_execution_start event (SDK emits this) + await handleToolExecutionStart( + ctx as never, + { type: "tool_execution_start", toolName: "read", toolCallId, args } as never, + ); + + // Step 2: Execute tool through the adapter wrapper (SDK calls this) + const extensionContext = {} as Parameters[4]; + await def.execute(toolCallId, args, undefined, undefined, extensionContext); + + // Step 3: Simulate tool_execution_end event (SDK emits this after execute returns) + await handleToolExecutionEnd( + ctx as never, + { + type: "tool_execution_end", + toolName: "read", + toolCallId, + isError: false, + result: { content: [{ type: "text", text: "ok" }] }, + } as never, + ); + + // The hook must fire exactly once — not zero, not two. + expect(hookMocks.runner.runAfterToolCall).toHaveBeenCalledTimes(1); + }); + + it("fires after_tool_call exactly once on error when both adapter and handler are active", async () => { + const tool = createFailingTool("exec"); + const defs = toToolDefinitions([tool]); + const def = defs[0]; + if (!def) { + throw new Error("missing tool definition"); + } + + const toolCallId = "integration-call-err"; + const args = { command: "fail" }; + const ctx = createToolHandlerCtx(); + + await handleToolExecutionStart( + ctx as never, + { type: "tool_execution_start", toolName: "exec", toolCallId, args } as never, + ); + + const extensionContext = {} as Parameters[4]; + await def.execute(toolCallId, args, undefined, undefined, extensionContext); + + await handleToolExecutionEnd( + ctx as never, + { + type: "tool_execution_end", + toolName: "exec", + toolCallId, + isError: true, + result: { status: "error", error: "tool failed" }, + } as never, + ); + + expect(hookMocks.runner.runAfterToolCall).toHaveBeenCalledTimes(1); + + const call = (hookMocks.runner.runAfterToolCall as ReturnType).mock.calls[0]; + const event = call?.[0] as { error?: unknown } | undefined; + expect(event?.error).toBeDefined(); + }); + + it("fires after_tool_call exactly once per tool across multiple sequential tool calls", async () => { + const tool = createTestTool("write"); + const defs = toToolDefinitions([tool]); + const def = defs[0]; + if (!def) { + throw new Error("missing tool definition"); + } + + const ctx = createToolHandlerCtx(); + const extensionContext = {} as Parameters[4]; + + for (let i = 0; i < 3; i++) { + const toolCallId = `sequential-call-${i}`; + const args = { path: `/tmp/file-${i}.txt`, content: "data" }; + + await handleToolExecutionStart( + ctx as never, + { type: "tool_execution_start", toolName: "write", toolCallId, args } as never, + ); + + await def.execute(toolCallId, args, undefined, undefined, extensionContext); + + await handleToolExecutionEnd( + ctx as never, + { + type: "tool_execution_end", + toolName: "write", + toolCallId, + isError: false, + result: { content: [{ type: "text", text: "written" }] }, + } as never, + ); + } + + expect(hookMocks.runner.runAfterToolCall).toHaveBeenCalledTimes(3); + }); +}); diff --git a/src/agents/pi-tool-definition-adapter.after-tool-call.test.ts b/src/agents/pi-tool-definition-adapter.after-tool-call.test.ts index 42784f1d726..ef621950adb 100644 --- a/src/agents/pi-tool-definition-adapter.after-tool-call.test.ts +++ b/src/agents/pi-tool-definition-adapter.after-tool-call.test.ts @@ -5,7 +5,7 @@ import { toToolDefinitions } from "./pi-tool-definition-adapter.js"; const hookMocks = vi.hoisted(() => ({ runner: { - hasHooks: vi.fn((_: string) => false), + hasHooks: vi.fn((_: string) => true), runAfterToolCall: vi.fn(async () => {}), }, isToolWrappedWithBeforeToolCallHook: vi.fn(() => false), @@ -39,31 +39,6 @@ function createReadTool() { type ToolExecute = ReturnType[number]["execute"]; const extensionContext = {} as Parameters[4]; -function enableAfterToolCallHook() { - hookMocks.runner.hasHooks.mockImplementation((name: string) => name === "after_tool_call"); -} - -async function executeReadTool(callId: string) { - const defs = toToolDefinitions([createReadTool()]); - const def = defs[0]; - if (!def) { - throw new Error("missing tool definition"); - } - const execute = (...args: Parameters<(typeof defs)[0]["execute"]>) => def.execute(...args); - return await execute(callId, { path: "/tmp/file" }, undefined, undefined, extensionContext); -} - -function expectReadAfterToolCallPayload(result: Awaited>) { - expect(hookMocks.runner.runAfterToolCall).toHaveBeenCalledWith( - { - toolName: "read", - params: { mode: "safe" }, - result, - }, - { toolName: "read" }, - ); -} - describe("pi tool definition adapter after_tool_call", () => { beforeEach(() => { hookMocks.runner.hasHooks.mockClear(); @@ -80,32 +55,21 @@ describe("pi tool definition adapter after_tool_call", () => { })); }); - it("dispatches after_tool_call once on successful adapter execution", async () => { - enableAfterToolCallHook(); - hookMocks.runBeforeToolCallHook.mockResolvedValue({ - blocked: false, - params: { mode: "safe" }, - }); - const result = await executeReadTool("call-ok"); + // Regression guard: after_tool_call is handled exclusively by + // handleToolExecutionEnd in the subscription handler to prevent + // duplicate invocations in embedded runs. + it("does not fire after_tool_call from the adapter (handled by subscription handler)", async () => { + const defs = toToolDefinitions([createReadTool()]); + const def = defs[0]; + if (!def) { + throw new Error("missing tool definition"); + } + await def.execute("call-ok", { path: "/tmp/file" }, undefined, undefined, extensionContext); - expect(result.details).toMatchObject({ ok: true }); - expect(hookMocks.runner.runAfterToolCall).toHaveBeenCalledTimes(1); - expectReadAfterToolCallPayload(result); + expect(hookMocks.runner.runAfterToolCall).not.toHaveBeenCalled(); }); - it("uses wrapped-tool adjusted params for after_tool_call payload", async () => { - enableAfterToolCallHook(); - hookMocks.isToolWrappedWithBeforeToolCallHook.mockReturnValue(true); - hookMocks.consumeAdjustedParamsForToolCall.mockReturnValue({ mode: "safe" } as unknown); - const result = await executeReadTool("call-ok-wrapped"); - - expect(result.details).toMatchObject({ ok: true }); - expect(hookMocks.runBeforeToolCallHook).not.toHaveBeenCalled(); - expectReadAfterToolCallPayload(result); - }); - - it("dispatches after_tool_call once on adapter error with normalized tool name", async () => { - enableAfterToolCallHook(); + it("does not fire after_tool_call from the adapter on error", async () => { const tool = { name: "bash", label: "Bash", @@ -121,31 +85,27 @@ describe("pi tool definition adapter after_tool_call", () => { if (!def) { throw new Error("missing tool definition"); } - const execute = (...args: Parameters<(typeof defs)[0]["execute"]>) => def.execute(...args); - const result = await execute("call-err", { cmd: "ls" }, undefined, undefined, extensionContext); + await def.execute("call-err", { cmd: "ls" }, undefined, undefined, extensionContext); - expect(result.details).toMatchObject({ - status: "error", - tool: "exec", - error: "boom", - }); - expect(hookMocks.runner.runAfterToolCall).toHaveBeenCalledTimes(1); - expect(hookMocks.runner.runAfterToolCall).toHaveBeenCalledWith( - { - toolName: "exec", - params: { cmd: "ls" }, - error: "boom", - }, - { toolName: "exec" }, - ); + expect(hookMocks.runner.runAfterToolCall).not.toHaveBeenCalled(); }); - it("does not break execution when after_tool_call hook throws", async () => { - enableAfterToolCallHook(); - hookMocks.runner.runAfterToolCall.mockRejectedValue(new Error("hook failed")); - const result = await executeReadTool("call-ok2"); + it("consumes adjusted params for wrapped tools to avoid leaks", async () => { + hookMocks.isToolWrappedWithBeforeToolCallHook.mockReturnValue(true); + const defs = toToolDefinitions([createReadTool()]); + const def = defs[0]; + if (!def) { + throw new Error("missing tool definition"); + } + await def.execute( + "call-wrapped", + { path: "/tmp/file" }, + undefined, + undefined, + extensionContext, + ); - expect(result.details).toMatchObject({ ok: true }); - expect(hookMocks.runner.runAfterToolCall).toHaveBeenCalledTimes(1); + expect(hookMocks.runBeforeToolCallHook).not.toHaveBeenCalled(); + expect(hookMocks.consumeAdjustedParamsForToolCall).toHaveBeenCalledWith("call-wrapped"); }); }); diff --git a/src/agents/pi-tool-definition-adapter.ts b/src/agents/pi-tool-definition-adapter.ts index a6221586242..b9e74d7946a 100644 --- a/src/agents/pi-tool-definition-adapter.ts +++ b/src/agents/pi-tool-definition-adapter.ts @@ -5,7 +5,6 @@ import type { } from "@mariozechner/pi-agent-core"; import type { ToolDefinition } from "@mariozechner/pi-coding-agent"; import { logDebug, logError } from "../logger.js"; -import { getGlobalHookRunner } from "../plugins/hook-runner-global.js"; import { isPlainObject } from "../utils.js"; import type { ClientToolDefinition } from "./pi-embedded-runner/run/params.js"; import type { HookContext } from "./pi-tools.before-tool-call.js"; @@ -166,27 +165,11 @@ export function toToolDefinitions(tools: AnyAgentTool[]): ToolDefinition[] { toolName: normalizedName, result: rawResult, }); - const afterParams = beforeHookWrapped - ? (consumeAdjustedParamsForToolCall(toolCallId) ?? executeParams) - : executeParams; - - // Call after_tool_call hook - const hookRunner = getGlobalHookRunner(); - if (hookRunner?.hasHooks("after_tool_call")) { - try { - await hookRunner.runAfterToolCall( - { - toolName: name, - params: isPlainObject(afterParams) ? afterParams : {}, - result, - }, - { toolName: name }, - ); - } catch (hookErr) { - logDebug( - `after_tool_call hook failed: tool=${normalizedName} error=${String(hookErr)}`, - ); - } + // Consume any adjusted params tracked by the before_tool_call hook to avoid leaks. + // after_tool_call is fired by handleToolExecutionEnd in the subscription handler + // to avoid duplicate invocations. + if (beforeHookWrapped) { + consumeAdjustedParamsForToolCall(toolCallId); } return result; @@ -210,32 +193,11 @@ export function toToolDefinitions(tools: AnyAgentTool[]): ToolDefinition[] { } logError(`[tools] ${normalizedName} failed: ${described.message}`); - const errorResult = jsonResult({ + return jsonResult({ status: "error", tool: normalizedName, error: described.message, }); - - // Call after_tool_call hook for errors too - const hookRunner = getGlobalHookRunner(); - if (hookRunner?.hasHooks("after_tool_call")) { - try { - await hookRunner.runAfterToolCall( - { - toolName: normalizedName, - params: isPlainObject(params) ? params : {}, - error: described.message, - }, - { toolName: normalizedName }, - ); - } catch (hookErr) { - logDebug( - `after_tool_call hook failed: tool=${normalizedName} error=${String(hookErr)}`, - ); - } - } - - return errorResult; } }, } satisfies ToolDefinition;