@mariozechner/pi-coding-agent 0.55.3 → 0.56.0

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. */
@@ -72,6 +73,7 @@ export class AgentSession {
     // Compaction state
     _compactionAbortController = undefined;
     _autoCompactionAbortController = undefined;
+    _overflowRecoveryAttempted = false;
     // Branch summarization state
     _branchSummaryAbortController = undefined;
     // Retry state
@@ -101,6 +103,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,10 +143,47 @@ 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") {
+            this._overflowRecoveryAttempted = false;
             const messageText = this._getUserMessageText(event.message);
             if (messageText) {
                 // Check steering queue first
@@ -180,9 +221,12 @@ export class AgentSession {
             // Track assistant message for auto-compaction (checked on agent_end)
             if (event.message.role === "assistant") {
                 this._lastAssistantMessage = event.message;
+                const assistantMsg = event.message;
+                if (assistantMsg.stopReason !== "error") {
+                    this._overflowRecoveryAttempted = false;
+                }
                 // Reset retry counter immediately on successful assistant response
                 // This prevents accumulation across multiple LLM calls within a turn
-                const assistantMsg = event.message;
                 if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) {
                     this._emit({
                         type: "auto_retry_end",
@@ -206,7 +250,7 @@ export class AgentSession {
             }
             await this._checkCompaction(msg);
         }
-    };
+    }
     /** Resolve the pending retry promise */
     _resolveRetry() {
         if (this._retryResolve) {
@@ -425,9 +469,11 @@ export class AgentSession {
         this._baseSystemPrompt = this._rebuildSystemPrompt(validToolNames);
         this.agent.setSystemPrompt(this._baseSystemPrompt);
     }
-    /** Whether auto-compaction is currently running */
+    /** Whether compaction or branch summarization is currently running */
     get isCompacting() {
-        return this._autoCompactionAbortController !== undefined || this._compactionAbortController !== undefined;
+        return (this._autoCompactionAbortController !== undefined ||
+            this._compactionAbortController !== undefined ||
+            this._branchSummaryAbortController !== undefined);
     }
     /** All messages including custom types like BashExecutionMessage */
     get messages() {
@@ -465,8 +511,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 +559,8 @@ export class AgentSession {
             customPrompt: loaderSystemPrompt,
             appendSystemPrompt,
             selectedTools: validToolNames,
+            toolSnippets,
+            promptGuidelines,
         });
     }
     // =========================================================================
@@ -984,8 +1066,11 @@ export class AgentSession {
         this.agent.setModel(next.model);
         this.sessionManager.appendModelChange(next.model.provider, next.model.id);
         this.settingsManager.setDefaultModelAndProvider(next.model.provider, next.model.id);
-        // Apply thinking level (setThinkingLevel clamps to model capabilities)
-        this.setThinkingLevel(next.thinkingLevel);
+        // Apply thinking level.
+        // - Explicit scoped model thinking level overrides current session level
+        // - Undefined scoped model thinking level inherits current session level
+        // setThinkingLevel clamps to model capabilities.
+        this.setThinkingLevel(next.thinkingLevel ?? this.thinkingLevel);
         await this._emitModelSelect(next.model, currentModel, "cycle");
         return { model: next.model, thinkingLevel: this.thinkingLevel, isScoped: true };
     }
@@ -1246,6 +1331,17 @@ export class AgentSession {
         const errorIsFromBeforeCompaction = compactionEntry !== null && assistantMessage.timestamp < new Date(compactionEntry.timestamp).getTime();
         // Case 1: Overflow - LLM returned context overflow error
         if (sameModel && !errorIsFromBeforeCompaction && isContextOverflow(assistantMessage, contextWindow)) {
+            if (this._overflowRecoveryAttempted) {
+                this._emit({
+                    type: "auto_compaction_end",
+                    result: undefined,
+                    aborted: false,
+                    willRetry: false,
+                    errorMessage: "Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.",
+                });
+                return;
+            }
+            this._overflowRecoveryAttempted = true;
             // Remove the error message from agent state (it IS saved to session for history,
             // but we don't want it in context for the retry)
             const messages = this.agent.state.messages;
@@ -1528,6 +1624,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 +1659,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 +1739,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 +1792,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({