@astralyn/pi 1.4.11 → 1.5.1

package/dist/core/agent-session.js

@@ -37,9 +37,11 @@ import { convertToLlm } from "./messages.js";
 import { ModelRegistry } from "./model-registry.js";
 import { NestedToolCallRunner } from "./nested-tool-calls.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
+import { QuiescenceLatch } from "./quiescence.js";
 import { exportSessionToJsonl } from "./session-export.js";
 import { getLatestCompactionEntry, SessionManager, } from "./session-manager.js";
 import { applyToolModifiers, DEFAULT_TOOL_NAMES, } from "./settings-manager.js";
+import { slashCommandName } from "./slash-command-name.js";
 import { BUILTIN_PATH_PREFIX, createSyntheticSourceInfo, isSyntheticPath } from "./source-info.js";
 import { buildSystemPrompt, buildSystemPromptSections, diffSystemPromptSections, normalizeBuildSystemPromptOptions, } from "./system-prompt.js";
 import { TaskSession } from "./tasks/session.js";
@@ -49,6 +51,16 @@ import { createLocalBashOperations } from "./tools/shell-operations.js";
 import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.js";
 import { addUsageToTotals, combineUsage, createUsageTotals, getAccountedUsages } from "./usage-totals.js";
 import { findLatestResponse, getBranchSelection, getVirtualModelState, isVirtualModel, VIRTUAL_MODEL_STATE_ENTRY, } from "./virtual-models.js";
+/**
+ * Stream functions that resolve auth through a given ModelRuntime themselves, such as the one
+ * createAgentSession installs. Summaries sent through them must not pass pre-resolved credentials:
+ * the runtime would take them as an explicit API key instead of the stored OAuth login (#11520).
+ */
+const runtimeStreamFunctions = new WeakMap();
+/** Mark `streamFn` as resolving auth through `modelRuntime` on its own. */
+export function markRuntimeStreamFunction(streamFn, modelRuntime) {
+    runtimeStreamFunctions.set(streamFn, modelRuntime);
+}
 /**
  * Parse a skill block from message text.
  * Returns null if the text doesn't contain a skill block.
@@ -114,11 +126,13 @@ export class AgentSession {
         this._isBeforeSettle = false;
         this._abortDuringBeforeSettle = false;
         this._isEmittingAgentSettled = false;
-        /** Last observed quiescence; a transition into true emits `session_quiescent` once. */
-        this._quiescent = true;
-        this._quiescenceDisposed = false;
-        this._quiescenceWaiters = new Set();
         this._deferredSettledActions = [];
+        /**
+         * Settled emissions in progress, with the actions they defer. A deferred prompt's input preflight
+         * is asynchronous and comes before its run marks the session busy.
+         */
+        this._settling = 0;
+        this._quiescence = new QuiescenceLatch(() => this._isBusy(), () => this._emit({ type: "session_quiescent" }));
         this._baseToolDefinitions = new Map();
         /**
          * Tools of the restored or reloaded loadout that are not registered yet, such as tools of MCP
@@ -141,13 +155,7 @@ export class AgentSession {
         this._toolPromptGuidelines = new Map();
         /** Internal handler for agent events - shared by subscribe and reconnect */
         this._handleAgentEvent = async (event) => {
-            if (event.type === "message_start")
-                this._tasksHost.delivery.messageStarted(event.message);
-            if (event.type === "message_end" &&
-                event.message.role === "assistant" &&
-                event.message.stopReason !== "error" &&
-                event.message.stopReason !== "aborted")
-                this._tasksHost.noteAssistantCompleted();
+            this._tasksHost.observe(event);
             // An abort requested while the run was starting, e.g. during before_agent_start, stops it
             // here. agent_start is awaited before the first request.
             if (event.type === "agent_start" && this._agentRunAbortRequested)
@@ -238,8 +246,6 @@ export class AgentSession {
             if (event.type === "turn_end") {
                 this._flushPendingCustomMessages();
             }
-            if (event.type === "agent_end")
-                this._tasksHost.delivery.agentEnded();
         };
         this.agent = config.agent;
         this.sessionManager = config.sessionManager;
@@ -268,7 +274,7 @@ export class AgentSession {
             onEntry: (entry) => this._emit({ type: "entry_appended", entry }),
             onError: (event, error) => this._extensionRunner?.emitError({ extensionPath: "<background>", event, error }),
         });
-        this._tasksHost.subscribe(() => this._updateQuiescence());
+        this._tasksHost.subscribe(() => this._quiescence.poke());
         this._scopedModels = config.scopedModels ?? [];
         this._resourceLoader = config.resourceLoader;
         this._customTools = config.customTools ?? [];
@@ -330,10 +336,10 @@ export class AgentSession {
         this._tasksHost.acknowledgeDelivery(taskId);
     }
     /**
-     * Run a lifecycle operation with delivery paused across its asynchronous gaps.
+     * Hold this session's completion delivery across an async operation, even one that replaces the session.
      * Lifecycle methods are thin shells over upstream bodies so those keep upstream's shape.
      */
-    async _withTasksPaused(operation) {
+    async withTaskNotificationsPaused(operation) {
         const resumeBackground = this.pauseTaskNotifications();
         try {
             return await operation();
@@ -342,6 +348,10 @@ export class AgentSession {
             resumeBackground();
         }
     }
+    /** Close task admission, abort the run and drain executors within `graceMs`, persisting their usage. */
+    async shutdownTasks(graceMs) {
+        await Promise.all([this._tasksHost.shutdown(graceMs), this.abort()]);
+    }
     get modelRuntime() {
         return this._modelRuntime;
     }
@@ -390,6 +400,9 @@ export class AgentSession {
         if (this.agent.streamFunction === streamSimple) {
             return { ...(await this._getRequiredRequestAuth(model, signal)), thinkingLevel };
         }
+        if (runtimeStreamFunctions.get(this.agent.streamFunction) === this._modelRuntime) {
+            return { model, thinkingLevel };
+        }
         try {
             const result = await this._modelRuntime.getAuth(model, { signal });
             if (!result)
@@ -783,48 +796,24 @@ export class AgentSession {
         for (const l of this._eventListeners) {
             l(event);
         }
-        // Compaction runs outside agent runs, so its start is what leaves quiescence.
-        if (event.type === "compaction_start")
-            this._updateQuiescence();
+        // Any event may follow a change quiescence reads, such as queued input or a compaction
+        // starting outside a run.
+        this._quiescence.poke();
     }
     /**
-     * Whether nothing more will happen without new input. Queued user input only matters during a
-     * run; left in the queue while idle it waits for the next prompt, and it also blocks automatic
-     * delivery, so completions that cannot start a turn do not count.
+     * Whether more will happen without new input. Queued user input only matters during a run; left
+     * in the queue while idle it waits for the next prompt and holds back automatic delivery.
      */
-    _computeQuiescent() {
-        if (this._isAgentRunActive ||
-            this._isEmittingAgentSettled ||
-            this._deferredSettledActions.length > 0 ||
+    _isBusy() {
+        return (!this.isIdle ||
+            this._settling > 0 ||
+            this.isRetrying ||
             this._pendingCustomMessages.length > 0 ||
-            this.isCompacting ||
-            this.isRetrying) {
-            return false;
-        }
-        if (this._tasksHost.runningBackground().length > 0)
-            return false;
-        const inputBlocksDelivery = this.pendingMessageCount > 0 || this.agent.hasQueuedMessages();
-        return !(this.model !== undefined && !inputBlocksDelivery && this._tasksHost.delivery.hasDeliverableWork());
+            !this._tasksHost.isQuiet());
     }
     /** Whether the session is quiescent now. See {@link AgentSessionEvent} `session_quiescent`. */
     get isQuiescent() {
-        return !this._quiescenceDisposed && this._computeQuiescent();
-    }
-    _updateQuiescence() {
-        if (this._quiescenceDisposed)
-            return;
-        const quiescent = this._computeQuiescent();
-        const entered = quiescent && !this._quiescent;
-        this._quiescent = quiescent;
-        if (!quiescent)
-            return;
-        if (entered)
-            this._emit({ type: "session_quiescent" });
-        // A listener may have started new work.
-        if (!this._quiescent)
-            return;
-        for (const waiter of [...this._quiescenceWaiters])
-            waiter.resolve();
+        return this._quiescence.isQuiescent;
     }
     /**
      * Resolve once the session is quiescent: no run, compaction, retry or branch summary, no
@@ -834,31 +823,7 @@ export class AgentSession {
      * the session is disposed.
      */
     waitForQuiescence(signal) {
-        if (this._quiescenceDisposed)
-            return Promise.reject(new Error("The session was disposed"));
-        if (signal?.aborted)
-            return Promise.reject(signal.reason);
-        if (this._computeQuiescent())
-            return Promise.resolve();
-        return new Promise((resolve, reject) => {
-            const waiter = {
-                resolve: () => {
-                    cleanup();
-                    resolve();
-                },
-                reject: (reason) => {
-                    cleanup();
-                    reject(reason);
-                },
-            };
-            const onAbort = () => waiter.reject(signal?.reason);
-            const cleanup = () => {
-                this._quiescenceWaiters.delete(waiter);
-                signal?.removeEventListener("abort", onAbort);
-            };
-            this._quiescenceWaiters.add(waiter);
-            signal?.addEventListener("abort", onAbort, { once: true });
-        });
+        return this._quiescence.wait(signal);
     }
     _emitQueueUpdate() {
         this._emit({
@@ -881,7 +846,7 @@ export class AgentSession {
         return this._idleWaitPromise;
     }
     _resolveIdleWaitIfIdle() {
-        this._updateQuiescence();
+        this._quiescence.poke();
         if (!this.isIdle || !this._resolveIdleWait) {
             return;
         }
@@ -895,7 +860,14 @@ export class AgentSession {
      * so a deferred prompt never races a delivery for the same idle gap.
      */
     async _emitAgentSettled() {
-        return this._withTasksPaused(() => this._emitAgentSettledBody());
+        this._settling++;
+        try {
+            await this.withTaskNotificationsPaused(() => this._emitAgentSettledBody());
+        }
+        finally {
+            this._settling--;
+            this._quiescence.poke();
+        }
     }
     async _emitAgentSettledBody() {
         this._cacheWarmer?.onAgentSettled();
@@ -1147,11 +1119,8 @@ export class AgentSession {
      * Call this when completely done with the session.
      */
     dispose() {
-        this._quiescenceDisposed = true;
-        const waiters = [...this._quiescenceWaiters];
+        this._quiescence.dispose();
         this._tasksHost.dispose();
-        for (const waiter of waiters)
-            waiter.reject(new Error("The session was disposed"));
         try {
             this.abortRetry();
             this.abortCompaction();
@@ -1314,9 +1283,14 @@ export class AgentSession {
      * the registered, non-hidden ones; they are declared to the model. Active tools with a
      * `prepareLoadout` hook can change the declared descriptions and hide declarations from
      * requests (see {@link _installHiddenDeclarationsProjection}).
+     *
+     * The active tools are a set: tools that stay active keep their current order and newly active
+     * ones are appended, so reordering the same names does not change the prompt (#11495).
      */
     _applyToolLoadout(toolNames) {
-        const tools = [...new Set(toolNames)].flatMap((name) => {
+        const requested = new Set(toolNames);
+        const kept = this.getActiveToolNames().filter((name) => requested.has(name));
+        const tools = [...new Set([...kept, ...requested])].flatMap((name) => {
             const tool = this._toolRegistry.get(name);
             return tool && this._getToolExposure(name) !== "hidden" && this._isActivatable(name) ? [tool] : [];
         });
@@ -1575,7 +1549,7 @@ export class AgentSession {
     async _startRun(trigger, onStarted) {
         this._isAgentRunActive = true;
         this._agentRunAbortRequested = false;
-        this._updateQuiescence();
+        this._quiescence.poke();
         let messages;
         try {
             messages = await this._prepareRun(trigger);
@@ -1807,10 +1781,10 @@ export class AgentSession {
         }
     }
     _isExtensionCommandPrompt(text, options) {
-        if (!(options?.expandPromptTemplates ?? true) || !text.startsWith("/"))
+        if (!(options?.expandPromptTemplates ?? true))
             return false;
-        const spaceIndex = text.indexOf(" ");
-        return (this._extensionRunner.getCommand(spaceIndex === -1 ? text.slice(1) : text.slice(1, spaceIndex)) !== undefined);
+        const name = slashCommandName(text);
+        return name !== undefined && this._extensionRunner.getCommand(name) !== undefined;
     }
     async _promptBody(text, options) {
         if (this._isEmittingAgentSettled) {
@@ -2414,7 +2388,7 @@ export class AgentSession {
      * @param customInstructions Optional instructions for the compaction summary
      */
     async compact(customInstructions) {
-        return this._withTasksPaused(() => this._compactBody(customInstructions));
+        return this.withTaskNotificationsPaused(() => this._compactBody(customInstructions));
     }
     async _compactBody(customInstructions) {
         await this.abort();
@@ -2709,7 +2683,7 @@ export class AgentSession {
      * @returns Whether the post-run loop should call `agent.continue()`
      */
     async _runAutoCompaction(reason, willRetry) {
-        return this._withTasksPaused(() => this._runAutoCompactionBody(reason, willRetry));
+        return this.withTaskNotificationsPaused(() => this._runAutoCompactionBody(reason, willRetry));
     }
     async _runAutoCompactionBody(reason, willRetry) {
         const model = this.model;
@@ -2852,7 +2826,7 @@ export class AgentSession {
         return this.settingsManager.getCompactionEnabled();
     }
     async bindExtensions(bindings) {
-        return this._withTasksPaused(() => this._bindExtensionsBody(bindings));
+        return this.withTaskNotificationsPaused(() => this._bindExtensionsBody(bindings));
     }
     async _bindExtensionsBody(bindings) {
         if (bindings.tasksEnabled !== undefined) {
@@ -3201,12 +3175,12 @@ export class AgentSession {
         });
     }
     async reload(options) {
-        return this._withTasksPaused(() => this._reloadBody(options));
+        return this.withTaskNotificationsPaused(() => this._reloadBody(options));
     }
     async _reloadBody(options) {
         const oldRunner = this._extensionRunner;
         const previousFlagValues = oldRunner.getFlagValues();
-        await Promise.all([this._tasksHost.shutdown(), this.abort()]);
+        await this.shutdownTasks();
         await emitSessionShutdownEvent(oldRunner, { type: "session_shutdown", reason: "reload" });
         oldRunner.invalidate();
         const getDefaultTools = () => this._usesDefaultTools
@@ -3508,7 +3482,7 @@ export class AgentSession {
         const resumeBackground = this.pauseTaskNotifications();
         // Set up abort controller for summarization
         this._branchSummaryAbortController = new AbortController();
-        this._updateQuiescence();
+        this._quiescence.poke();
         try {
             let extensionSummary;
             let fromExtension = false;