@sheason/pi-coding-agent 0.74.1-sheason.0 → 0.78.0-sheason.0.6.0-alpha.2

package/dist/core/agent-session.js

@@ -13,10 +13,11 @@
  * Modes use this class and add their own I/O layer on top.
  */
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
-import { basename, dirname, resolve } from "node:path";
-import { clampThinkingLevel, cleanupSessionResources, getSupportedThinkingLevels, isContextOverflow, modelsAreEqual, resetApiProviders, } from "@sheason/pi-ai";
+import { basename, dirname } from "node:path";
+import { clampThinkingLevel, cleanupSessionResources, getSupportedThinkingLevels, isContextOverflow, modelsAreEqual, resetApiProviders, streamSimple, } from "@sheason/pi-ai";
 import { theme } from "../modes/interactive/theme/theme.js";
 import { stripFrontmatter } from "../utils/frontmatter.js";
+import { resolvePath } from "../utils/paths.js";
 import { sleep } from "../utils/sleep.js";
 import { formatNoApiKeyFoundMessage, formatNoModelSelectedMessage } from "./auth-guidance.js";
 import { executeBashWithOperations } from "./bash-executor.js";
@@ -64,7 +65,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,14 +80,14 @@ export class AgentSession {
     // Retry state
     _retryAbortController = undefined;
     _retryAttempt = 0;
-    _retryPromise = undefined;
-    _retryResolve = undefined;
     // Bash execution state
     _bashAbortController = undefined;
     _pendingBashMessages = [];
     // Extension system
     _extensionRunner;
     _turnIndex = 0;
+    /** Timestamp when the current agent run started, for per-turn TPS calculation */
+    _agentStartTime = undefined;
     _resourceLoader;
     _customTools;
     _baseToolDefinitions = new Map();
@@ -95,10 +95,12 @@ export class AgentSession {
     _extensionRunnerRef;
     _initialActiveToolNames;
     _allowedToolNames;
+    _excludedToolNames;
     _baseToolsOverride;
     _sessionStartEvent;
     _extensionUIContext;
     _extensionCommandContextActions;
+    _extensionAbortHandler;
     _extensionShutdownHandler;
     _extensionErrorListener;
     _extensionErrorUnsubscriber;
@@ -124,6 +126,7 @@ export class AgentSession {
         this._extensionRunnerRef = config.extensionRunnerRef;
         this._initialActiveToolNames = config.initialActiveToolNames;
         this._allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined;
+        this._excludedToolNames = config.excludedToolNames ? new Set(config.excludedToolNames) : undefined;
         this._baseToolsOverride = config.baseToolsOverride;
         this._sessionStartEvent = config.sessionStartEvent ?? { type: "session_start", reason: "startup" };
         // Always subscribe to agent events for internal handling
@@ -158,6 +161,13 @@ export class AgentSession {
         }
         throw new Error(formatNoApiKeyFoundMessage(model.provider));
     }
+    async _getCompactionRequestAuth(model) {
+        if (this.agent.streamFn === streamSimple) {
+            return this._getRequiredRequestAuth(model);
+        }
+        const result = await this._modelRegistry.getApiKeyAndHeaders(model);
+        return result.ok ? { apiKey: result.apiKey, headers: result.headers } : {};
+    }
     /**
      * Install tool hooks once on the Agent instance.
      *
@@ -172,7 +182,6 @@ export class AgentSession {
             if (!runner.hasHandlers("tool_call")) {
                 return undefined;
             }
-            await this._agentEventQueue;
             try {
                 return await runner.emitToolCall({
                     type: "tool_call",
@@ -228,72 +237,81 @@ export class AgentSession {
             followUp: [...this._followUpMessages],
         });
     }
+    /** Emit a state_update event with the current token/context/model snapshot */
+    _emitStateUpdate() {
+        const model = this.model;
+        const state = this.state;
+        // Compute cumulative token usage
+        let totalInput = 0;
+        let totalOutput = 0;
+        let totalCacheRead = 0;
+        let totalCacheWrite = 0;
+        let totalCost = 0;
+        for (const entry of this.sessionManager.getEntries()) {
+            if (entry.type === "message" && entry.message.role === "assistant") {
+                totalInput += entry.message.usage.input;
+                totalOutput += entry.message.usage.output;
+                totalCacheRead += entry.message.usage.cacheRead;
+                totalCacheWrite += entry.message.usage.cacheWrite;
+                totalCost += entry.message.usage.cost.total;
+            }
+        }
+        const usingSubscription = model ? this.modelRegistry.isUsingOAuth(model) : false;
+        const rawContextUsage = this.getContextUsage();
+        const snapshot = {
+            model: model?.id ?? "",
+            thinkingLevel: state.thinkingLevel,
+            // Note: isStreaming and isCompacting are NOT included here.
+            // They are managed by agent_start/agent_end and compaction_start/compaction_end events
+            // to avoid race conditions where a late state_update overwrites the correct state.
+            tokenUsage: {
+                input: totalInput,
+                output: totalOutput,
+                cacheRead: totalCacheRead,
+                cacheWrite: totalCacheWrite,
+                cost: totalCost,
+                usingSubscription,
+            },
+            contextUsage: rawContextUsage
+                ? {
+                    tokens: rawContextUsage.tokens,
+                    contextWindow: rawContextUsage.contextWindow,
+                    percent: rawContextUsage.percent,
+                }
+                : { tokens: null, contextWindow: model?.contextWindow ?? 0, percent: null },
+            modelInfo: model
+                ? { id: model.id, provider: model.provider, reasoning: model.reasoning, contextWindow: model.contextWindow }
+                : { id: "", provider: "", reasoning: false, contextWindow: 0 },
+            autoCompactEnabled: this.autoCompactionEnabled,
+            cwd: this.sessionManager.getCwd(),
+            availableProviderCount: this.modelRegistry.getAvailable().length,
+        };
+        this._emit({ type: "state_update", snapshot });
+    }
     // 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) {
-        // 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);
+    _handleAgentEvent = async (event) => {
+        // When a queued message starts, remove it BEFORE emitting so the UI sees the updated queue state.
+        if (event.type === "message_start" && (event.message.role === "user" || event.message.role === "custom")) {
+            if (event.message.role === "user") {
+                this._overflowRecoveryAttempted = false;
+            }
+            const messageText = event.message.role === "user"
+                ? this._getUserMessageText(event.message)
+                : this._getCustomMessageText(event.message);
             if (messageText) {
-                // Check steering queue first
-                const steeringIndex = this._steeringMessages.indexOf(messageText);
-                if (steeringIndex !== -1) {
-                    this._steeringMessages.splice(steeringIndex, 1);
-                    this._emitQueueUpdate();
-                }
-                else {
-                    // Check follow-up queue
-                    const followUpIndex = this._followUpMessages.indexOf(messageText);
-                    if (followUpIndex !== -1) {
-                        this._followUpMessages.splice(followUpIndex, 1);
-                        this._emitQueueUpdate();
-                    }
-                }
+                this._removeQueuedMessageText(messageText);
             }
         }
+        // Track agent start time for per-turn TPS calculation
+        if (event.type === "agent_start") {
+            this._agentStartTime = Date.now();
+        }
         // 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
@@ -325,40 +343,87 @@ export class AgentSession {
                     });
                     this._retryAttempt = 0;
                 }
+                // Push state update (token usage changed)
+                this._emitStateUpdate();
             }
         }
-        // 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
+        else if (event.type === "agent_end") {
+            // Push state update (streaming ended)
+            this._emitStateUpdate();
+            // Emit per-turn stats
+            if (this._agentStartTime !== undefined) {
+                const duration = (Date.now() - this._agentStartTime) / 1000;
+                this._agentStartTime = undefined;
+                // Sum token usage from all assistant messages in this agent run
+                let turnInput = 0;
+                let turnOutput = 0;
+                let turnCacheRead = 0;
+                let turnCacheWrite = 0;
+                for (const msg of event.messages) {
+                    if (msg.role === "assistant") {
+                        const usage = msg.usage;
+                        turnInput += usage.input;
+                        turnOutput += usage.output;
+                        turnCacheRead += usage.cacheRead;
+                        turnCacheWrite += usage.cacheWrite;
+                    }
+                }
+                const total = turnInput + turnOutput + turnCacheRead + turnCacheWrite;
+                const tps = duration > 0 ? turnOutput / duration : 0;
+                this._emit({
+                    type: "turn_stats",
+                    tps,
+                    output: turnOutput,
+                    input: turnInput,
+                    cacheRead: turnCacheRead,
+                    cacheWrite: turnCacheWrite,
+                    total,
+                    duration,
+                });
             }
-            this._resolveRetry();
-            await this._checkCompaction(msg);
         }
-    }
-    /** Resolve the pending retry promise */
-    _resolveRetry() {
-        if (this._retryResolve) {
-            this._retryResolve();
-            this._retryResolve = undefined;
-            this._retryPromise = undefined;
+    };
+    _willRetryAfterAgentEnd(event) {
+        const settings = this.settingsManager.getRetrySettings();
+        if (!settings.enabled || this._retryAttempt >= settings.maxRetries) {
+            return false;
         }
+        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) {
         if (message.role !== "user")
             return "";
-        const content = message.content;
+        return this._getContentText(message.content);
+    }
+    _getCustomMessageText(message) {
+        return this._getContentText(message.content);
+    }
+    _getContentText(content) {
         if (typeof content === "string")
             return content;
         const textBlocks = content.filter((c) => c.type === "text");
         return textBlocks.map((c) => c.text).join("");
     }
+    _removeQueuedMessageText(messageText) {
+        const steeringIndex = this._steeringMessages.indexOf(messageText);
+        if (steeringIndex !== -1) {
+            this._steeringMessages.splice(steeringIndex, 1);
+            this._emitQueueUpdate();
+            return;
+        }
+        const followUpIndex = this._followUpMessages.indexOf(messageText);
+        if (followUpIndex !== -1) {
+            this._followUpMessages.splice(followUpIndex, 1);
+            this._emitQueueUpdate();
+        }
+    }
     /** Find the last assistant message in agent state (including aborted ones) */
     _findLastAssistantMessage() {
         const messages = this.agent.state.messages;
@@ -372,7 +437,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) {
@@ -506,6 +571,16 @@ export class AgentSession {
      * Call this when completely done with the session.
      */
     dispose() {
+        try {
+            this.abortRetry();
+            this.abortCompaction();
+            this.abortBranchSummary();
+            this.abortBash();
+            this.agent.abort();
+        }
+        catch {
+            // Dispose must succeed even if an abort hook throws.
+        }
         this._extensionRunner.invalidate("This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().");
         this._disconnectFromAgent();
         this._eventListeners = [];
@@ -546,13 +621,14 @@ export class AgentSession {
         return this.agent.state.tools.map((t) => t.name);
     }
     /**
-     * Get all configured tools with name, description, parameter schema, and source metadata.
+     * Get all configured tools with name, description, parameter schema, prompt guidelines, and source metadata.
      */
     getAllTools() {
         return Array.from(this._toolDefinitions.values()).map(({ definition, sourceInfo }) => ({
             name: definition.name,
             description: definition.description,
             parameters: definition.parameters,
+            promptGuidelines: definition.promptGuidelines,
             sourceInfo,
         }));
     }
@@ -678,6 +754,42 @@ 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;
+        }
+        if (await this._checkCompaction(msg)) {
+            return true;
+        }
+        // The agent loop drains both queues before emitting agent_end. Any messages
+        // here were queued by agent_end extension handlers and need a continuation.
+        return this.agent.hasQueuedMessages();
+    }
     /**
      * Send a prompt to the agent.
      * - Handles extension commands (registered via pi.registerCommand) immediately, even during streaming
@@ -706,7 +818,7 @@ export class AgentSession {
             let currentText = text;
             let currentImages = options?.images;
             if (this._extensionRunner.hasHandlers("input")) {
-                const inputResult = await this._extensionRunner.emitInput(currentText, currentImages, options?.source ?? "interactive");
+                const inputResult = await this._extensionRunner.emitInput(currentText, currentImages, options?.source ?? "interactive", this.isStreaming ? options?.streamingBehavior : undefined);
                 if (inputResult.action === "handled") {
                     preflightResult?.(true);
                     return;
@@ -753,8 +865,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 = [];
@@ -805,8 +925,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.
@@ -968,15 +1087,20 @@ export class AgentSession {
             this._pendingNextTurnMessages.push(appMessage);
         }
         else if (this.isStreaming) {
+            const queuedText = this._getCustomMessageText(appMessage);
             if (options?.deliverAs === "followUp") {
+                this._followUpMessages.push(queuedText);
+                this._emitQueueUpdate();
                 this.agent.followUp(appMessage);
             }
             else {
+                this._steeringMessages.push(queuedText);
+                this._emitQueueUpdate();
                 this.agent.steer(appMessage);
             }
         }
         else if (options?.triggerTurn) {
-            await this.agent.prompt(appMessage);
+            await this._runAgentPrompt(appMessage);
         }
         else {
             this.agent.state.messages.push(appMessage);
@@ -1077,7 +1201,7 @@ export class AgentSession {
      * Validates that auth is configured, saves to session and settings.
      * @throws Error if no auth is configured for the model
      */
-    async setModel(model, options = {}) {
+    async setModel(model) {
         if (!this._modelRegistry.hasConfiguredAuth(model)) {
             throw new Error(`No API key for ${model.provider}/${model.id}`);
         }
@@ -1085,12 +1209,11 @@ export class AgentSession {
         const thinkingLevel = this._getThinkingLevelForModelSwitch();
         this.agent.state.model = model;
         this.sessionManager.appendModelChange(model.provider, model.id);
-        if (options.persistDefault !== false) {
-            this.settingsManager.setDefaultModelAndProvider(model.provider, model.id);
-        }
+        this.settingsManager.setDefaultModelAndProvider(model.provider, model.id);
         // Re-clamp thinking level for new model's capabilities
         this.setThinkingLevel(thinkingLevel);
         await this._emitModelSelect(model, previousModel, "set");
+        this._emitStateUpdate();
     }
     /**
      * Cycle to next/previous model.
@@ -1126,6 +1249,7 @@ export class AgentSession {
         // setThinkingLevel clamps to model capabilities.
         this.setThinkingLevel(thinkingLevel);
         await this._emitModelSelect(next.model, currentModel, "cycle");
+        this._emitStateUpdate();
         return { model: next.model, thinkingLevel: this.thinkingLevel, isScoped: true };
     }
     async _cycleAvailableModel(direction) {
@@ -1146,6 +1270,7 @@ export class AgentSession {
         // Re-clamp thinking level for new model's capabilities
         this.setThinkingLevel(thinkingLevel);
         await this._emitModelSelect(nextModel, currentModel, "cycle");
+        this._emitStateUpdate();
         return { model: nextModel, thinkingLevel: this.thinkingLevel, isScoped: false };
     }
     // =========================================================================
@@ -1174,6 +1299,7 @@ export class AgentSession {
                 level: effectiveLevel,
                 previousLevel,
             });
+            this._emitStateUpdate();
         }
     }
     /**
@@ -1227,6 +1353,7 @@ export class AgentSession {
     setSteeringMode(mode) {
         this.agent.steeringMode = mode;
         this.settingsManager.setSteeringMode(mode);
+        this._emitStateUpdate();
     }
     /**
      * Set follow-up message mode.
@@ -1235,6 +1362,7 @@ export class AgentSession {
     setFollowUpMode(mode) {
         this.agent.followUpMode = mode;
         this.settingsManager.setFollowUpMode(mode);
+        this._emitStateUpdate();
     }
     // =========================================================================
     // Compaction
@@ -1253,7 +1381,7 @@ export class AgentSession {
             if (!this.model) {
                 throw new Error(formatNoModelSelectedMessage());
             }
-            const { apiKey, headers } = await this._getRequiredRequestAuth(this.model);
+            const { apiKey, headers } = await this._getCompactionRequestAuth(this.model);
             const pathEntries = this.sessionManager.getBranch();
             const settings = this.settingsManager.getCompactionSettings();
             const preparation = prepareCompaction(pathEntries, settings);
@@ -1296,7 +1424,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                const result = await compact(preparation, this.model, apiKey, headers, customInstructions, this._compactionAbortController.signal, this.thinkingLevel);
+                const result = await compact(preparation, this.model, apiKey, headers, customInstructions, this._compactionAbortController.signal, this.thinkingLevel, this.agent.streamFn);
                 summary = result.summary;
                 firstKeptEntryId = result.firstKeptEntryId;
                 tokensBefore = result.tokensBefore;
@@ -1309,6 +1437,8 @@ export class AgentSession {
             const newEntries = this.sessionManager.getEntries();
             const sessionContext = this.sessionManager.buildSessionContext();
             this.agent.state.messages = sessionContext.messages;
+            // Prune old entries from memory and disk to prevent unbounded growth
+            this.sessionManager.pruneAfterCompaction();
             // Get the saved compaction entry for the extension event
             const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary);
             if (this._extensionRunner && savedCompactionEntry) {
@@ -1331,6 +1461,7 @@ export class AgentSession {
                 aborted: false,
                 willRetry: false,
             });
+            this._emitStateUpdate();
             return compactionResult;
         }
         catch (error) {
@@ -1344,6 +1475,7 @@ export class AgentSession {
                 willRetry: false,
                 errorMessage: aborted ? undefined : `Compaction failed: ${message}`,
             });
+            this._emitStateUpdate();
             throw error;
         }
         finally {
@@ -1378,10 +1510,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)
@@ -1394,7 +1526,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)) {
@@ -1407,7 +1539,8 @@ 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;
+                this._emitStateUpdate();
+                return false;
             }
             this._overflowRecoveryAttempted = true;
             // Remove the error message from agent state (it IS saved to session for history,
@@ -1416,8 +1549,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.
@@ -1427,7 +1559,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.
@@ -1435,7 +1567,7 @@ export class AgentSession {
             if (compactionEntry &&
                 usageMsg.role === "assistant" &&
                 usageMsg.timestamp <= new Date(compactionEntry.timestamp).getTime()) {
-                return;
+                return false;
             }
             contextTokens = estimate.tokens;
         }
@@ -1443,8 +1575,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.
@@ -1462,20 +1595,30 @@ export class AgentSession {
                     aborted: false,
                     willRetry: false,
                 });
-                return;
+                this._emitStateUpdate();
+                return false;
             }
-            const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model);
-            if (!authResult.ok || !authResult.apiKey) {
-                this._emit({
-                    type: "compaction_end",
-                    reason,
-                    result: undefined,
-                    aborted: false,
-                    willRetry: false,
-                });
-                return;
+            let apiKey;
+            let headers;
+            if (this.agent.streamFn === streamSimple) {
+                const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model);
+                if (!authResult.ok || !authResult.apiKey) {
+                    this._emit({
+                        type: "compaction_end",
+                        reason,
+                        result: undefined,
+                        aborted: false,
+                        willRetry: false,
+                    });
+                    this._emitStateUpdate();
+                    return false;
+                }
+                apiKey = authResult.apiKey;
+                headers = authResult.headers;
+            }
+            else {
+                ({ apiKey, headers } = await this._getCompactionRequestAuth(this.model));
             }
-            const { apiKey, headers } = authResult;
             const pathEntries = this.sessionManager.getBranch();
             const preparation = prepareCompaction(pathEntries, settings);
             if (!preparation) {
@@ -1486,7 +1629,8 @@ export class AgentSession {
                     aborted: false,
                     willRetry: false,
                 });
-                return;
+                this._emitStateUpdate();
+                return false;
             }
             let extensionCompaction;
             let fromExtension = false;
@@ -1506,7 +1650,8 @@ export class AgentSession {
                         aborted: true,
                         willRetry: false,
                     });
-                    return;
+                    this._emitStateUpdate();
+                    return false;
                 }
                 if (extensionResult?.compaction) {
                     extensionCompaction = extensionResult.compaction;
@@ -1526,7 +1671,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                const compactResult = await compact(preparation, this.model, apiKey, headers, undefined, this._autoCompactionAbortController.signal, this.thinkingLevel);
+                const compactResult = await compact(preparation, this.model, apiKey, headers, undefined, this._autoCompactionAbortController.signal, this.thinkingLevel, this.agent.streamFn);
                 summary = compactResult.summary;
                 firstKeptEntryId = compactResult.firstKeptEntryId;
                 tokensBefore = compactResult.tokensBefore;
@@ -1540,12 +1685,14 @@ export class AgentSession {
                     aborted: true,
                     willRetry: false,
                 });
-                return;
+                return false;
             }
             this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension);
             const newEntries = this.sessionManager.getEntries();
             const sessionContext = this.sessionManager.buildSessionContext();
             this.agent.state.messages = sessionContext.messages;
+            // Prune old entries from memory and disk to prevent unbounded growth
+            this.sessionManager.pruneAfterCompaction();
             // Get the saved compaction entry for the extension event
             const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary);
             if (this._extensionRunner && savedCompactionEntry) {
@@ -1562,23 +1709,18 @@ export class AgentSession {
                 details,
             };
             this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry });
+            this._emitStateUpdate();
             if (willRetry) {
                 const messages = this.agent.state.messages;
                 const lastMsg = messages[messages.length - 1];
                 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";
@@ -1592,6 +1734,8 @@ export class AgentSession {
                     ? `Context overflow recovery failed: ${errorMessage}`
                     : `Auto-compaction failed: ${errorMessage}`,
             });
+            this._emitStateUpdate();
+            return false;
         }
         finally {
             this._autoCompactionAbortController = undefined;
@@ -1602,6 +1746,7 @@ export class AgentSession {
      */
     setAutoCompactionEnabled(enabled) {
         this.settingsManager.setCompactionEnabled(enabled);
+        this._emitStateUpdate();
     }
     /** Whether auto-compaction is enabled */
     get autoCompactionEnabled() {
@@ -1614,6 +1759,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;
         }
@@ -1753,7 +1901,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?.();
@@ -1787,7 +1941,8 @@ export class AgentSession {
         const previousRegistryNames = new Set(this._toolRegistry.keys());
         const previousActiveToolNames = this.getActiveToolNames();
         const allowedToolNames = this._allowedToolNames;
-        const isAllowedTool = (name) => !allowedToolNames || allowedToolNames.has(name);
+        const excludedToolNames = this._excludedToolNames;
+        const isAllowedTool = (name) => (!allowedToolNames || allowedToolNames.has(name)) && !excludedToolNames?.has(name);
         const registeredTools = this._extensionRunner.getAllRegisteredTools();
         const allCustomTools = [
             ...registeredTools,
@@ -1917,6 +2072,9 @@ export class AgentSession {
     // =========================================================================
     // Auto-Retry
     // =========================================================================
+    _isNonRetryableProviderLimitError(errorMessage) {
+        return /GoUsageLimitError|FreeUsageLimitError|Monthly usage limit reached|available balance|insufficient_quota|out of budget|quota exceeded|billing/i.test(errorMessage);
+    }
     /**
      * Check if an error is retryable (overloaded, rate limit, server errors).
      * Context overflow errors are NOT retryable (handled by compaction instead).
@@ -1929,37 +2087,24 @@ export class AgentSession {
         if (isContextOverflow(message, contextWindow))
             return false;
         const err = message.errorMessage;
+        if (this._isNonRetryableProviderLimitError(err))
+            return false;
         // Match: overloaded_error, provider returned error, rate limit, 429, 500, 502, 503, 504, service unavailable, network/connection errors (including connection lost), WebSocket transport closes/errors, fetch failed, premature stream endings, HTTP/2 closed before response, terminated, retry delay exceeded
         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);
@@ -1984,23 +2129,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;
     }
     /**
@@ -2008,23 +2147,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() {
@@ -2443,7 +2569,7 @@ export class AgentSession {
      * @returns The resolved output file path.
      */
     exportToJsonl(outputPath) {
-        const filePath = resolve(outputPath ?? `session-${new Date().toISOString().replace(/[:.]/g, "-")}.jsonl`);
+        const filePath = resolvePath(outputPath ?? `session-${new Date().toISOString().replace(/[:.]/g, "-")}.jsonl`, process.cwd());
         const dir = dirname(filePath);
         if (!existsSync(dir)) {
             mkdirSync(dir, { recursive: true });