gsd-pi 2.11.0 → 2.12.0

package/packages/pi-coding-agent/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 { FallbackResolver } from "./fallback-resolver.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
 import { getLatestCompactionEntry } from "./session-manager.js";
@@ -119,6 +119,10 @@ 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);
+        // Install tool hooks that await the event queue before emitting extension events.
+        // This ensures extensions always see settled state (e.g., assistant message appended)
+        // even when tools execute in parallel.
+        this._installAgentToolHooks();
         this._buildRuntime({
             activeToolNames: this._initialActiveToolNames,
             includeAllExtensionTools: true,
@@ -246,6 +250,66 @@ export class AgentSession {
             this._retryPromise = undefined;
         }
     }
+    /**
+     * Install beforeToolCall/afterToolCall hooks on the Agent.
+     *
+     * These hooks await `_agentEventQueue` before emitting extension events,
+     * ensuring that all prior events (including `message_end` which appends
+     * the assistant message) have fully settled. This prevents a race condition
+     * in parallel tool execution where extension `tool_call` handlers could
+     * see stale agent state.
+     */
+    _installAgentToolHooks() {
+        this.agent.setBeforeToolCall(async ({ toolCall, args }) => {
+            // Wait for all queued agent events to settle before emitting to extensions
+            await this._agentEventQueue;
+            if (!this._extensionRunner?.hasHandlers("tool_call"))
+                return undefined;
+            try {
+                const callResult = await this._extensionRunner.emitToolCall({
+                    type: "tool_call",
+                    toolName: toolCall.name,
+                    toolCallId: toolCall.id,
+                    input: args,
+                });
+                if (callResult?.block) {
+                    return {
+                        block: true,
+                        reason: callResult.reason || "Tool execution was blocked by an extension",
+                    };
+                }
+            }
+            catch (err) {
+                if (err instanceof Error) {
+                    return { block: true, reason: err.message };
+                }
+                return { block: true, reason: `Extension failed, blocking execution: ${String(err)}` };
+            }
+            return undefined;
+        });
+        this.agent.setAfterToolCall(async ({ toolCall, args, result, isError }) => {
+            // Wait for all queued agent events to settle
+            await this._agentEventQueue;
+            if (!this._extensionRunner?.hasHandlers("tool_result"))
+                return undefined;
+            const resultResult = await this._extensionRunner.emitToolResult({
+                type: "tool_result",
+                toolName: toolCall.name,
+                toolCallId: toolCall.id,
+                input: args,
+                content: result.content,
+                details: result.details,
+                isError,
+            });
+            if (resultResult) {
+                return {
+                    content: resultResult.content ?? undefined,
+                    details: resultResult.details ?? undefined,
+                };
+            }
+            return undefined;
+        });
+    }
     /** Extract text content from a message */
     _getUserMessageText(message) {
         if (message.role !== "user")
@@ -1724,13 +1788,10 @@ 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;
-        }
+        // Tool interception (tool_call/tool_result extension events) is handled by
+        // beforeToolCall/afterToolCall hooks installed in _installAgentToolHooks(),
+        // which await _agentEventQueue for safe parallel execution.
+        this._toolRegistry = toolRegistry;
         const nextActiveToolNames = options?.activeToolNames
             ? [...options.activeToolNames]
             : [...previousActiveToolNames];