@astralyn/pi 0.83.0 → 0.84.0

package/dist/core/agent-session.js

@@ -15,11 +15,12 @@
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { basename, dirname } from "node:path";
 import { contentText } from "@earendil-works/pi-ai";
-import { clampThinkingLevel, cleanupSessionResources, getSupportedThinkingLevels, isContextOverflow, isRetryableAssistantError, modelsAreEqual, resetApiProviders, streamSimple, } from "@earendil-works/pi-ai/compat";
+import { clampThinkingLevel, cleanupSessionResources, getSupportedThinkingLevels, isContextOverflow, isRecoverableLength, isRetryableAssistantError, modelsAreEqual, resetApiProviders, streamSimple, } from "@earendil-works/pi-ai/compat";
 import { getThemeByName, 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 { normalizeToolResultImages } from "../utils/tool-result-images.js";
 import { formatNoApiKeyFoundMessage, formatNoModelSelectedMessage } from "./auth-guidance.js";
 import { executeBashWithOperations } from "./bash-executor.js";
 import { calculateContextTokens, collectEntriesForBranchSummary, compact, estimateContextTokens, estimateTokens, generateBranchSummary, prepareCompaction, shouldCompact, } from "./compaction/index.js";
@@ -79,68 +80,104 @@ const THINKING_LEVELS = ["off", "minimal", "low", "medium", "high"];
 // AgentSession Class
 // ============================================================================
 export class AgentSession {
-    agent;
-    sessionManager;
-    settingsManager;
-    _scopedModels;
-    // Event subscription state
-    _unsubscribeAgent;
-    _eventListeners = [];
-    _isAgentRunActive = false;
-    _idleWaitPromise;
-    _resolveIdleWait;
-    /** Tracks pending steering messages for UI display. Removed when delivered. */
-    _steeringMessages = [];
-    /** Tracks pending follow-up messages for UI display. Removed when delivered. */
-    _followUpMessages = [];
-    /** Messages queued to be included with the next user prompt as context ("asides"). */
-    _pendingNextTurnMessages = [];
-    // Compaction state
-    _compactionAbortController = undefined;
-    _autoCompactionAbortController = undefined;
-    _stopAfterTurnRequested = false;
-    /** An extension cancelled a mid-turn compaction; skip further mid-turn checks this run. */
-    _midTurnCompactionDeclined = false;
-    _overflowRecoveryAttempted = false;
-    // Branch summarization state
-    _branchSummaryAbortController = undefined;
-    // Retry state
-    _retryAbortController = undefined;
-    _retryAttempt = 0;
-    // Bash execution state
-    _bashAbortControllers = new Set();
-    _pendingBashMessages = [];
-    // Extension system
-    _extensionRunner;
-    _turnIndex = 0;
-    _resourceLoader;
-    _customTools;
-    _baseToolDefinitions = new Map();
-    _cwd;
-    _extensionRunnerRef;
-    _initialActiveToolNames;
-    _allowedToolNames;
-    _excludedToolNames;
-    _baseToolsOverride;
-    _sessionStartEvent;
-    _extensionUIContext;
-    _extensionMode = "print";
-    _extensionCommandContextActions;
-    _extensionAbortHandler;
-    _extensionShutdownHandler;
-    _extensionErrorListener;
-    _extensionErrorUnsubscriber;
-    _modelRuntime;
-    // Tool registry for extension getTools/setTools
-    _toolRegistry = new Map();
-    _toolDefinitions = new Map();
-    _toolPromptSnippets = new Map();
-    _toolPromptGuidelines = new Map();
-    // Base system prompt (without extension appends) - used to apply fresh appends each turn
-    _baseSystemPrompt = "";
-    _baseSystemPromptOptions;
-    _systemPromptOverride;
     constructor(config) {
+        this._eventListeners = [];
+        this._isAgentRunActive = false;
+        /** Tracks pending steering messages for UI display. Removed when delivered. */
+        this._steeringMessages = [];
+        /** Tracks pending follow-up messages for UI display. Removed when delivered. */
+        this._followUpMessages = [];
+        /** Messages queued to be included with the next user prompt as context ("asides"). */
+        this._pendingNextTurnMessages = [];
+        // Compaction state
+        this._compactionAbortController = undefined;
+        this._autoCompactionAbortController = undefined;
+        this._stopAfterTurnRequested = false;
+        /** An extension cancelled a mid-turn compaction; skip further mid-turn checks this run. */
+        this._midTurnCompactionDeclined = false;
+        this._overflowRecoveryAttempted = false;
+        // Branch summarization state
+        this._branchSummaryAbortController = undefined;
+        // Retry state
+        this._retryAbortController = undefined;
+        this._retryAttempt = 0;
+        // Bash execution state
+        this._bashAbortControllers = new Set();
+        this._pendingBashMessages = [];
+        this._turnIndex = 0;
+        this._baseToolDefinitions = new Map();
+        this._extensionMode = "print";
+        // Tool registry for extension getTools/setTools
+        this._toolRegistry = new Map();
+        this._toolDefinitions = new Map();
+        this._toolPromptSnippets = new Map();
+        this._toolPromptGuidelines = new Map();
+        // Base system prompt (without extension appends) - used to apply fresh appends each turn
+        this._baseSystemPrompt = "";
+        // Track last assistant message for auto-compaction check
+        this._lastAssistantMessage = undefined;
+        /** Internal handler for agent events - shared by subscribe and reconnect */
+        this._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") {
+                this._overflowRecoveryAttempted = false;
+                const messageText = contentText(event.message.content, "");
+                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();
+                        }
+                    }
+                }
+            }
+            // Emit to extensions first
+            await this._emitExtensionEvent(event);
+            // Notify all listeners
+            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
+                if (event.message.role === "custom") {
+                    // Persist as CustomMessageEntry
+                    this.sessionManager.appendCustomMessageEntry(event.message.customType, event.message.content, event.message.display, event.message.details);
+                }
+                else if (event.message.role === "user" ||
+                    event.message.role === "assistant" ||
+                    event.message.role === "toolResult") {
+                    // Regular LLM message - persist as SessionMessageEntry
+                    this.sessionManager.appendMessage(event.message);
+                }
+                // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere
+                // 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" && assistantMsg.stopReason !== "length") {
+                        this._overflowRecoveryAttempted = false;
+                    }
+                    // Reset retry counter immediately on successful assistant response
+                    // This prevents accumulation across multiple LLM calls within a turn
+                    if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) {
+                        this._emit({
+                            type: "auto_retry_end",
+                            success: true,
+                            attempt: this._retryAttempt,
+                        });
+                        this._retryAttempt = 0;
+                    }
+                }
+            }
+        };
         this.agent = config.agent;
         this.sessionManager = config.sessionManager;
         this.settingsManager = config.settingsManager;
@@ -181,7 +218,9 @@ export class AgentSession {
             throw error;
         }
         if (result && (result.auth.apiKey || result.auth.headers)) {
+            const requestModel = result.auth.baseUrl ? { ...model, baseUrl: result.auth.baseUrl } : model;
             return {
+                model: requestModel,
                 apiKey: result.auth.apiKey,
                 headers: withoutDeletedHeaders(result.auth.headers),
                 env: result.env,
@@ -201,12 +240,18 @@ export class AgentSession {
         }
         try {
             const result = await this._modelRuntime.getAuth(model);
-            return result
-                ? { apiKey: result.auth.apiKey, headers: withoutDeletedHeaders(result.auth.headers), env: result.env }
-                : {};
+            if (!result)
+                return { model };
+            const requestModel = result.auth.baseUrl ? { ...model, baseUrl: result.auth.baseUrl } : model;
+            return {
+                model: requestModel,
+                apiKey: result.auth.apiKey,
+                headers: withoutDeletedHeaders(result.auth.headers),
+                env: result.env,
+            };
         }
         catch {
-            return {};
+            return { model };
         }
     }
     /**
@@ -240,27 +285,31 @@ export class AgentSession {
         };
         this.agent.afterToolCall = async ({ toolCall, args, result, isError }) => {
             const runner = this._extensionRunner;
-            if (!runner.hasHandlers("tool_result")) {
-                return undefined;
-            }
-            const hookResult = await runner.emitToolResult({
-                type: "tool_result",
-                toolName: toolCall.name,
-                toolCallId: toolCall.id,
-                input: args,
-                content: result.content,
-                details: result.details,
-                isError,
-                usage: result.usage,
+            const hookResult = runner.hasHandlers("tool_result")
+                ? await runner.emitToolResult({
+                    type: "tool_result",
+                    toolName: toolCall.name,
+                    toolCallId: toolCall.id,
+                    input: args,
+                    content: result.content,
+                    details: result.details,
+                    isError,
+                    usage: result.usage,
+                })
+                : undefined;
+            const content = hookResult?.content ?? result.content ?? [];
+            // Runs after the extension hook so images injected or replaced by extensions are normalized too.
+            const normalizedContent = await normalizeToolResultImages(content, {
+                autoResizeImages: this.settingsManager.getImageAutoResize(),
             });
-            if (!hookResult) {
+            if (!hookResult && normalizedContent === content) {
                 return undefined;
             }
             return {
-                content: hookResult.content,
-                details: hookResult.details,
-                isError: hookResult.isError ?? isError,
-                usage: hookResult.usage,
+                content: normalizedContent,
+                details: hookResult?.details,
+                isError: hookResult?.isError ?? isError,
+                usage: hookResult?.usage,
             };
         };
     }
@@ -344,70 +393,6 @@ export class AgentSession {
             this._resolveIdleWaitIfIdle();
         }
     }
-    // Track last assistant message for auto-compaction check
-    _lastAssistantMessage = undefined;
-    /** Internal handler for agent events - shared by subscribe and reconnect */
-    _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") {
-            this._overflowRecoveryAttempted = false;
-            const messageText = contentText(event.message.content, "");
-            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();
-                    }
-                }
-            }
-        }
-        // Emit to extensions first
-        await this._emitExtensionEvent(event);
-        // Notify all listeners
-        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
-            if (event.message.role === "custom") {
-                // Persist as CustomMessageEntry
-                this.sessionManager.appendCustomMessageEntry(event.message.customType, event.message.content, event.message.display, event.message.details);
-            }
-            else if (event.message.role === "user" ||
-                event.message.role === "assistant" ||
-                event.message.role === "toolResult") {
-                // Regular LLM message - persist as SessionMessageEntry
-                this.sessionManager.appendMessage(event.message);
-            }
-            // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere
-            // 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
-                if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) {
-                    this._emit({
-                        type: "auto_retry_end",
-                        success: true,
-                        attempt: this._retryAttempt,
-                    });
-                    this._retryAttempt = 0;
-                }
-            }
-        }
-    };
     _willRetryAfterAgentEnd(event) {
         const settings = this.settingsManager.getRetrySettings();
         if (!settings.enabled || this._retryAttempt >= settings.maxRetries) {
@@ -552,26 +537,13 @@ export class AgentSession {
             }
         };
     }
-    /**
-     * Temporarily disconnect from agent events.
-     * User listeners are preserved and will receive events again after resubscribe().
-     * Used internally during operations that need to pause event processing.
-     */
+    /** Disconnect from agent events during disposal. */
     _disconnectFromAgent() {
         if (this._unsubscribeAgent) {
             this._unsubscribeAgent();
             this._unsubscribeAgent = undefined;
         }
     }
-    /**
-     * Reconnect to agent events after _disconnectFromAgent().
-     * Preserves all existing listeners.
-     */
-    _reconnectToAgent() {
-        if (this._unsubscribeAgent)
-            return; // Already connected
-        this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
-    }
     /**
      * Remove all listeners and disconnect from agent.
      * Call this when completely done with the session.
@@ -839,6 +811,9 @@ export class AgentSession {
                     return;
                 }
             }
+            if (this._compactionAbortController !== undefined) {
+                throw new Error("Cannot submit a prompt while compaction is in progress. Wait for compaction to finish and retry.");
+            }
             // Emit input event for extension interception (before skill/template expansion)
             let currentText = text;
             let currentImages = options?.images;
@@ -1260,11 +1235,8 @@ export class AgentSession {
         return this._cycleAvailableModel(direction);
     }
     async _cycleScopedModel(direction) {
-        const checks = await Promise.all(this._scopedModels.map(async (scoped) => ({
-            scoped,
-            auth: await this._modelRuntime.checkAuth(scoped.model.provider),
-        })));
-        const scopedModels = checks.filter(({ auth }) => auth !== undefined).map(({ scoped }) => scoped);
+        const availableIds = new Set(this._modelRuntime.getAvailableSnapshot().map((model) => `${model.provider}\0${model.id}`));
+        const scopedModels = this._scopedModels.filter((scoped) => availableIds.has(`${scoped.model.provider}\0${scoped.model.id}`));
         if (scopedModels.length <= 1)
             return undefined;
         const currentModel = this.model;
@@ -1288,7 +1260,7 @@ export class AgentSession {
         return { model: next.model, thinkingLevel: this.thinkingLevel, isScoped: true };
     }
     async _cycleAvailableModel(direction) {
-        const availableModels = await this._modelRuntime.getAvailable();
+        const availableModels = this._modelRuntime.getAvailableSnapshot();
         if (availableModels.length <= 1)
             return undefined;
         const currentModel = this.model;
@@ -1408,7 +1380,6 @@ export class AgentSession {
      * @param customInstructions Optional instructions for the compaction summary
      */
     async compact(customInstructions) {
-        this._disconnectFromAgent();
         await this.abort();
         this._compactionAbortController = new AbortController();
         this._emit({ type: "compaction_start", reason: "manual" });
@@ -1416,7 +1387,7 @@ export class AgentSession {
             if (!this.model) {
                 throw new Error(formatNoModelSelectedMessage());
             }
-            const { apiKey, headers, env } = await this._getSummarizationRequestAuth(this.model);
+            const { model: requestModel, apiKey, headers, env } = await this._getSummarizationRequestAuth(this.model);
             const pathEntries = this.sessionManager.getBranch();
             const settings = this.settingsManager.getCompactionSettings();
             const preparation = prepareCompaction(pathEntries, settings);
@@ -1464,7 +1435,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                const result = await compact(preparation, this.model, apiKey, headers, customInstructions, this._compactionAbortController.signal, this.thinkingLevel, this.agent.streamFunction, env, this.settingsManager.getRetrySettings(), this._summarizationRetryCallbacks({ source: "compaction", reason: "manual" }));
+                const result = await compact(preparation, requestModel, apiKey, headers, customInstructions, this._compactionAbortController.signal, this.thinkingLevel, this.agent.streamFunction, env, this.settingsManager.getRetrySettings(), this._summarizationRetryCallbacks({ source: "compaction", reason: "manual" }));
                 summary = result.summary;
                 firstKeptEntryId = result.firstKeptEntryId;
                 tokensBefore = result.tokensBefore;
@@ -1498,6 +1469,8 @@ export class AgentSession {
                 usage,
                 details,
             };
+            // compaction_end listeners may submit queued prompts, so expose idle state before notifying them.
+            this._compactionAbortController = undefined;
             this._emit({
                 type: "compaction_end",
                 reason: "manual",
@@ -1510,6 +1483,7 @@ export class AgentSession {
         catch (error) {
             const message = error instanceof Error ? error.message : String(error);
             const aborted = message === "Compaction cancelled" || (error instanceof Error && error.name === "AbortError");
+            this._compactionAbortController = undefined;
             this._emit({
                 type: "compaction_end",
                 reason: "manual",
@@ -1522,7 +1496,6 @@ export class AgentSession {
         }
         finally {
             this._compactionAbortController = undefined;
-            this._reconnectToAgent();
         }
     }
     /**
@@ -1597,7 +1570,8 @@ export class AgentSession {
      * Called after agent_end and before prompt submission.
      *
      * Two cases:
-     * 1. Overflow: LLM returned context overflow error, remove error message from agent state, compact, auto-retry
+     * 1. Recoverable failure: LLM returned context overflow or stopped below its desired output limit;
+     *    remove the assistant message, compact, and auto-retry once
      * 2. Threshold: Context over threshold, compact, NO auto-retry (user continues manually)
      *
      * @param assistantMessage The assistant message to check
@@ -1624,11 +1598,13 @@ export class AgentSession {
         if (assistantIsFromBeforeCompaction) {
             return false;
         }
-        // Case 1: Overflow - LLM returned context overflow error, or reported usage exceeded
-        // the configured window. A successful response over the configured window should compact
-        // but must not retry: the assistant answer already completed and agent.continue() cannot
-        // continue from an assistant message.
-        if (sameModel && isContextOverflow(assistantMessage, contextWindow)) {
+        // Case 1: Recoverable failure. Explicit/silent context overflow still uses context metadata.
+        // A length stop is recoverable when output ended below the model's original desired limit,
+        // independent of the configured context size or any context-clamped provider request limit.
+        // A successful response over the configured window should compact but must not retry: the
+        // assistant answer already completed and agent.continue() cannot continue from an assistant.
+        const recoverableLength = sameModel && isRecoverableLength(assistantMessage, this.model?.maxTokens ?? 0);
+        if (sameModel && (isContextOverflow(assistantMessage, contextWindow) || recoverableLength)) {
             const willRetry = assistantMessage.stopReason !== "stop";
             if (!willRetry) {
                 return await this._runAutoCompaction("overflow", false, timing);
@@ -1646,8 +1622,8 @@ export class AgentSession {
                 return false;
             }
             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)
+            // Remove the failed or truncated message from agent state. It remains in session history,
+            // but must not be included in the compact-and-retry context.
             const messages = this.agent.state.messages;
             if (messages.length > 0 && messages[messages.length - 1].role === "assistant") {
                 this.agent.state.messages = messages.slice(0, -1);
@@ -1781,14 +1757,20 @@ export class AgentSession {
                 details = extensionCompaction.details;
             }
             else {
+                let requestModel = model;
                 let apiKey;
                 let headers;
                 let env;
                 if (streamFunction === streamSimple) {
-                    ({ apiKey, headers, env } = await this._getRequiredRequestAuth(model));
+                    ({ model: requestModel, apiKey, headers, env } = await this._getRequiredRequestAuth(model));
                 }
                 else {
-                    ({ apiKey, headers, env } = await this._getSummarizationRequestAuth(model, streamFunction));
+                    ({
+                        model: requestModel,
+                        apiKey,
+                        headers,
+                        env,
+                    } = await this._getSummarizationRequestAuth(model, streamFunction));
                 }
                 if (controller.signal.aborted) {
                     this._emit({
@@ -1801,7 +1783,7 @@ export class AgentSession {
                     return { compacted: false, shouldContinue: false };
                 }
                 // Generate compaction result
-                const compactResult = await compact(preparation, model, apiKey, headers, undefined, controller.signal, thinkingLevel, streamFunction, env, retrySettings, this._summarizationRetryCallbacks({ source: "compaction", reason }));
+                const compactResult = await compact(preparation, requestModel, apiKey, headers, undefined, controller.signal, thinkingLevel, streamFunction, env, retrySettings, this._summarizationRetryCallbacks({ source: "compaction", reason }));
                 summary = compactResult.summary;
                 firstKeptEntryId = compactResult.firstKeptEntryId;
                 tokensBefore = compactResult.tokensBefore;
@@ -1825,7 +1807,13 @@ export class AgentSession {
             if (willRetry) {
                 const messages = this.agent.state.messages;
                 const lastMsg = messages[messages.length - 1];
-                if (lastMsg?.role === "assistant" && lastMsg.stopReason === "error") {
+                // The overflow response was persisted on message_end before _checkCompaction() removed it
+                // from agent state. Rebuilding state from the new compaction can restore that kept entry,
+                // leaving an assistant as the final message. agent.continue() rejects that state, so remove
+                // the retriable error or truncated-length response again before continuing the interrupted turn.
+                if (lastMsg?.role === "assistant" &&
+                    (lastMsg.stopReason === "error" ||
+                        lastMsg.stopReason === "length")) {
                     this.agent.state.messages = messages.slice(0, -1);
                 }
             }
@@ -2219,8 +2207,10 @@ export class AgentSession {
         });
     }
     async reload(options) {
-        const previousFlagValues = this._extensionRunner.getFlagValues();
-        await emitSessionShutdownEvent(this._extensionRunner, { type: "session_shutdown", reason: "reload" });
+        const oldRunner = this._extensionRunner;
+        const previousFlagValues = oldRunner.getFlagValues();
+        await emitSessionShutdownEvent(oldRunner, { type: "session_shutdown", reason: "reload" });
+        oldRunner.invalidate();
         await this.settingsManager.reload();
         this.syncQueueModesFromSettings();
         resetApiProviders();
@@ -2539,10 +2529,10 @@ export class AgentSession {
             let summaryUsage;
             if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) {
                 const model = this.model;
-                const { apiKey, headers, env } = await this._getSummarizationRequestAuth(model);
+                const { model: requestModel, apiKey, headers, env } = await this._getSummarizationRequestAuth(model);
                 const branchSummarySettings = this.settingsManager.getBranchSummarySettings();
                 const result = await generateBranchSummary(entriesToSummarize, {
-                    model,
+                    model: requestModel,
                     apiKey,
                     headers,
                     env,