@bastani/atomic 0.8.21 → 0.8.22-0

package/dist/core/agent-session.js

@@ -71,11 +71,47 @@ export function parseSkillBlock(text) {
         userMessage: afterClosing ? afterClosing.slice(2).trim() || undefined : undefined,
     };
 }
+function drainAgentMessageQueue(queue) {
+    if (!queue)
+        return [];
+    const drained = [];
+    while (queue.hasItems()) {
+        drained.push(...queue.drain());
+    }
+    return drained;
+}
+function normalizeInterruptAbortMessage(value) {
+    const trimmed = value?.trim();
+    return trimmed && trimmed.length > 0 ? trimmed : undefined;
+}
+function isGenericAbortText(value) {
+    const normalized = value.trim().toLowerCase().replace(/[.!]+$/, "");
+    return (normalized === "operation aborted" ||
+        normalized === "the operation was aborted" ||
+        normalized === "request was aborted" ||
+        normalized === "this operation was aborted" ||
+        normalized === "extension custom ui aborted");
+}
+function isSingleGenericAbortTextContent(content) {
+    return (Array.isArray(content) &&
+        content.length === 1 &&
+        typeof content[0] === "object" &&
+        content[0] !== null &&
+        content[0].type === "text" &&
+        typeof content[0].text === "string" &&
+        isGenericAbortText(content[0].text));
+}
+function replacementAbortContent(text) {
+    return [{ type: "text", text }];
+}
 // ============================================================================
 // Constants
 // ============================================================================
 /** Standard thinking levels */
 const THINKING_LEVELS = ["off", "minimal", "low", "medium", "high"];
+function customMessageExcludesContext(message) {
+    return message.excludeFromContext === true;
+}
 // ============================================================================
 // AgentSession Class
 // ============================================================================
@@ -87,6 +123,14 @@ export class AgentSession {
         this._steeringMessages = [];
         /** Tracks pending follow-up messages for UI display. Removed when delivered. */
         this._followUpMessages = [];
+        /** Serializes interrupt custom-message delivery so only one immediate prompt runs at a time. */
+        this._interruptDeliveryQueue = Promise.resolve();
+        /** Number of interrupt custom messages enqueued or currently delivering. */
+        this._pendingInterruptDeliveries = 0;
+        /** Queues held out of pi-agent-core while an interrupt sequence is active. */
+        this._activeInterruptQueueHold = undefined;
+        /** Replacement text for generic abort results produced by an interrupt-delivered custom message. */
+        this._activeInterruptAbortMessage = undefined;
         /** Messages queued to be included with the next user prompt as context ("asides"). */
         this._pendingNextTurnMessages = [];
         // Compaction state
@@ -150,6 +194,10 @@ export class AgentSession {
             includeAllExtensionTools: true,
         });
     }
+    /** Orchestration context for this session, when owned by a workflow/subagent runtime. */
+    get orchestrationContext() {
+        return this._orchestrationContext;
+    }
     /** Model registry for API key resolution and model discovery */
     get modelRegistry() {
         return this._modelRegistry;
@@ -291,6 +339,7 @@ export class AgentSession {
                 }
             }
         }
+        this._applyInterruptAbortMessage(event);
         // Emit to extensions first
         await this._emitExtensionEvent(event);
         // Notify all listeners
@@ -300,7 +349,7 @@ export class AgentSession {
             // Check if this is a custom message from extensions
             if (event.message.role === "custom") {
                 // Persist as CustomMessageEntry
-                this.sessionManager.appendCustomMessageEntry(event.message.customType, event.message.content, event.message.display, event.message.details);
+                this.sessionManager.appendCustomMessageEntry(event.message.customType, event.message.content, event.message.display, event.message.details, customMessageExcludesContext(event.message));
             }
             else if (event.message.role === "user" ||
                 event.message.role === "assistant" ||
@@ -342,6 +391,27 @@ export class AgentSession {
             await this._checkCompaction(msg);
         }
     }
+    _applyInterruptAbortMessage(event) {
+        const abortMessage = this._activeInterruptAbortMessage;
+        if (!abortMessage)
+            return;
+        if (event.type === "tool_execution_end" && event.isError && isSingleGenericAbortTextContent(event.result.content)) {
+            event.result.content = replacementAbortContent(abortMessage);
+            return;
+        }
+        if (event.type !== "message_start" && event.type !== "message_end")
+            return;
+        if (event.message.role === "toolResult" && event.message.isError && isSingleGenericAbortTextContent(event.message.content)) {
+            event.message.content = replacementAbortContent(abortMessage);
+            return;
+        }
+        if (event.message.role === "assistant") {
+            const assistantMessage = event.message;
+            if (assistantMessage.stopReason === "aborted") {
+                assistantMessage.errorMessage = abortMessage;
+            }
+        }
+    }
     /** Resolve the pending retry promise */
     _resolveRetry() {
         if (this._retryResolve) {
@@ -952,11 +1022,11 @@ export class AgentSession {
         if (images) {
             content.push(...images);
         }
-        this.agent.steer({
+        this._queueAgentMessage({
             role: "user",
             content,
             timestamp: Date.now(),
-        });
+        }, "steer");
     }
     /**
      * Internal: Queue a follow-up message (already expanded, no extension command check).
@@ -968,11 +1038,11 @@ export class AgentSession {
         if (images) {
             content.push(...images);
         }
-        this.agent.followUp({
+        this._queueAgentMessage({
             role: "user",
             content,
             timestamp: Date.now(),
-        });
+        }, "followUp");
     }
     /**
      * Throw an error if the text is an extension command.
@@ -988,14 +1058,16 @@ export class AgentSession {
     /**
      * Send a custom message to the session. Creates a CustomMessageEntry.
      *
-     * Handles three cases:
-     * - Streaming: queues message, processed when loop pulls from queue
+     * Handles five cases:
+     * - Streaming + interrupt trigger: aborts the active run and starts an immediate custom-message turn
+     * - Streaming + explicit display-only context exclusion: appends to state/session, no turn and no queue
+     * - Streaming otherwise: queues message, processed when loop pulls from queue
      * - Not streaming + triggerTurn: appends to state/session, starts new turn
      * - Not streaming + no trigger: appends to state/session, no turn
      *
      * @param message Custom message with customType, content, display, details
      * @param options.triggerTurn If true and not streaming, triggers a new LLM turn
-     * @param options.deliverAs Delivery mode: "steer", "followUp", or "nextTurn"
+     * @param options.deliverAs Delivery mode: "steer", "followUp", "nextTurn", or "interrupt"
      */
     async sendCustomMessage(message, options) {
         const appMessage = {
@@ -1005,26 +1077,129 @@ export class AgentSession {
             display: message.display,
             details: message.details,
             timestamp: Date.now(),
+            ...(options?.excludeFromContext === true ? { excludeFromContext: true } : {}),
         };
         if (options?.deliverAs === "nextTurn") {
             this._pendingNextTurnMessages.push(appMessage);
         }
+        else if (options?.deliverAs === "interrupt" && options.triggerTurn) {
+            await this._enqueueInterruptCustomMessage(appMessage, options);
+        }
+        else if (this.isStreaming && options?.excludeFromContext === true && options.triggerTurn !== true && options.deliverAs === undefined) {
+            this._appendCustomMessage(appMessage);
+        }
         else if (this.isStreaming) {
-            if (options?.deliverAs === "followUp") {
-                this.agent.followUp(appMessage);
+            this._queueAgentMessage(appMessage, options?.deliverAs === "followUp" ? "followUp" : "steer");
+        }
+        else if (options?.triggerTurn) {
+            await this.agent.prompt(appMessage);
+        }
+        else {
+            this._appendCustomMessage(appMessage);
+        }
+    }
+    _appendCustomMessage(message) {
+        this.agent.state.messages.push(message);
+        this.sessionManager.appendCustomMessageEntry(message.customType, message.content, message.display, message.details, customMessageExcludesContext(message));
+        this._emit({ type: "message_start", message });
+        this._emit({ type: "message_end", message });
+    }
+    _enqueueInterruptCustomMessage(message, options) {
+        this._pendingInterruptDeliveries += 1;
+        // Establish the hold synchronously when the interrupt is enqueued, not when
+        // the serialized delivery callback later starts. Callers commonly fire and
+        // forget sendCustomMessage(), then queue additional steer/follow-up messages
+        // before the promise chain gets a microtask; those messages must be captured
+        // in the active interrupt hold instead of pi-agent-core's live queues.
+        this._ensureActiveInterruptQueueHold();
+        const delivery = this._interruptDeliveryQueue.then(async () => {
+            try {
+                await this._sendInterruptCustomMessageNow(message, options);
+            }
+            finally {
+                this._pendingInterruptDeliveries -= 1;
+                if (this._pendingInterruptDeliveries === 0) {
+                    this._restoreAndClearActiveInterruptQueueHold();
+                }
+            }
+        });
+        this._interruptDeliveryQueue = delivery.catch(() => undefined);
+        return delivery;
+    }
+    async _sendInterruptCustomMessageNow(message, options) {
+        this.abortRetry();
+        this._ensureActiveInterruptQueueHold();
+        if (this.isStreaming) {
+            const previousAbortMessage = this._activeInterruptAbortMessage;
+            this._activeInterruptAbortMessage = normalizeInterruptAbortMessage(options?.interruptAbortMessage);
+            try {
+                this.agent.abort();
+                await this.agent.waitForIdle();
+                await this._agentEventQueue;
+            }
+            finally {
+                this._activeInterruptAbortMessage = previousAbortMessage;
+            }
+        }
+        await this.agent.prompt(message);
+    }
+    _ensureActiveInterruptQueueHold() {
+        if (this._activeInterruptQueueHold !== undefined) {
+            return this._activeInterruptQueueHold;
+        }
+        const drained = this._drainQueuedAgentMessages();
+        this._activeInterruptQueueHold = {
+            steering: [...drained.steering],
+            followUp: [...drained.followUp],
+        };
+        return this._activeInterruptQueueHold;
+    }
+    _restoreAndClearActiveInterruptQueueHold() {
+        const hold = this._activeInterruptQueueHold;
+        if (hold === undefined) {
+            return;
+        }
+        const currentCoreQueues = this._drainQueuedAgentMessages();
+        this._restoreQueuedAgentMessages({
+            steering: [...hold.steering, ...currentCoreQueues.steering],
+            followUp: [...hold.followUp, ...currentCoreQueues.followUp],
+        });
+        this._activeInterruptQueueHold = undefined;
+    }
+    _queueAgentMessage(message, delivery) {
+        const hold = this._activeInterruptQueueHold;
+        if (hold !== undefined) {
+            if (delivery === "followUp") {
+                hold.followUp.push(message);
             }
             else {
-                this.agent.steer(appMessage);
+                hold.steering.push(message);
             }
+            return;
         }
-        else if (options?.triggerTurn) {
-            await this.agent.prompt(appMessage);
+        if (delivery === "followUp") {
+            this.agent.followUp(message);
         }
         else {
-            this.agent.state.messages.push(appMessage);
-            this.sessionManager.appendCustomMessageEntry(message.customType, message.content, message.display, message.details);
-            this._emit({ type: "message_start", message: appMessage });
-            this._emit({ type: "message_end", message: appMessage });
+            this.agent.steer(message);
+        }
+    }
+    _drainQueuedAgentMessages() {
+        // pi-agent-core exposes public clear methods but no public drain/restore pair.
+        // Interrupts need to prevent the aborting run from consuming queued steer/follow-up
+        // messages while still preserving those queues for a later turn.
+        const agentWithQueues = this.agent;
+        return {
+            steering: drainAgentMessageQueue(agentWithQueues.steeringQueue),
+            followUp: drainAgentMessageQueue(agentWithQueues.followUpQueue),
+        };
+    }
+    _restoreQueuedAgentMessages(queues) {
+        for (const message of queues.steering) {
+            this.agent.steer(message);
+        }
+        for (const message of queues.followUp) {
+            this.agent.followUp(message);
         }
     }
     /**
@@ -1075,6 +1250,10 @@ export class AgentSession {
         this._steeringMessages = [];
         this._followUpMessages = [];
         this.agent.clearAllQueues();
+        if (this._activeInterruptQueueHold !== undefined) {
+            this._activeInterruptQueueHold.steering.length = 0;
+            this._activeInterruptQueueHold.followUp.length = 0;
+        }
         this._emitQueueUpdate();
         return { steering, followUp };
     }