@mariozechner/pi-coding-agent 0.57.1 → 0.58.0

package/dist/core/agent-session.js

@@ -24,7 +24,7 @@ import { calculateContextTokens, collectEntriesForBranchSummary, compact, estima
 import { DEFAULT_THINKING_LEVEL } from "./defaults.js";
 import { exportSessionToHtml } from "./export-html/index.js";
 import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
-import { ExtensionRunner, wrapRegisteredTools, wrapToolsWithExtensions, } from "./extensions/index.js";
+import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
 import { getLatestCompactionEntry } from "./session-manager.js";
 import { BUILTIN_SLASH_COMMANDS } from "./slash-commands.js";
@@ -122,6 +122,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,
             includeAllExtensionTools: true,
@@ -131,6 +132,59 @@ export class AgentSession {
     get modelRegistry() {
         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.
+     */
+    _installAgentToolHooks() {
+        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,
+                });
+            }
+            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,
+                content: result.content,
+                details: isError ? undefined : result.details,
+                isError,
+            });
+            if (!hookResult || isError) {
+                return undefined;
+            }
+            return {
+                content: hookResult.content,
+                details: hookResult.details,
+            };
+        });
+    }
     // =========================================================================
     // Event Subscription
     // =========================================================================
@@ -1719,13 +1773,7 @@ export class AgentSession {
         for (const tool of wrappedExtensionTools) {
             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]
             : [...previousActiveToolNames];
@@ -1817,7 +1865,7 @@ export class AgentSession {
             return false;
         const err = message.errorMessage;
         // Match: overloaded_error, rate limit, 429, 500, 502, 503, 504, service unavailable, connection errors, fetch failed, terminated, retry delay exceeded
-        return /overloaded|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server error|internal error|connection.?error|connection.?refused|other side closed|fetch failed|upstream.?connect|reset before headers|terminated|retry delay/i.test(err);
+        return /overloaded|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server.?error|internal.?error|connection.?error|connection.?refused|other side closed|fetch failed|upstream.?connect|reset before headers|terminated|retry delay/i.test(err);
     }
     /**
      * Handle retryable errors with exponential backoff.