@averyyy/pi-coding-agent 0.85.1-piclient.2 → 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";
@@ -152,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. */
@@ -163,15 +167,22 @@ 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 = [];
     // Extension system
     _extensionRunner;
@@ -743,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;
@@ -896,11 +915,16 @@ export class AgentSession {
     // =========================================================================
     // Prompting
     // =========================================================================
-    async _runAgentPrompt(messages) {
+    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();
                 }
@@ -918,7 +942,9 @@ export class AgentSession {
                     throw error;
                 }
             }
-            await this._maybeGenerateSessionName();
+            if (!this._isAborting) {
+                await this._maybeGenerateSessionName();
+            }
         }
         finally {
             this._systemPromptOverride = undefined;
@@ -927,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");
     }
@@ -960,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,
@@ -978,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.
@@ -992,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;
@@ -999,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",
@@ -1045,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)
@@ -1069,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.");
                 }
@@ -1102,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();
@@ -1270,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).
@@ -1286,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.
@@ -1324,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) {
@@ -1431,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) {
@@ -1706,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());
             }
@@ -1849,6 +1999,11 @@ export class AgentSession {
         }
         finally {
             this._clearManualCompactionState();
+            this._compactionAbortController = undefined;
+            if (this._compactionCompletionPromise === compactionCompletion) {
+                this._compactionCompletionPromise = undefined;
+            }
+            resolveCompactionCompletion?.();
         }
     }
     /**
@@ -1876,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;
@@ -1909,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
@@ -1940,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.
@@ -1968,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, {
@@ -2004,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")) {
@@ -2071,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,
@@ -2154,6 +2327,7 @@ export class AgentSession {
             return false;
         }
         finally {
+            externalSignal?.removeEventListener("abort", abortExternal);
             this._autoCompactionAbortController = undefined;
             this._resolveIdleWaitIfIdle();
         }
@@ -2714,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();
@@ -2731,6 +2910,8 @@ export class AgentSession {
         }
         finally {
             this._bashAbortControllers.delete(abortController);
+            this._bashCompletionPromises.delete(bashCompletion);
+            resolveBashCompletion?.();
         }
     }
     /**
@@ -2853,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;
@@ -2977,6 +3163,10 @@ export class AgentSession {
         }
         finally {
             this._branchSummaryAbortController = undefined;
+            if (this._branchSummaryCompletionPromise === branchSummaryCompletion) {
+                this._branchSummaryCompletionPromise = undefined;
+            }
+            resolveBranchSummaryCompletion?.();
             this._resolveIdleWaitIfIdle();
         }
     }