diff --git a/src/agents/pi-tools.hooks.test.ts b/src/agents/pi-tools.hooks.test.ts new file mode 100644 index 000000000..c9d58cd15 --- /dev/null +++ b/src/agents/pi-tools.hooks.test.ts @@ -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" }, + ); + }); +}); diff --git a/src/agents/pi-tools.hooks.ts b/src/agents/pi-tools.hooks.ts new file mode 100644 index 000000000..0ee1f05df --- /dev/null +++ b/src/agents/pi-tools.hooks.ts @@ -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 }, + { 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>>; + 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, + result: error ? undefined : result!, + error, + durationMs, + }, + { agentId: ctx.agentId, sessionKey: ctx.sessionKey, toolName }, + ) + .catch(() => { + // swallow — after_tool_call is fire-and-forget + }); + } + } + + return result; + }, + }; +} diff --git a/src/agents/pi-tools.ts b/src/agents/pi-tools.ts index d763393a4..5e1dd77ca 100644 --- a/src/agents/pi-tools.ts +++ b/src/agents/pi-tools.ts @@ -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; }