diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index e83c3ae4a..28031460f 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"; @@ -432,6 +433,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, @@ -447,6 +452,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, @@ -455,8 +469,8 @@ export async function runEmbeddedAttempt( model: params.model, thinkingLevel: mapThinkingLevel(params.thinkLevel), systemPrompt, - tools: builtInTools, - customTools: allCustomTools, + tools: hookedBuiltInTools, + customTools: hookedCustomTools, sessionManager, settingsManager, skills: [], @@ -676,9 +690,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(); @@ -693,7 +704,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, @@ -821,7 +832,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)); +}