From bdf04ab4266a14c22a643210dc88c20382efe9b9 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 26 Jan 2026 19:12:04 +0000 Subject: [PATCH] Wire up before_tool_call and after_tool_call hooks in agent execution loop - Add pi-tools.hooks.ts: wraps tool execute methods to fire plugin hooks - before_tool_call runs before each tool, supports blocking and param modification - after_tool_call fires as fire-and-forget after tool returns with timing data - Hook errors are caught and logged, never breaking tool execution - Tools are only wrapped when hooks are actually registered (zero overhead otherwise) - Move hookRunner initialization earlier in attempt.ts for tool wrapping access - Extract hookAgentId to avoid repeated string splitting --- src/agents/pi-embedded-runner/run/attempt.ts | 25 ++-- src/agents/pi-tools.hooks.ts | 115 +++++++++++++++++++ 2 files changed, 133 insertions(+), 7 deletions(-) create mode 100644 src/agents/pi-tools.hooks.ts diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index f1c487470..323371492 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -75,6 +75,7 @@ import { prewarmSessionFile, trackSessionManagerAccess } from "../session-manage import { prepareSessionManagerForRun } from "../session-manager-init.js"; import { buildEmbeddedSystemPrompt, createSystemPromptOverride } from "../system-prompt.js"; import { splitSdkTools } from "../tool-split.js"; +import { wrapToolsWithHooks } from "../../pi-tools.hooks.js"; import { toClientToolDefinitions } from "../../pi-tool-definition-adapter.js"; import { buildSystemPromptParams } from "../../system-prompt-params.js"; import { describeUnknownError, mapThinkingLevel } from "../utils.js"; @@ -428,6 +429,10 @@ export async function runEmbeddedAttempt( model: params.model, }); + // Get hook runner once for tool-call hooks, before_agent_start, and agent_end + const hookRunner = getGlobalHookRunner(); + const hookAgentId = params.sessionKey?.split(":")[0] ?? "main"; + const { builtInTools, customTools } = splitSdkTools({ tools, sandboxEnabled: !!sandbox?.enabled, @@ -443,6 +448,15 @@ export async function runEmbeddedAttempt( const allCustomTools = [...customTools, ...clientToolDefs]; + // Wrap tools with before_tool_call / after_tool_call plugin hooks + const hookCtx = { agentId: hookAgentId, sessionKey: params.sessionKey }; + const hookedBuiltInTools = hookRunner + ? wrapToolsWithHooks(builtInTools, hookRunner, hookCtx) + : builtInTools; + const hookedCustomTools = hookRunner + ? wrapToolsWithHooks(allCustomTools, hookRunner, hookCtx) + : allCustomTools; + ({ session } = await createAgentSession({ cwd: resolvedWorkspace, agentDir, @@ -451,8 +465,8 @@ export async function runEmbeddedAttempt( model: params.model, thinkingLevel: mapThinkingLevel(params.thinkLevel), systemPrompt, - tools: builtInTools, - customTools: allCustomTools, + tools: hookedBuiltInTools, + customTools: hookedCustomTools, sessionManager, settingsManager, skills: [], @@ -672,9 +686,6 @@ export async function runEmbeddedAttempt( } } - // Get hook runner once for both before_agent_start and agent_end hooks - const hookRunner = getGlobalHookRunner(); - let promptError: unknown = null; try { const promptStartedAt = Date.now(); @@ -689,7 +700,7 @@ export async function runEmbeddedAttempt( messages: activeSession.messages, }, { - agentId: params.sessionKey?.split(":")[0] ?? "main", + agentId: hookAgentId, sessionKey: params.sessionKey, workspaceDir: params.workspaceDir, messageProvider: params.messageProvider ?? undefined, @@ -817,7 +828,7 @@ export async function runEmbeddedAttempt( durationMs: Date.now() - promptStartedAt, }, { - agentId: params.sessionKey?.split(":")[0] ?? "main", + agentId: hookAgentId, sessionKey: params.sessionKey, workspaceDir: params.workspaceDir, messageProvider: params.messageProvider ?? undefined, diff --git a/src/agents/pi-tools.hooks.ts b/src/agents/pi-tools.hooks.ts new file mode 100644 index 000000000..73a3df776 --- /dev/null +++ b/src/agents/pi-tools.hooks.ts @@ -0,0 +1,115 @@ +/** + * Tool Hook Wrappers + * + * Wraps tool execute methods to fire before_tool_call and after_tool_call + * plugin hooks around every tool invocation. + */ + +import type { HookRunner } from "../plugins/hooks.js"; +import { createSubsystemLogger } from "../logging/subsystem.js"; +import type { AnyAgentTool } from "./pi-tools.types.js"; + +const log = createSubsystemLogger("tools/hooks"); + +/** + * Wrap a single tool's execute method to fire before_tool_call and after_tool_call hooks. + */ +export function wrapToolWithHooks( + tool: AnyAgentTool, + hookRunner: HookRunner, + ctx: { agentId?: string; sessionKey?: string }, +): AnyAgentTool { + const originalExecute = tool.execute; + if (!originalExecute) return tool; + + const toolName = tool.name; + + return { + ...tool, + execute: async ( + toolCallId: string, + params: unknown, + signal?: AbortSignal, + onUpdate?: (data: unknown) => void, + ) => { + const hookCtx = { + agentId: ctx.agentId, + sessionKey: ctx.sessionKey, + toolName, + }; + + // --- before_tool_call --- + let effectiveParams = params; + if (hookRunner.hasHooks("before_tool_call")) { + try { + const beforeResult = await hookRunner.runBeforeToolCall( + { + toolName, + params: (params ?? {}) as Record, + }, + hookCtx, + ); + if (beforeResult?.block) { + const reason = beforeResult.blockReason ?? "Blocked by plugin hook"; + log.debug(`before_tool_call: blocked ${toolName} — ${reason}`); + return `[Tool call blocked] ${reason}`; + } + if (beforeResult?.params) { + effectiveParams = beforeResult.params; + } + } catch (err) { + log.debug(`before_tool_call hook error for ${toolName}: ${String(err)}`); + // Hook errors must not break tool execution + } + } + + // --- execute --- + const startMs = Date.now(); + let result: unknown; + let error: string | undefined; + try { + result = await originalExecute(toolCallId, effectiveParams, signal, onUpdate); + return result; + } catch (err) { + error = String(err); + throw err; + } finally { + // --- after_tool_call (fire-and-forget) --- + if (hookRunner.hasHooks("after_tool_call")) { + hookRunner + .runAfterToolCall( + { + toolName, + params: (effectiveParams ?? {}) as Record, + result, + error, + durationMs: Date.now() - startMs, + }, + hookCtx, + ) + .catch((hookErr) => { + log.debug(`after_tool_call hook error for ${toolName}: ${String(hookErr)}`); + }); + } + } + }, + }; +} + +/** + * Wrap all tools in an array with before/after tool call hooks. + * Returns the original array unchanged if no tool call hooks are registered. + */ +export function wrapToolsWithHooks( + tools: AnyAgentTool[], + hookRunner: HookRunner, + ctx: { agentId?: string; sessionKey?: string }, +): AnyAgentTool[] { + if ( + !hookRunner.hasHooks("before_tool_call") && + !hookRunner.hasHooks("after_tool_call") + ) { + return tools; + } + return tools.map((tool) => wrapToolWithHooks(tool, hookRunner, ctx)); +}