@selesai/code 0.5.13 → 0.5.15

package/dist/core/agent-session.js

@@ -14,6 +14,7 @@
  */
 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 { getThemeByName, theme } from "../modes/interactive/theme/theme.js";
 import { stripFrontmatter } from "../utils/frontmatter.js";
@@ -28,6 +29,7 @@ import { exportSessionToHtml } from "./export-html/index.js";
 import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
 import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { emitSessionShutdownEvent } from "./extensions/runner.js";
+import { ModelRegistry } from "./model-registry.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
 import { CURRENT_SESSION_VERSION, getLatestCompactionEntry } from "./session-manager.js";
 import { createSyntheticSourceInfo } from "./source-info.js";
@@ -35,6 +37,7 @@ import { buildSystemPrompt } from "./system-prompt.js";
 import { createLocalBashOperations } from "./tools/bash.js";
 import { createAllToolDefinitions } from "./tools/index.js";
 import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.js";
+import { addUsageToTotals, createUsageTotals } from "./usage-totals.js";
 /**
  * Parse a skill block from message text.
  * Returns null if the text doesn't contain a skill block.
@@ -50,6 +53,14 @@ export function parseSkillBlock(text) {
         userMessage: match[4]?.trim() || undefined,
     };
 }
+// ============================================================================
+// Types
+// ============================================================================
+function withoutDeletedHeaders(headers) {
+    return headers
+        ? Object.fromEntries(Object.entries(headers).filter((entry) => entry[1] !== null))
+        : undefined;
+}
 function estimateMessagesTokens(messages) {
     let tokens = 0;
     for (const message of messages) {
@@ -114,8 +125,7 @@ export class AgentSession {
     _extensionShutdownHandler;
     _extensionErrorListener;
     _extensionErrorUnsubscriber;
-    // Model registry for API key resolution
-    _modelRegistry;
+    _modelRuntime;
     // Tool registry for extension getTools/setTools
     _toolRegistry = new Map();
     _toolDefinitions = new Map();
@@ -133,7 +143,7 @@ export class AgentSession {
         this._resourceLoader = config.resourceLoader;
         this._customTools = config.customTools ?? [];
         this._cwd = config.cwd;
-        this._modelRegistry = config.modelRegistry;
+        this._modelRuntime = config.modelRuntime;
         this._extensionRunnerRef = config.extensionRunnerRef;
         this._initialActiveToolNames = config.initialActiveToolNames;
         this._allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined;
@@ -150,22 +160,29 @@ export class AgentSession {
             includeAllExtensionTools: true,
         });
     }
-    /** Model registry for API key resolution and model discovery */
-    get modelRegistry() {
-        return this._modelRegistry;
+    get modelRuntime() {
+        return this._modelRuntime;
     }
     async _getRequiredRequestAuth(model) {
-        const result = await this._modelRegistry.getApiKeyAndHeaders(model);
-        if (!result.ok) {
-            if (result.error.startsWith("No API key found")) {
+        let result;
+        try {
+            result = await this._modelRuntime.getAuth(model);
+        }
+        catch (error) {
+            const cause = error instanceof Error ? error.cause : undefined;
+            if (cause instanceof Error && cause.message === "authHeader requires a resolved API key") {
                 throw new Error(formatNoApiKeyFoundMessage(model.provider));
             }
-            throw new Error(result.error);
+            throw error;
         }
-        if (result.apiKey) {
-            return { apiKey: result.apiKey, headers: result.headers, env: result.env };
+        if (result?.auth.apiKey) {
+            return {
+                apiKey: result.auth.apiKey,
+                headers: withoutDeletedHeaders(result.auth.headers),
+                env: result.env,
+            };
         }
-        const isOAuth = this._modelRegistry.isUsingOAuth(model);
+        const isOAuth = this._modelRuntime.isUsingOAuth(model.provider);
         if (isOAuth) {
             throw new Error(`Authentication failed for "${model.provider}". ` +
                 `Credentials may have expired or network is unavailable. ` +
@@ -173,12 +190,19 @@ export class AgentSession {
         }
         throw new Error(formatNoApiKeyFoundMessage(model.provider));
     }
-    async _getCompactionRequestAuth(model) {
-        if (this.agent.streamFn === streamSimple) {
+    async _getSummarizationRequestAuth(model) {
+        if (this.agent.streamFunction === streamSimple) {
             return this._getRequiredRequestAuth(model);
         }
-        const result = await this._modelRegistry.getApiKeyAndHeaders(model);
-        return result.ok ? { apiKey: result.apiKey, headers: result.headers, env: result.env } : {};
+        try {
+            const result = await this._modelRuntime.getAuth(model);
+            return result
+                ? { apiKey: result.auth.apiKey, headers: withoutDeletedHeaders(result.auth.headers), env: result.env }
+                : {};
+        }
+        catch {
+            return {};
+        }
     }
     /**
      * Install tool hooks once on the Agent instance.
@@ -222,6 +246,7 @@ export class AgentSession {
                 content: result.content,
                 details: result.details,
                 isError,
+                usage: result.usage,
             });
             if (!hookResult) {
                 return undefined;
@@ -230,6 +255,7 @@ export class AgentSession {
                 content: hookResult.content,
                 details: hookResult.details,
                 isError: hookResult.isError ?? isError,
+                usage: hookResult.usage,
             };
         };
     }
@@ -341,7 +367,7 @@ export class AgentSession {
         // 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);
+            const messageText = contentText(event.message.content, "");
             if (messageText) {
                 // Check steering queue first
                 const steeringIndex = this._steeringMessages.indexOf(messageText);
@@ -410,16 +436,6 @@ export class AgentSession {
         }
         return false;
     }
-    /** Extract text content from a message */
-    _getUserMessageText(message) {
-        if (message.role !== "user")
-            return "";
-        const content = message.content;
-        if (typeof content === "string")
-            return content;
-        const textBlocks = content.filter((c) => c.type === "text");
-        return textBlocks.map((c) => c.text).join("");
-    }
     /** Find the last assistant message in agent state (including aborted ones) */
     _findLastAssistantMessage() {
         const messages = this.agent.state.messages;
@@ -875,8 +891,10 @@ export class AgentSession {
             if (!this.model) {
                 throw new Error(formatNoModelSelectedMessage());
             }
-            if (!this._modelRegistry.hasConfiguredAuth(this.model)) {
-                const isOAuth = this._modelRegistry.isUsingOAuth(this.model);
+            const hasConfiguredAuth = this._modelRuntime.hasConfiguredAuth(this.model.provider) ||
+                (await this._modelRuntime.checkAuth(this.model.provider)) !== undefined;
+            if (!hasConfiguredAuth) {
+                const isOAuth = this._modelRuntime.isUsingOAuth(this.model.provider);
                 if (isOAuth) {
                     throw new Error(`Authentication failed for "${this.model.provider}". ` +
                         `Credentials may have expired or network is unavailable. ` +
@@ -1221,7 +1239,7 @@ export class AgentSession {
      * @throws Error if no auth is configured for the model
      */
     async setModel(model) {
-        if (!this._modelRegistry.hasConfiguredAuth(model)) {
+        if (!(await this._modelRuntime.checkAuth(model.provider))) {
             throw new Error(`No API key for ${model.provider}/${model.id}`);
         }
         const previousModel = this.model;
@@ -1246,7 +1264,11 @@ export class AgentSession {
         return this._cycleAvailableModel(direction);
     }
     async _cycleScopedModel(direction) {
-        const scopedModels = this._scopedModels.filter((scoped) => this._modelRegistry.hasConfiguredAuth(scoped.model));
+        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);
         if (scopedModels.length <= 1)
             return undefined;
         const currentModel = this.model;
@@ -1270,7 +1292,7 @@ export class AgentSession {
         return { model: next.model, thinkingLevel: this.thinkingLevel, isScoped: true };
     }
     async _cycleAvailableModel(direction) {
-        const availableModels = await this._modelRegistry.getAvailable();
+        const availableModels = await this._modelRuntime.getAvailable();
         if (availableModels.length <= 1)
             return undefined;
         const currentModel = this.model;
@@ -1398,7 +1420,7 @@ export class AgentSession {
             if (!this.model) {
                 throw new Error(formatNoModelSelectedMessage());
             }
-            const { apiKey, headers, env } = await this._getCompactionRequestAuth(this.model);
+            const { apiKey, headers, env } = await this._getSummarizationRequestAuth(this.model);
             const pathEntries = this.sessionManager.getBranch();
             const settings = this.settingsManager.getCompactionSettings();
             const preparation = prepareCompaction(pathEntries, settings);
@@ -1433,26 +1455,29 @@ export class AgentSession {
             let summary;
             let firstKeptEntryId;
             let tokensBefore;
+            let usage;
             let details;
             if (extensionCompaction) {
                 // Extension provided compaction content
                 summary = extensionCompaction.summary;
                 firstKeptEntryId = extensionCompaction.firstKeptEntryId;
                 tokensBefore = extensionCompaction.tokensBefore;
+                usage = extensionCompaction.usage;
                 details = extensionCompaction.details;
             }
             else {
                 // Generate compaction result
-                const result = await compact(preparation, this.model, apiKey, headers, customInstructions, this._compactionAbortController.signal, this.thinkingLevel, this.agent.streamFn, env);
+                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" }));
                 summary = result.summary;
                 firstKeptEntryId = result.firstKeptEntryId;
                 tokensBefore = result.tokensBefore;
+                usage = result.usage;
                 details = result.details;
             }
             if (this._compactionAbortController.signal.aborted) {
                 throw new Error("Compaction cancelled");
             }
-            this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension);
+            this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage);
             const newEntries = this.sessionManager.getEntries();
             const sessionContext = this.sessionManager.buildSessionContext();
             this.agent.state.messages = sessionContext.messages;
@@ -1473,6 +1498,7 @@ export class AgentSession {
                 firstKeptEntryId,
                 tokensBefore,
                 estimatedTokensAfter,
+                usage,
                 details,
             };
             this._emit({
@@ -1619,17 +1645,16 @@ export class AgentSession {
             let apiKey;
             let headers;
             let env;
-            if (this.agent.streamFn === streamSimple) {
-                const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model);
-                if (!authResult.ok || !authResult.apiKey) {
+            if (this.agent.streamFunction === streamSimple) {
+                const authResult = await this._modelRuntime.getAuth(this.model);
+                if (!authResult?.auth.apiKey)
                     return false;
-                }
-                apiKey = authResult.apiKey;
-                headers = authResult.headers;
+                apiKey = authResult.auth.apiKey;
+                headers = withoutDeletedHeaders(authResult.auth.headers);
                 env = authResult.env;
             }
             else {
-                ({ apiKey, headers, env } = await this._getCompactionRequestAuth(this.model));
+                ({ apiKey, headers, env } = await this._getSummarizationRequestAuth(this.model));
             }
             const pathEntries = this.sessionManager.getBranch();
             const preparation = prepareCompaction(pathEntries, settings);
@@ -1669,20 +1694,23 @@ export class AgentSession {
             let summary;
             let firstKeptEntryId;
             let tokensBefore;
+            let usage;
             let details;
             if (extensionCompaction) {
                 // Extension provided compaction content
                 summary = extensionCompaction.summary;
                 firstKeptEntryId = extensionCompaction.firstKeptEntryId;
                 tokensBefore = extensionCompaction.tokensBefore;
+                usage = extensionCompaction.usage;
                 details = extensionCompaction.details;
             }
             else {
                 // Generate compaction result
-                const compactResult = await compact(preparation, this.model, apiKey, headers, undefined, this._autoCompactionAbortController.signal, this.thinkingLevel, this.agent.streamFn, env);
+                const compactResult = await compact(preparation, this.model, apiKey, headers, undefined, this._autoCompactionAbortController.signal, this.thinkingLevel, this.agent.streamFunction, env, this.settingsManager.getRetrySettings(), this._summarizationRetryCallbacks({ source: "compaction", reason }));
                 summary = compactResult.summary;
                 firstKeptEntryId = compactResult.firstKeptEntryId;
                 tokensBefore = compactResult.tokensBefore;
+                usage = compactResult.usage;
                 details = compactResult.details;
             }
             if (this._autoCompactionAbortController.signal.aborted) {
@@ -1695,7 +1723,7 @@ export class AgentSession {
                 });
                 return false;
             }
-            this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension);
+            this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage);
             const newEntries = this.sessionManager.getEntries();
             const sessionContext = this.sessionManager.buildSessionContext();
             this.agent.state.messages = sessionContext.messages;
@@ -1716,6 +1744,7 @@ export class AgentSession {
                 firstKeptEntryId,
                 tokensBefore,
                 estimatedTokensAfter,
+                usage,
                 details,
             };
             this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry });
@@ -1843,7 +1872,7 @@ export class AgentSession {
         if (!currentModel) {
             return;
         }
-        const refreshedModel = this._modelRegistry.find(currentModel.provider, currentModel.id);
+        const refreshedModel = this._modelRuntime.getModel(currentModel.provider, currentModel.id);
         if (!refreshedModel || refreshedModel === currentModel) {
             return;
         }
@@ -1912,7 +1941,7 @@ export class AgentSession {
             refreshTools: () => this._refreshToolRegistry(),
             getCommands,
             setModel: async (model) => {
-                if (!this.modelRegistry.hasConfiguredAuth(model))
+                if (!this._modelRuntime.hasConfiguredAuth(model.provider))
                     return false;
                 await this.setModel(model);
                 return true;
@@ -1952,11 +1981,15 @@ export class AgentSession {
             getSystemPromptOptions: () => this._baseSystemPromptOptions,
         }, {
             registerProvider: (name, config) => {
-                this._modelRegistry.registerProvider(name, config);
+                this._modelRuntime.registerProvider(name, config);
+                this._refreshCurrentModelFromRegistry();
+            },
+            registerNativeProvider: (provider) => {
+                this._modelRuntime.registerNativeProvider(provider);
                 this._refreshCurrentModelFromRegistry();
             },
             unregisterProvider: (name) => {
-                this._modelRegistry.unregisterProvider(name);
+                this._modelRuntime.unregisterProvider(name);
                 this._refreshCurrentModelFromRegistry();
             },
         });
@@ -2058,7 +2091,7 @@ export class AgentSession {
                 extensionsResult.runtime.flagValues.set(name, value);
             }
         }
-        this._extensionRunner = new ExtensionRunner(extensionsResult.extensions, extensionsResult.runtime, this._cwd, this.sessionManager, this._modelRegistry);
+        this._extensionRunner = new ExtensionRunner(extensionsResult.extensions, extensionsResult.runtime, this._cwd, this.sessionManager, new ModelRegistry(this._modelRuntime));
         if (this._extensionRunnerRef) {
             this._extensionRunnerRef.current = this._extensionRunner;
         }
@@ -2108,6 +2141,34 @@ export class AgentSession {
             return false;
         return isRetryableAssistantError(message);
     }
+    /**
+     * Retry policy + callbacks shared by compaction and branch-summary summarization calls.
+     * Uses the same `settings.retry` budget/backoff as agent-turn retries so a single transient
+     * stream drop no longer fails the whole operation. `source` carries the context
+     * the TUI needs to render the retry and recreate the underlying indicator.
+     */
+    _summarizationRetryCallbacks(source) {
+        return {
+            onRetryScheduled: (attempt, maxAttempts, delayMs, errorMessage) => {
+                this._emit({
+                    type: "summarization_retry_scheduled",
+                    attempt,
+                    maxAttempts,
+                    delayMs,
+                    errorMessage,
+                });
+            },
+            onRetryAttemptStart: () => {
+                this._emit({
+                    type: "summarization_retry_attempt_start",
+                    ...source,
+                });
+            },
+            onRetryFinished: () => {
+                this._emit({ type: "summarization_retry_finished" });
+            },
+        };
+    }
     /**
      * Prepare a retryable error for continuation with exponential backoff.
      * @returns true if the caller should continue the agent, false otherwise
@@ -2353,9 +2414,10 @@ export class AgentSession {
             // Run default summarizer if needed
             let summaryText;
             let summaryDetails;
+            let summaryUsage;
             if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) {
                 const model = this.model;
-                const { apiKey, headers, env } = await this._getRequiredRequestAuth(model);
+                const { apiKey, headers, env } = await this._getSummarizationRequestAuth(model);
                 const branchSummarySettings = this.settingsManager.getBranchSummarySettings();
                 const result = await generateBranchSummary(entriesToSummarize, {
                     model,
@@ -2366,7 +2428,9 @@ export class AgentSession {
                     customInstructions,
                     replaceInstructions,
                     reserveTokens: branchSummarySettings.reserveTokens,
-                    streamFn: this.agent.streamFn,
+                    streamFn: this.agent.streamFunction,
+                    retry: this.settingsManager.getRetrySettings(),
+                    callbacks: this._summarizationRetryCallbacks({ source: "branchSummary" }),
                 });
                 if (result.aborted) {
                     return { cancelled: true, aborted: true };
@@ -2375,6 +2439,7 @@ export class AgentSession {
                     throw new Error(result.error);
                 }
                 summaryText = result.summary;
+                summaryUsage = result.usage;
                 summaryDetails = {
                     readFiles: result.readFiles || [],
                     modifiedFiles: result.modifiedFiles || [],
@@ -2383,6 +2448,7 @@ export class AgentSession {
             else if (extensionSummary) {
                 summaryText = extensionSummary.summary;
                 summaryDetails = extensionSummary.details;
+                summaryUsage = extensionSummary.usage;
             }
             // Determine the new leaf position based on target type
             let newLeafId;
@@ -2390,18 +2456,12 @@ export class AgentSession {
             if (targetEntry.type === "message" && targetEntry.message.role === "user") {
                 // User message: leaf = parent (null if root), text goes to editor
                 newLeafId = targetEntry.parentId;
-                editorText = this._extractUserMessageText(targetEntry.message.content);
+                editorText = contentText(targetEntry.message.content, "");
             }
             else if (targetEntry.type === "custom_message") {
                 // Custom message: leaf = parent (null if root), text goes to editor
                 newLeafId = targetEntry.parentId;
-                editorText =
-                    typeof targetEntry.content === "string"
-                        ? targetEntry.content
-                        : targetEntry.content
-                            .filter((c) => c.type === "text")
-                            .map((c) => c.text)
-                            .join("");
+                editorText = contentText(targetEntry.content, "");
             }
             else {
                 // Non-user message: leaf = selected node
@@ -2412,7 +2472,7 @@ export class AgentSession {
             let summaryEntry;
             if (summaryText) {
                 // Create summary at target position (can be null for root)
-                const summaryId = this.sessionManager.branchWithSummary(newLeafId, summaryText, summaryDetails, fromExtension);
+                const summaryId = this.sessionManager.branchWithSummary(newLeafId, summaryText, summaryDetails, fromExtension, summaryUsage);
                 summaryEntry = this.sessionManager.getEntry(summaryId);
                 // Attach label to the summary entry
                 if (label) {
@@ -2460,24 +2520,13 @@ export class AgentSession {
                 continue;
             if (entry.message.role !== "user")
                 continue;
-            const text = this._extractUserMessageText(entry.message.content);
+            const text = contentText(entry.message.content, "");
             if (text) {
                 result.push({ entryId: entry.id, text });
             }
         }
         return result;
     }
-    _extractUserMessageText(content) {
-        if (typeof content === "string")
-            return content;
-        if (Array.isArray(content)) {
-            return content
-                .filter((c) => c.type === "text")
-                .map((c) => c.text)
-                .join("");
-        }
-        return "";
-    }
     /**
      * Get session statistics. Aggregates over ALL session entries (including
      * history that was compacted away), so token/cost totals reflect what was
@@ -2489,12 +2538,11 @@ export class AgentSession {
         let toolResults = 0;
         let totalMessages = 0;
         let toolCalls = 0;
-        let totalInput = 0;
-        let totalOutput = 0;
-        let totalCacheRead = 0;
-        let totalCacheWrite = 0;
-        let totalCost = 0;
+        const usageTotals = createUsageTotals();
         for (const entry of this.sessionManager.getEntries()) {
+            if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) {
+                addUsageToTotals(usageTotals, entry.usage);
+            }
             if (entry.type !== "message")
                 continue;
             totalMessages++;
@@ -2504,6 +2552,9 @@ export class AgentSession {
             }
             else if (message.role === "toolResult") {
                 toolResults++;
+                if (message.usage) {
+                    addUsageToTotals(usageTotals, message.usage);
+                }
             }
             else if (message.role === "assistant") {
                 assistantMessages++;
@@ -2511,12 +2562,7 @@ export class AgentSession {
                 if (Array.isArray(assistantMsg.content)) {
                     toolCalls += assistantMsg.content.filter((c) => c.type === "toolCall").length;
                 }
-                const usage = assistantMsg.usage;
-                totalInput += usage.input;
-                totalOutput += usage.output;
-                totalCacheRead += usage.cacheRead;
-                totalCacheWrite += usage.cacheWrite;
-                totalCost += usage.cost.total;
+                addUsageToTotals(usageTotals, assistantMsg.usage);
             }
         }
         return {
@@ -2528,13 +2574,13 @@ export class AgentSession {
             toolResults,
             totalMessages,
             tokens: {
-                input: totalInput,
-                output: totalOutput,
-                cacheRead: totalCacheRead,
-                cacheWrite: totalCacheWrite,
-                total: totalInput + totalOutput + totalCacheRead + totalCacheWrite,
+                input: usageTotals.input,
+                output: usageTotals.output,
+                cacheRead: usageTotals.cacheRead,
+                cacheWrite: usageTotals.cacheWrite,
+                total: usageTotals.input + usageTotals.output + usageTotals.cacheRead + usageTotals.cacheWrite,
             },
-            cost: totalCost,
+            cost: usageTotals.cost,
             contextUsage: this.getContextUsage(),
         };
     }