@averyyy/pi-coding-agent 0.85.1-piclient.1 → 0.85.1-piclient.3

package/dist/core/agent-session.js

@@ -30,7 +30,7 @@ import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
 import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { emitSessionShutdownEvent } from "./extensions/runner.js";
 import { ModelRegistry } from "./model-registry.js";
-import { compactPiServer } from "./pi-server-client.js";
+import { cancelPiServerOperations, compactPiServer, } from "./pi-server-client.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
 import { CURRENT_SESSION_VERSION, getLatestCompactionEntry } from "./session-manager.js";
 import { createSyntheticSourceInfo } from "./source-info.js";
@@ -74,9 +74,6 @@ function estimateMessagesTokens(messages) {
 // ============================================================================
 const INTRA_TURN_COMPACTION_MARKER_TYPE = "pi:intra-turn-compaction";
 const INTRA_TURN_COMPACTION_MARKER_TEXT = "Continue from the compaction summary above. The previous tool loop was compacted before the next provider request; use the summary's Operational State, file lists, failures, last command, and next steps.";
-const VALIDATION_HINT_MESSAGE_TYPE = "pi:validation-hint";
-const VALIDATION_HINT_FILE_LIMIT = 12;
-const VALIDATION_HINT_FAILURE_CHAR_LIMIT = 1200;
 const AUTO_SESSION_NAME_INPUT_CHAR_LIMIT = 4000;
 const AUTO_SESSION_NAME_MAX_CHARS = 80;
 const AUTO_SESSION_NAME_PROMPT = "Name this coding session from the transcript below. Return only a concise title, in the user's language when clear. No quotes, no markdown, no explanation. Do not call tools.";
@@ -86,12 +83,6 @@ function isPiServerMode() {
 function toProviderReasoning(thinkingLevel) {
     return thinkingLevel === "off" ? undefined : thinkingLevel;
 }
-function getStringField(input, key) {
-    if (!input || typeof input !== "object" || Array.isArray(input))
-        return undefined;
-    const value = input[key];
-    return typeof value === "string" && value.trim().length > 0 ? value : undefined;
-}
 function getPiServerFailurePhase(message) {
     for (const diagnostic of message.diagnostics ?? []) {
         if (diagnostic.type !== "pi_server_failure")
@@ -112,19 +103,6 @@ function getPiServerFailureRetryable(message) {
     }
     return undefined;
 }
-function toolTextContent(content) {
-    return content
-        .filter((part) => part.type === "text")
-        .map((part) => part.text)
-        .join("\n")
-        .trim();
-}
-function compactValidationSummary(text) {
-    const trimmed = text.trim();
-    if (trimmed.length <= VALIDATION_HINT_FAILURE_CHAR_LIMIT)
-        return trimmed;
-    return `...${trimmed.slice(trimmed.length - VALIDATION_HINT_FAILURE_CHAR_LIMIT + 3)}`;
-}
 function messageTextContent(message) {
     if (!("content" in message))
         return "";
@@ -174,6 +152,10 @@ export class AgentSession {
     _isAgentRunActive = false;
     _idleWaitPromise;
     _resolveIdleWait;
+    _abortPromise;
+    _isAborting = false;
+    _abortError;
+    _continueAfterAbortRequested = false;
     /** Tracks pending steering messages for UI display. Removed when delivered. */
     _steeringMessages = [];
     /** Tracks pending follow-up messages for UI display. Removed when delivered. */
@@ -185,18 +167,23 @@ export class AgentSession {
     // Compaction state
     _compactionAbortController = undefined;
     _autoCompactionAbortController = undefined;
+    _compactionCompletionPromise = undefined;
+    _prePromptAbortController = undefined;
+    _prePromptCompletionPromise = undefined;
     _overflowRecoveryAttempted = false;
     _lastServerCompactionTimestamp = undefined;
     // Branch summarization state
     _branchSummaryAbortController = undefined;
+    _branchSummaryCompletionPromise = undefined;
     // Retry state
     _retryAbortController = undefined;
     _retryAttempt = 0;
+    // Post-turn session naming is a provider request too, so cancellation must reach it.
+    _sessionNameAbortController = undefined;
     // Bash execution state
     _bashAbortControllers = new Set();
+    _bashCompletionPromises = new Set();
     _pendingBashMessages = [];
-    _validationModifiedFiles = new Set();
-    _validationFailure = undefined;
     // Extension system
     _extensionRunner;
     _turnIndex = 0;
@@ -320,71 +307,6 @@ export class AgentSession {
             return { model };
         }
     }
-    _resetValidationHintState() {
-        this._validationModifiedFiles.clear();
-        this._validationFailure = undefined;
-    }
-    _recordValidationToolResult(toolName, input, content, isError) {
-        if ((toolName === "edit" || toolName === "write") && !isError) {
-            const path = getStringField(input, "path") ?? getStringField(input, "file_path");
-            if (path) {
-                this._validationModifiedFiles.add(path);
-            }
-            return;
-        }
-        if (toolName !== "bash") {
-            return;
-        }
-        const command = getStringField(input, "command");
-        if (!command) {
-            return;
-        }
-        if (isError) {
-            this._validationFailure = {
-                command,
-                summary: compactValidationSummary(toolTextContent(content)),
-            };
-        }
-        else if (this._validationFailure?.command === command) {
-            this._resetValidationHintState();
-        }
-        else if (!this._validationFailure) {
-            this._validationModifiedFiles.clear();
-        }
-    }
-    _createValidationHintMessage() {
-        const modifiedFiles = Array.from(this._validationModifiedFiles).slice(-VALIDATION_HINT_FILE_LIMIT);
-        if (modifiedFiles.length === 0 && !this._validationFailure) {
-            return undefined;
-        }
-        const lines = ["<validation-hint>"];
-        if (modifiedFiles.length > 0) {
-            lines.push("Recent modified files:", ...modifiedFiles.map((file) => `- ${file}`));
-        }
-        if (this._validationFailure) {
-            lines.push("Recent failed command:", this._validationFailure.command, "Failure summary:", this._validationFailure.summary || "(no output)", "Suggested next validation command:", this._validationFailure.command);
-        }
-        else {
-            lines.push("Suggested next step:", "Run the smallest project validation that covers the modified files.");
-        }
-        lines.push("Inspect failures, fix the root cause, then rerun validation before finalizing.", "</validation-hint>");
-        return {
-            role: "custom",
-            customType: VALIDATION_HINT_MESSAGE_TYPE,
-            content: [{ type: "text", text: lines.join("\n") }],
-            display: false,
-            details: { modifiedFiles, failedCommand: this._validationFailure?.command },
-            timestamp: Date.now(),
-        };
-    }
-    _withValidationHint(context) {
-        const messages = context.messages.filter((message) => !(message.role === "custom" && message.customType === VALIDATION_HINT_MESSAGE_TYPE));
-        const hint = this._createValidationHintMessage();
-        return {
-            ...context,
-            messages: hint ? [...messages, hint] : messages,
-        };
-    }
     /**
      * Install tool hooks once on the Agent instance.
      *
@@ -433,8 +355,6 @@ export class AgentSession {
             const normalizedContent = await normalizeToolResultImages(content, {
                 autoResizeImages: this.settingsManager.getImageAutoResize(),
             });
-            const finalIsError = hookResult?.isError ?? isError;
-            this._recordValidationToolResult(toolCall.name, args, normalizedContent, finalIsError);
             if (!hookResult && normalizedContent === content) {
                 return undefined;
             }
@@ -458,11 +378,11 @@ export class AgentSession {
             const previousContext = previousSnapshot?.context ?? nextTurn.context;
             return {
                 ...previousSnapshot,
-                context: this._withValidationHint({
+                context: {
                     ...previousContext,
                     systemPrompt: this._systemPromptOverride ?? this._baseSystemPrompt,
                     tools: this.agent.state.tools.slice(),
-                }),
+                },
                 model: this.agent.state.model,
                 thinkingLevel: this.agent.state.thinkingLevel,
             };
@@ -834,6 +754,14 @@ export class AgentSession {
     get isStreaming() {
         return this._isAgentRunActive;
     }
+    /** Whether the current run is being cancelled. New queued intent is drained after cancellation. */
+    get isAborting() {
+        return this._isAborting;
+    }
+    /** Error from the most recent cancellation attempt, if it failed. */
+    get abortError() {
+        return this._abortError;
+    }
     /** Whether the session has no active agent run, compaction, branch summary, retry, or queued continuation. */
     get isIdle() {
         return !this._isAgentRunActive && !this.isCompacting;
@@ -987,12 +915,16 @@ export class AgentSession {
     // =========================================================================
     // Prompting
     // =========================================================================
-    async _runAgentPrompt(messages) {
-        this._resetValidationHintState();
+    async _runAgentOperation(start) {
         this._isAgentRunActive = true;
         try {
-            await this.agent.prompt(messages);
+            await start();
             while (await this._handlePostAgentRun()) {
+                // Esc cancels the current agent run. Keep messages queued during that
+                // cancellation for the single deferred continuation in abort().
+                if (this._isAborting) {
+                    break;
+                }
                 try {
                     await this.agent.continue();
                 }
@@ -1010,7 +942,9 @@ export class AgentSession {
                     throw error;
                 }
             }
-            await this._maybeGenerateSessionName();
+            if (!this._isAborting) {
+                await this._maybeGenerateSessionName();
+            }
         }
         finally {
             this._systemPromptOverride = undefined;
@@ -1019,6 +953,9 @@ export class AgentSession {
             await this._emitAgentSettled();
         }
     }
+    async _runAgentPrompt(messages) {
+        await this._runAgentOperation(() => this.agent.prompt(messages));
+    }
     _sessionHasInfoEntry() {
         return this.sessionManager.getEntries().some((entry) => entry.type === "session_info");
     }
@@ -1052,6 +989,7 @@ export class AgentSession {
         const transcript = this._buildSessionNameTranscript();
         if (!model || !transcript)
             return;
+        this._sessionNameAbortController = new AbortController();
         try {
             const stream = await this.agent.streamFunction(model, {
                 systemPrompt: this.systemPrompt,
@@ -1070,6 +1008,7 @@ export class AgentSession {
                 transport: this.agent.transport,
                 thinkingBudgets: this.agent.thinkingBudgets,
                 maxRetryDelayMs: this.agent.maxRetryDelayMs,
+                signal: this._sessionNameAbortController.signal,
             });
             for await (const _event of stream) {
                 // Drain the title stream without surfacing it as a chat message.
@@ -1084,6 +1023,9 @@ export class AgentSession {
         catch {
             // Session naming is a post-turn convenience; the completed assistant turn stays authoritative.
         }
+        finally {
+            this._sessionNameAbortController = undefined;
+        }
     }
     async _handlePostAgentRun() {
         const msg = this._lastAssistantMessage;
@@ -1091,9 +1033,18 @@ export class AgentSession {
         if (!msg) {
             return false;
         }
+        if (this._isAborting) {
+            if (isPiServerMode() && this._isTerminalAssistantFailure(msg)) {
+                this._detachPiServerTerminalAssistantFailure();
+            }
+            return false;
+        }
         if (this._isRetryableError(msg) && (await this._prepareRetry(msg))) {
             return true;
         }
+        if (this._isAborting) {
+            return false;
+        }
         if (msg.stopReason === "error" && this._retryAttempt > 0) {
             this._emit({
                 type: "auto_retry_end",
@@ -1137,7 +1088,7 @@ export class AgentSession {
                     return;
                 }
             }
-            if (this._compactionAbortController !== undefined) {
+            if (this._compactionAbortController !== undefined && !this._isAborting) {
                 throw new Error("Cannot submit a prompt while compaction is in progress. Wait for compaction to finish and retry.");
             }
             // Emit input event for extension interception (before skill/template expansion)
@@ -1161,7 +1112,7 @@ export class AgentSession {
                 expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]);
             }
             // If streaming, queue via steer() or followUp() based on option
-            if (this.isStreaming) {
+            if (this.isStreaming || this._isAborting) {
                 if (!options?.streamingBehavior) {
                     throw new Error("Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.");
                 }
@@ -1194,9 +1145,35 @@ export class AgentSession {
             }
             // Check if we need to compact before sending (catches aborted responses).
             // The user's new prompt is sent below, so do not call agent.continue() here.
+            if (this._isAborting) {
+                preflightResult?.(false);
+                return;
+            }
             const lastAssistant = this._findLastAssistantMessage();
             if (lastAssistant) {
-                await this._checkCompaction(lastAssistant, false);
+                const prePromptAbortController = new AbortController();
+                let resolvePrePromptCompletion;
+                const prePromptCompletion = new Promise((resolve) => {
+                    resolvePrePromptCompletion = resolve;
+                });
+                this._prePromptAbortController = prePromptAbortController;
+                this._prePromptCompletionPromise = prePromptCompletion;
+                try {
+                    await this._checkCompaction(lastAssistant, false, prePromptAbortController.signal);
+                }
+                finally {
+                    if (this._prePromptAbortController === prePromptAbortController) {
+                        this._prePromptAbortController = undefined;
+                    }
+                    if (this._prePromptCompletionPromise === prePromptCompletion) {
+                        this._prePromptCompletionPromise = undefined;
+                    }
+                    resolvePrePromptCompletion?.();
+                }
+                if (prePromptAbortController.signal.aborted || this._isAborting) {
+                    preflightResult?.(false);
+                    return;
+                }
             }
             if (isPiServerMode()) {
                 this._detachPiServerTerminalAssistantFailure();
@@ -1362,6 +1339,9 @@ export class AgentSession {
             content,
             timestamp: Date.now(),
         });
+        if (this._isAborting) {
+            this._continueAfterAbortRequested = true;
+        }
     }
     /**
      * Internal: Queue a follow-up message (already expanded, no extension command check).
@@ -1378,6 +1358,9 @@ export class AgentSession {
             content,
             timestamp: Date.now(),
         });
+        if (this._isAborting) {
+            this._continueAfterAbortRequested = true;
+        }
     }
     /**
      * Throw an error if the text is an extension command.
@@ -1416,12 +1399,18 @@ export class AgentSession {
         if (options?.deliverAs === "nextTurn") {
             this._pendingNextTurnMessages.push(appMessage);
         }
-        else if (this.isStreaming && options?.triggerTurn !== false) {
+        else if ((this.isStreaming || this._isAborting) && options?.triggerTurn !== false) {
             if (options?.deliverAs === "followUp") {
                 this.agent.followUp(appMessage);
+                if (this._isAborting) {
+                    this._continueAfterAbortRequested = true;
+                }
             }
             else {
                 this.agent.steer(appMessage);
+                if (this._isAborting) {
+                    this._continueAfterAbortRequested = true;
+                }
             }
         }
         else if (options?.triggerTurn) {
@@ -1523,15 +1512,79 @@ export class AgentSession {
     get resourceLoader() {
         return this._resourceLoader;
     }
-    /**
-     * Abort current operation and wait for agent to become idle.
-     */
-    async abort() {
+    _startQueuedAgentRunAfterAbort() {
+        if (this._isAgentRunActive || this._isAborting || !this.agent.hasQueuedMessages()) {
+            return;
+        }
+        void this._runAgentOperation(() => this.agent.continue()).catch((error) => {
+            this._extensionRunner.emitError({
+                extensionPath: "<runtime>",
+                event: "abort_continuation",
+                error: error instanceof Error ? error.message : String(error),
+            });
+        });
+    }
+    async _performAbort() {
         this.abortRetry();
         this.abortCompaction();
         this.abortBranchSummary();
+        this.abortBash();
+        this._prePromptAbortController?.abort();
+        this._sessionNameAbortController?.abort();
         this.agent.abort();
-        await this.waitForIdle();
+        // _runAgentOperation owns the settled event and the idle waiter. Do not
+        // start a queued continuation until that cleanup has completed.
+        try {
+            const cleanupPromises = [
+                this.waitForIdle(),
+                ...(this._compactionCompletionPromise ? [this._compactionCompletionPromise] : []),
+                ...(this._prePromptCompletionPromise ? [this._prePromptCompletionPromise] : []),
+                ...(this._branchSummaryCompletionPromise ? [this._branchSummaryCompletionPromise] : []),
+                ...this._bashCompletionPromises,
+            ];
+            if (isPiServerMode()) {
+                cleanupPromises.push(cancelPiServerOperations(this.sessionId));
+            }
+            const cleanupResults = await Promise.allSettled(cleanupPromises);
+            const cleanupFailure = cleanupResults.find((result) => result.status === "rejected");
+            if (cleanupFailure?.status === "rejected") {
+                throw cleanupFailure.reason;
+            }
+        }
+        catch (error) {
+            this._abortError = error instanceof Error ? error.message : String(error);
+            this._abortPromise = undefined;
+            this._emit({ type: "abort_error", errorMessage: this._abortError });
+            throw error;
+        }
+        this._isAborting = false;
+        const continueAfterAbort = this._continueAfterAbortRequested && this.agent.hasQueuedMessages();
+        this._continueAfterAbortRequested = false;
+        this._abortError = undefined;
+        this._abortPromise = undefined;
+        this._emit({ type: "abort_end" });
+        if (continueAfterAbort) {
+            this._startQueuedAgentRunAfterAbort();
+        }
+    }
+    /**
+     * Abort the current operation and wait for cancellation cleanup to finish.
+     * Messages queued while cancellation is in progress are sent in one deferred continuation.
+     */
+    abort() {
+        if (this._abortPromise) {
+            return this._abortPromise;
+        }
+        if (this._isAborting) {
+            this._abortError = undefined;
+        }
+        else {
+            this._isAborting = true;
+        }
+        const abortPromise = this._performAbort();
+        this._abortPromise = abortPromise;
+        this._emit({ type: "abort_start" });
+        return abortPromise;
     }
     async waitForIdle() {
         if (this.isIdle) {
@@ -1798,8 +1851,13 @@ export class AgentSession {
     async compact(customInstructions) {
         await this.abort();
         this._compactionAbortController = new AbortController();
-        this._emit({ type: "compaction_start", reason: "manual" });
+        let resolveCompactionCompletion;
+        const compactionCompletion = new Promise((resolve) => {
+            resolveCompactionCompletion = resolve;
+        });
+        this._compactionCompletionPromise = compactionCompletion;
         try {
+            this._emit({ type: "compaction_start", reason: "manual" });
             if (!this.model) {
                 throw new Error(formatNoModelSelectedMessage());
             }
@@ -1941,6 +1999,11 @@ export class AgentSession {
         }
         finally {
             this._clearManualCompactionState();
+            this._compactionAbortController = undefined;
+            if (this._compactionCompletionPromise === compactionCompletion) {
+                this._compactionCompletionPromise = undefined;
+            }
+            resolveCompactionCompletion?.();
         }
     }
     /**
@@ -1968,7 +2031,9 @@ export class AgentSession {
      * @param assistantMessage The assistant message to check
      * @param skipAbortedCheck If false, include aborted messages (for pre-prompt check). Default: true
      */
-    async _checkCompaction(assistantMessage, skipAbortedCheck = true) {
+    async _checkCompaction(assistantMessage, skipAbortedCheck = true, externalSignal) {
+        if (externalSignal?.aborted)
+            return false;
         const settings = this.settingsManager.getCompactionSettings();
         if (!settings.enabled)
             return false;
@@ -2001,7 +2066,7 @@ export class AgentSession {
         if (contextOverflow || recoverableLength) {
             const willRetry = assistantMessage.stopReason !== "stop";
             if (!willRetry) {
-                return await this._runAutoCompaction("overflow", false);
+                return await this._runAutoCompaction("overflow", false, undefined, undefined, externalSignal);
             }
             if (this._overflowRecoveryAttempted) {
                 const errorMessage = contextOverflow
@@ -2032,7 +2097,7 @@ export class AgentSession {
                     this.agent.state.messages = messages.slice(0, -1);
                 }
             }
-            return await this._runAutoCompaction("overflow", willRetry, terminalFailurePreparation);
+            return await this._runAutoCompaction("overflow", willRetry, terminalFailurePreparation, undefined, externalSignal);
         }
         // Case 2: Threshold - context is getting large
         // For error messages or all-zero usage messages, estimate from the last valid response.
@@ -2060,25 +2125,33 @@ export class AgentSession {
             contextTokens = directContextTokens;
         }
         if (shouldCompact(contextTokens, contextWindow, settings)) {
-            return await this._runAutoCompaction("threshold", false);
+            return await this._runAutoCompaction("threshold", false, undefined, undefined, externalSignal);
         }
         return false;
     }
     /**
      * Internal: Run auto-compaction with events.
      */
-    async _runAutoCompaction(reason, willRetry, preparationOverride, preparationOptions) {
+    async _runAutoCompaction(reason, willRetry, preparationOverride, preparationOptions, externalSignal) {
         const settings = this.settingsManager.getCompactionSettings();
         let started = false;
         let fromExtension = false;
+        if (externalSignal?.aborted)
+            return false;
+        const abortExternal = () => this._autoCompactionAbortController?.abort();
+        externalSignal?.addEventListener("abort", abortExternal, { once: true });
         try {
             if (!this.model) {
                 return false;
             }
             const { model: requestModel, apiKey, headers, env } = await this._getSummarizationRequestAuth(this.model);
+            if (externalSignal?.aborted)
+                return false;
             if (isPiServerMode()) {
                 this._emit({ type: "compaction_start", reason });
                 this._autoCompactionAbortController = new AbortController();
+                if (externalSignal?.aborted)
+                    this._autoCompactionAbortController.abort();
                 started = true;
                 const signal = this._autoCompactionAbortController.signal;
                 const piServerResult = await compactPiServer(requestModel, {
@@ -2096,18 +2169,26 @@ export class AgentSession {
                     reasoning: toProviderReasoning(this.thinkingLevel),
                     onHistoryReconciled: (snapshot) => this.reconcilePiServerHistory(snapshot),
                 });
+                if (externalSignal?.aborted || signal.aborted) {
+                    this._emit({ type: "compaction_end", reason, result: undefined, aborted: true, willRetry: false });
+                    return false;
+                }
                 const result = await this._applyPiServerCompactionResult(piServerResult, reason, willRetry);
                 this._lastServerCompactionTimestamp = Date.now();
                 this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry });
                 return willRetry ? true : this.agent.hasQueuedMessages();
             }
             const pathEntries = this.sessionManager.getBranch();
+            if (externalSignal?.aborted)
+                return false;
             const preparation = preparationOverride ?? prepareCompaction(pathEntries, settings);
             if (!preparation) {
                 return false;
             }
             this._emit({ type: "compaction_start", reason });
             this._autoCompactionAbortController = new AbortController();
+            if (externalSignal?.aborted)
+                this._autoCompactionAbortController.abort();
             started = true;
             let extensionCompaction;
             if (this._extensionRunner.hasHandlers("session_before_compact")) {
@@ -2163,7 +2244,7 @@ export class AgentSession {
                 usage = compactResult.usage;
                 details = compactResult.details;
             }
-            if (this._autoCompactionAbortController.signal.aborted) {
+            if (externalSignal?.aborted || this._autoCompactionAbortController.signal.aborted) {
                 this._emit({
                     type: "compaction_end",
                     reason,
@@ -2246,6 +2327,7 @@ export class AgentSession {
             return false;
         }
         finally {
+            externalSignal?.removeEventListener("abort", abortExternal);
             this._autoCompactionAbortController = undefined;
             this._resolveIdleWaitIfIdle();
         }
@@ -2806,6 +2888,11 @@ export class AgentSession {
     async executeBash(command, onChunk, options) {
         const abortController = new AbortController();
         this._bashAbortControllers.add(abortController);
+        let resolveBashCompletion;
+        const bashCompletion = new Promise((resolve) => {
+            resolveBashCompletion = resolve;
+        });
+        this._bashCompletionPromises.add(bashCompletion);
         // Apply command prefix if configured (e.g., "shopt -s expand_aliases" for alias support)
         const prefix = this.settingsManager.getShellCommandPrefix();
         const shellPath = this.settingsManager.getShellPath();
@@ -2823,6 +2910,8 @@ export class AgentSession {
         }
         finally {
             this._bashAbortControllers.delete(abortController);
+            this._bashCompletionPromises.delete(bashCompletion);
+            resolveBashCompletion?.();
         }
     }
     /**
@@ -2945,6 +3034,11 @@ export class AgentSession {
         };
         // Set up abort controller for summarization
         this._branchSummaryAbortController = new AbortController();
+        let resolveBranchSummaryCompletion;
+        const branchSummaryCompletion = new Promise((resolve) => {
+            resolveBranchSummaryCompletion = resolve;
+        });
+        this._branchSummaryCompletionPromise = branchSummaryCompletion;
         try {
             let extensionSummary;
             let fromExtension = false;
@@ -3069,6 +3163,10 @@ export class AgentSession {
         }
         finally {
             this._branchSummaryAbortController = undefined;
+            if (this._branchSummaryCompletionPromise === branchSummaryCompletion) {
+                this._branchSummaryCompletionPromise = undefined;
+            }
+            resolveBranchSummaryCompletion?.();
             this._resolveIdleWaitIfIdle();
         }
     }