@mariozechner/pi-coding-agent 0.55.2 → 0.55.4

package/dist/core/agent-session.js

@@ -63,6 +63,7 @@ export class AgentSession {
     // Event subscription state
     _unsubscribeAgent;
     _eventListeners = [];
+    _agentEventQueue = Promise.resolve();
     /** Tracks pending steering messages for UI display. Removed when delivered. */
     _steeringMessages = [];
     /** Tracks pending follow-up messages for UI display. Removed when delivered. */
@@ -101,6 +102,8 @@ export class AgentSession {
     _modelRegistry;
     // Tool registry for extension getTools/setTools
     _toolRegistry = new Map();
+    _toolPromptSnippets = new Map();
+    _toolPromptGuidelines = new Map();
     // Base system prompt (without extension appends) - used to apply fresh appends each turn
     _baseSystemPrompt = "";
     constructor(config) {
@@ -139,7 +142,43 @@ export class AgentSession {
     // Track last assistant message for auto-compaction check
     _lastAssistantMessage = undefined;
     /** Internal handler for agent events - shared by subscribe and reconnect */
-    _handleAgentEvent = async (event) => {
+    _handleAgentEvent = (event) => {
+        // Create retry promise synchronously before queueing async processing.
+        // Agent.emit() calls this handler synchronously, and prompt() calls waitForRetry()
+        // as soon as agent.prompt() resolves. If _retryPromise is created only inside
+        // _processAgentEvent, slow earlier queued events can delay agent_end processing
+        // and waitForRetry() can miss the in-flight retry.
+        this._createRetryPromiseForAgentEnd(event);
+        this._agentEventQueue = this._agentEventQueue.then(() => this._processAgentEvent(event), () => this._processAgentEvent(event));
+        // Keep queue alive if an event handler fails
+        this._agentEventQueue.catch(() => { });
+    };
+    _createRetryPromiseForAgentEnd(event) {
+        if (event.type !== "agent_end" || this._retryPromise) {
+            return;
+        }
+        const settings = this.settingsManager.getRetrySettings();
+        if (!settings.enabled) {
+            return;
+        }
+        const lastAssistant = this._findLastAssistantInMessages(event.messages);
+        if (!lastAssistant || !this._isRetryableError(lastAssistant)) {
+            return;
+        }
+        this._retryPromise = new Promise((resolve) => {
+            this._retryResolve = resolve;
+        });
+    }
+    _findLastAssistantInMessages(messages) {
+        for (let i = messages.length - 1; i >= 0; i--) {
+            const message = messages[i];
+            if (message.role === "assistant") {
+                return message;
+            }
+        }
+        return undefined;
+    }
+    async _processAgentEvent(event) {
         // When a user message starts, check if it's from either queue and remove it BEFORE emitting
         // This ensures the UI sees the updated queue state
         if (event.type === "message_start" && event.message.role === "user") {
@@ -206,7 +245,7 @@ export class AgentSession {
             }
             await this._checkCompaction(msg);
         }
-    };
+    }
     /** Resolve the pending retry promise */
     _resolveRetry() {
         if (this._retryResolve) {
@@ -465,8 +504,42 @@ export class AgentSession {
     get promptTemplates() {
         return this._resourceLoader.getPrompts().prompts;
     }
+    _normalizePromptSnippet(text) {
+        if (!text)
+            return undefined;
+        const oneLine = text
+            .replace(/[\r\n]+/g, " ")
+            .replace(/\s+/g, " ")
+            .trim();
+        return oneLine.length > 0 ? oneLine : undefined;
+    }
+    _normalizePromptGuidelines(guidelines) {
+        if (!guidelines || guidelines.length === 0) {
+            return [];
+        }
+        const unique = new Set();
+        for (const guideline of guidelines) {
+            const normalized = guideline.trim();
+            if (normalized.length > 0) {
+                unique.add(normalized);
+            }
+        }
+        return Array.from(unique);
+    }
     _rebuildSystemPrompt(toolNames) {
-        const validToolNames = toolNames.filter((name) => this._baseToolRegistry.has(name));
+        const validToolNames = toolNames.filter((name) => this._toolRegistry.has(name));
+        const toolSnippets = {};
+        const promptGuidelines = [];
+        for (const name of validToolNames) {
+            const snippet = this._toolPromptSnippets.get(name);
+            if (snippet) {
+                toolSnippets[name] = snippet;
+            }
+            const toolGuidelines = this._toolPromptGuidelines.get(name);
+            if (toolGuidelines) {
+                promptGuidelines.push(...toolGuidelines);
+            }
+        }
         const loaderSystemPrompt = this._resourceLoader.getSystemPrompt();
         const loaderAppendSystemPrompt = this._resourceLoader.getAppendSystemPrompt();
         const appendSystemPrompt = loaderAppendSystemPrompt.length > 0 ? loaderAppendSystemPrompt.join("\n\n") : undefined;
@@ -479,6 +552,8 @@ export class AgentSession {
             customPrompt: loaderSystemPrompt,
             appendSystemPrompt,
             selectedTools: validToolNames,
+            toolSnippets,
+            promptGuidelines,
         });
     }
     // =========================================================================
@@ -1528,6 +1603,7 @@ export class AgentSession {
             getActiveTools: () => this.getActiveToolNames(),
             getAllTools: () => this.getAllTools(),
             setActiveTools: (toolNames) => this.setActiveToolsByName(toolNames),
+            refreshTools: () => this._refreshToolRegistry(),
             getCommands,
             setModel: async (model) => {
                 const key = await this.modelRegistry.getApiKey(model);
@@ -1562,6 +1638,57 @@ export class AgentSession {
             getSystemPrompt: () => this.systemPrompt,
         });
     }
+    _refreshToolRegistry(options) {
+        const previousRegistryNames = new Set(this._toolRegistry.keys());
+        const previousActiveToolNames = this.getActiveToolNames();
+        const registeredTools = this._extensionRunner?.getAllRegisteredTools() ?? [];
+        const allCustomTools = [
+            ...registeredTools,
+            ...this._customTools.map((def) => ({ definition: def, extensionPath: "<sdk>" })),
+        ];
+        this._toolPromptSnippets = new Map(allCustomTools
+            .map((registeredTool) => {
+            const snippet = this._normalizePromptSnippet(registeredTool.definition.promptSnippet ?? registeredTool.definition.description);
+            return snippet ? [registeredTool.definition.name, snippet] : undefined;
+        })
+            .filter((entry) => entry !== undefined));
+        this._toolPromptGuidelines = new Map(allCustomTools
+            .map((registeredTool) => {
+            const guidelines = this._normalizePromptGuidelines(registeredTool.definition.promptGuidelines);
+            return guidelines.length > 0 ? [registeredTool.definition.name, guidelines] : undefined;
+        })
+            .filter((entry) => entry !== undefined));
+        const wrappedExtensionTools = this._extensionRunner
+            ? wrapRegisteredTools(allCustomTools, this._extensionRunner)
+            : [];
+        const toolRegistry = new Map(this._baseToolRegistry);
+        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;
+        }
+        const nextActiveToolNames = options?.activeToolNames
+            ? [...options.activeToolNames]
+            : [...previousActiveToolNames];
+        if (options?.includeAllExtensionTools) {
+            for (const tool of wrappedExtensionTools) {
+                nextActiveToolNames.push(tool.name);
+            }
+        }
+        else if (!options?.activeToolNames) {
+            for (const toolName of this._toolRegistry.keys()) {
+                if (!previousRegistryNames.has(toolName)) {
+                    nextActiveToolNames.push(toolName);
+                }
+            }
+        }
+        this.setActiveToolsByName([...new Set(nextActiveToolNames)]);
+    }
     _buildRuntime(options) {
         const autoResizeImages = this.settingsManager.getImageAutoResize();
         const shellCommandPrefix = this.settingsManager.getShellCommandPrefix();
@@ -1591,47 +1718,14 @@ export class AgentSession {
             this._bindExtensionCore(this._extensionRunner);
             this._applyExtensionBindings(this._extensionRunner);
         }
-        const registeredTools = this._extensionRunner?.getAllRegisteredTools() ?? [];
-        const allCustomTools = [
-            ...registeredTools,
-            ...this._customTools.map((def) => ({ definition: def, extensionPath: "<sdk>" })),
-        ];
-        const wrappedExtensionTools = this._extensionRunner
-            ? wrapRegisteredTools(allCustomTools, this._extensionRunner)
-            : [];
-        const toolRegistry = new Map(this._baseToolRegistry);
-        for (const tool of wrappedExtensionTools) {
-            toolRegistry.set(tool.name, tool);
-        }
         const defaultActiveToolNames = this._baseToolsOverride
             ? Object.keys(this._baseToolsOverride)
             : ["read", "bash", "edit", "write"];
         const baseActiveToolNames = options.activeToolNames ?? defaultActiveToolNames;
-        const activeToolNameSet = new Set(baseActiveToolNames);
-        if (options.includeAllExtensionTools) {
-            for (const tool of wrappedExtensionTools) {
-                activeToolNameSet.add(tool.name);
-            }
-        }
-        const extensionToolNames = new Set(wrappedExtensionTools.map((tool) => tool.name));
-        const activeBaseTools = Array.from(activeToolNameSet)
-            .filter((name) => this._baseToolRegistry.has(name) && !extensionToolNames.has(name))
-            .map((name) => this._baseToolRegistry.get(name));
-        const activeExtensionTools = wrappedExtensionTools.filter((tool) => activeToolNameSet.has(tool.name));
-        const activeToolsArray = [...activeBaseTools, ...activeExtensionTools];
-        if (this._extensionRunner) {
-            const wrappedActiveTools = wrapToolsWithExtensions(activeToolsArray, this._extensionRunner);
-            this.agent.setTools(wrappedActiveTools);
-            const wrappedAllTools = wrapToolsWithExtensions(Array.from(toolRegistry.values()), this._extensionRunner);
-            this._toolRegistry = new Map(wrappedAllTools.map((tool) => [tool.name, tool]));
-        }
-        else {
-            this.agent.setTools(activeToolsArray);
-            this._toolRegistry = toolRegistry;
-        }
-        const systemPromptToolNames = Array.from(activeToolNameSet).filter((name) => this._baseToolRegistry.has(name));
-        this._baseSystemPrompt = this._rebuildSystemPrompt(systemPromptToolNames);
-        this.agent.setSystemPrompt(this._baseSystemPrompt);
+        this._refreshToolRegistry({
+            activeToolNames: baseActiveToolNames,
+            includeAllExtensionTools: options.includeAllExtensionTools,
+        });
     }
     async reload() {
         const previousFlagValues = this._extensionRunner?.getFlagValues();
@@ -1677,15 +1771,18 @@ export class AgentSession {
      */
     async _handleRetryableError(message) {
         const settings = this.settingsManager.getRetrySettings();
-        if (!settings.enabled)
+        if (!settings.enabled) {
+            this._resolveRetry();
             return false;
-        this._retryAttempt++;
-        // Create retry promise on first attempt so waitForRetry() can await it
-        if (this._retryAttempt === 1 && !this._retryPromise) {
+        }
+        // Retry promise is created synchronously in _handleAgentEvent for agent_end.
+        // Keep a defensive fallback here in case a future refactor bypasses that path.
+        if (!this._retryPromise) {
             this._retryPromise = new Promise((resolve) => {
                 this._retryResolve = resolve;
             });
         }
+        this._retryAttempt++;
         if (this._retryAttempt > settings.maxRetries) {
             // Max retries exceeded, emit final failure and reset
             this._emit({