@astralyn/pi 1.3.1 → 1.4.1

package/dist/core/agent-session.js

@@ -16,7 +16,7 @@ import { createTasksToolDefinition } from "./tools/tasks.js";
 import { readFileSync } from "node:fs";
 import { basename, dirname } from "node:path";
 import { runToolCall, } from "@earendil-works/pi-agent-core";
-import { contentText, getCurrentSystemMessage, retryDelayMs } from "@earendil-works/pi-ai";
+import { contentText, getCurrentSystemMessage, getCurrentTools, retryDelayMs } from "@earendil-works/pi-ai";
 import { clampThinkingLevel, cleanupSessionResources, getSupportedThinkingLevels, isContextOverflow, isRecoverableLength, isRetryableAssistantError, modelsAreEqual, resetApiProviders, streamSimple, } from "@earendil-works/pi-ai/compat";
 import { getThemeByName, theme } from "../modes/interactive/theme/theme.js";
 import { stripFrontmatter } from "../utils/frontmatter.js";
@@ -110,7 +110,6 @@ export class AgentSession {
         this._turnIndex = 0;
         this._entryIdsByMessage = new WeakMap();
         this._boundaryDispatchedMessages = new WeakSet();
-        this._lastAssistantToolResults = [];
         this._lastActivityOutcome = "completed";
         this._isBeforeSettle = false;
         this._abortDuringBeforeSettle = false;
@@ -140,6 +139,11 @@ export class AgentSession {
         this._handleAgentEvent = async (event) => {
             if (event.type === "message_start")
                 this._tasksHost.delivery.messageStarted(event.message);
+            if (event.type === "message_end" &&
+                event.message.role === "assistant" &&
+                event.message.stopReason !== "error" &&
+                event.message.stopReason !== "aborted")
+                this._tasksHost.noteAssistantCompleted();
             // An abort requested while the run was starting, e.g. during before_agent_start, stops it
             // here. agent_start is awaited before the first request.
             if (event.type === "agent_start" && this._agentRunAbortRequested)
@@ -228,7 +232,6 @@ export class AgentSession {
             // extension and listener dispatch above also picks up messages that turn_end
             // handlers queued.
             if (event.type === "turn_end") {
-                this._lastAssistantToolResults = event.toolResults;
                 this._flushPendingCustomMessages();
             }
             if (event.type === "agent_end")
@@ -251,6 +254,7 @@ export class AgentSession {
             deliver: (message, waitForPersistence) => {
                 if (this.isStreaming) {
                     this.agent.steer(message);
+                    this._tasksHost.interruptWaits();
                     return waitForPersistence();
                 }
                 // An idle delivery starts its own custom-message run, which consumes the
@@ -291,6 +295,7 @@ export class AgentSession {
         this._installAgentBoundaryHooks();
         this._installHiddenDeclarationsProjection();
         this._installAgentForcedPromptProjection();
+        this._installTaskContextReminderProjection();
         this._buildRuntime({
             activeToolNames: this._initialActiveToolNames,
             includeAllExtensionTools: true,
@@ -883,7 +888,24 @@ export class AgentSession {
         }
         return undefined;
     }
-    _omitRecoveryAttempt(message, toolResults = []) {
+    /**
+     * Persistently omit a failed attempt and its tool results from model context. Results are
+     * found in the projected context, so the pair stays together whether the attempt comes from
+     * the current run or a resumed session. An unpaired result fails the next request (#11227).
+     */
+    _omitRecoveryAttempt(message) {
+        const toolCallIds = new Set(message.content.flatMap((block) => (block.type === "toolCall" ? [block.id] : [])));
+        // A context edit replaces the projected copy, so match it by source entry as well.
+        const entryId = this._findPersistedMessageEntryId(message);
+        const messages = this.agent.state.messages;
+        const index = messages.findIndex((m) => m === message || (entryId !== undefined && this._entryIdsByMessage.get(m) === entryId));
+        const toolResults = [];
+        for (const m of index < 0 ? [] : messages.slice(index + 1)) {
+            if (m.role === "assistant")
+                break;
+            if (m.role === "toolResult" && toolCallIds.has(m.toolCallId))
+                toolResults.push(m);
+        }
         const targets = [message, ...toolResults];
         const targetIds = targets.map((target) => this._findPersistedMessageEntryId(target));
         const unresolvedProjectedTarget = targets.some((target, index) => targetIds[index] === undefined && this.agent.state.messages.includes(target));
@@ -1412,6 +1434,14 @@ export class AgentSession {
             return [head, ...transformed.filter((message) => message.role !== "system")];
         };
     }
+    /** Reminders live only in the request projection and survive failed requests until observed. */
+    _installTaskContextReminderProjection() {
+        const previousTransformContext = this.agent.transformContext;
+        this.agent.transformContext = async (messages, signal) => {
+            const transformed = previousTransformContext ? await previousTransformContext(messages, signal) : messages;
+            return this._tasksHost.projectContextReminder(transformed);
+        };
+    }
     /**
      * Restore the active tool loadout declared by the session transcript, if it declares one.
      * Tools reachable only from other tools are never declared, but they do not depend on the active
@@ -1420,11 +1450,27 @@ export class AgentSession {
      */
     _restoreToolsFromTranscript(modifiers = []) {
         this._pendingToolNames.clear();
-        const current = getCurrentSystemMessage(this.sessionManager.buildSessionContext().messages);
-        if (!current)
+        // All system messages on the branch, including those before a compaction. Compaction snapshots
+        // are derived from these, so replaying them gives the current loadout.
+        const systemMessages = this.sessionManager
+            .getBranch()
+            .flatMap((entry) => (entry.type === "message" && entry.message.role === "system" ? [entry.message] : []));
+        // No recorded loadout: keep the current tools.
+        if (systemMessages.length === 0)
             return;
-        const names = applyToolModifiers((current.toolsAdded ?? []).map((tool) => tool.name), modifiers);
+        // Tools the branch ever added or removed. Removed tools stay inactive.
+        const mentioned = new Set(systemMessages
+            .flatMap((message) => [...(message.toolsAdded ?? []), ...(message.toolsRemoved ?? [])])
+            .map((tool) => tool.name));
+        // Extension tools the branch never mentioned were registered after it recorded its loadout, so
+        // they are activated like newly registered tools. Built-in tools follow `defaultTools` instead.
+        const newTools = [...this._toolRegistry.keys()].filter((name) => !mentioned.has(name) && !this._baseToolDefinitions.has(name) && this._isActivatedOnRegistration(name));
+        // The recorded loadout plus the new tools, with `--tools +name`/`-name` applied.
+        const names = applyToolModifiers([...getCurrentTools(systemMessages).map((tool) => tool.name), ...newTools], modifiers);
+        // Recorded tools that are not registered yet, such as MCP tools of servers still connecting, are
+        // activated when they register. Excluded tools never register, so they are not kept pending.
         this._pendingToolNames = new Set(names.filter((name) => this._isAllowedTool(name)));
+        // Ignores names that are not registered.
         this._setActiveTools(names);
     }
     // =========================================================================
@@ -1547,9 +1593,7 @@ export class AgentSession {
     }
     async _handlePostAgentRun() {
         const message = this._lastAssistantMessage;
-        const toolResults = this._lastAssistantToolResults;
         this._lastAssistantMessage = undefined;
-        this._lastAssistantToolResults = [];
         if (this._agentRunAbortRequested) {
             this._finishCancelledRetry();
             return false;
@@ -1575,7 +1619,7 @@ export class AgentSession {
             });
             this._retryAttempt = 0;
         }
-        if (await this._checkCompaction(message, true, toolResults)) {
+        if (await this._checkCompaction(message)) {
             return !this._agentRunAbortRequested;
         }
         // The low-level loop drains both queues before agent_end. Messages queued by
@@ -1853,6 +1897,7 @@ export class AgentSession {
             content,
             timestamp: Date.now(),
         });
+        this._tasksHost.interruptWaits();
     }
     /**
      * Internal: Queue a follow-up message (already expanded, no extension command check).
@@ -1909,6 +1954,7 @@ export class AgentSession {
             }
             else {
                 this.agent.steer(appMessage);
+                this._tasksHost.interruptWaits();
             }
         }
         else if (options?.triggerTurn) {
@@ -2349,6 +2395,7 @@ export class AgentSession {
                 throw new Error("Compaction cancelled");
             }
             this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage);
+            this._tasksHost.noteCompaction();
             const newEntries = this.sessionManager.getEntries();
             this._refreshFinalizedContext();
             const estimatedTokensAfter = estimateMessagesTokens(this.sessionManager.buildSessionProjection().messages);
@@ -2442,7 +2489,7 @@ export class AgentSession {
      * @param skipAbortedCheck If false, include aborted messages (for pre-prompt check). Default: true
      * @returns Whether the post-run loop should call `agent.continue()` for overflow recovery or queued messages
      */
-    async _checkCompaction(assistantMessage, skipAbortedCheck = true, toolResults = []) {
+    async _checkCompaction(assistantMessage, skipAbortedCheck = true) {
         const settings = this.settingsManager.getCompactionSettings(this._modelForMessage(assistantMessage) ?? this.model);
         if (!settings.enabled)
             return false;
@@ -2519,7 +2566,7 @@ export class AgentSession {
             }
             // Persistently omit the selected final attempt before post-run recovery compaction.
             this._overflowRecoveryAttempted = true;
-            this._omitRecoveryAttempt(assistantMessage, toolResults);
+            this._omitRecoveryAttempt(assistantMessage);
             const retry = await this._runAutoCompaction("overflow", willRetry);
             if (retry)
                 this._failedResponse = assistantMessage;
@@ -2641,6 +2688,7 @@ export class AgentSession {
             }
             abortController.signal.throwIfAborted();
             this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage);
+            this._tasksHost.noteCompaction();
             const newEntries = this.sessionManager.getEntries();
             this._refreshFinalizedContext();
             const estimatedTokensAfter = estimateMessagesTokens(this.sessionManager.buildSessionProjection().messages);
@@ -3045,7 +3093,7 @@ export class AgentSession {
                 bash: { commandPrefix: shellCommandPrefix, shellPath },
             });
         this._baseToolDefinitions = new Map(Object.entries(baseToolDefinitions).map(([name, tool]) => [name, tool]));
-        this._baseToolDefinitions.set("tasks", createTasksToolDefinition((callId, taskId) => this._tasksHost.delivery.prepareWait(callId, taskId)));
+        this._baseToolDefinitions.set("tasks", createTasksToolDefinition((callId, taskId) => this._tasksHost.delivery.prepareWait(callId, taskId), () => this._tasksHost.waitInterruptSignal));
         const extensionsResult = this._resourceLoader.getExtensions();
         if (options.flagValues) {
             for (const [name, value] of options.flagValues) {