@mariozechner/pi-coding-agent 0.49.3 → 0.50.1

package/dist/core/agent-session.js

@@ -13,15 +13,35 @@
  * Modes use this class and add their own I/O layer on top.
  */
 import { readFileSync } from "node:fs";
-import { isContextOverflow, modelsAreEqual, supportsXhigh } from "@mariozechner/pi-ai";
-import { getAuthPath } from "../config.js";
+import { join } from "node:path";
+import { isContextOverflow, modelsAreEqual, resetApiProviders, supportsXhigh } from "@mariozechner/pi-ai";
+import { getDocsPath } from "../config.js";
 import { theme } from "../modes/interactive/theme/theme.js";
 import { stripFrontmatter } from "../utils/frontmatter.js";
+import { sleep } from "../utils/sleep.js";
 import { executeBash as executeBashCommand, executeBashWithOperations } from "./bash-executor.js";
 import { calculateContextTokens, collectEntriesForBranchSummary, compact, estimateContextTokens, generateBranchSummary, prepareCompaction, shouldCompact, } from "./compaction/index.js";
 import { exportSessionToHtml } from "./export-html/index.js";
 import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
+import { ExtensionRunner, wrapRegisteredTools, wrapToolsWithExtensions, } from "./extensions/index.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
+import { buildSystemPrompt } from "./system-prompt.js";
+import { createAllTools } from "./tools/index.js";
+/**
+ * Parse a skill block from message text.
+ * Returns null if the text doesn't contain a skill block.
+ */
+export function parseSkillBlock(text) {
+    const match = text.match(/^<skill name="([^"]+)" location="([^"]+)">\n([\s\S]*?)\n<\/skill>(?:\n\n([\s\S]+))?$/);
+    if (!match)
+        return null;
+    return {
+        name: match[1],
+        location: match[2],
+        content: match[3],
+        userMessage: match[4]?.trim() || undefined,
+    };
+}
 // ============================================================================
 // Constants
 // ============================================================================
@@ -37,7 +57,6 @@ export class AgentSession {
     sessionManager;
     settingsManager;
     _scopedModels;
-    _promptTemplates;
     // Event subscription state
     _unsubscribeAgent;
     _eventListeners = [];
@@ -63,34 +82,43 @@ export class AgentSession {
     // Extension system
     _extensionRunner = undefined;
     _turnIndex = 0;
-    _skills;
-    _skillWarnings;
-    _skillsSettings;
+    _resourceLoader;
+    _customTools;
+    _baseToolRegistry = new Map();
+    _cwd;
+    _extensionRunnerRef;
+    _initialActiveToolNames;
+    _baseToolsOverride;
+    _extensionUIContext;
+    _extensionCommandContextActions;
+    _extensionShutdownHandler;
+    _extensionErrorListener;
+    _extensionErrorUnsubscriber;
     // Model registry for API key resolution
     _modelRegistry;
     // Tool registry for extension getTools/setTools
-    _toolRegistry;
-    // Function to rebuild system prompt when tools change
-    _rebuildSystemPrompt;
+    _toolRegistry = new Map();
     // Base system prompt (without extension appends) - used to apply fresh appends each turn
-    _baseSystemPrompt;
+    _baseSystemPrompt = "";
     constructor(config) {
         this.agent = config.agent;
         this.sessionManager = config.sessionManager;
         this.settingsManager = config.settingsManager;
         this._scopedModels = config.scopedModels ?? [];
-        this._promptTemplates = config.promptTemplates ?? [];
-        this._extensionRunner = config.extensionRunner;
-        this._skills = config.skills ?? [];
-        this._skillWarnings = config.skillWarnings ?? [];
-        this._skillsSettings = config.skillsSettings;
+        this._resourceLoader = config.resourceLoader;
+        this._customTools = config.customTools ?? [];
+        this._cwd = config.cwd;
         this._modelRegistry = config.modelRegistry;
-        this._toolRegistry = config.toolRegistry ?? new Map();
-        this._rebuildSystemPrompt = config.rebuildSystemPrompt;
-        this._baseSystemPrompt = config.agent.state.systemPrompt;
+        this._extensionRunnerRef = config.extensionRunnerRef;
+        this._initialActiveToolNames = config.initialActiveToolNames;
+        this._baseToolsOverride = config.baseToolsOverride;
         // Always subscribe to agent events for internal handling
         // (session persistence, extensions, auto-compaction, retry logic)
         this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
+        this._buildRuntime({
+            activeToolNames: this._initialActiveToolNames,
+            includeAllExtensionTools: true,
+        });
     }
     /** Model registry for API key resolution and model discovery */
     get modelRegistry() {
@@ -334,10 +362,8 @@ export class AgentSession {
         }
         this.agent.setTools(tools);
         // Rebuild base system prompt with new tool set
-        if (this._rebuildSystemPrompt) {
-            this._baseSystemPrompt = this._rebuildSystemPrompt(validToolNames);
-            this.agent.setSystemPrompt(this._baseSystemPrompt);
-        }
+        this._baseSystemPrompt = this._rebuildSystemPrompt(validToolNames);
+        this.agent.setSystemPrompt(this._baseSystemPrompt);
     }
     /** Whether auto-compaction is currently running */
     get isCompacting() {
@@ -373,7 +399,23 @@ export class AgentSession {
     }
     /** File-based prompt templates */
     get promptTemplates() {
-        return this._promptTemplates;
+        return this._resourceLoader.getPrompts().prompts;
+    }
+    _rebuildSystemPrompt(toolNames) {
+        const validToolNames = toolNames.filter((name) => this._baseToolRegistry.has(name));
+        const loaderSystemPrompt = this._resourceLoader.getSystemPrompt();
+        const loaderAppendSystemPrompt = this._resourceLoader.getAppendSystemPrompt();
+        const appendSystemPrompt = loaderAppendSystemPrompt.length > 0 ? loaderAppendSystemPrompt.join("\n\n") : undefined;
+        const loadedSkills = this._resourceLoader.getSkills().skills;
+        const loadedContextFiles = this._resourceLoader.getAgentsFiles().agentsFiles;
+        return buildSystemPrompt({
+            cwd: this._cwd,
+            skills: loadedSkills,
+            contextFiles: loadedContextFiles,
+            customPrompt: loaderSystemPrompt,
+            appendSystemPrompt,
+            selectedTools: validToolNames,
+        });
     }
     // =========================================================================
     // Prompting
@@ -415,7 +457,7 @@ export class AgentSession {
         let expandedText = currentText;
         if (expandPromptTemplates) {
             expandedText = this._expandSkillCommand(expandedText);
-            expandedText = expandPromptTemplate(expandedText, [...this._promptTemplates]);
+            expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]);
         }
         // If streaming, queue via steer() or followUp() based on option
         if (this.isStreaming) {
@@ -435,7 +477,7 @@ export class AgentSession {
         // Validate model
         if (!this.model) {
             throw new Error("No model selected.\n\n" +
-                `Use /login, set an API key environment variable, or create ${getAuthPath()}\n\n` +
+                `Use /login or set an API key environment variable. See ${join(getDocsPath(), "authentication.md")}\n\n` +
                 "Then use /model to select a model.");
         }
         // Validate API key
@@ -448,7 +490,7 @@ export class AgentSession {
                     `Run '/login ${this.model.provider}' to re-authenticate.`);
             }
             throw new Error(`No API key found for ${this.model.provider}.\n\n` +
-                `Use /login, set an API key environment variable, or create ${getAuthPath()}`);
+                `Use /login or set an API key environment variable. See ${join(getDocsPath(), "authentication.md")}`);
         }
         // Check if we need to compact before sending (catches aborted responses)
         const lastAssistant = this._findLastAssistantMessage();
@@ -540,15 +582,14 @@ export class AgentSession {
         const spaceIndex = text.indexOf(" ");
         const skillName = spaceIndex === -1 ? text.slice(7) : text.slice(7, spaceIndex);
         const args = spaceIndex === -1 ? "" : text.slice(spaceIndex + 1).trim();
-        const skill = this._skills.find((s) => s.name === skillName);
+        const skill = this.resourceLoader.getSkills().skills.find((s) => s.name === skillName);
         if (!skill)
             return text; // Unknown skill, pass through
         try {
             const content = readFileSync(skill.filePath, "utf-8");
             const body = stripFrontmatter(content).trim();
-            const header = `Skill location: ${skill.filePath}\nReferences are relative to ${skill.baseDir}.`;
-            const skillMessage = `${header}\n\n${body}`;
-            return args ? `${skillMessage}\n\n---\n\nUser: ${args}` : skillMessage;
+            const skillBlock = `<skill name="${skill.name}" location="${skill.filePath}">\nReferences are relative to ${skill.baseDir}.\n\n${body}\n</skill>`;
+            return args ? `${skillBlock}\n\n${args}` : skillBlock;
         }
         catch (err) {
             // Emit error like extension commands do
@@ -573,7 +614,7 @@ export class AgentSession {
         }
         // Expand skill commands and prompt templates
         let expandedText = this._expandSkillCommand(text);
-        expandedText = expandPromptTemplate(expandedText, [...this._promptTemplates]);
+        expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]);
         await this._queueSteer(expandedText);
     }
     /**
@@ -589,7 +630,7 @@ export class AgentSession {
         }
         // Expand skill commands and prompt templates
         let expandedText = this._expandSkillCommand(text);
-        expandedText = expandPromptTemplate(expandedText, [...this._promptTemplates]);
+        expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]);
         await this._queueFollowUp(expandedText);
     }
     /**
@@ -665,6 +706,8 @@ export class AgentSession {
         else {
             this.agent.appendMessage(appMessage);
             this.sessionManager.appendCustomMessageEntry(message.customType, message.content, message.display, message.details);
+            this._emit({ type: "message_start", message: appMessage });
+            this._emit({ type: "message_end", message: appMessage });
         }
     }
     /**
@@ -729,16 +772,8 @@ export class AgentSession {
     getFollowUpMessages() {
         return this._followUpMessages;
     }
-    get skillsSettings() {
-        return this._skillsSettings;
-    }
-    /** Skills loaded by SDK (empty if --no-skills or skills: [] was passed) */
-    get skills() {
-        return this._skills;
-    }
-    /** Skill loading warnings captured by SDK */
-    get skillWarnings() {
-        return this._skillWarnings;
+    get resourceLoader() {
+        return this._resourceLoader;
     }
     /**
      * Abort current operation and wait for agent to become idle.
@@ -752,7 +787,8 @@ export class AgentSession {
      * Start a new session, optionally with initial messages and parent tracking.
      * Clears all messages and starts a new session.
      * Listeners are preserved and will continue receiving events.
-     * @param options - Optional initial messages and parent session path
+     * @param options.parentSession - Optional parent session path for tracking
+     * @param options.setup - Optional callback to initialize session (e.g., append messages)
      * @returns true if completed, false if cancelled by extension
      */
     async newSession(options) {
@@ -770,11 +806,18 @@ export class AgentSession {
         this._disconnectFromAgent();
         await this.abort();
         this.agent.reset();
-        this.sessionManager.newSession(options);
+        this.sessionManager.newSession({ parentSession: options?.parentSession });
         this.agent.sessionId = this.sessionManager.getSessionId();
         this._steeringMessages = [];
         this._followUpMessages = [];
         this._pendingNextTurnMessages = [];
+        // Run setup callback if provided (e.g., to append initial messages)
+        if (options?.setup) {
+            await options.setup(this.sessionManager);
+            // Sync agent state with session manager after setup
+            const sessionContext = this.sessionManager.buildSessionContext();
+            this.agent.replaceMessages(sessionContext.messages);
+        }
         this._reconnectToAgent();
         // Emit session_switch event with reason "new" to extensions
         if (this._extensionRunner) {
@@ -879,12 +922,6 @@ export class AgentSession {
         await this._emitModelSelect(nextModel, currentModel, "cycle");
         return { model: nextModel, thinkingLevel: this.thinkingLevel, isScoped: false };
     }
-    /**
-     * Get all available models with valid API keys.
-     */
-    async getAvailableModels() {
-        return this._modelRegistry.getAvailable();
-    }
     // =========================================================================
     // Thinking Level Management
     // =========================================================================
@@ -1255,6 +1292,188 @@ export class AgentSession {
     get autoCompactionEnabled() {
         return this.settingsManager.getCompactionEnabled();
     }
+    async bindExtensions(bindings) {
+        if (bindings.uiContext !== undefined) {
+            this._extensionUIContext = bindings.uiContext;
+        }
+        if (bindings.commandContextActions !== undefined) {
+            this._extensionCommandContextActions = bindings.commandContextActions;
+        }
+        if (bindings.shutdownHandler !== undefined) {
+            this._extensionShutdownHandler = bindings.shutdownHandler;
+        }
+        if (bindings.onError !== undefined) {
+            this._extensionErrorListener = bindings.onError;
+        }
+        if (this._extensionRunner) {
+            this._applyExtensionBindings(this._extensionRunner);
+            await this._extensionRunner.emit({ type: "session_start" });
+        }
+    }
+    _applyExtensionBindings(runner) {
+        runner.setUIContext(this._extensionUIContext);
+        runner.bindCommandContext(this._extensionCommandContextActions);
+        this._extensionErrorUnsubscriber?.();
+        this._extensionErrorUnsubscriber = this._extensionErrorListener
+            ? runner.onError(this._extensionErrorListener)
+            : undefined;
+    }
+    _bindExtensionCore(runner) {
+        runner.bindCore({
+            sendMessage: (message, options) => {
+                this.sendCustomMessage(message, options).catch((err) => {
+                    runner.emitError({
+                        extensionPath: "<runtime>",
+                        event: "send_message",
+                        error: err instanceof Error ? err.message : String(err),
+                    });
+                });
+            },
+            sendUserMessage: (content, options) => {
+                this.sendUserMessage(content, options).catch((err) => {
+                    runner.emitError({
+                        extensionPath: "<runtime>",
+                        event: "send_user_message",
+                        error: err instanceof Error ? err.message : String(err),
+                    });
+                });
+            },
+            appendEntry: (customType, data) => {
+                this.sessionManager.appendCustomEntry(customType, data);
+            },
+            setSessionName: (name) => {
+                this.sessionManager.appendSessionInfo(name);
+            },
+            getSessionName: () => {
+                return this.sessionManager.getSessionName();
+            },
+            setLabel: (entryId, label) => {
+                this.sessionManager.appendLabelChange(entryId, label);
+            },
+            getActiveTools: () => this.getActiveToolNames(),
+            getAllTools: () => this.getAllTools(),
+            setActiveTools: (toolNames) => this.setActiveToolsByName(toolNames),
+            setModel: async (model) => {
+                const key = await this.modelRegistry.getApiKey(model);
+                if (!key)
+                    return false;
+                await this.setModel(model);
+                return true;
+            },
+            getThinkingLevel: () => this.thinkingLevel,
+            setThinkingLevel: (level) => this.setThinkingLevel(level),
+        }, {
+            getModel: () => this.model,
+            isIdle: () => !this.isStreaming,
+            abort: () => this.abort(),
+            hasPendingMessages: () => this.pendingMessageCount > 0,
+            shutdown: () => {
+                this._extensionShutdownHandler?.();
+            },
+            getContextUsage: () => this.getContextUsage(),
+            compact: (options) => {
+                void (async () => {
+                    try {
+                        const result = await this.compact(options?.customInstructions);
+                        options?.onComplete?.(result);
+                    }
+                    catch (error) {
+                        const err = error instanceof Error ? error : new Error(String(error));
+                        options?.onError?.(err);
+                    }
+                })();
+            },
+        });
+    }
+    _buildRuntime(options) {
+        const autoResizeImages = this.settingsManager.getImageAutoResize();
+        const shellCommandPrefix = this.settingsManager.getShellCommandPrefix();
+        const baseTools = this._baseToolsOverride
+            ? this._baseToolsOverride
+            : createAllTools(this._cwd, {
+                read: { autoResizeImages },
+                bash: { commandPrefix: shellCommandPrefix },
+            });
+        this._baseToolRegistry = new Map(Object.entries(baseTools).map(([name, tool]) => [name, tool]));
+        const extensionsResult = this._resourceLoader.getExtensions();
+        if (options.flagValues) {
+            for (const [name, value] of options.flagValues) {
+                extensionsResult.runtime.flagValues.set(name, value);
+            }
+        }
+        const hasExtensions = extensionsResult.extensions.length > 0;
+        const hasCustomTools = this._customTools.length > 0;
+        this._extensionRunner =
+            hasExtensions || hasCustomTools
+                ? new ExtensionRunner(extensionsResult.extensions, extensionsResult.runtime, this._cwd, this.sessionManager, this._modelRegistry)
+                : undefined;
+        if (this._extensionRunnerRef) {
+            this._extensionRunnerRef.current = this._extensionRunner;
+        }
+        if (this._extensionRunner) {
+            this._bindExtensionCore(this._extensionRunner);
+            this._applyExtensionBindings(this._extensionRunner);
+        }
+        const registeredTools = this._extensionRunner?.getAllRegisteredTools() ?? [];
+        const allCustomTools = [
+            ...registeredTools,
+            ...this._customTools.map((def) => ({ definition: def, extensionPath: "<sdk>" })),
+        ];
+        const wrappedExtensionTools = this._extensionRunner
+            ? wrapRegisteredTools(allCustomTools, this._extensionRunner)
+            : [];
+        const toolRegistry = new Map(this._baseToolRegistry);
+        for (const tool of wrappedExtensionTools) {
+            toolRegistry.set(tool.name, tool);
+        }
+        const defaultActiveToolNames = this._baseToolsOverride
+            ? Object.keys(this._baseToolsOverride)
+            : ["read", "bash", "edit", "write"];
+        const baseActiveToolNames = options.activeToolNames ?? defaultActiveToolNames;
+        const activeToolNameSet = new Set(baseActiveToolNames);
+        if (options.includeAllExtensionTools) {
+            for (const tool of wrappedExtensionTools) {
+                activeToolNameSet.add(tool.name);
+            }
+        }
+        const extensionToolNames = new Set(wrappedExtensionTools.map((tool) => tool.name));
+        const activeBaseTools = Array.from(activeToolNameSet)
+            .filter((name) => this._baseToolRegistry.has(name) && !extensionToolNames.has(name))
+            .map((name) => this._baseToolRegistry.get(name));
+        const activeExtensionTools = wrappedExtensionTools.filter((tool) => activeToolNameSet.has(tool.name));
+        const activeToolsArray = [...activeBaseTools, ...activeExtensionTools];
+        if (this._extensionRunner) {
+            const wrappedActiveTools = wrapToolsWithExtensions(activeToolsArray, this._extensionRunner);
+            this.agent.setTools(wrappedActiveTools);
+            const wrappedAllTools = wrapToolsWithExtensions(Array.from(toolRegistry.values()), this._extensionRunner);
+            this._toolRegistry = new Map(wrappedAllTools.map((tool) => [tool.name, tool]));
+        }
+        else {
+            this.agent.setTools(activeToolsArray);
+            this._toolRegistry = toolRegistry;
+        }
+        const systemPromptToolNames = Array.from(activeToolNameSet).filter((name) => this._baseToolRegistry.has(name));
+        this._baseSystemPrompt = this._rebuildSystemPrompt(systemPromptToolNames);
+        this.agent.setSystemPrompt(this._baseSystemPrompt);
+    }
+    async reload() {
+        const previousFlagValues = this._extensionRunner?.getFlagValues();
+        await this._extensionRunner?.emit({ type: "session_shutdown" });
+        resetApiProviders();
+        await this._resourceLoader.reload();
+        this._buildRuntime({
+            activeToolNames: this.getActiveToolNames(),
+            flagValues: previousFlagValues,
+            includeAllExtensionTools: true,
+        });
+        const hasBindings = this._extensionUIContext ||
+            this._extensionCommandContextActions ||
+            this._extensionShutdownHandler ||
+            this._extensionErrorListener;
+        if (this._extensionRunner && hasBindings) {
+            await this._extensionRunner.emit({ type: "session_start" });
+        }
+    }
     // =========================================================================
     // Auto-Retry
     // =========================================================================
@@ -1270,8 +1489,8 @@ export class AgentSession {
         if (isContextOverflow(message, contextWindow))
             return false;
         const err = message.errorMessage;
-        // Match: overloaded_error, rate limit, 429, 500, 502, 503, 504, service unavailable, connection errors, fetch failed
-        return /overloaded|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server error|internal error|connection.?error|connection.?refused|other side closed|fetch failed|upstream.?connect|reset before headers/i.test(err);
+        // Match: overloaded_error, rate limit, 429, 500, 502, 503, 504, service unavailable, connection errors, fetch failed, terminated
+        return /overloaded|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server error|internal error|connection.?error|connection.?refused|other side closed|fetch failed|upstream.?connect|reset before headers|terminated/i.test(err);
     }
     /**
      * Handle retryable errors with exponential backoff.
@@ -1316,7 +1535,7 @@ export class AgentSession {
         // Wait with exponential backoff (abortable)
         this._retryAbortController = new AbortController();
         try {
-            await this._sleep(delayMs, this._retryAbortController.signal);
+            await sleep(delayMs, this._retryAbortController.signal);
         }
         catch {
             // Aborted during sleep - emit end event so UI can clean up
@@ -1341,22 +1560,6 @@ export class AgentSession {
         }, 0);
         return true;
     }
-    /**
-     * Sleep helper that respects abort signal.
-     */
-    _sleep(ms, signal) {
-        return new Promise((resolve, reject) => {
-            if (signal?.aborted) {
-                reject(new Error("Aborted"));
-                return;
-            }
-            const timeout = setTimeout(resolve, ms);
-            signal?.addEventListener("abort", () => {
-                clearTimeout(timeout);
-                reject(new Error("Aborted"));
-            });
-        });
-    }
     /**
      * Cancel in-progress retry.
      */