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:
parent
3fe4b2595a
commit
dbbca8e040
236
src/agents/pi-tools.hooks.test.ts
Normal file
236
src/agents/pi-tools.hooks.test.ts
Normal 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" },
|
||||
);
|
||||
});
|
||||
});
|
||||
88
src/agents/pi-tools.hooks.ts
Normal file
88
src/agents/pi-tools.hooks.ts
Normal 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;
|
||||
},
|
||||
};
|
||||
}
|
||||
@ -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;
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user