@mariozechner/pi-coding-agent 0.61.1 → 0.63.0

package/dist/core/agent-session.js

@@ -27,9 +27,10 @@ import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
 import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
 import { CURRENT_SESSION_VERSION, getLatestCompactionEntry } from "./session-manager.js";
-import { BUILTIN_SLASH_COMMANDS } from "./slash-commands.js";
+import { createSyntheticSourceInfo } from "./source-info.js";
 import { buildSystemPrompt } from "./system-prompt.js";
-import { createAllTools } from "./tools/index.js";
+import { createAllToolDefinitions } from "./tools/index.js";
+import { createToolDefinitionFromAgentTool, wrapToolDefinition } from "./tools/tool-definition-wrapper.js";
 /**
  * Parse a skill block from message text.
  * Returns null if the text doesn't contain a skill block.
@@ -89,7 +90,7 @@ export class AgentSession {
     _turnIndex = 0;
     _resourceLoader;
     _customTools;
-    _baseToolRegistry = new Map();
+    _baseToolDefinitions = new Map();
     _cwd;
     _extensionRunnerRef;
     _initialActiveToolNames;
@@ -103,6 +104,7 @@ export class AgentSession {
     _modelRegistry;
     // 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
@@ -132,6 +134,23 @@ export class AgentSession {
     get modelRegistry() {
         return this._modelRegistry;
     }
+    async _getRequiredRequestAuth(model) {
+        const result = await this._modelRegistry.getApiKeyAndHeaders(model);
+        if (!result.ok) {
+            throw new Error(result.error);
+        }
+        if (result.apiKey) {
+            return { apiKey: result.apiKey, headers: result.headers };
+        }
+        const isOAuth = this._modelRegistry.isUsingOAuth(model);
+        if (isOAuth) {
+            throw new Error(`Authentication failed for "${model.provider}". ` +
+                `Credentials may have expired or network is unavailable. ` +
+                `Run '/login ${model.provider}' to re-authenticate.`);
+        }
+        throw new Error(`No API key found for ${model.provider}.\n\n` +
+            `Use /login or set an API key environment variable. See ${join(getDocsPath(), "providers.md")}`);
+    }
     /**
      * Install tool hooks once on the Agent instance.
      *
@@ -493,15 +512,19 @@ export class AgentSession {
         return this.agent.state.tools.map((t) => t.name);
     }
     /**
-     * Get all configured tools with name, description, and parameter schema.
+     * Get all configured tools with name, description, parameter schema, and source metadata.
      */
     getAllTools() {
-        return Array.from(this._toolRegistry.values()).map((t) => ({
-            name: t.name,
-            description: t.description,
-            parameters: t.parameters,
+        return Array.from(this._toolDefinitions.values()).map(({ definition, sourceInfo }) => ({
+            name: definition.name,
+            description: definition.description,
+            parameters: definition.parameters,
+            sourceInfo,
         }));
     }
+    getToolDefinition(name) {
+        return this._toolDefinitions.get(name)?.definition;
+    }
     /**
      * Set active tools by name.
      * Only tools in the registry can be enabled. Unknown tool names are ignored.
@@ -680,9 +703,7 @@ export class AgentSession {
                 `Use /login or set an API key environment variable. See ${join(getDocsPath(), "providers.md")}\n\n` +
                 "Then use /model to select a model.");
         }
-        // Validate API key
-        const apiKey = await this._modelRegistry.getApiKey(this.model);
-        if (!apiKey) {
+        if (!this._modelRegistry.hasConfiguredAuth(this.model)) {
             const isOAuth = this._modelRegistry.isUsingOAuth(this.model);
             if (isOAuth) {
                 throw new Error(`Authentication failed for "${this.model.provider}". ` +
@@ -1059,12 +1080,11 @@ export class AgentSession {
     }
     /**
      * Set model directly.
-     * Validates API key, saves to session and settings.
-     * @throws Error if no API key available for the model
+     * Validates that auth is configured, saves to session and settings.
+     * @throws Error if no auth is configured for the model
      */
     async setModel(model) {
-        const apiKey = await this._modelRegistry.getApiKey(model);
-        if (!apiKey) {
+        if (!this._modelRegistry.hasConfiguredAuth(model)) {
             throw new Error(`No API key for ${model.provider}/${model.id}`);
         }
         const previousModel = this.model;
@@ -1088,27 +1108,11 @@ export class AgentSession {
         }
         return this._cycleAvailableModel(direction);
     }
-    async _getScopedModelsWithApiKey() {
-        const apiKeysByProvider = new Map();
-        const result = [];
-        for (const scoped of this._scopedModels) {
-            const provider = scoped.model.provider;
-            let apiKey;
-            if (apiKeysByProvider.has(provider)) {
-                apiKey = apiKeysByProvider.get(provider);
-            }
-            else {
-                apiKey = await this._modelRegistry.getApiKeyForProvider(provider);
-                apiKeysByProvider.set(provider, apiKey);
-            }
-            if (apiKey) {
-                result.push(scoped);
-            }
-        }
-        return result;
+    _getScopedModelsWithAuth() {
+        return this._scopedModels.filter((scoped) => this._modelRegistry.hasConfiguredAuth(scoped.model));
     }
     async _cycleScopedModel(direction) {
-        const scopedModels = await this._getScopedModelsWithApiKey();
+        const scopedModels = this._getScopedModelsWithAuth();
         if (scopedModels.length <= 1)
             return undefined;
         const currentModel = this.model;
@@ -1142,10 +1146,6 @@ export class AgentSession {
         const len = availableModels.length;
         const nextIndex = direction === "forward" ? (currentIndex + 1) % len : (currentIndex - 1 + len) % len;
         const nextModel = availableModels[nextIndex];
-        const apiKey = await this._modelRegistry.getApiKey(nextModel);
-        if (!apiKey) {
-            throw new Error(`No API key for ${nextModel.provider}/${nextModel.id}`);
-        }
         const thinkingLevel = this._getThinkingLevelForModelSwitch();
         this.agent.setModel(nextModel);
         this.sessionManager.appendModelChange(nextModel.provider, nextModel.id);
@@ -1274,10 +1274,7 @@ export class AgentSession {
             if (!this.model) {
                 throw new Error("No model selected");
             }
-            const apiKey = await this._modelRegistry.getApiKey(this.model);
-            if (!apiKey) {
-                throw new Error(`No API key for ${this.model.provider}`);
-            }
+            const { apiKey, headers } = await this._getRequiredRequestAuth(this.model);
             const pathEntries = this.sessionManager.getBranch();
             const settings = this.settingsManager.getCompactionSettings();
             const preparation = prepareCompaction(pathEntries, settings);
@@ -1320,7 +1317,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                const result = await compact(preparation, this.model, apiKey, customInstructions, this._compactionAbortController.signal);
+                const result = await compact(preparation, this.model, apiKey, headers, customInstructions, this._compactionAbortController.signal);
                 summary = result.summary;
                 firstKeptEntryId = result.firstKeptEntryId;
                 tokensBefore = result.tokensBefore;
@@ -1460,11 +1457,12 @@ export class AgentSession {
                 this._emit({ type: "auto_compaction_end", result: undefined, aborted: false, willRetry: false });
                 return;
             }
-            const apiKey = await this._modelRegistry.getApiKey(this.model);
-            if (!apiKey) {
+            const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model);
+            if (!authResult.ok || !authResult.apiKey) {
                 this._emit({ type: "auto_compaction_end", result: undefined, aborted: false, willRetry: false });
                 return;
             }
+            const { apiKey, headers } = authResult;
             const pathEntries = this.sessionManager.getBranch();
             const preparation = prepareCompaction(pathEntries, settings);
             if (!preparation) {
@@ -1503,7 +1501,7 @@ export class AgentSession {
             }
             else {
                 // Generate compaction result
-                const compactResult = await compact(preparation, this.model, apiKey, undefined, this._autoCompactionAbortController.signal);
+                const compactResult = await compact(preparation, this.model, apiKey, headers, undefined, this._autoCompactionAbortController.signal);
                 summary = compactResult.summary;
                 firstKeptEntryId = compactResult.firstKeptEntryId;
                 tokensBefore = compactResult.tokensBefore;
@@ -1656,36 +1654,24 @@ export class AgentSession {
         this.agent.setModel(refreshedModel);
     }
     _bindExtensionCore(runner) {
-        const normalizeLocation = (source) => {
-            if (source === "user" || source === "project" || source === "path") {
-                return source;
-            }
-            return undefined;
-        };
-        const reservedBuiltins = new Set(BUILTIN_SLASH_COMMANDS.map((command) => command.name));
         const getCommands = () => {
-            const extensionCommands = runner
-                .getRegisteredCommandsWithPaths()
-                .filter(({ command }) => !reservedBuiltins.has(command.name))
-                .map(({ command, extensionPath }) => ({
-                name: command.name,
+            const extensionCommands = runner.getRegisteredCommands().map((command) => ({
+                name: command.invocationName,
                 description: command.description,
                 source: "extension",
-                path: extensionPath,
+                sourceInfo: command.sourceInfo,
             }));
             const templates = this.promptTemplates.map((template) => ({
                 name: template.name,
                 description: template.description,
                 source: "prompt",
-                location: normalizeLocation(template.source),
-                path: template.filePath,
+                sourceInfo: template.sourceInfo,
             }));
             const skills = this._resourceLoader.getSkills().skills.map((skill) => ({
                 name: `skill:${skill.name}`,
                 description: skill.description,
                 source: "skill",
-                location: normalizeLocation(skill.source),
-                path: skill.filePath,
+                sourceInfo: skill.sourceInfo,
             }));
             return [...extensionCommands, ...templates, ...skills];
         };
@@ -1726,8 +1712,7 @@ export class AgentSession {
             refreshTools: () => this._refreshToolRegistry(),
             getCommands,
             setModel: async (model) => {
-                const key = await this.modelRegistry.getApiKey(model);
-                if (!key)
+                if (!this.modelRegistry.hasConfiguredAuth(model))
                     return false;
                 await this.setModel(model);
                 return true;
@@ -1773,24 +1758,44 @@ export class AgentSession {
         const registeredTools = this._extensionRunner?.getAllRegisteredTools() ?? [];
         const allCustomTools = [
             ...registeredTools,
-            ...this._customTools.map((def) => ({ definition: def, extensionPath: "<sdk>" })),
+            ...this._customTools.map((definition) => ({
+                definition,
+                sourceInfo: createSyntheticSourceInfo(`<sdk:${definition.name}>`, { source: "sdk" }),
+            })),
         ];
-        this._toolPromptSnippets = new Map(allCustomTools
-            .map((registeredTool) => {
-            const snippet = this._normalizePromptSnippet(registeredTool.definition.promptSnippet);
-            return snippet ? [registeredTool.definition.name, snippet] : undefined;
+        const definitionRegistry = new Map(Array.from(this._baseToolDefinitions.entries()).map(([name, definition]) => [
+            name,
+            {
+                definition,
+                sourceInfo: createSyntheticSourceInfo(`<builtin:${name}>`, { source: "builtin" }),
+            },
+        ]));
+        for (const tool of allCustomTools) {
+            definitionRegistry.set(tool.definition.name, {
+                definition: tool.definition,
+                sourceInfo: tool.sourceInfo,
+            });
+        }
+        this._toolDefinitions = definitionRegistry;
+        this._toolPromptSnippets = new Map(Array.from(definitionRegistry.values())
+            .map(({ definition }) => {
+            const snippet = this._normalizePromptSnippet(definition.promptSnippet);
+            return snippet ? [definition.name, snippet] : undefined;
         })
             .filter((entry) => entry !== undefined));
-        this._toolPromptGuidelines = new Map(allCustomTools
-            .map((registeredTool) => {
-            const guidelines = this._normalizePromptGuidelines(registeredTool.definition.promptGuidelines);
-            return guidelines.length > 0 ? [registeredTool.definition.name, guidelines] : undefined;
+        this._toolPromptGuidelines = new Map(Array.from(definitionRegistry.values())
+            .map(({ definition }) => {
+            const guidelines = this._normalizePromptGuidelines(definition.promptGuidelines);
+            return guidelines.length > 0 ? [definition.name, guidelines] : undefined;
         })
             .filter((entry) => entry !== undefined));
         const wrappedExtensionTools = this._extensionRunner
             ? wrapRegisteredTools(allCustomTools, this._extensionRunner)
             : [];
-        const toolRegistry = new Map(this._baseToolRegistry);
+        const toolRegistry = new Map(Array.from(this._baseToolDefinitions.values()).map((definition) => [
+            definition.name,
+            wrapToolDefinition(definition),
+        ]));
         for (const tool of wrappedExtensionTools) {
             toolRegistry.set(tool.name, tool);
         }
@@ -1815,13 +1820,16 @@ export class AgentSession {
     _buildRuntime(options) {
         const autoResizeImages = this.settingsManager.getImageAutoResize();
         const shellCommandPrefix = this.settingsManager.getShellCommandPrefix();
-        const baseTools = this._baseToolsOverride
-            ? this._baseToolsOverride
-            : createAllTools(this._cwd, {
+        const baseToolDefinitions = this._baseToolsOverride
+            ? Object.fromEntries(Object.entries(this._baseToolsOverride).map(([name, tool]) => [
+                name,
+                createToolDefinitionFromAgentTool(tool),
+            ]))
+            : createAllToolDefinitions(this._cwd, {
                 read: { autoResizeImages },
                 bash: { commandPrefix: shellCommandPrefix },
             });
-        this._baseToolRegistry = new Map(Object.entries(baseTools).map(([name, tool]) => [name, tool]));
+        this._baseToolDefinitions = new Map(Object.entries(baseToolDefinitions).map(([name, tool]) => [name, tool]));
         const extensionsResult = this._resourceLoader.getExtensions();
         if (options.flagValues) {
             for (const [name, value] of options.flagValues) {
@@ -2285,14 +2293,12 @@ export class AgentSession {
         let summaryDetails;
         if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) {
             const model = this.model;
-            const apiKey = await this._modelRegistry.getApiKey(model);
-            if (!apiKey) {
-                throw new Error(`No API key for ${model.provider}`);
-            }
+            const { apiKey, headers } = await this._getRequiredRequestAuth(model);
             const branchSummarySettings = this.settingsManager.getBranchSummarySettings();
             const result = await generateBranchSummary(entriesToSummarize, {
                 model,
                 apiKey,
+                headers,
                 signal: this._branchSummaryAbortController.signal,
                 customInstructions,
                 replaceInstructions,
@@ -2449,6 +2455,7 @@ export class AgentSession {
                 total: totalInput + totalOutput + totalCacheRead + totalCacheWrite,
             },
             cost: totalCost,
+            contextUsage: this.getContextUsage(),
         };
     }
     getContextUsage() {
@@ -2500,13 +2507,10 @@ export class AgentSession {
     async exportToHtml(outputPath) {
         const themeName = this.settingsManager.getTheme();
         // Create tool renderer if we have an extension runner (for custom tool HTML rendering)
-        let toolRenderer;
-        if (this._extensionRunner) {
-            toolRenderer = createToolHtmlRenderer({
-                getToolDefinition: (name) => this._extensionRunner.getToolDefinition(name),
-                theme,
-            });
-        }
+        const toolRenderer = createToolHtmlRenderer({
+            getToolDefinition: (name) => this.getToolDefinition(name),
+            theme,
+        });
         return await exportSessionToHtml(this.sessionManager, this.state, {
             outputPath,
             themeName,