@averyyy/pi-coding-agent 0.99.1-piclient.1 → 1.0.0-piclient.1

package/dist/core/agent-session.js

@@ -39,6 +39,7 @@ import { cancelPiServerOperations, compactPiServer, hashStaticContext, } from ".
 import { hashPiServerProviderMessages } from "./pi-server-protocol.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
 import { CURRENT_SESSION_VERSION, getLatestCompactionEntry, SessionManager, } from "./session-manager.js";
+import { DEFAULT_TOOL_NAMES } from "./settings-manager.js";
 import { BUILTIN_PATH_PREFIX, createSyntheticSourceInfo, isSyntheticPath } from "./source-info.js";
 import { buildSystemPrompt, buildSystemPromptSections, diffSystemPromptSections, normalizeBuildSystemPromptOptions, } from "./system-prompt.js";
 import { createLocalBashOperations } from "./tools/bash.js";
@@ -235,6 +236,13 @@ export class AgentSession {
     _cwd;
     _extensionRunnerRef;
     _initialActiveToolNames;
+    /**
+     * Tools of the restored or reloaded loadout that are not registered yet, such as tools of MCP
+     * servers that are still connecting. They are activated when they are registered, and dropped when
+     * `setActiveToolsByName()` deactivates a tool or the next agent run starts.
+     */
+    _pendingToolNames = new Set();
+    _usesDefaultTools;
     _allowedToolNames;
     _excludedToolNames;
     _baseToolsOverride;
@@ -277,6 +285,7 @@ export class AgentSession {
         }
         this._extensionRunnerRef = config.extensionRunnerRef;
         this._initialActiveToolNames = config.initialActiveToolNames;
+        this._usesDefaultTools = config.usesDefaultTools ?? false;
         this._allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined;
         this._excludedToolNames = config.excludedToolNames ? new Set(config.excludedToolNames) : undefined;
         this._baseToolsOverride = config.baseToolsOverride;
@@ -1299,9 +1308,23 @@ export class AgentSession {
      * Changes take effect on the next agent turn.
      */
     setActiveToolsByName(toolNames) {
+        const previous = this.getActiveToolNames();
+        this._setActiveTools(toolNames);
+        // A loadout that deactivates a tool replaces the restored one, whose pending tools are dropped.
+        // One that only adds tools, like activating tool_search, keeps them.
+        const active = new Set(this.getActiveToolNames());
+        if (previous.some((name) => !active.has(name)))
+            this._pendingToolNames.clear();
+    }
+    _setActiveTools(toolNames) {
         const tools = this._applyToolLoadout(toolNames);
+        for (const tool of tools)
+            this._pendingToolNames.delete(tool.name);
         this._rebuildSystemPrompt(tools.map((tool) => tool.name));
     }
+    _isAllowedTool(name) {
+        return (!this._allowedToolNames || this._allowedToolNames.has(name)) && !this._excludedToolNames?.has(name);
+    }
     _getToolExposure(name) {
         return this._toolDefinitions.get(name)?.definition.exposure ?? "direct";
     }
@@ -1437,7 +1460,8 @@ export class AgentSession {
         const toolSnippets = {};
         for (const name of this._toolRegistry.keys()) {
             const snippet = this._toolPromptSnippets.get(name);
-            if (snippet)
+            // Tools without a snippet are not listed. Hidden tools are only callable through another tool.
+            if (snippet && !this._hiddenDeclarations.has(name))
                 toolSnippets[name] = snippet;
         }
         const loaderSystemPrompt = this._resourceLoader.getSystemPrompt();
@@ -1468,6 +1492,8 @@ export class AgentSession {
      */
     _preparePromptAndToolLoadout(options, messages = this.agent.state.messages) {
         options.selectedTools = this._applyToolLoadout(options.selectedTools).map((tool) => tool.name);
+        // The tool list must match the declarations the request carries, so hidden tools are not listed.
+        options.toolSnippets = Object.fromEntries(Object.entries(options.toolSnippets).filter(([name]) => !this._hiddenDeclarations.has(name)));
         const sections = diffSystemPromptSections(getCurrentSystemMessage(messages)?.sections ?? {}, buildSystemPromptSections(options));
         return sections ? { role: "system", content: "", sections, timestamp: Date.now() } : undefined;
     }
@@ -1532,10 +1558,13 @@ export class AgentSession {
      * set, so the transcript's declarations are the whole loadout.
      */
     _restoreToolsFromTranscript() {
+        this._pendingToolNames.clear();
         const current = getCurrentSystemMessage(this.sessionManager.buildSessionContext().messages);
         if (!current)
             return;
-        this.setActiveToolsByName((current.toolsAdded ?? []).map((tool) => tool.name));
+        const names = (current.toolsAdded ?? []).map((tool) => tool.name);
+        this._pendingToolNames = new Set(names.filter((name) => this._isAllowedTool(name)));
+        this._setActiveTools(names);
     }
     // =========================================================================
     // Prompting
@@ -1545,6 +1574,9 @@ export class AgentSession {
         // Compaction before the prompt may have scheduled a retry; the new prompt replaces it.
         this._failedResponse = undefined;
         this._recordSelection();
+        // The run records the loadout in the transcript; restored tools that did not register by now
+        // are dropped, so a tool that never registers does not stay pending.
+        this._pendingToolNames.clear();
         this._isAgentRunActive = true;
         try {
             await start();
@@ -3433,8 +3465,6 @@ export class AgentSession {
         const previousActivatedOnRegistration = new Set([...this._toolRegistry.keys()].filter((name) => this._isActivatedOnRegistration(name)));
         const previousActiveToolNames = this.getActiveToolNames();
         const allowedToolNames = this._allowedToolNames;
-        const excludedToolNames = this._excludedToolNames;
-        const isAllowedTool = (name) => (!allowedToolNames || allowedToolNames.has(name)) && !excludedToolNames?.has(name);
         const registeredTools = this._extensionRunner.getAllRegisteredTools();
         const allCustomTools = [
             ...registeredTools,
@@ -3442,9 +3472,9 @@ export class AgentSession {
                 definition,
                 sourceInfo: createSyntheticSourceInfo(`<sdk:${definition.name}>`, { source: "sdk" }),
             })),
-        ].filter((tool) => isAllowedTool(tool.definition.name));
+        ].filter((tool) => this._isAllowedTool(tool.definition.name));
         const definitionRegistry = new Map(Array.from(this._baseToolDefinitions.entries())
-            .filter(([name]) => isAllowedTool(name))
+            .filter(([name]) => this._isAllowedTool(name))
             .map(([name, definition]) => [
             name,
             {
@@ -3474,7 +3504,7 @@ export class AgentSession {
         const runner = this._extensionRunner;
         const wrappedExtensionTools = wrapRegisteredTools(allCustomTools, runner);
         const wrappedBuiltInTools = wrapRegisteredTools(Array.from(this._baseToolDefinitions.values())
-            .filter((definition) => isAllowedTool(definition.name))
+            .filter((definition) => this._isAllowedTool(definition.name))
             .map((definition) => ({
             definition,
             sourceInfo: createSyntheticSourceInfo(`${BUILTIN_PATH_PREFIX}${definition.name}`, {
@@ -3486,7 +3516,7 @@ export class AgentSession {
             toolRegistry.set(tool.name, tool);
         }
         this._toolRegistry = toolRegistry;
-        const nextActiveToolNames = (options?.activeToolNames ? [...options.activeToolNames] : [...previousActiveToolNames]).filter((name) => isAllowedTool(name));
+        const nextActiveToolNames = (options?.activeToolNames ? [...options.activeToolNames] : [...previousActiveToolNames]).filter((name) => this._isAllowedTool(name));
         if (allowedToolNames) {
             for (const toolName of this._toolRegistry.keys()) {
                 // Naming a tool activates it even when it is not active by default.
@@ -3508,7 +3538,9 @@ export class AgentSession {
                 }
             }
         }
-        this.setActiveToolsByName([...new Set(nextActiveToolNames)]);
+        // Pending tools that are registered now become active.
+        nextActiveToolNames.push(...this._pendingToolNames);
+        this._setActiveTools([...new Set(nextActiveToolNames)]);
     }
     /** Whether activating the tool declares it to the model. */
     _isDeclarable(name) {
@@ -3559,12 +3591,21 @@ export class AgentSession {
         const previousFlagValues = oldRunner.getFlagValues();
         await emitSessionShutdownEvent(oldRunner, { type: "session_shutdown", reason: "reload" });
         oldRunner.invalidate();
+        const previousDefaultTools = new Set(this._usesDefaultTools ? (this.settingsManager.getDefaultTools() ?? DEFAULT_TOOL_NAMES) : []);
         await this.settingsManager.reload();
         this.syncQueueModesFromSettings();
         resetApiProviders();
         await this._resourceLoader.reload();
+        // Activate tools newly added to defaultTools. Removed ones stay active, and tools disabled
+        // during the session stay disabled unless the setting newly adds them.
+        const addedDefaultTools = this._usesDefaultTools
+            ? (this.settingsManager.getDefaultTools() ?? DEFAULT_TOOL_NAMES).filter((name) => !previousDefaultTools.has(name))
+            : [];
+        // Tools the new extensions register later, such as MCP tools, are pending until then.
+        for (const name of this.getActiveToolNames())
+            this._pendingToolNames.add(name);
         this._buildRuntime({
-            activeToolNames: this.getActiveToolNames(),
+            activeToolNames: [...this.getActiveToolNames(), ...addedDefaultTools],
             flagValues: previousFlagValues,
             includeAllExtensionTools: true,
         });