@astralyn/pi 0.82.0 → 0.82.2

package/dist/core/agent-session.js

@@ -70,6 +70,8 @@ function estimateMessagesTokens(messages) {
 // ============================================================================
 // Constants
 // ============================================================================
+const CONTEXT_REMAINS_OVER_COMPACTION_THRESHOLD = "Estimated retained context remains above the configured auto-compaction threshold. Reduce compaction.keepRecentTokens, remove large retained content, or switch to a model with a larger context window before continuing.";
+const AUTO_COMPACTION_STOPPED_BEFORE_NEXT_PROVIDER = "Stopped before the next provider request because auto-compaction did not produce a safe context. Review the compaction warning before continuing.";
 /** Standard thinking levels */
 const THINKING_LEVELS = ["off", "minimal", "low", "medium", "high"];
 // ============================================================================
@@ -95,6 +97,7 @@ export class AgentSession {
     // Compaction state
     _compactionAbortController = undefined;
     _autoCompactionAbortController = undefined;
+    _stopAfterTurnRequested = false;
     _overflowRecoveryAttempted = false;
     // Branch summarization state
     _branchSummaryAbortController = undefined;
@@ -189,8 +192,8 @@ export class AgentSession {
         }
         throw new Error(formatNoApiKeyFoundMessage(model.provider));
     }
-    async _getSummarizationRequestAuth(model) {
-        if (this.agent.streamFunction === streamSimple) {
+    async _getSummarizationRequestAuth(model, streamFunction = this.agent.streamFunction) {
+        if (streamFunction === streamSimple) {
             return this._getRequiredRequestAuth(model);
         }
         try {
@@ -264,8 +267,25 @@ export class AgentSession {
                 ? async (_turn, signal) => await this.agent.prepareNextTurn?.(signal)
                 : undefined);
         this.agent.prepareNextTurnWithContext = async (turn, signal) => {
-            const previousSnapshot = await previousPrepareNextTurnWithContext?.(turn, signal);
-            const previousContext = previousSnapshot?.context ?? turn.context;
+            const compaction = await this._maybeCompactBeforeNextToolTurn(turn, signal);
+            if (compaction.stopErrorMessage) {
+                // Agent does not expose the low-level graceful turn-stop hook. Throwing here
+                // exits before queue polling or another provider request; the run records a
+                // normal error/aborted lifecycle instead of risking an oversized context.
+                this._stopAfterTurnRequested = true;
+                throw new Error(compaction.stopErrorMessage);
+            }
+            const effectiveTurn = compaction.messages === undefined
+                ? turn
+                : {
+                    ...turn,
+                    context: {
+                        ...turn.context,
+                        messages: compaction.messages,
+                    },
+                };
+            const previousSnapshot = await previousPrepareNextTurnWithContext?.(effectiveTurn, signal);
+            const previousContext = previousSnapshot?.context ?? effectiveTurn.context;
             return {
                 ...previousSnapshot,
                 context: {
@@ -743,6 +763,7 @@ export class AgentSession {
     // =========================================================================
     async _runAgentPrompt(messages) {
         this._isAgentRunActive = true;
+        this._stopAfterTurnRequested = false;
         try {
             await this.agent.prompt(messages);
             while (await this._handlePostAgentRun()) {
@@ -759,6 +780,12 @@ export class AgentSession {
         const msg = this._lastAssistantMessage;
         this._lastAssistantMessage = undefined;
         if (!msg) {
+            this._stopAfterTurnRequested = false;
+            return false;
+        }
+        const stopAfterTurnRequested = this._stopAfterTurnRequested;
+        this._stopAfterTurnRequested = false;
+        if (stopAfterTurnRequested) {
             return false;
         }
         if (this._isRetryableError(msg) && (await this._prepareRetry(msg))) {
@@ -776,6 +803,10 @@ export class AgentSession {
         if (await this._checkCompaction(msg)) {
             return true;
         }
+        if (this._stopAfterTurnRequested) {
+            this._stopAfterTurnRequested = false;
+            return false;
+        }
         // 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();
@@ -861,6 +892,13 @@ export class AgentSession {
             if (lastAssistant) {
                 await this._checkCompaction(lastAssistant, false);
             }
+            const compactionSettings = this.settingsManager.getCompactionSettings();
+            // A compaction entry can be valid while its retained suffix is still too large.
+            // Refuse the existing context here rather than discovering that via provider overflow.
+            if (compactionSettings.enabled &&
+                shouldCompact(this._estimateContextTokensForCompaction(this.agent.state.messages), this.model.contextWindow, compactionSettings)) {
+                throw new Error(CONTEXT_REMAINS_OVER_COMPACTION_THRESHOLD);
+            }
             // Build messages array (custom message if any, then user message)
             messages = [];
             // Add user message
@@ -1164,6 +1202,7 @@ export class AgentSession {
      */
     async abort() {
         this.abortRetry();
+        this.abortCompaction();
         this.agent.abort();
         await this.waitForIdle();
     }
@@ -1494,6 +1533,49 @@ export class AgentSession {
     abortBranchSummary() {
         this._branchSummaryAbortController?.abort();
     }
+    _estimateContextTokensForCompaction(messages) {
+        const estimate = estimateContextTokens(messages);
+        const compactionEntry = getLatestCompactionEntry(this.sessionManager.getBranch());
+        if (compactionEntry && estimate.lastUsageIndex !== null) {
+            const usageMessage = messages[estimate.lastUsageIndex];
+            if (usageMessage?.role === "assistant" &&
+                usageMessage.timestamp <= new Date(compactionEntry.timestamp).getTime()) {
+                // A retained pre-compaction usage describes the old, larger context. Fall back
+                // to message estimates rather than immediately compacting again from stale data.
+                return estimateMessagesTokens(messages);
+            }
+        }
+        return estimate.tokens;
+    }
+    async _maybeCompactBeforeNextToolTurn(turn, signal) {
+        if (signal?.aborted) {
+            return { stopErrorMessage: AUTO_COMPACTION_STOPPED_BEFORE_NEXT_PROVIDER };
+        }
+        if (turn.toolResults.length === 0) {
+            return {};
+        }
+        const settings = this.settingsManager.getCompactionSettings();
+        if (!settings.enabled) {
+            return {};
+        }
+        const contextTokens = this._estimateContextTokensForCompaction(turn.context.messages);
+        const contextWindow = this.model?.contextWindow ?? 0;
+        if (!shouldCompact(contextTokens, contextWindow, settings)) {
+            return {};
+        }
+        const outcome = await this._runAutoCompactionWithOutcome("threshold", false, signal);
+        if (!outcome.compacted) {
+            return { stopErrorMessage: AUTO_COMPACTION_STOPPED_BEFORE_NEXT_PROVIDER };
+        }
+        return {
+            messages: this.agent.state.messages.slice(),
+            ...(signal?.aborted === true
+                ? { stopErrorMessage: AUTO_COMPACTION_STOPPED_BEFORE_NEXT_PROVIDER }
+                : outcome.mustStop === true
+                    ? { stopErrorMessage: CONTEXT_REMAINS_OVER_COMPACTION_THRESHOLD }
+                    : {}),
+        };
+    }
     /**
      * Check if compaction is needed and run it.
      * Called after agent_end and before prompt submission.
@@ -1536,6 +1618,7 @@ export class AgentSession {
                 return await this._runAutoCompaction("overflow", false);
             }
             if (this._overflowRecoveryAttempted) {
+                this._stopAfterTurnRequested = true;
                 this._emit({
                     type: "compaction_end",
                     reason: "overflow",
@@ -1589,31 +1672,57 @@ export class AgentSession {
      * Internal: Run auto-compaction with events.
      */
     async _runAutoCompaction(reason, willRetry) {
+        const outcome = await this._runAutoCompactionWithOutcome(reason, willRetry);
+        if (!outcome.compacted || outcome.mustStop) {
+            this._stopAfterTurnRequested = true;
+        }
+        return outcome.shouldContinue;
+    }
+    async _runAutoCompactionWithOutcome(reason, willRetry, parentSignal) {
         const settings = this.settingsManager.getCompactionSettings();
         let started = false;
+        let controller;
+        let removeParentAbortListener;
         try {
-            if (!this.model) {
-                return false;
-            }
-            let apiKey;
-            let headers;
-            let env;
-            if (this.agent.streamFunction === streamSimple) {
-                ({ apiKey, headers, env } = await this._getRequiredRequestAuth(this.model));
-            }
-            else {
-                ({ apiKey, headers, env } = await this._getSummarizationRequestAuth(this.model));
+            if (!this.model ||
+                this._autoCompactionAbortController ||
+                this._compactionAbortController ||
+                parentSignal?.aborted) {
+                return { compacted: false, shouldContinue: false };
+            }
+            const model = this.model;
+            const streamFunction = this.agent.streamFunction;
+            const thinkingLevel = this.thinkingLevel;
+            const retrySettings = this.settingsManager.getRetrySettings();
+            controller = new AbortController();
+            if (parentSignal) {
+                const abortFromParent = () => controller?.abort();
+                parentSignal.addEventListener("abort", abortFromParent, { once: true });
+                removeParentAbortListener = () => parentSignal.removeEventListener("abort", abortFromParent);
+            }
+            if (controller.signal.aborted) {
+                return { compacted: false, shouldContinue: false };
+            }
+            this._autoCompactionAbortController = controller;
+            started = true;
+            this._emit({ type: "compaction_start", reason });
+            if (controller.signal.aborted) {
+                this._emit({
+                    type: "compaction_end",
+                    reason,
+                    result: undefined,
+                    aborted: true,
+                    willRetry: false,
+                });
+                return { compacted: false, shouldContinue: false };
             }
+            let extensionCompaction;
+            let fromExtension = false;
             const pathEntries = this.sessionManager.getBranch();
             const preparation = prepareCompaction(pathEntries, settings);
             if (!preparation) {
-                return false;
+                throw new Error("Nothing to compact while preserving the configured recent context.");
             }
-            this._emit({ type: "compaction_start", reason });
-            this._autoCompactionAbortController = new AbortController();
-            started = true;
-            let extensionCompaction;
-            let fromExtension = false;
             if (this._extensionRunner.hasHandlers("session_before_compact")) {
                 const extensionResult = (await this._extensionRunner.emit({
                     type: "session_before_compact",
@@ -1622,7 +1731,7 @@ export class AgentSession {
                     customInstructions: undefined,
                     reason,
                     willRetry,
-                    signal: this._autoCompactionAbortController.signal,
+                    signal: controller.signal,
                 }));
                 if (extensionResult?.cancel) {
                     this._emit({
@@ -1632,7 +1741,7 @@ export class AgentSession {
                         aborted: true,
                         willRetry: false,
                     });
-                    return false;
+                    return { compacted: false, shouldContinue: false };
                 }
                 if (extensionResult?.compaction) {
                     extensionCompaction = extensionResult.compaction;
@@ -1653,15 +1762,34 @@ export class AgentSession {
                 details = extensionCompaction.details;
             }
             else {
+                let apiKey;
+                let headers;
+                let env;
+                if (streamFunction === streamSimple) {
+                    ({ apiKey, headers, env } = await this._getRequiredRequestAuth(model));
+                }
+                else {
+                    ({ apiKey, headers, env } = await this._getSummarizationRequestAuth(model, streamFunction));
+                }
+                if (controller.signal.aborted) {
+                    this._emit({
+                        type: "compaction_end",
+                        reason,
+                        result: undefined,
+                        aborted: true,
+                        willRetry: false,
+                    });
+                    return { compacted: false, shouldContinue: false };
+                }
                 // Generate compaction result
-                const compactResult = await compact(preparation, this.model, apiKey, headers, undefined, this._autoCompactionAbortController.signal, this.thinkingLevel, this.agent.streamFunction, env, this.settingsManager.getRetrySettings(), this._summarizationRetryCallbacks({ source: "compaction", reason }));
+                const compactResult = await compact(preparation, model, apiKey, headers, undefined, controller.signal, thinkingLevel, streamFunction, env, retrySettings, this._summarizationRetryCallbacks({ source: "compaction", reason }));
                 summary = compactResult.summary;
                 firstKeptEntryId = compactResult.firstKeptEntryId;
                 tokensBefore = compactResult.tokensBefore;
                 usage = compactResult.usage;
                 details = compactResult.details;
             }
-            if (this._autoCompactionAbortController.signal.aborted) {
+            if (controller.signal.aborted) {
                 this._emit({
                     type: "compaction_end",
                     reason,
@@ -1669,13 +1797,22 @@ export class AgentSession {
                     aborted: true,
                     willRetry: false,
                 });
-                return false;
+                return { compacted: false, shouldContinue: false };
             }
             this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage);
             const newEntries = this.sessionManager.getEntries();
             const sessionContext = this.sessionManager.buildSessionContext();
             this.agent.state.messages = sessionContext.messages;
-            const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages);
+            if (willRetry) {
+                const messages = this.agent.state.messages;
+                const lastMsg = messages[messages.length - 1];
+                if (lastMsg?.role === "assistant" && lastMsg.stopReason === "error") {
+                    this.agent.state.messages = messages.slice(0, -1);
+                }
+            }
+            const estimatedTokensAfter = this._estimateContextTokensForCompaction(this.agent.state.messages);
+            const contextStillOverThreshold = shouldCompact(estimatedTokensAfter, model.contextWindow, settings);
+            const effectiveWillRetry = willRetry && !contextStillOverThreshold;
             // Get the saved compaction entry for the extension event
             const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary);
             if (this._extensionRunner && savedCompactionEntry) {
@@ -1684,7 +1821,7 @@ export class AgentSession {
                     compactionEntry: savedCompactionEntry,
                     fromExtension,
                     reason,
-                    willRetry,
+                    willRetry: effectiveWillRetry,
                 });
             }
             const result = {
@@ -1695,37 +1832,50 @@ export class AgentSession {
                 usage,
                 details,
             };
-            this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry });
-            if (willRetry) {
-                const messages = this.agent.state.messages;
-                const lastMsg = messages[messages.length - 1];
-                if (lastMsg?.role === "assistant" && lastMsg.stopReason === "error") {
-                    this.agent.state.messages = messages.slice(0, -1);
-                }
-                return true;
+            this._emit({
+                type: "compaction_end",
+                reason,
+                result,
+                aborted: false,
+                willRetry: effectiveWillRetry,
+                ...(contextStillOverThreshold ? { errorMessage: CONTEXT_REMAINS_OVER_COMPACTION_THRESHOLD } : {}),
+            });
+            if (contextStillOverThreshold) {
+                return { compacted: true, shouldContinue: false, mustStop: true };
+            }
+            if (effectiveWillRetry) {
+                return { compacted: true, shouldContinue: true };
             }
             // Auto-compaction can complete while follow-up/steering/custom messages are waiting.
             // Continue once so queued messages are delivered.
-            return this.agent.hasQueuedMessages();
+            return { compacted: true, shouldContinue: this.agent.hasQueuedMessages() };
         }
         catch (error) {
+            const aborted = controller?.signal.aborted || (error instanceof Error && error.name === "AbortError");
             const errorMessage = error instanceof Error ? error.message : "compaction failed";
             if (started) {
                 this._emit({
                     type: "compaction_end",
                     reason,
                     result: undefined,
-                    aborted: false,
+                    aborted,
                     willRetry: false,
-                    errorMessage: reason === "overflow"
-                        ? `Context overflow recovery failed: ${errorMessage}`
-                        : `Auto-compaction failed: ${errorMessage}`,
+                    ...(aborted
+                        ? {}
+                        : {
+                            errorMessage: reason === "overflow"
+                                ? `Context overflow recovery failed: ${errorMessage}`
+                                : `Auto-compaction failed: ${errorMessage}`,
+                        }),
                 });
             }
-            return false;
+            return { compacted: false, shouldContinue: false };
         }
         finally {
-            this._autoCompactionAbortController = undefined;
+            removeParentAbortListener?.();
+            if (this._autoCompactionAbortController === controller) {
+                this._autoCompactionAbortController = undefined;
+            }
         }
     }
     /**