@ashx-j/lunr 0.2.26 → 0.2.28

package/dist/core/agent-session.js

@@ -1,3 +1,4 @@
+import { collectChildUsage, collectUsageRequests, totalChildUsage, totalRequestUsage, } from "./usage-accounting.js";
 /**
  * AgentSession - Core abstraction for agent lifecycle and session management.
  *
@@ -14,7 +15,7 @@
  */
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { basename, dirname, join, resolve } from "node:path";
-import { clampThinkingLevel, cleanupSessionResources, getSupportedThinkingLevels, isContextOverflow, isRetryableAssistantError, modelsAreEqual, resetApiProviders, streamSimple, } from "@ashx-j/lunr-ai/compat";
+import { clampThinkingLevel, cleanupSessionResources, estimateContextTokens as estimateRequestContextTokens, getSupportedThinkingLevels, isContextOverflow, isRetryableAssistantError, modelsAreEqual, resetApiProviders, streamSimple, } from "@ashx-j/lunr-ai/compat";
 import { getAgentDir } from "../config.js";
 import { getThemeByName, theme } from "../modes/interactive/theme/theme.js";
 import { stripFrontmatter } from "../utils/frontmatter.js";
@@ -27,13 +28,16 @@ import { DEFAULT_THINKING_LEVEL } from "./defaults.js";
 import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { emitSessionShutdownEvent } from "./extensions/runner.js";
 import { MEMORY_TOOL_NAMES } from "./memory-cap.js";
+import { convertToLlm } from "./messages.js";
 import { isUserInstructionsPath, loadSelectedUserInstructions } from "./model-instructions.js";
 import { ModelRegistry } from "./model-registry.js";
-import { gateToolCall } from "./permissions.js";
+import { AUTO_MODE_ADDENDUM, deletePermissionContext, gateToolCall, getPermissionMode, initializePermissionContext, setPermissionMode, } from "./permissions.js";
+import { READ_ONLY_MODE_ADDENDUM } from "./plan-mode.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
 import { rollbackSnapshotBeforeWrite } from "./rollback.js";
 import { CURRENT_SESSION_VERSION, getLatestCompactionEntry } from "./session-manager.js";
 import { createSyntheticSourceInfo } from "./source-info.js";
+import { isSubagentChildProcess, resolveChildRuntimePermissionMode, SUBAGENT_CHILD_PERMISSION_ENV, } from "./subagent-permission-inherit.js";
 import { registerSubagentWaitInterruptionOwner, } from "./subagent-wait-interruption.js";
 import { buildSystemPrompt } from "./system-prompt.js";
 import { measureStartup } from "./timings.js";
@@ -113,7 +117,12 @@ export class AgentSession {
     _retryAbortController = undefined;
     _retryAttempt = 0;
     // Bash execution state
-    _bashAbortController = undefined;
+    _bashOperations = new Map();
+    _operations = new Set();
+    _closing = false;
+    _disposed = false;
+    _drainPromise;
+    _shutdownPromise;
     _pendingBashMessages = [];
     // Extension system
     _extensionRunner;
@@ -146,7 +155,7 @@ export class AgentSession {
     _baseSystemPrompt = "";
     _baseSystemPromptOptions;
     _systemPromptOverride;
-    // lunr: core-owned system-prompt append (read-only mode addendum) applied on top of base/override
+    // Optional host-owned prompt append, separate from session permission guidance.
     _systemPromptAppend;
     // lunr: core-owned tool-call gates (read-only mode) checked before extension tool_call handlers
     _toolCallGates = [];
@@ -154,6 +163,13 @@ export class AgentSession {
         this.agent = config.agent;
         this.sessionManager = config.sessionManager;
         this.settingsManager = config.settingsManager;
+        const permissionMode = config.permissionMode ??
+            (isSubagentChildProcess()
+                ? resolveChildRuntimePermissionMode(process.env[SUBAGENT_CHILD_PERMISSION_ENV])
+                : undefined) ??
+            this.sessionManager.getPermissionMode() ??
+            this.settingsManager.getDefaultPermissionMode();
+        initializePermissionContext(this.sessionId, permissionMode, config.permissionMode);
         this._scopedModels = config.scopedModels ?? [];
         this._resourceLoader = config.resourceLoader;
         this._customTools = config.customTools ?? [];
@@ -341,12 +357,17 @@ export class AgentSession {
                 : undefined);
         this.agent.prepareNextTurnWithContext = async (turn, signal) => {
             const previousSnapshot = await previousPrepareNextTurnWithContext?.(turn, signal);
-            let previousContext = previousSnapshot?.context ?? turn.context;
+            let previousContext = {
+                ...(previousSnapshot?.context ?? turn.context),
+                systemPrompt: this._withSystemPromptAppend(this._systemPromptOverride ?? this._baseSystemPrompt),
+                tools: this.agent.state.tools.slice(),
+            };
             const model = this.model;
             const settings = this.settingsManager.getCompactionSettings();
             if (model?.provider === "openai-codex" &&
                 settings.enabled &&
-                shouldCompact(estimateContextTokens(previousContext.messages).tokens, model.contextWindow, settings)) {
+                shouldCompact(estimateRequestContextTokens({ ...previousContext, messages: convertToLlm(previousContext.messages) })
+                    .tokens, model.contextWindow, settings)) {
                 const previousCompactionId = getLatestCompactionEntry(this.sessionManager.getBranch())?.id;
                 await this._runAutoCompaction("threshold", true);
                 const currentCompaction = getLatestCompactionEntry(this.sessionManager.getBranch());
@@ -435,6 +456,8 @@ export class AgentSession {
                 }
             }
         }
+        if (event.type === "message_end")
+            this._activePromptRun?.completion?.messages.push(event.message);
         // Emit to extensions first
         await this._emitExtensionEvent(event);
         // Notify all listeners
@@ -659,27 +682,80 @@ export class AgentSession {
     }
     assertCanStartWork() {
         this.sessionManager.assertWritable();
+        if (this._closing)
+            throw new Error("Session is closing. New work is not accepted.");
         if (this.transferring)
             throw new Error("Session transfer is in progress. Wait or reclaim the session.");
     }
-    dispose() {
+    /** Stop admission synchronously, retaining listeners and session ownership. */
+    stopAdmission() {
+        this._closing = true;
         this._invalidateWaitPromptHandoffs("Session closed before the pending prompt could run.");
-        for (const unregister of this._unregisterSubagentWaitInterruptionOwners.splice(0))
-            unregister();
+    }
+    async _trackOperation(operation) {
+        let settle = () => { };
+        const completion = new Promise((resolve) => {
+            settle = resolve;
+        });
+        this._operations.add(completion);
         try {
-            this.abortRetry();
-            this.abortCompaction();
-            this.abortBranchSummary();
-            this.abortBash();
-            this.agent.abort();
+            return await operation();
         }
-        catch {
-            // Dispose must succeed even if an abort hook throws.
+        finally {
+            this._operations.delete(completion);
+            settle();
+        }
+    }
+    /** Cancel owned work and wait for final persistence before extension shutdown. */
+    drain() {
+        if (this._drainPromise)
+            return this._drainPromise;
+        this.stopAdmission();
+        if (this._activePromptRun)
+            this._activePromptRun.gracefulStopRequested = true;
+        this.abortRetry();
+        this.abortCompaction();
+        this.abortBranchSummary();
+        this.abortBash();
+        this.clearQueue();
+        this.agent.abort();
+        this._drainPromise = (async () => {
+            await Promise.all([...this._operations]);
+            await this.waitForIdle();
+            await this.agent.waitForIdle();
+            this._flushPendingBashMessages();
+        })();
+        return this._drainPromise;
+    }
+    /** Await cancellation and persistence, then release listeners and ownership. */
+    shutdown() {
+        if (!this._shutdownPromise) {
+            this._shutdownPromise = this.drain().then(() => this._disposeSettled());
         }
+        return this._shutdownPromise;
+    }
+    /** Compatibility wrapper. Use await shutdown() when release must be complete. */
+    dispose() {
+        if (this._operations.size === 0 && !this._isAgentRunActive && !this.agent.state.isStreaming) {
+            this.stopAdmission();
+            this._disposeSettled();
+            return;
+        }
+        void this.shutdown().catch(() => {
+            /* Awaitable callers receive shutdown failures. */
+        });
+    }
+    _disposeSettled() {
+        if (this._disposed)
+            return;
+        this._disposed = true;
+        for (const unregister of this._unregisterSubagentWaitInterruptionOwners.splice(0))
+            unregister();
         this._extensionRunner.invalidate("This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().");
         this._disconnectFromAgent();
         this._eventListeners = [];
         cleanupSessionResources(this.sessionId);
+        deletePermissionContext(this.sessionId);
         this.sessionManager.dispose();
     }
     // =========================================================================
@@ -737,8 +813,23 @@ export class AgentSession {
         this._systemPromptAppend = next;
         this.agent.state.systemPrompt = this._withSystemPromptAppend(this._systemPromptOverride ?? this._baseSystemPrompt);
     }
+    get permissionMode() {
+        return getPermissionMode(this.sessionId);
+    }
+    _permissionAddendum(mode) {
+        return mode === "read-only" ? READ_ONLY_MODE_ADDENDUM : mode === "auto" ? AUTO_MODE_ADDENDUM : undefined;
+    }
+    /** Update the same permission context consulted by headless and interactive tools. */
+    setPermissionMode(mode) {
+        this.assertCanStartWork();
+        setPermissionMode(mode, this.sessionId);
+        this.sessionManager.setPermissionMode(mode);
+        this.agent.state.systemPrompt = this._withSystemPromptAppend(this._systemPromptOverride ?? this._baseSystemPrompt);
+    }
     _withSystemPromptAppend(prompt) {
-        return this._systemPromptAppend ? `${prompt}\n\n${this._systemPromptAppend}` : prompt;
+        return [prompt, this._permissionAddendum(this.permissionMode), this._systemPromptAppend]
+            .filter(Boolean)
+            .join("\n\n");
     }
     /** Current retry attempt (0 if not retrying) */
     get retryAttempt() {
@@ -944,28 +1035,67 @@ export class AgentSession {
         for (const reservation of handoff.reservations.splice(0))
             reservation.reject(error);
     }
-    async _runAgentPrompt(messages) {
+    _reservePromptRun() {
+        this.assertCanStartWork();
+        if (this._isAgentRunActive || this._manualCompactionPromise) {
+            throw new Error("Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.");
+        }
         const run = { waits: new Set(), gracefulStopRequested: false };
         this._activePromptRun = run;
         this._isAgentRunActive = true;
-        try {
-            await this.agent.prompt(messages);
-            while (await this._handlePostAgentRun(run)) {
-                await this.agent.continue();
+        return run;
+    }
+    async _releasePromptRun(run) {
+        if (this._activePromptRun !== run)
+            return;
+        this._activePromptRun = undefined;
+        this._systemPromptOverride = undefined;
+        this._flushPendingBashMessages();
+        await this._emitAgentSettled();
+    }
+    async _runAgentPrompt(messages, owner) {
+        const run = owner ?? this._reservePromptRun();
+        await this._trackOperation(async () => {
+            try {
+                this.assertCanStartWork();
+                if (run.gracefulStopRequested)
+                    throw new Error("Prompt was cancelled during preparation.");
+                await this.agent.prompt(Array.isArray(messages) ? messages : [messages], {
+                    consumeFollowUps: !run.completion,
+                });
+                while (await this._handlePostAgentRun(run)) {
+                    if (run.gracefulStopRequested || this._closing || run.completion?.cancelled)
+                        break;
+                    await this.agent.continue({ consumeFollowUps: !run.completion });
+                }
+                if (run.completion) {
+                    const completion = run.completion;
+                    run.completion = undefined;
+                    completion.detach?.();
+                    const hasQueuedWork = !run.gracefulStopRequested && !this._closing && this.agent.hasQueuedMessages();
+                    if (!hasQueuedWork)
+                        await this._releasePromptRun(run);
+                    completion.settle();
+                    // Later queued work retains normal session lifecycle, with a fresh queue admission.
+                    if (hasQueuedWork) {
+                        await this.agent.continueQueued();
+                        while (await this._handlePostAgentRun(run)) {
+                            if (run.gracefulStopRequested || this._closing)
+                                break;
+                            await this.agent.continue();
+                        }
+                    }
+                }
             }
-        }
-        finally {
-            if (this._activePromptRun === run)
-                this._activePromptRun = undefined;
-            this._systemPromptOverride = undefined;
-            this._flushPendingBashMessages();
-            await this._emitAgentSettled();
-        }
+            finally {
+                await this._releasePromptRun(run);
+            }
+        });
     }
     async _handlePostAgentRun(run) {
         const msg = this._lastAssistantMessage;
         this._lastAssistantMessage = undefined;
-        if (!msg || run.gracefulStopRequested) {
+        if (!msg || run.gracefulStopRequested || run.completion?.signal?.aborted) {
             return false;
         }
         // lunr: quota-dead or auth-invalid keys rotate immediately — same-key backoff
@@ -992,9 +1122,9 @@ export class AgentSession {
         if (await this._checkCompaction(msg)) {
             return true;
         }
-        // The agent loop drains both queues before emitting agent_end. Any messages
-        // here were queued by agent_end extension handlers and need a continuation.
-        return this.agent.hasQueuedMessages();
+        // Correlated requests leave follow-ups queued. Ordinary runs also continue
+        // messages queued by agent_end extension handlers.
+        return !run.completion && this.agent.hasQueuedMessages();
     }
     /**
      * Send a prompt to the agent.
@@ -1007,36 +1137,116 @@ export class AgentSession {
      */
     async prompt(text, options) {
         this.assertCanStartWork();
-        const preflightResult = options?.preflightResult;
-        let prepared;
-        let messages;
+        const commandName = text.startsWith("/") ? text.slice(1).split(" ")[0] : undefined;
+        // Lifecycle commands may replace the session; they cannot own a run they must drain.
+        if ((options?.expandPromptTemplates ?? true) && commandName && this._extensionRunner.getCommand(commandName)) {
+            try {
+                await this._tryExecuteExtensionCommand(text);
+                options?.preflightResult?.(true);
+            }
+            catch (error) {
+                options?.preflightResult?.(false);
+                throw error;
+            }
+            return;
+        }
+        return this._trackOperation(() => this._prompt(text, options));
+    }
+    async _prompt(text, options, completion) {
+        let run;
+        let accepted = false;
         try {
-            prepared = await this._preparePromptInput(text, options, this.isStreaming ? options?.streamingBehavior : undefined);
+            this.assertCanStartWork();
+            const wasStreaming = this.isStreaming;
+            if (!wasStreaming) {
+                run = this._reservePromptRun();
+                run.completion = completion;
+                if (completion?.signal) {
+                    const owner = run;
+                    const cancel = () => {
+                        if (this._activePromptRun !== owner || owner.completion !== completion)
+                            return;
+                        completion.cancelled = true;
+                        this.agent.abort();
+                        this.abortRetry();
+                        this._autoCompactionAbortController?.abort();
+                    };
+                    completion.signal.addEventListener("abort", cancel, { once: true });
+                    completion.detach = () => completion.signal?.removeEventListener("abort", cancel);
+                    completion.signal.throwIfAborted();
+                }
+            }
+            else if (!options?.streamingBehavior) {
+                throw new Error("Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.");
+            }
+            const prepared = await this._preparePromptInput(text, options, wasStreaming ? options?.streamingBehavior : undefined);
+            this.assertCanStartWork();
+            run?.completion?.signal?.throwIfAborted();
+            if (run?.gracefulStopRequested)
+                throw new Error("Prompt was cancelled during preparation.");
             if (!prepared) {
-                preflightResult?.(true);
+                accepted = true;
+                options?.preflightResult?.(true);
                 return;
             }
-            if (this.isStreaming) {
-                if (!options?.streamingBehavior) {
-                    throw new Error("Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.");
-                }
-                if (options.streamingBehavior === "followUp") {
+            if (wasStreaming && this.isStreaming) {
+                if (options?.streamingBehavior === "followUp")
                     await this._queueFollowUp(prepared.text, prepared.images);
-                }
-                else {
+                else
                     await this._queueSteer(prepared.text, prepared.images);
-                }
-                preflightResult?.(true);
+                accepted = true;
+                options?.preflightResult?.(true);
                 return;
             }
-            messages = await this._createPromptMessages(prepared);
+            run ??= this._reservePromptRun();
+            run.completion = completion;
+            const messages = await this._createPromptMessages(prepared);
+            this.assertCanStartWork();
+            run.completion?.signal?.throwIfAborted();
+            if (run.gracefulStopRequested)
+                throw new Error("Prompt was cancelled during preparation.");
+            accepted = true;
+            options?.preflightResult?.(true);
+            await this._runAgentPrompt(messages, run);
         }
         catch (error) {
-            preflightResult?.(false);
+            if (!accepted)
+                options?.preflightResult?.(false);
             throw error;
         }
-        preflightResult?.(true);
-        await this._runAgentPrompt(messages);
+        finally {
+            completion?.detach?.();
+            if (run)
+                await this._releasePromptRun(run);
+        }
+    }
+    /**
+     * Admit one scheduled request, including steering, retries and compaction.
+     * Completes before later follow-ups; signal cancellation belongs only to this request.
+     */
+    async promptWithCompletion(text, options) {
+        this.assertCanStartWork();
+        if (!this.isIdle || this.isBashRunning || this.isCompacting || this.pendingMessageCount > 0) {
+            throw new Error("Session is busy. Retry the scheduled prompt after its work settles.");
+        }
+        if (text.startsWith("/") && (options?.expandPromptTemplates ?? true))
+            this._throwIfExtensionCommand(text);
+        options?.signal?.throwIfAborted();
+        const messages = [];
+        return new Promise((resolve, reject) => {
+            const completion = {
+                messages,
+                signal: options?.signal,
+                settle: () => {
+                    if (completion.cancelled)
+                        reject(completion.signal?.reason);
+                    else
+                        resolve({ messages });
+                },
+            };
+            // Track the full session operation while returning only the admitted request's completion.
+            void this._trackOperation(() => this._prompt(text, options, completion)).then(completion.settle, reject);
+        });
     }
     async _preparePromptInput(text, options, streamingBehavior) {
         const expandPromptTemplates = options?.expandPromptTemplates ?? true;
@@ -1113,7 +1323,9 @@ export class AgentSession {
         return messages;
     }
     async interruptSubagentWaitWithPrompt(text, options) {
-        const acceptance = this._waitPromptAcceptance.then(() => this._reserveWaitPrompt(text, options));
+        this.assertCanStartWork();
+        const previousAcceptance = this._waitPromptAcceptance;
+        const acceptance = this._trackOperation(() => previousAcceptance.then(() => this._reserveWaitPrompt(text, options)));
         this._waitPromptAcceptance = acceptance.then(() => { }, () => { });
         return acceptance;
     }
@@ -1147,6 +1359,9 @@ export class AgentSession {
             }
             throw error;
         }
+        this.assertCanStartWork();
+        if (handoff.generation !== this._waitPromptGeneration)
+            throw new Error("Pending prompt was cancelled during preparation.");
         if (!prepared) {
             if (handoff.reservations.length === 0 && this._waitPromptHandoff === handoff) {
                 this._waitPromptHandoff = undefined;
@@ -1176,17 +1391,21 @@ export class AgentSession {
             handoff.generation === this._waitPromptGeneration &&
             handoff.reservations.length > 0) {
             const reservation = handoff.reservations[0];
+            let run;
             try {
                 this.assertCanStartWork();
+                run = this._reservePromptRun();
                 const messages = await this._createPromptMessages(reservation.input);
                 handoff.accepting = false;
-                await this._runAgentPrompt(messages);
+                await this._runAgentPrompt(messages, run);
                 reservation.resolve();
             }
             catch (error) {
                 reservation.reject(error);
             }
             finally {
+                if (run)
+                    await this._releasePromptRun(run);
                 if (handoff.reservations[0] === reservation)
                     handoff.reservations.shift();
             }
@@ -1448,8 +1667,13 @@ export class AgentSession {
     async abort() {
         this._invalidateWaitPromptHandoffs("The active run was aborted before the pending prompt could run.");
         this.abortRetry();
+        if (this._activePromptRun && !this.agent.signal)
+            this._activePromptRun.gracefulStopRequested = true;
+        const bash = [...this._bashOperations.values()];
+        this.abortBash();
         this.agent.abort();
         await this.waitForIdle();
+        await Promise.all(bash);
     }
     async waitForIdle() {
         if (this.isIdle) {
@@ -1690,12 +1914,13 @@ export class AgentSession {
      * @param customInstructions Optional instructions for the compaction summary
      */
     async compact(customInstructions) {
+        this.assertCanStartWork();
         const existing = this._latestCompactionResult();
         if (existing)
             return existing;
         if (this._manualCompactionPromise)
             return this._manualCompactionPromise;
-        const promise = this._compact(customInstructions);
+        const promise = this._trackOperation(() => this._compact(customInstructions));
         this._manualCompactionPromise = promise;
         try {
             return await promise;
@@ -1706,16 +1931,15 @@ export class AgentSession {
         }
     }
     async _compact(customInstructions) {
-        this._disconnectFromAgent();
         await this.abort();
+        this.assertCanStartWork();
         const existing = this._latestCompactionResult();
-        if (existing) {
-            this._reconnectToAgent();
+        if (existing)
             return existing;
-        }
+        this._disconnectFromAgent();
         this._compactionAbortController = new AbortController();
-        this._emit({ type: "compaction_start", reason: "manual" });
         try {
+            this._emit({ type: "compaction_start", reason: "manual" });
             if (!this.model) {
                 throw new Error(formatNoModelSelectedMessage());
             }
@@ -1764,7 +1988,7 @@ export class AgentSession {
             }
             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.streamFn, env, (response) => this.sessionManager.appendRequestUsage("compaction", response));
                 summary = result.summary;
                 firstKeptEntryId = result.firstKeptEntryId;
                 tokensBefore = result.tokensBefore;
@@ -1998,7 +2222,7 @@ export class AgentSession {
             }
             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.streamFn, env, (response) => this.sessionManager.appendRequestUsage("compaction", response));
                 summary = compactResult.summary;
                 firstKeptEntryId = compactResult.firstKeptEntryId;
                 tokensBefore = compactResult.tokensBefore;
@@ -2046,9 +2270,9 @@ export class AgentSession {
                 }
                 return true;
             }
-            // Auto-compaction can complete while follow-up/steering/custom messages are waiting.
-            // Continue once so queued messages are delivered.
-            return this.agent.hasQueuedMessages();
+            // Only ordinary runs continue queued requests after successful compaction.
+            // Correlated completion must settle before those requests begin.
+            return !this._activePromptRun?.completion && this.agent.hasQueuedMessages();
         }
         catch (error) {
             const errorMessage = error instanceof Error ? error.message : "compaction failed";
@@ -2310,6 +2534,7 @@ export class AgentSession {
         }, {
             getModel: () => this.resolvedCatalogModel(),
             isIdle: () => this.isIdle,
+            promptWithCompletion: (text, options) => this.promptWithCompletion(text, options),
             isProjectTrusted: () => this.settingsManager.isProjectTrusted(),
             getSignal: () => this.agent.signal,
             abort: () => {
@@ -2569,19 +2794,17 @@ export class AgentSession {
      * false when rotation is impossible, leaving the error as the final message.
      */
     async _tryRotateSubscriptionKey(message) {
-        const model = this.model;
-        if (!model)
+        const requestKey = this._modelRuntime.getRequestSubscriptionKey(message);
+        if (!requestKey)
             return false;
-        const providerId = model.provider;
+        const providerId = requestKey.providerId;
         // A runtime --api-key override wins over auth.json in RuntimeCredentials.read(),
         // so mirroring a rotated key into the store would have no effect — skip rotation.
         if (this._modelRuntime.hasRuntimeApiKey(providerId))
             return false;
         const manager = this._modelRuntime.subscriptionManager;
-        if ((await manager.list(providerId)).length < 2)
-            return false;
         const errorMessage = message.errorMessage ?? "Unknown error";
-        const entry = await manager.rotateOnFailure(providerId, errorMessage);
+        const entry = await manager.rotateOnFailure(providerId, errorMessage, requestKey.fingerprint);
         if (!entry)
             return false;
         // Close out any in-flight retry indicator; a fresh key gets a fresh retry budget.
@@ -2630,21 +2853,29 @@ export class AgentSession {
      */
     async executeBash(command, onChunk, options) {
         this.assertCanStartWork();
-        this._bashAbortController = new AbortController();
-        // Apply command prefix if configured (e.g., "shopt -s expand_aliases" for alias support)
-        const prefix = this.settingsManager.getShellCommandPrefix();
-        const shellPath = this.settingsManager.getShellPath();
-        const resolvedCommand = prefix ? `${prefix}\n${command}` : command;
+        const controller = new AbortController();
+        let settle = () => { };
+        const completion = new Promise((resolve) => {
+            settle = resolve;
+        });
+        this._bashOperations.set(controller, completion);
+        this._operations.add(completion);
         try {
+            // Apply the configured shell prefix within this operation's cleanup boundary.
+            const prefix = this.settingsManager.getShellCommandPrefix();
+            const shellPath = this.settingsManager.getShellPath();
+            const resolvedCommand = prefix ? `${prefix}\n${command}` : command;
             const result = await executeBashWithOperations(resolvedCommand, this.sessionManager.getCwd(), options?.operations ?? createLocalBashOperations({ shellPath, sessionId: this.sessionId }), {
                 onChunk,
-                signal: this._bashAbortController.signal,
+                signal: controller.signal,
             });
             this.recordBashResult(command, result, options);
             return result;
         }
         finally {
-            this._bashAbortController = undefined;
+            this._bashOperations.delete(controller);
+            this._operations.delete(completion);
+            settle();
         }
     }
     /**
@@ -2676,14 +2907,15 @@ export class AgentSession {
         }
     }
     /**
-     * Cancel running bash command.
+     * Cancel every running direct bash operation.
      */
     abortBash() {
-        this._bashAbortController?.abort();
+        for (const controller of this._bashOperations.keys())
+            controller.abort();
     }
     /** Whether a bash command is currently running */
     get isBashRunning() {
-        return this._bashAbortController !== undefined;
+        return this._bashOperations.size > 0;
     }
     /** Whether there are pending bash messages waiting to be flushed */
     get hasPendingBashMessages() {
@@ -2731,6 +2963,12 @@ export class AgentSession {
      * @returns Result with editorText (if user message) and cancelled status
      */
     async navigateTree(targetId, options = {}) {
+        this.assertCanStartWork();
+        if (this._branchSummaryAbortController)
+            throw new Error("Tree navigation is already in progress.");
+        return this._trackOperation(() => this._navigateTree(targetId, options));
+    }
+    async _navigateTree(targetId, options = {}) {
         const oldLeafId = this.sessionManager.getLeafId();
         // No-op if already at target
         if (targetId === oldLeafId) {
@@ -2807,6 +3045,7 @@ export class AgentSession {
                     replaceInstructions,
                     reserveTokens: branchSummarySettings.reserveTokens,
                     streamFn: this.agent.streamFn,
+                    onUsage: (response) => this.sessionManager.appendRequestUsage("branch-summary", response),
                 });
                 if (result.aborted) {
                     return { cancelled: true, aborted: true };
@@ -2824,6 +3063,8 @@ export class AgentSession {
                 summaryText = extensionSummary.summary;
                 summaryDetails = extensionSummary.details;
             }
+            if (this._branchSummaryAbortController.signal.aborted)
+                return { cancelled: true, aborted: true };
             // Determine the new leaf position based on target type
             let newLeafId;
             let editorText;
@@ -2929,7 +3170,7 @@ export class AgentSession {
     /**
      * Get session statistics. Aggregates over ALL session entries (including
      * history that was compacted away), so token/cost totals reflect what was
-     * actually billed across the session.
+     * reported or estimated across requests incurred in this session.
      */
     getSessionStats() {
         let userMessages = 0;
@@ -2937,11 +3178,12 @@ 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 entries = this.sessionManager.getEntries();
+        const parentUsage = totalRequestUsage(collectUsageRequests(entries));
+        const childUsage = totalChildUsage(collectChildUsage(entries));
+        const combinedUsage = { ...parentUsage };
+        for (const key of Object.keys(combinedUsage))
+            combinedUsage[key] += childUsage[key];
         for (const entry of this.sessionManager.getEntries()) {
             if (entry.type !== "message")
                 continue;
@@ -2959,12 +3201,6 @@ 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;
             }
         }
         return {
@@ -2976,13 +3212,16 @@ export class AgentSession {
             toolResults,
             totalMessages,
             tokens: {
-                input: totalInput,
-                output: totalOutput,
-                cacheRead: totalCacheRead,
-                cacheWrite: totalCacheWrite,
-                total: totalInput + totalOutput + totalCacheRead + totalCacheWrite,
+                input: parentUsage.input,
+                output: parentUsage.output,
+                cacheRead: parentUsage.cacheRead,
+                cacheWrite: parentUsage.cacheWrite,
+                total: parentUsage.total,
             },
-            cost: totalCost,
+            cost: parentUsage.cost,
+            parentUsage,
+            childUsage,
+            combinedUsage,
             contextUsage: this.getContextUsage(),
         };
     }
@@ -2993,38 +3232,17 @@ export class AgentSession {
         const contextWindow = model.contextWindow ?? 0;
         if (contextWindow <= 0)
             return undefined;
-        // After compaction, the last assistant usage reflects pre-compaction context size.
-        // We can only trust usage from an assistant that responded after the latest compaction.
-        // If no such assistant exists, context token count is unknown until the next LLM response.
-        const branchEntries = this.sessionManager.getBranch();
-        const latestCompaction = getLatestCompactionEntry(branchEntries);
-        if (latestCompaction) {
-            // Check if there's a valid assistant usage after the compaction boundary
-            const compactionIndex = branchEntries.lastIndexOf(latestCompaction);
-            let hasPostCompactionUsage = false;
-            for (let i = branchEntries.length - 1; i > compactionIndex; i--) {
-                const entry = branchEntries[i];
-                if (entry.type === "message" && entry.message.role === "assistant") {
-                    const assistant = entry.message;
-                    if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error") {
-                        const contextTokens = calculateContextTokens(assistant.usage);
-                        if (contextTokens > 0) {
-                            hasPostCompactionUsage = true;
-                            break;
-                        }
-                    }
-                }
-            }
-            if (!hasPostCompactionUsage) {
-                return { tokens: null, contextWindow, percent: null };
-            }
-        }
-        const estimate = estimateContextTokens(this.messages);
+        const estimate = estimateRequestContextTokens({
+            systemPrompt: this.systemPrompt,
+            tools: this.agent.state.tools,
+            messages: convertToLlm(this.messages),
+        });
         const percent = (estimate.tokens / contextWindow) * 100;
         return {
             tokens: estimate.tokens,
             contextWindow,
             percent,
+            estimated: estimate.lastUsageIndex === null || estimate.trailingTokens > 0,
         };
     }
     /**