Merge pull request #1 from dp-web4/feat/tool-call-hooks

feat: wire before_tool_call / after_tool_call plugin hooks into tool pipeline

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" },
+    );
+  });
+});

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<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;
+    },
+  };
+}

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;
 }