@earendil-works/pi-coding-agent 0.75.2 → 0.75.4

package/dist/core/agent-session.js

@@ -64,7 +64,6 @@ 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. */
@@ -80,8 +79,6 @@ export class AgentSession {
     // Retry state
     _retryAbortController = undefined;
     _retryAttempt = 0;
-    _retryPromise = undefined;
-    _retryResolve = undefined;
     // Bash execution state
     _bashAbortController = undefined;
     _pendingBashMessages = [];
@@ -99,6 +96,7 @@ export class AgentSession {
     _sessionStartEvent;
     _extensionUIContext;
     _extensionCommandContextActions;
+    _extensionAbortHandler;
     _extensionShutdownHandler;
     _extensionErrorListener;
     _extensionErrorUnsubscriber;
@@ -179,7 +177,6 @@ export class AgentSession {
             if (!runner.hasHandlers("tool_call")) {
                 return undefined;
             }
-            await this._agentEventQueue;
             try {
                 return await runner.emitToolCall({
                     type: "tool_call",
@@ -238,43 +235,7 @@ export class AgentSession {
     // Track last assistant message for auto-compaction check
     _lastAssistantMessage = undefined;
     /** Internal handler for agent events - shared by subscribe and reconnect */
-    _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) {
+    _handleAgentEvent = async (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") {
@@ -300,7 +261,7 @@ export class AgentSession {
         // Emit to extensions first
         await this._emitExtensionEvent(event);
         // Notify all listeners
-        this._emit(event);
+        this._emit(event.type === "agent_end" ? { ...event, willRetry: this._willRetryAfterAgentEnd(event) } : event);
         // Handle session persistence
         if (event.type === "message_end") {
             // Check if this is a custom message from extensions
@@ -334,27 +295,19 @@ export class AgentSession {
                 }
             }
         }
-        // Check auto-retry and auto-compaction after agent completes
-        if (event.type === "agent_end" && this._lastAssistantMessage) {
-            const msg = this._lastAssistantMessage;
-            this._lastAssistantMessage = undefined;
-            // Check for retryable errors first (overloaded, rate limit, server errors)
-            if (this._isRetryableError(msg)) {
-                const didRetry = await this._handleRetryableError(msg);
-                if (didRetry)
-                    return; // Retry was initiated, don't proceed to compaction
-            }
-            this._resolveRetry();
-            await this._checkCompaction(msg);
+    };
+    _willRetryAfterAgentEnd(event) {
+        const settings = this.settingsManager.getRetrySettings();
+        if (!settings.enabled || this._retryAttempt >= settings.maxRetries) {
+            return false;
         }
-    }
-    /** Resolve the pending retry promise */
-    _resolveRetry() {
-        if (this._retryResolve) {
-            this._retryResolve();
-            this._retryResolve = undefined;
-            this._retryPromise = undefined;
+        for (let i = event.messages.length - 1; i >= 0; i--) {
+            const message = event.messages[i];
+            if (message.role === "assistant") {
+                return this._isRetryableError(message);
+            }
         }
+        return false;
     }
     /** Extract text content from a message */
     _getUserMessageText(message) {
@@ -379,7 +332,7 @@ export class AgentSession {
     }
     _replaceMessageInPlace(target, replacement) {
         // Agent-core stores the finalized message object in its state before emitting message_end.
-        // SessionManager persistence happens later in _processAgentEvent() with event.message.
+        // SessionManager persistence happens later in _handleAgentEvent() with event.message.
         // Mutating this object in place keeps agent state, later turn/agent events, listeners,
         // and the eventual SessionManager.appendMessage(event.message) persistence in sync.
         if (target === replacement) {
@@ -685,6 +638,37 @@ export class AgentSession {
     // =========================================================================
     // Prompting
     // =========================================================================
+    async _runAgentPrompt(messages) {
+        try {
+            await this.agent.prompt(messages);
+            while (await this._handlePostAgentRun()) {
+                await this.agent.continue();
+            }
+        }
+        finally {
+            this._flushPendingBashMessages();
+        }
+    }
+    async _handlePostAgentRun() {
+        const msg = this._lastAssistantMessage;
+        this._lastAssistantMessage = undefined;
+        if (!msg) {
+            return false;
+        }
+        if (this._isRetryableError(msg) && (await this._prepareRetry(msg))) {
+            return true;
+        }
+        if (msg.stopReason === "error" && this._retryAttempt > 0) {
+            this._emit({
+                type: "auto_retry_end",
+                success: false,
+                attempt: this._retryAttempt,
+                finalError: msg.errorMessage,
+            });
+            this._retryAttempt = 0;
+        }
+        return await this._checkCompaction(msg);
+    }
     /**
      * Send a prompt to the agent.
      * - Handles extension commands (registered via pi.registerCommand) immediately, even during streaming
@@ -760,8 +744,16 @@ export class AgentSession {
             }
             // Check if we need to compact before sending (catches aborted responses)
             const lastAssistant = this._findLastAssistantMessage();
-            if (lastAssistant) {
-                await this._checkCompaction(lastAssistant, false);
+            if (lastAssistant && (await this._checkCompaction(lastAssistant, false))) {
+                try {
+                    await this.agent.continue();
+                    while (await this._handlePostAgentRun()) {
+                        await this.agent.continue();
+                    }
+                }
+                finally {
+                    this._flushPendingBashMessages();
+                }
             }
             // Build messages array (custom message if any, then user message)
             messages = [];
@@ -812,8 +804,7 @@ export class AgentSession {
             return;
         }
         preflightResult?.(true);
-        await this.agent.prompt(messages);
-        await this.waitForRetry();
+        await this._runAgentPrompt(messages);
     }
     /**
      * Try to execute an extension command. Returns true if command was found and executed.
@@ -983,7 +974,7 @@ export class AgentSession {
             }
         }
         else if (options?.triggerTurn) {
-            await this.agent.prompt(appMessage);
+            await this._runAgentPrompt(appMessage);
         }
         else {
             this.agent.state.messages.push(appMessage);
@@ -1383,10 +1374,10 @@ export class AgentSession {
     async _checkCompaction(assistantMessage, skipAbortedCheck = true) {
         const settings = this.settingsManager.getCompactionSettings();
         if (!settings.enabled)
-            return;
+            return false;
         // Skip if message was aborted (user cancelled) - unless skipAbortedCheck is false
         if (skipAbortedCheck && assistantMessage.stopReason === "aborted")
-            return;
+            return false;
         const contextWindow = this.model?.contextWindow ?? 0;
         // Skip overflow check if the message came from a different model.
         // This handles the case where user switched from a smaller-context model (e.g. opus)
@@ -1399,7 +1390,7 @@ export class AgentSession {
         const compactionEntry = getLatestCompactionEntry(this.sessionManager.getBranch());
         const assistantIsFromBeforeCompaction = compactionEntry !== null && assistantMessage.timestamp <= new Date(compactionEntry.timestamp).getTime();
         if (assistantIsFromBeforeCompaction) {
-            return;
+            return false;
         }
         // Case 1: Overflow - LLM returned context overflow error
         if (sameModel && isContextOverflow(assistantMessage, contextWindow)) {
@@ -1412,7 +1403,7 @@ export class AgentSession {
                     willRetry: false,
                     errorMessage: "Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.",
                 });
-                return;
+                return false;
             }
             this._overflowRecoveryAttempted = true;
             // Remove the error message from agent state (it IS saved to session for history,
@@ -1421,8 +1412,7 @@ export class AgentSession {
             if (messages.length > 0 && messages[messages.length - 1].role === "assistant") {
                 this.agent.state.messages = messages.slice(0, -1);
             }
-            await this._runAutoCompaction("overflow", true);
-            return;
+            return await this._runAutoCompaction("overflow", true);
         }
         // Case 2: Threshold - context is getting large
         // For error messages (no usage data), estimate from last successful response.
@@ -1432,7 +1422,7 @@ export class AgentSession {
             const messages = this.agent.state.messages;
             const estimate = estimateContextTokens(messages);
             if (estimate.lastUsageIndex === null)
-                return; // No usage data at all
+                return false; // No usage data at all
             // Verify the usage source is post-compaction. Kept pre-compaction messages
             // have stale usage reflecting the old (larger) context and would falsely
             // trigger compaction right after one just finished.
@@ -1440,7 +1430,7 @@ export class AgentSession {
             if (compactionEntry &&
                 usageMsg.role === "assistant" &&
                 usageMsg.timestamp <= new Date(compactionEntry.timestamp).getTime()) {
-                return;
+                return false;
             }
             contextTokens = estimate.tokens;
         }
@@ -1448,8 +1438,9 @@ export class AgentSession {
             contextTokens = calculateContextTokens(assistantMessage.usage);
         }
         if (shouldCompact(contextTokens, contextWindow, settings)) {
-            await this._runAutoCompaction("threshold", false);
+            return await this._runAutoCompaction("threshold", false);
         }
+        return false;
     }
     /**
      * Internal: Run auto-compaction with events.
@@ -1467,7 +1458,7 @@ export class AgentSession {
                     aborted: false,
                     willRetry: false,
                 });
-                return;
+                return false;
             }
             let apiKey;
             let headers;
@@ -1481,7 +1472,7 @@ export class AgentSession {
                         aborted: false,
                         willRetry: false,
                     });
-                    return;
+                    return false;
                 }
                 apiKey = authResult.apiKey;
                 headers = authResult.headers;
@@ -1499,7 +1490,7 @@ export class AgentSession {
                     aborted: false,
                     willRetry: false,
                 });
-                return;
+                return false;
             }
             let extensionCompaction;
             let fromExtension = false;
@@ -1519,7 +1510,7 @@ export class AgentSession {
                         aborted: true,
                         willRetry: false,
                     });
-                    return;
+                    return false;
                 }
                 if (extensionResult?.compaction) {
                     extensionCompaction = extensionResult.compaction;
@@ -1553,7 +1544,7 @@ export class AgentSession {
                     aborted: true,
                     willRetry: false,
                 });
-                return;
+                return false;
             }
             this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension);
             const newEntries = this.sessionManager.getEntries();
@@ -1581,17 +1572,11 @@ export class AgentSession {
                 if (lastMsg?.role === "assistant" && lastMsg.stopReason === "error") {
                     this.agent.state.messages = messages.slice(0, -1);
                 }
-                setTimeout(() => {
-                    this.agent.continue().catch(() => { });
-                }, 100);
-            }
-            else if (this.agent.hasQueuedMessages()) {
-                // Auto-compaction can complete while follow-up/steering/custom messages are waiting.
-                // Kick the loop so queued messages are actually delivered.
-                setTimeout(() => {
-                    this.agent.continue().catch(() => { });
-                }, 100);
+                return true;
             }
+            // Auto-compaction can complete while follow-up/steering/custom messages are waiting.
+            // Continue once so queued messages are delivered.
+            return this.agent.hasQueuedMessages();
         }
         catch (error) {
             const errorMessage = error instanceof Error ? error.message : "compaction failed";
@@ -1605,6 +1590,7 @@ export class AgentSession {
                     ? `Context overflow recovery failed: ${errorMessage}`
                     : `Auto-compaction failed: ${errorMessage}`,
             });
+            return false;
         }
         finally {
             this._autoCompactionAbortController = undefined;
@@ -1627,6 +1613,9 @@ export class AgentSession {
         if (bindings.commandContextActions !== undefined) {
             this._extensionCommandContextActions = bindings.commandContextActions;
         }
+        if (bindings.abortHandler !== undefined) {
+            this._extensionAbortHandler = bindings.abortHandler;
+        }
         if (bindings.shutdownHandler !== undefined) {
             this._extensionShutdownHandler = bindings.shutdownHandler;
         }
@@ -1766,7 +1755,13 @@ export class AgentSession {
             getModel: () => this.model,
             isIdle: () => !this.isStreaming,
             getSignal: () => this.agent.signal,
-            abort: () => this.abort(),
+            abort: () => {
+                if (this._extensionAbortHandler) {
+                    this._extensionAbortHandler();
+                    return;
+                }
+                void this.abort();
+            },
             hasPendingMessages: () => this.pendingMessageCount > 0,
             shutdown: () => {
                 this._extensionShutdownHandler?.();
@@ -1946,33 +1941,18 @@ export class AgentSession {
         return /overloaded|provider.?returned.?error|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server.?error|internal.?error|network.?error|connection.?error|connection.?refused|connection.?lost|websocket.?closed|websocket.?error|other side closed|fetch failed|upstream.?connect|reset before headers|socket hang up|ended without|stream ended before message_stop|http2 request did not get a response|timed? out|timeout|terminated|retry delay/i.test(err);
     }
     /**
-     * Handle retryable errors with exponential backoff.
-     * @returns true if retry was initiated, false if max retries exceeded or disabled
+     * Prepare a retryable error for continuation with exponential backoff.
+     * @returns true if the caller should continue the agent, false otherwise
      */
-    async _handleRetryableError(message) {
+    async _prepareRetry(message) {
         const settings = this.settingsManager.getRetrySettings();
         if (!settings.enabled) {
-            this._resolveRetry();
             return false;
         }
-        // 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({
-                type: "auto_retry_end",
-                success: false,
-                attempt: this._retryAttempt - 1,
-                finalError: message.errorMessage,
-            });
-            this._retryAttempt = 0;
-            this._resolveRetry(); // Resolve so waitForRetry() completes
+            // Preserve the completed attempt count so post-run handling can emit the final failure.
+            this._retryAttempt--;
             return false;
         }
         const delayMs = settings.baseDelayMs * 2 ** (this._retryAttempt - 1);
@@ -1997,23 +1977,17 @@ export class AgentSession {
             // Aborted during sleep - emit end event so UI can clean up
             const attempt = this._retryAttempt;
             this._retryAttempt = 0;
-            this._retryAbortController = undefined;
             this._emit({
                 type: "auto_retry_end",
                 success: false,
                 attempt,
                 finalError: "Retry cancelled",
             });
-            this._resolveRetry();
             return false;
         }
-        this._retryAbortController = undefined;
-        // Retry via continue() - use setTimeout to break out of event handler chain
-        setTimeout(() => {
-            this.agent.continue().catch(() => {
-                // Retry failed - will be caught by next agent_end
-            });
-        }, 0);
+        finally {
+            this._retryAbortController = undefined;
+        }
         return true;
     }
     /**
@@ -2021,23 +1995,10 @@ export class AgentSession {
      */
     abortRetry() {
         this._retryAbortController?.abort();
-        // Note: _retryAttempt is reset in the catch block of _autoRetry
-        this._resolveRetry();
-    }
-    /**
-     * Wait for any in-progress retry to complete.
-     * Returns immediately if no retry is in progress.
-     */
-    async waitForRetry() {
-        if (!this._retryPromise) {
-            return;
-        }
-        await this._retryPromise;
-        await this.agent.waitForIdle();
     }
     /** Whether auto-retry is currently in progress */
     get isRetrying() {
-        return this._retryPromise !== undefined;
+        return this._retryAbortController !== undefined;
     }
     /** Whether auto-retry is enabled */
     get autoRetryEnabled() {