fix(coding-agent): sync tool hooks with agent event processing closes #2113

packages/coding-agent/src/core/agent-session.ts

@@ -67,7 +67,6 @@ import {
 	type TurnEndEvent,
 	type TurnStartEvent,
 	wrapRegisteredTools,
-	wrapToolsWithExtensions,
 } from "./extensions/index.js";
 import type { BashExecutionMessage, CustomMessage } from "./messages.js";
 import type { ModelRegistry } from "./model-registry.js";
@@ -291,6 +290,7 @@ export class AgentSession {
 		// Always subscribe to agent events for internal handling
 		// (session persistence, extensions, auto-compaction, retry logic)
 		this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
+		this._installAgentToolHooks();
 
 		this._buildRuntime({
 			activeToolNames: this._initialActiveToolNames,
@@ -303,6 +303,65 @@ export class AgentSession {
 		return this._modelRegistry;
 	}
 
+	/**
+	 * Install tool hooks once on the Agent instance.
+	 *
+	 * The callbacks read `this._extensionRunner` at execution time, so extension reload swaps in the
+	 * new runner without reinstalling hooks. Extension-specific tool wrappers are still used to adapt
+	 * registered tool execution to the extension context. Tool call and tool result interception now
+	 * happens here instead of in wrappers.
+	 */
+	private _installAgentToolHooks(): void {
+		this.agent.setBeforeToolCall(async ({ toolCall, args }) => {
+			const runner = this._extensionRunner;
+			if (!runner?.hasHandlers("tool_call")) {
+				return undefined;
+			}
+
+			await this._agentEventQueue;
+
+			try {
+				return await runner.emitToolCall({
+					type: "tool_call",
+					toolName: toolCall.name,
+					toolCallId: toolCall.id,
+					input: args as Record<string, unknown>,
+				});
+			} catch (err) {
+				if (err instanceof Error) {
+					throw err;
+				}
+				throw new Error(`Extension failed, blocking execution: ${String(err)}`);
+			}
+		});
+
+		this.agent.setAfterToolCall(async ({ toolCall, args, result, isError }) => {
+			const runner = this._extensionRunner;
+			if (!runner?.hasHandlers("tool_result")) {
+				return undefined;
+			}
+
+			const hookResult = await runner.emitToolResult({
+				type: "tool_result",
+				toolName: toolCall.name,
+				toolCallId: toolCall.id,
+				input: args as Record<string, unknown>,
+				content: result.content,
+				details: isError ? undefined : result.details,
+				isError,
+			});
+
+			if (!hookResult || isError) {
+				return undefined;
+			}
+
+			return {
+				content: hookResult.content,
+				details: hookResult.details,
+			};
+		});
+	}
+
 	// =========================================================================
 	// Event Subscription
 	// =========================================================================
@@ -2144,13 +2203,7 @@ export class AgentSession {
 		for (const tool of wrappedExtensionTools as AgentTool[]) {
 			toolRegistry.set(tool.name, tool);
 		}
-
-		if (this._extensionRunner) {
-			const wrappedAllTools = wrapToolsWithExtensions(Array.from(toolRegistry.values()), this._extensionRunner);
-			this._toolRegistry = new Map(wrappedAllTools.map((tool) => [tool.name, tool]));
-		} else {
-			this._toolRegistry = toolRegistry;
-		}
+		this._toolRegistry = toolRegistry;
 
 		const nextActiveToolNames = options?.activeToolNames
 			? [...options.activeToolNames]

packages/coding-agent/src/core/extensions/index.ts

@@ -163,9 +163,4 @@ export {
 	isToolCallEventType,
 	isWriteToolResult,
 } from "./types.js";
-export {
-	wrapRegisteredTool,
-	wrapRegisteredTools,
-	wrapToolsWithExtensions,
-	wrapToolWithExtensions,
-} from "./wrapper.js";
+export { wrapRegisteredTool, wrapRegisteredTools } from "./wrapper.js";

packages/coding-agent/src/core/extensions/wrapper.ts

@@ -1,10 +1,13 @@
 /**
- * Tool wrappers for extensions.
+ * Tool wrappers for extension-registered tools.
+ *
+ * These wrappers only adapt tool execution so extension tools receive the runner context.
+ * Tool call and tool result interception is handled by AgentSession via agent-core hooks.
  */
 
-import type { AgentTool, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core";
+import type { AgentTool } from "@mariozechner/pi-agent-core";
 import type { ExtensionRunner } from "./runner.js";
-import type { RegisteredTool, ToolCallEventResult } from "./types.js";
+import type { RegisteredTool } from "./types.js";
 
 /**
  * Wrap a RegisteredTool into an AgentTool.
@@ -29,90 +32,3 @@ export function wrapRegisteredTool(registeredTool: RegisteredTool, runner: Exten
 export function wrapRegisteredTools(registeredTools: RegisteredTool[], runner: ExtensionRunner): AgentTool[] {
 	return registeredTools.map((rt) => wrapRegisteredTool(rt, runner));
 }
-
-/**
- * Wrap a tool with extension callbacks for interception.
- * - Emits tool_call event before execution (can block)
- * - Emits tool_result event after execution (can modify result)
- */
-export function wrapToolWithExtensions<T>(tool: AgentTool<any, T>, runner: ExtensionRunner): AgentTool<any, T> {
-	return {
-		...tool,
-		execute: async (
-			toolCallId: string,
-			params: Record<string, unknown>,
-			signal?: AbortSignal,
-			onUpdate?: AgentToolUpdateCallback<T>,
-		) => {
-			// Emit tool_call event - extensions can block execution
-			if (runner.hasHandlers("tool_call")) {
-				try {
-					const callResult = (await runner.emitToolCall({
-						type: "tool_call",
-						toolName: tool.name,
-						toolCallId,
-						input: params,
-					})) as ToolCallEventResult | undefined;
-
-					if (callResult?.block) {
-						const reason = callResult.reason || "Tool execution was blocked by an extension";
-						throw new Error(reason);
-					}
-				} catch (err) {
-					if (err instanceof Error) {
-						throw err;
-					}
-					throw new Error(`Extension failed, blocking execution: ${String(err)}`);
-				}
-			}
-
-			// Execute the actual tool
-			try {
-				const result = await tool.execute(toolCallId, params, signal, onUpdate);
-
-				// Emit tool_result event - extensions can modify the result
-				if (runner.hasHandlers("tool_result")) {
-					const resultResult = await runner.emitToolResult({
-						type: "tool_result",
-						toolName: tool.name,
-						toolCallId,
-						input: params,
-						content: result.content,
-						details: result.details,
-						isError: false,
-					});
-
-					if (resultResult) {
-						return {
-							content: resultResult.content ?? result.content,
-							details: (resultResult.details ?? result.details) as T,
-						};
-					}
-				}
-
-				return result;
-			} catch (err) {
-				// Emit tool_result event for errors
-				if (runner.hasHandlers("tool_result")) {
-					await runner.emitToolResult({
-						type: "tool_result",
-						toolName: tool.name,
-						toolCallId,
-						input: params,
-						content: [{ type: "text", text: err instanceof Error ? err.message : String(err) }],
-						details: undefined,
-						isError: true,
-					});
-				}
-				throw err;
-			}
-		},
-	};
-}
-
-/**
- * Wrap all tools with extension callbacks.
- */
-export function wrapToolsWithExtensions<T>(tools: AgentTool<any, T>[], runner: ExtensionRunner): AgentTool<any, T>[] {
-	return tools.map((tool) => wrapToolWithExtensions(tool, runner));
-}

packages/coding-agent/src/core/index.ts

@@ -57,5 +57,4 @@ export {
 	type ToolResultEvent,
 	type TurnEndEvent,
 	type TurnStartEvent,
-	wrapToolsWithExtensions,
 } from "./extensions/index.js";

packages/coding-agent/src/index.ts

@@ -137,8 +137,6 @@ export {
 	isWriteToolResult,
 	wrapRegisteredTool,
 	wrapRegisteredTools,
-	wrapToolsWithExtensions,
-	wrapToolWithExtensions,
 } from "./core/extensions/index.js";
 // Footer data provider (git branch + extension statuses - data not otherwise available to extensions)
 export type { ReadonlyFooterDataProvider } from "./core/footer-data-provider.js";