gsd-pi 2.78.1-dev.0fdacd524 → 2.78.1-dev.2ff77b0ce

package/packages/pi-coding-agent/dist/core/agent-session.js

@@ -84,6 +84,12 @@ export class AgentSession {
         // Extension system
         this._extensionRunner = undefined;
         this._turnIndex = 0;
+        this._processingAgentEnd = false;
+        /** True while newSession()/switchSession() is in progress; signals agent_end
+         * post-handlers to bail rather than corrupt new-session state. */
+        this._sessionSwitchPending = false;
+        this._processingQueuedAgentEnd = false;
+        this._sessionTransitionStartedDuringAgentEnd = false;
         this._baseToolRegistry = new Map();
         // Tool registry for extension getTools/setTools
         this._toolRegistry = new Map();
@@ -201,7 +207,24 @@ export class AgentSession {
             }
         }
         // Emit to extensions first
-        await this._emitExtensionEvent(event);
+        let skipAgentEndPostHandlers = false;
+        if (event.type === "agent_end") {
+            this._processingQueuedAgentEnd = true;
+            try {
+                await this._emitExtensionEvent(event);
+            }
+            finally {
+                this._processingQueuedAgentEnd = false;
+                skipAgentEndPostHandlers = this._sessionTransitionStartedDuringAgentEnd;
+                this._sessionTransitionStartedDuringAgentEnd = false;
+            }
+            if (skipAgentEndPostHandlers) {
+                return;
+            }
+        }
+        else {
+            await this._emitExtensionEvent(event);
+        }
         // Notify all listeners
         this._emit(event);
         // Handle session persistence
@@ -245,6 +268,12 @@ export class AgentSession {
         }
         // Check auto-retry and auto-compaction after agent completes
         if (event.type === "agent_end" && this._lastAssistantMessage) {
+            // A session transition started during agent_end handler execution -
+            // bail to avoid running retry/compaction against new-session state.
+            if (this._sessionSwitchPending) {
+                this._lastAssistantMessage = undefined;
+                return;
+            }
             const msg = this._lastAssistantMessage;
             this._lastAssistantMessage = undefined;
             // Check for retryable errors first (overloaded, rate limit, server errors)
@@ -361,26 +390,33 @@ export class AgentSession {
     }
     /** Emit extension events based on agent events */
     async _emitExtensionEvent(event) {
-        if (!this._extensionRunner)
+        const extensionRunner = this._extensionRunner;
+        if (!extensionRunner)
             return;
         if (event.type === "agent_start") {
             this._turnIndex = 0;
-            await this._extensionRunner.emit({ type: "agent_start" });
+            await extensionRunner.emit({ type: "agent_start" });
         }
         else if (event.type === "agent_end") {
-            await this._extensionRunner.emit({ type: "agent_end", messages: event.messages });
-            // `stop` fires on true quiescence: the agent cleanly completed and is now
-            // waiting for the user. Use the last assistant message's stopReason to
-            // distinguish clean completion from error/cancellation.
-            const last = event.messages[event.messages.length - 1];
-            const stopReason = last?.role === "assistant"
-                ? last.stopReason === "aborted"
-                    ? "cancelled"
-                    : last.stopReason === "error"
-                        ? "error"
-                        : "completed"
-                : "completed";
-            await this._extensionRunner.emitStop({ reason: stopReason, lastMessage: last });
+            this._processingAgentEnd = true;
+            try {
+                await extensionRunner.emit({ type: "agent_end", messages: event.messages });
+                // `stop` fires on true quiescence: the agent cleanly completed and is now
+                // waiting for the user. Use the last assistant message's stopReason to
+                // distinguish clean completion from error/cancellation.
+                const last = event.messages[event.messages.length - 1];
+                const stopReason = last?.role === "assistant"
+                    ? last.stopReason === "aborted"
+                        ? "cancelled"
+                        : last.stopReason === "error"
+                            ? "error"
+                            : "completed"
+                    : "completed";
+                await extensionRunner.emitStop({ reason: stopReason, lastMessage: last });
+            }
+            finally {
+                this._processingAgentEnd = false;
+            }
         }
         else if (event.type === "turn_start") {
             const extensionEvent = {
@@ -388,7 +424,7 @@ export class AgentSession {
                 turnIndex: this._turnIndex,
                 timestamp: Date.now(),
             };
-            await this._extensionRunner.emit(extensionEvent);
+            await extensionRunner.emit(extensionEvent);
         }
         else if (event.type === "turn_end") {
             const extensionEvent = {
@@ -397,7 +433,7 @@ export class AgentSession {
                 message: event.message,
                 toolResults: event.toolResults,
             };
-            await this._extensionRunner.emit(extensionEvent);
+            await extensionRunner.emit(extensionEvent);
             this._turnIndex++;
         }
         else if (event.type === "message_start") {
@@ -405,7 +441,7 @@ export class AgentSession {
                 type: "message_start",
                 message: event.message,
             };
-            await this._extensionRunner.emit(extensionEvent);
+            await extensionRunner.emit(extensionEvent);
         }
         else if (event.type === "message_update") {
             const extensionEvent = {
@@ -413,14 +449,14 @@ export class AgentSession {
                 message: event.message,
                 assistantMessageEvent: event.assistantMessageEvent,
             };
-            await this._extensionRunner.emit(extensionEvent);
+            await extensionRunner.emit(extensionEvent);
         }
         else if (event.type === "message_end") {
             const extensionEvent = {
                 type: "message_end",
                 message: event.message,
             };
-            await this._extensionRunner.emit(extensionEvent);
+            await extensionRunner.emit(extensionEvent);
         }
         else if (event.type === "tool_execution_start") {
             const extensionEvent = {
@@ -429,7 +465,7 @@ export class AgentSession {
                 toolName: event.toolName,
                 args: event.args,
             };
-            await this._extensionRunner.emit(extensionEvent);
+            await extensionRunner.emit(extensionEvent);
         }
         else if (event.type === "tool_execution_update") {
             const extensionEvent = {
@@ -439,7 +475,7 @@ export class AgentSession {
                 args: event.args,
                 partialResult: event.partialResult,
             };
-            await this._extensionRunner.emit(extensionEvent);
+            await extensionRunner.emit(extensionEvent);
         }
         else if (event.type === "tool_execution_end") {
             const extensionEvent = {
@@ -449,7 +485,7 @@ export class AgentSession {
                 result: event.result,
                 isError: event.isError,
             };
-            await this._extensionRunner.emit(extensionEvent);
+            await extensionRunner.emit(extensionEvent);
         }
     }
     /**
@@ -1175,21 +1211,49 @@ export class AgentSession {
         // between tool execution and response processing. Also fire Stop so
         // Layer 0 hooks see a consistent view of session quiescence.
         if (!this.isStreaming && this._extensionRunner) {
-            const messages = this.agent.state.messages;
-            await this._extensionRunner.emit({
-                type: "agent_end",
-                messages,
-            });
-            const last = messages[messages.length - 1];
-            const stopReason = last?.role === "assistant"
-                ? last.stopReason === "aborted"
-                    ? "cancelled"
-                    : last.stopReason === "error"
-                        ? "error"
-                        : "completed"
-                : "cancelled";
-            await this._extensionRunner.emitStop({ reason: stopReason, lastMessage: last });
+            const wasProcessingAgentEnd = this._processingAgentEnd;
+            this._processingAgentEnd = true;
+            try {
+                const messages = this.agent.state.messages;
+                await this._extensionRunner.emit({
+                    type: "agent_end",
+                    messages,
+                });
+                const last = messages[messages.length - 1];
+                const stopReason = last?.role === "assistant"
+                    ? last.stopReason === "aborted"
+                        ? "cancelled"
+                        : last.stopReason === "error"
+                            ? "error"
+                            : "completed"
+                    : "cancelled";
+                await this._extensionRunner.emitStop({ reason: stopReason, lastMessage: last });
+            }
+            finally {
+                this._processingAgentEnd = wasProcessingAgentEnd;
+            }
+        }
+    }
+    async _settleCurrentTurnForSessionTransition() {
+        if (this._processingAgentEnd) {
+            // Wait for the agent to fully settle. When called from inside an
+            // agent_end extension handler, the agent may already be idle - but
+            // _processAgentEvent still has retry/compaction tail work to run after
+            // _emitExtensionEvent returns. waitForIdle() is effectively a no-op when
+            // already idle, so awaiting it unconditionally is safe and ensures we
+            // don't proceed into the session reset while that tail is still on the stack.
+            await this.agent.waitForIdle();
+            if (this._processingQueuedAgentEnd) {
+                this._sessionTransitionStartedDuringAgentEnd = true;
+                this._lastAssistantMessage = undefined;
+            }
+            return;
         }
+        // #4243: Normal session transitions must abort before disconnecting so
+        // message_end/agent_end events fire while listeners are still connected.
+        // During agent_end handling the turn is already ending; aborting there can
+        // convert a successful auto-mode handoff into an aborted provider message.
+        await this.abort();
     }
     /**
      * Start a new session, optionally with initial messages and parent tracking.
@@ -1211,19 +1275,22 @@ export class AgentSession {
                 return false;
             }
         }
-        // #4243: Must call abort() BEFORE _disconnectFromAgent() so that
-        // message_end/agent_end events fire and the #4216 finalization code
-        // can run before we unsubscribe from the event bus.
-        await this.abort();
-        // #3731: If the caller aborted (e.g. runUnit() timed out and restored cwd to
-        // project root), discard this session before capturing process.cwd() and
-        // rebuilding the tool runtime. Without this check, the late newSession()
-        // would rebuild tools with root cwd, breaking worktree isolation.
-        if (options?.abortSignal?.aborted) {
-            return false;
+        this._sessionSwitchPending = true;
+        try {
+            await this._settleCurrentTurnForSessionTransition();
+            // #3731: If the caller aborted (e.g. runUnit() timed out and restored cwd to
+            // project root), discard this session before capturing process.cwd() and
+            // rebuilding the tool runtime. Without this check, the late newSession()
+            // would rebuild tools with root cwd, breaking worktree isolation.
+            if (options?.abortSignal?.aborted) {
+                return false;
+            }
+            this._disconnectFromAgent();
+            this.agent.reset();
+        }
+        finally {
+            this._sessionSwitchPending = false;
         }
-        this._disconnectFromAgent();
-        this.agent.reset();
         // Update cwd to current process directory — auto-mode may have chdir'd
         // into a worktree since the original session was created.
         const previousCwd = this._cwd;
@@ -1949,11 +2016,14 @@ export class AgentSession {
                 return false;
             }
         }
-        // #4243: Must call abort() BEFORE _disconnectFromAgent() so that
-        // message_end/agent_end events fire and the #4216 finalization code
-        // can run before we unsubscribe from the event bus.
-        await this.abort();
-        this._disconnectFromAgent();
+        this._sessionSwitchPending = true;
+        try {
+            await this._settleCurrentTurnForSessionTransition();
+            this._disconnectFromAgent();
+        }
+        finally {
+            this._sessionSwitchPending = false;
+        }
         this._steeringMessages = [];
         this._followUpMessages = [];
         this._pendingNextTurnMessages = [];