feat: wire before_tool_call / after_tool_call plugin hooks into tool pipeline

The typed plugin hooks `before_tool_call` and `after_tool_call` were fully
implemented in the hook runner but never invoked during agent tool execution.

This adds a tool wrapper (`pi-tools.hooks.ts`) that intercepts every tool's
`execute()` to fire these hooks, following the same pattern as the existing
`wrapToolWithAbortSignal` wrapper:

- `before_tool_call` runs sequentially before execution and can modify params
  or block the call entirely.
- `after_tool_call` runs fire-and-forget after execution with result, error,
  and duration.

Includes 12 unit tests covering the full lifecycle: pass-through, blocking,
param modification, error propagation, and hook isolation.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
dp-web4 2026-01-27 19:42:51 -08:00
parent 3fe4b2595a
commit dbbca8e040
3 changed files with 330 additions and 1 deletions

View File

@ -0,0 +1,236 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { wrapToolWithHooks, type ToolHookContext } from "./pi-tools.hooks.js";
import type { AnyAgentTool } from "./pi-tools.types.js";
// Mock the global hook runner
const mockRunBeforeToolCall = vi.fn();
const mockRunAfterToolCall = vi.fn();
const mockHasHooks = vi.fn();
vi.mock("../plugins/hook-runner-global.js", () => ({
getGlobalHookRunner: () => ({
hasHooks: mockHasHooks,
runBeforeToolCall: mockRunBeforeToolCall,
runAfterToolCall: mockRunAfterToolCall,
}),
}));
function createMockTool(name: string, executeFn?: AnyAgentTool["execute"]): AnyAgentTool {
return {
name,
description: `Mock ${name} tool`,
schema: {} as AnyAgentTool["schema"],
execute:
executeFn ??
(async (_id: string, params: unknown) => {
return `result from ${name}`;
}),
};
}
const ctx: ToolHookContext = { agentId: "test-agent", sessionKey: "test-session" };
describe("wrapToolWithHooks", () => {
beforeEach(() => {
vi.clearAllMocks();
mockHasHooks.mockReturnValue(false);
mockRunBeforeToolCall.mockResolvedValue(undefined);
mockRunAfterToolCall.mockResolvedValue(undefined);
});
it("should preserve tool name, description, and schema", () => {
const tool = createMockTool("read");
const wrapped = wrapToolWithHooks(tool, ctx);
expect(wrapped.name).toBe("read");
expect(wrapped.description).toBe("Mock read tool");
expect(wrapped.schema).toBe(tool.schema);
});
it("should return tool unchanged if it has no execute", () => {
const tool = { name: "noop", description: "no-op", schema: {} } as unknown as AnyAgentTool;
const wrapped = wrapToolWithHooks(tool, ctx);
expect(wrapped).toBe(tool);
});
it("should call original execute when no hooks are registered", async () => {
const executeFn = vi.fn().mockResolvedValue("ok");
const tool = createMockTool("read", executeFn);
const wrapped = wrapToolWithHooks(tool, ctx);
const result = await wrapped.execute!(
"call-1",
{ path: "/foo" },
undefined as any,
undefined as any,
);
expect(result).toBe("ok");
expect(executeFn).toHaveBeenCalledWith("call-1", { path: "/foo" }, undefined, undefined);
});
it("should call before_tool_call hook before execution", async () => {
mockHasHooks.mockImplementation((name: string) => name === "before_tool_call");
const callOrder: string[] = [];
mockRunBeforeToolCall.mockImplementation(async () => {
callOrder.push("before");
return undefined;
});
const executeFn = vi.fn().mockImplementation(async () => {
callOrder.push("execute");
return "ok";
});
const tool = createMockTool("exec", executeFn);
const wrapped = wrapToolWithHooks(tool, ctx);
await wrapped.execute!("call-1", {}, undefined as any, undefined as any);
expect(callOrder).toEqual(["before", "execute"]);
expect(mockRunBeforeToolCall).toHaveBeenCalledWith(
{ toolName: "exec", params: {} },
{ agentId: "test-agent", sessionKey: "test-session", toolName: "exec" },
);
});
it("should block execution when before_tool_call returns block: true", async () => {
mockHasHooks.mockImplementation((name: string) => name === "before_tool_call");
mockRunBeforeToolCall.mockResolvedValue({ block: true, blockReason: "Not allowed" });
const executeFn = vi.fn();
const tool = createMockTool("exec", executeFn);
const wrapped = wrapToolWithHooks(tool, ctx);
const result = await wrapped.execute!("call-1", {}, undefined as any, undefined as any);
expect(result).toEqual({
content: [{ type: "text", text: "[blocked] Not allowed" }],
details: undefined,
});
expect(executeFn).not.toHaveBeenCalled();
});
it("should use default block reason when none provided", async () => {
mockHasHooks.mockImplementation((name: string) => name === "before_tool_call");
mockRunBeforeToolCall.mockResolvedValue({ block: true });
const tool = createMockTool("exec", vi.fn());
const wrapped = wrapToolWithHooks(tool, ctx);
const result = await wrapped.execute!("call-1", {}, undefined as any, undefined as any);
expect(result).toEqual({
content: [{ type: "text", text: "[blocked] Blocked by plugin hook" }],
details: undefined,
});
});
it("should pass modified params from before_tool_call to execute", async () => {
mockHasHooks.mockImplementation((name: string) => name === "before_tool_call");
mockRunBeforeToolCall.mockResolvedValue({ params: { path: "/modified" } });
const executeFn = vi.fn().mockResolvedValue("ok");
const tool = createMockTool("read", executeFn);
const wrapped = wrapToolWithHooks(tool, ctx);
await wrapped.execute!("call-1", { path: "/original" }, undefined as any, undefined as any);
expect(executeFn).toHaveBeenCalledWith("call-1", { path: "/modified" }, undefined, undefined);
});
it("should call after_tool_call hook after successful execution", async () => {
mockHasHooks.mockImplementation((name: string) => name === "after_tool_call");
const executeFn = vi.fn().mockResolvedValue("tool output");
const tool = createMockTool("read", executeFn);
const wrapped = wrapToolWithHooks(tool, ctx);
await wrapped.execute!("call-1", { path: "/foo" }, undefined as any, undefined as any);
// Allow microtask for fire-and-forget
await new Promise((r) => setTimeout(r, 10));
expect(mockRunAfterToolCall).toHaveBeenCalledWith(
expect.objectContaining({
toolName: "read",
params: { path: "/foo" },
result: "tool output",
error: undefined,
durationMs: expect.any(Number),
}),
{ agentId: "test-agent", sessionKey: "test-session", toolName: "read" },
);
});
it("should call after_tool_call with error when execution fails", async () => {
mockHasHooks.mockImplementation((name: string) => name === "after_tool_call");
const executeFn = vi.fn().mockRejectedValue(new Error("boom"));
const tool = createMockTool("exec", executeFn);
const wrapped = wrapToolWithHooks(tool, ctx);
await expect(
wrapped.execute!("call-1", {}, undefined as any, undefined as any),
).rejects.toThrow("boom");
// Allow microtask for fire-and-forget
await new Promise((r) => setTimeout(r, 10));
expect(mockRunAfterToolCall).toHaveBeenCalledWith(
expect.objectContaining({
toolName: "exec",
params: {},
result: undefined,
error: "boom",
durationMs: expect.any(Number),
}),
{ agentId: "test-agent", sessionKey: "test-session", toolName: "exec" },
);
});
it("should call both before and after hooks in correct order", async () => {
mockHasHooks.mockReturnValue(true);
const callOrder: string[] = [];
mockRunBeforeToolCall.mockImplementation(async () => {
callOrder.push("before");
return undefined;
});
mockRunAfterToolCall.mockImplementation(async () => {
callOrder.push("after");
});
const executeFn = vi.fn().mockImplementation(async () => {
callOrder.push("execute");
return "ok";
});
const tool = createMockTool("write", executeFn);
const wrapped = wrapToolWithHooks(tool, ctx);
await wrapped.execute!("call-1", {}, undefined as any, undefined as any);
// Allow microtask for fire-and-forget
await new Promise((r) => setTimeout(r, 10));
expect(callOrder).toEqual(["before", "execute", "after"]);
});
it("should not break when after_tool_call hook throws", async () => {
mockHasHooks.mockImplementation((name: string) => name === "after_tool_call");
mockRunAfterToolCall.mockRejectedValue(new Error("hook error"));
const executeFn = vi.fn().mockResolvedValue("ok");
const tool = createMockTool("read", executeFn);
const wrapped = wrapToolWithHooks(tool, ctx);
// Should not throw despite hook error
const result = await wrapped.execute!("call-1", {}, undefined as any, undefined as any);
expect(result).toBe("ok");
});
it("should work with empty context", async () => {
mockHasHooks.mockReturnValue(true);
mockRunBeforeToolCall.mockResolvedValue(undefined);
const executeFn = vi.fn().mockResolvedValue("ok");
const tool = createMockTool("read", executeFn);
const wrapped = wrapToolWithHooks(tool, {});
await wrapped.execute!("call-1", {}, undefined as any, undefined as any);
expect(mockRunBeforeToolCall).toHaveBeenCalledWith(
{ toolName: "read", params: {} },
{ agentId: undefined, sessionKey: undefined, toolName: "read" },
);
});
});

View File

@ -0,0 +1,88 @@
/**
* Tool hook wrapper wires plugin before_tool_call / after_tool_call hooks
* into the tool execution pipeline.
*
* Follows the same wrapping pattern as pi-tools.abort.ts.
*/
import { getGlobalHookRunner } from "../plugins/hook-runner-global.js";
import type { AnyAgentTool } from "./pi-tools.types.js";
export type ToolHookContext = {
agentId?: string;
sessionKey?: string;
};
/**
* Wrap a tool so that plugin before_tool_call and after_tool_call hooks fire
* around every invocation.
*
* - before_tool_call runs sequentially and may modify params or block the call.
* - after_tool_call runs in parallel (fire-and-forget) with result/error/duration.
*/
export function wrapToolWithHooks(tool: AnyAgentTool, ctx: ToolHookContext): AnyAgentTool {
const originalExecute = tool.execute;
if (!originalExecute) return tool;
return {
...tool,
execute: async (toolCallId, params, signal, onUpdate) => {
const hookRunner = getGlobalHookRunner();
const toolName = tool.name;
// --- before_tool_call ---
if (hookRunner?.hasHooks("before_tool_call")) {
const beforeResult = await hookRunner.runBeforeToolCall(
{ toolName, params: (params ?? {}) as Record<string, unknown> },
{ agentId: ctx.agentId, sessionKey: ctx.sessionKey, toolName },
);
if (beforeResult?.block) {
const reason = beforeResult.blockReason ?? "Blocked by plugin hook";
return {
content: [{ type: "text" as const, text: `[blocked] ${reason}` }],
details: undefined,
};
}
// Allow hooks to modify params
if (beforeResult?.params) {
params = beforeResult.params;
}
}
// --- execute ---
const start = Date.now();
let result: Awaited<ReturnType<NonNullable<AnyAgentTool["execute"]>>>;
let error: string | undefined;
try {
result = await originalExecute(toolCallId, params, signal, onUpdate);
} catch (err) {
error = err instanceof Error ? err.message : String(err);
throw err;
} finally {
const durationMs = Date.now() - start;
// --- after_tool_call (fire-and-forget) ---
if (hookRunner?.hasHooks("after_tool_call")) {
hookRunner
.runAfterToolCall(
{
toolName,
params: (params ?? {}) as Record<string, unknown>,
result: error ? undefined : result!,
error,
durationMs,
},
{ agentId: ctx.agentId, sessionKey: ctx.sessionKey, toolName },
)
.catch(() => {
// swallow — after_tool_call is fire-and-forget
});
}
}
return result;
},
};
}

View File

@ -19,6 +19,7 @@ import { listChannelAgentTools } from "./channel-tools.js";
import { createMoltbotTools } from "./moltbot-tools.js";
import type { ModelAuthMode } from "./model-auth.js";
import { wrapToolWithAbortSignal } from "./pi-tools.abort.js";
import { wrapToolWithHooks } from "./pi-tools.hooks.js";
import {
filterToolsByPolicy,
isToolAllowedByPolicies,
@ -413,8 +414,12 @@ export function createMoltbotCodingTools(options?: {
? normalized.map((tool) => wrapToolWithAbortSignal(tool, options.abortSignal))
: normalized;
// Wrap tools with plugin before_tool_call / after_tool_call hooks.
const hookCtx = { agentId, sessionKey: options?.sessionKey };
const withHooks = withAbort.map((tool) => wrapToolWithHooks(tool, hookCtx));
// NOTE: Keep canonical (lowercase) tool names here.
// pi-ai's Anthropic OAuth transport remaps tool names to Claude Code-style names
// on the wire and maps them back for tool dispatch.
return withAbort;
return withHooks;
}