gsd-pi 2.78.1-dev.8a893322c → 2.78.1-dev.a7b6e59b7

package/packages/pi-coding-agent/src/core/agent-session-abort-order.test.ts

@@ -77,6 +77,38 @@ function recordCallOrder<O extends object>(
 	return order;
 }
 
+function makeAssistantMessage(text: string) {
+	return {
+		role: "assistant",
+		content: [{ type: "text", text }],
+		usage: {
+			input: 1,
+			output: 1,
+			cacheRead: 0,
+			cacheWrite: 0,
+			total: 2,
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+		},
+		stopReason: "stop",
+		timestamp: Date.now(),
+	} as any;
+}
+
+function installAgentEndSessionTransition(
+	session: AgentSession,
+	transition: () => Promise<unknown>,
+): void {
+	(session as any)._extensionRunner = {
+		hasHandlers: () => false,
+		emit: async (event: any) => {
+			if (event.type === "agent_end") {
+				await transition();
+			}
+		},
+		emitStop: async () => {},
+	};
+}
+
 describe("#4243 — abort() must run before _disconnectFromAgent()", () => {
 	beforeEach(() => {
 		testDir = mkdtempSync(join(tmpdir(), "agent-session-abort-"));
@@ -106,6 +138,293 @@ describe("#4243 — abort() must run before _disconnectFromAgent()", () => {
 		);
 	});
 
+	it("newSession() waits instead of aborting while agent_end processing is still streaming", async () => {
+		const session = await createSession();
+		const order: string[] = [];
+		let releaseIdle!: () => void;
+		const idle = new Promise<void>((resolve) => {
+			releaseIdle = resolve;
+		});
+
+		(session as any)._processingAgentEnd = true;
+		(session as any).agent.state.isStreaming = true;
+		(session as any).agent.waitForIdle = () => {
+			order.push("waitForIdle");
+			return idle;
+		};
+		(session as any).abort = async () => {
+			order.push("abort");
+		};
+		const originalDisconnect = (session as any)._disconnectFromAgent.bind(session);
+		(session as any)._disconnectFromAgent = () => {
+			order.push("_disconnectFromAgent");
+			originalDisconnect();
+		};
+
+		const pendingNewSession = session.newSession();
+		await Promise.resolve();
+		assert.deepEqual(order, ["waitForIdle"]);
+		assert.equal(order.includes("abort"), false);
+
+		releaseIdle();
+		const ok = await pendingNewSession;
+		assert.equal(ok, true);
+		assert.deepEqual(order, ["waitForIdle", "_disconnectFromAgent"]);
+		assert.equal(order.includes("abort"), false);
+	});
+
+	it("newSession() waits during agent_end processing even once already idle", async () => {
+		const session = await createSession();
+		const order: string[] = [];
+
+		(session as any)._processingAgentEnd = true;
+		(session as any).agent.state.isStreaming = false;
+		(session as any).agent.waitForIdle = async () => {
+			order.push("waitForIdle");
+		};
+		(session as any).abort = async () => {
+			order.push("abort");
+		};
+		const originalDisconnect = (session as any)._disconnectFromAgent.bind(session);
+		(session as any)._disconnectFromAgent = () => {
+			order.push("_disconnectFromAgent");
+			originalDisconnect();
+		};
+
+		const ok = await session.newSession();
+		assert.equal(ok, true);
+		assert.deepEqual(order, ["waitForIdle", "_disconnectFromAgent"]);
+		assert.equal(order.includes("abort"), false);
+	});
+
+	it("abort() marks synthetic agent_end processing while extension handlers run", async () => {
+		const session = await createSession();
+		const observedProcessingStates: boolean[] = [];
+
+		(session as any).agent.abort = () => {};
+		(session as any).agent.waitForIdle = async () => {};
+		(session as any)._extensionRunner = {
+			emit: async (event: any) => {
+				if (event.type === "agent_end") {
+					observedProcessingStates.push((session as any)._processingAgentEnd);
+				}
+			},
+			emitStop: async () => {
+				observedProcessingStates.push((session as any)._processingAgentEnd);
+			},
+		};
+
+		await session.abort();
+
+		assert.deepEqual(observedProcessingStates, [true, true]);
+		assert.equal((session as any)._processingAgentEnd, false);
+	});
+
+	it("newSession() during agent_end preserves the previous session for resume", async () => {
+		const session = await createSession({ persistSessions: true });
+		const previousSessionFile = session.sessionFile;
+		assert.ok(previousSessionFile, "need a persisted session file");
+
+		session.sessionManager.appendMessage({
+			role: "user",
+			content: [{ type: "text", text: "persisted prompt" }],
+		} as any);
+		session.sessionManager.appendMessage({
+			role: "assistant",
+			content: [{ type: "text", text: "persisted response" }],
+			usage: {
+				input: 1,
+				output: 1,
+				cacheRead: 0,
+				cacheWrite: 0,
+				total: 2,
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+			},
+			stopReason: "stop",
+			timestamp: Date.now(),
+		} as any);
+		session.agent.replaceMessages(session.sessionManager.buildSessionContext().messages);
+
+		(session as any)._processingAgentEnd = true;
+		(session as any).agent.waitForIdle = async () => {};
+
+		const ok = await session.newSession();
+		assert.equal(ok, true);
+		assert.notEqual(session.sessionFile, previousSessionFile);
+		assert.deepEqual(session.messages, []);
+
+		(session as any)._processingAgentEnd = false;
+		const switched = await session.switchSession(previousSessionFile);
+		assert.equal(switched, true);
+
+		const restoredText = session.messages
+			.flatMap((message: any) => message.content ?? [])
+			.filter((part: any) => part.type === "text")
+			.map((part: any) => part.text);
+		assert.deepEqual(restoredText, ["persisted prompt", "persisted response"]);
+	});
+
+	it("switchSession() waits instead of aborting while agent_end processing is still streaming", async () => {
+		const session = await createSession({ persistSessions: true });
+		const previousSessionFile = session.sessionFile;
+		assert.ok(previousSessionFile, "need a persisted session file");
+
+		session.sessionManager.appendMessage({
+			role: "user",
+			content: [{ type: "text", text: "switch persisted prompt" }],
+		} as any);
+		session.sessionManager.appendMessage({
+			role: "assistant",
+			content: [{ type: "text", text: "switch persisted response" }],
+			usage: {
+				input: 1,
+				output: 1,
+				cacheRead: 0,
+				cacheWrite: 0,
+				total: 2,
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+			},
+			stopReason: "stop",
+			timestamp: Date.now(),
+		} as any);
+		session.agent.replaceMessages(session.sessionManager.buildSessionContext().messages);
+
+		const ok = await session.newSession();
+		assert.equal(ok, true);
+		const activeSessionFile = session.sessionFile;
+		assert.ok(activeSessionFile, "need an active session file");
+		assert.notEqual(activeSessionFile, previousSessionFile);
+		assert.deepEqual(session.messages, []);
+
+		const order: string[] = [];
+		let releaseIdle!: () => void;
+		const idle = new Promise<void>((resolve) => {
+			releaseIdle = resolve;
+		});
+
+		(session as any)._processingAgentEnd = true;
+		(session as any).agent.state.isStreaming = true;
+		(session as any).agent.waitForIdle = () => {
+			order.push("waitForIdle");
+			return idle;
+		};
+		(session as any).abort = async () => {
+			order.push("abort");
+		};
+		const originalDisconnect = (session as any)._disconnectFromAgent.bind(session);
+		(session as any)._disconnectFromAgent = () => {
+			order.push("_disconnectFromAgent");
+			originalDisconnect();
+		};
+
+		const pendingSwitch = session.switchSession(previousSessionFile);
+		await Promise.resolve();
+		assert.deepEqual(order, ["waitForIdle"]);
+		assert.equal(order.includes("abort"), false);
+		assert.equal(session.sessionFile, activeSessionFile);
+		assert.deepEqual(session.messages, []);
+
+		releaseIdle();
+		const switched = await pendingSwitch;
+		assert.equal(switched, true);
+		assert.deepEqual(order, ["waitForIdle", "_disconnectFromAgent"]);
+		assert.equal(order.includes("abort"), false);
+		assert.equal(session.sessionFile, previousSessionFile);
+
+		const restoredText = session.messages
+			.flatMap((message: any) => message.content ?? [])
+			.filter((part: any) => part.type === "text")
+			.map((part: any) => part.text);
+		assert.deepEqual(restoredText, ["switch persisted prompt", "switch persisted response"]);
+	});
+
+	it("newSession() during agent_end skips stale post-handlers after the transition starts", async () => {
+		const session = await createSession();
+		const assistantMessage = makeAssistantMessage("old response");
+		let compactionChecks = 0;
+		let listenerAgentEnds = 0;
+
+		(session as any)._lastAssistantMessage = assistantMessage;
+		(session as any)._compactionOrchestrator.checkCompaction = async () => {
+			compactionChecks++;
+		};
+		session.subscribe((event: any) => {
+			if (event.type === "agent_end") listenerAgentEnds++;
+		});
+		installAgentEndSessionTransition(session, () => session.newSession());
+
+		await (session as any)._processAgentEvent({
+			type: "agent_end",
+			messages: [assistantMessage],
+		});
+
+		assert.equal(compactionChecks, 0);
+		assert.equal(listenerAgentEnds, 0);
+		assert.equal((session as any)._lastAssistantMessage, undefined);
+		assert.equal((session as any)._sessionSwitchPending, false);
+		assert.equal((session as any)._sessionTransitionStartedDuringAgentEnd, false);
+	});
+
+	it("switchSession() during agent_end skips stale post-handlers after the transition starts", async () => {
+		const session = await createSession({ persistSessions: true });
+		const previousSessionFile = session.sessionFile;
+		assert.ok(previousSessionFile, "need a persisted session file");
+
+		const ok = await session.newSession();
+		assert.equal(ok, true);
+		assert.notEqual(session.sessionFile, previousSessionFile);
+
+		const assistantMessage = makeAssistantMessage("old switch response");
+		let compactionChecks = 0;
+		let listenerAgentEnds = 0;
+
+		(session as any)._lastAssistantMessage = assistantMessage;
+		(session as any)._compactionOrchestrator.checkCompaction = async () => {
+			compactionChecks++;
+		};
+		session.subscribe((event: any) => {
+			if (event.type === "agent_end") listenerAgentEnds++;
+		});
+		installAgentEndSessionTransition(session, () => session.switchSession(previousSessionFile));
+
+		await (session as any)._processAgentEvent({
+			type: "agent_end",
+			messages: [assistantMessage],
+		});
+
+		assert.equal(session.sessionFile, previousSessionFile);
+		assert.equal(compactionChecks, 0);
+		assert.equal(listenerAgentEnds, 0);
+		assert.equal((session as any)._lastAssistantMessage, undefined);
+		assert.equal((session as any)._sessionSwitchPending, false);
+		assert.equal((session as any)._sessionTransitionStartedDuringAgentEnd, false);
+	});
+
+	it("agent_end post-handlers bail while a session switch is pending", async () => {
+		const session = await createSession();
+		const assistantMessage = makeAssistantMessage("old pending response");
+		let compactionChecks = 0;
+		let listenerAgentEnds = 0;
+
+		(session as any)._lastAssistantMessage = assistantMessage;
+		(session as any)._sessionSwitchPending = true;
+		(session as any)._compactionOrchestrator.checkCompaction = async () => {
+			compactionChecks++;
+		};
+		session.subscribe((event: any) => {
+			if (event.type === "agent_end") listenerAgentEnds++;
+		});
+
+		await (session as any)._processAgentEvent({
+			type: "agent_end",
+			messages: [assistantMessage],
+		});
+
+		assert.equal(compactionChecks, 0);
+		assert.equal(listenerAgentEnds, 1);
+		assert.equal((session as any)._lastAssistantMessage, undefined);
+	});
+
 	it("switchSession() invokes abort() before _disconnectFromAgent()", async () => {
 		const session = await createSession({ persistSessions: true });
 		// Seed a session file to switch to (switchSession reads from the session manager).

package/packages/pi-coding-agent/src/core/agent-session.ts

@@ -272,6 +272,12 @@ export class AgentSession {
 	// Extension system
 	private _extensionRunner: ExtensionRunner | undefined = undefined;
 	private _turnIndex = 0;
+	private _processingAgentEnd = false;
+	/** True while newSession()/switchSession() is in progress; signals agent_end
+	 * post-handlers to bail rather than corrupt new-session state. */
+	private _sessionSwitchPending = false;
+	private _processingQueuedAgentEnd = false;
+	private _sessionTransitionStartedDuringAgentEnd = false;
 
 	private _resourceLoader: ResourceLoader;
 	private _customTools: ToolDefinition[];
@@ -433,7 +439,23 @@ 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);
@@ -491,6 +513,13 @@ 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;
 
@@ -621,33 +650,39 @@ export class AgentSession {
 
 	/** Emit extension events based on agent events */
 	private async _emitExtensionEvent(event: AgentEvent): Promise<void> {
-		if (!this._extensionRunner) return;
+		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: "completed" | "cancelled" | "error" | "blocked" =
-				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: "completed" | "cancelled" | "error" | "blocked" =
+					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: TurnStartEvent = {
 				type: "turn_start",
 				turnIndex: this._turnIndex,
 				timestamp: Date.now(),
 			};
-			await this._extensionRunner.emit(extensionEvent);
+			await extensionRunner.emit(extensionEvent);
 		} else if (event.type === "turn_end") {
 			const extensionEvent: TurnEndEvent = {
 				type: "turn_end",
@@ -655,27 +690,27 @@ 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") {
 			const extensionEvent: MessageStartEvent = {
 				type: "message_start",
 				message: event.message,
 			};
-			await this._extensionRunner.emit(extensionEvent);
+			await extensionRunner.emit(extensionEvent);
 		} else if (event.type === "message_update") {
 			const extensionEvent: MessageUpdateEvent = {
 				type: "message_update",
 				message: event.message,
 				assistantMessageEvent: event.assistantMessageEvent,
 			};
-			await this._extensionRunner.emit(extensionEvent);
+			await extensionRunner.emit(extensionEvent);
 		} else if (event.type === "message_end") {
 			const extensionEvent: MessageEndEvent = {
 				type: "message_end",
 				message: event.message,
 			};
-			await this._extensionRunner.emit(extensionEvent);
+			await extensionRunner.emit(extensionEvent);
 		} else if (event.type === "tool_execution_start") {
 			const extensionEvent: ToolExecutionStartEvent = {
 				type: "tool_execution_start",
@@ -683,7 +718,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: ToolExecutionUpdateEvent = {
 				type: "tool_execution_update",
@@ -692,7 +727,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: ToolExecutionEndEvent = {
 				type: "tool_execution_end",
@@ -701,7 +736,7 @@ export class AgentSession {
 				result: event.result,
 				isError: event.isError,
 			};
-			await this._extensionRunner.emit(extensionEvent);
+			await extensionRunner.emit(extensionEvent);
 		}
 	}
 
@@ -1546,24 +1581,54 @@ 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: "completed" | "cancelled" | "error" | "blocked" =
-				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: "completed" | "cancelled" | "error" | "blocked" =
+					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;
+			}
 		}
 	}
 
+	private async _settleCurrentTurnForSessionTransition(): Promise<void> {
+		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.
 	 * Clears all messages and starts a new session.
@@ -1592,21 +1657,23 @@ export class AgentSession {
 			}
 		}
 
-	// #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._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;
+			}
 
-		// #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;
@@ -2457,12 +2524,14 @@ export class AgentSession {
 			}
 		}
 
-	// #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._steeringMessages = [];
+		this._sessionSwitchPending = true;
+		try {
+			await this._settleCurrentTurnForSessionTransition();
+			this._disconnectFromAgent();
+		} finally {
+			this._sessionSwitchPending = false;
+		}
+		this._steeringMessages = [];
 		this._followUpMessages = [];
 		this._pendingNextTurnMessages = [];