fix(coding-agent): preserve startup extension UI

closes #5943

packages/coding-agent/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Fixed
 
+- Fixed transient extension UI and session-start messages emitted during session replacement or reload so they remain visible, and kept reload input blocked until reload completes ([#5943](https://github.com/earendil-works/pi/issues/5943)).
 - Fixed the plan-mode example to preserve active custom tools, skip the action prompt when no plan is found, and queue refinement/execution follow-ups correctly from `agent_end` ([#5940](https://github.com/earendil-works/pi/issues/5940)).
 - Fixed `pi update` to install the exact version returned by the Pi update check, make `--force` reinstall that checked version, fail instead of falling back to an unversioned reinstall when no version is available, and report both the old and updated versions.
 

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

@@ -2443,7 +2443,7 @@ export class AgentSession {
 		});
 	}
 
-	async reload(): Promise<void> {
+	async reload(options?: { beforeSessionStart?: () => void | Promise<void> }): Promise<void> {
 		const previousFlagValues = this._extensionRunner.getFlagValues();
 		await emitSessionShutdownEvent(this._extensionRunner, { type: "session_shutdown", reason: "reload" });
 		await this.settingsManager.reload();
@@ -2462,6 +2462,7 @@ export class AgentSession {
 			this._extensionShutdownHandler ||
 			this._extensionErrorListener;
 		if (hasBindings) {
+			await options?.beforeSessionStart?.();
 			await this._extensionRunner.emit({ type: "session_start", reason: "reload" });
 			await this.extendResourcesFromExtensions("reload");
 		}

packages/coding-agent/src/modes/interactive/interactive-mode.ts

@@ -395,7 +395,7 @@ export class InteractiveMode {
 			this.resetExtensionUI();
 		});
 		this.runtimeHost.setRebindSession(async () => {
-			await this.rebindCurrentSession();
+			await this.rebindCurrentSession({ renderBeforeBind: true });
 		});
 		this.version = VERSION;
 		this.ui = new TUI(new ProcessTerminal(), this.settingsManager.getShowHardwareCursor());
@@ -1543,12 +1543,7 @@ export class InteractiveMode {
 					}
 					this.statusContainer.clear();
 					try {
-						const result = await this.runtimeHost.newSession(options);
+						return await this.runtimeHost.newSession(options);
-						if (!result.cancelled) {
-							this.renderCurrentSessionState();
-							this.ui.requestRender();
-						}
-						return result;
 					} catch (error: unknown) {
 						return this.handleFatalRuntimeError("Failed to create session", error);
 					}
@@ -1557,7 +1552,6 @@ export class InteractiveMode {
 					try {
 						const result = await this.runtimeHost.fork(entryId, options);
 						if (!result.cancelled) {
-							this.renderCurrentSessionState();
 							this.editor.setText(result.selectedText ?? "");
 							this.showStatus("Forked to new session");
 						}
@@ -1631,12 +1625,18 @@ export class InteractiveMode {
 		}
 	}
 
-	private async rebindCurrentSession(): Promise<void> {
+	private async rebindCurrentSession(options: { renderBeforeBind?: boolean } = {}): Promise<void> {
 		this.unsubscribe?.();
 		this.unsubscribe = undefined;
 		this.applyRuntimeSettings();
-		await this.bindCurrentSessionExtensions();
+		if (options.renderBeforeBind) {
-		this.subscribeToAgent();
+			this.renderCurrentSessionState();
+			this.subscribeToAgent();
+			await this.bindCurrentSessionExtensions();
+		} else {
+			await this.bindCurrentSessionExtensions();
+			this.subscribeToAgent();
+		}
 		await this.updateAvailableProviderCount();
 		this.updateEditorBorderColor();
 		this.updateTerminalTitle();
@@ -4376,7 +4376,6 @@ export class InteractiveMode {
 							return;
 						}
 
-						this.renderCurrentSessionState();
 						this.editor.setText(result.selectedText ?? "");
 						done();
 						this.showStatus("Forked to new session");
@@ -4409,7 +4408,6 @@ export class InteractiveMode {
 				return;
 			}
 
-			this.renderCurrentSessionState();
 			this.editor.setText("");
 			this.showStatus("Cloned to new session");
 		} catch (error: unknown) {
@@ -4601,7 +4599,6 @@ export class InteractiveMode {
 			if (result.cancelled) {
 				return result;
 			}
-			this.renderCurrentSessionState();
 			this.showStatus("Resumed session");
 			return result;
 		} catch (error: unknown) {
@@ -4619,7 +4616,6 @@ export class InteractiveMode {
 				if (result.cancelled) {
 					return result;
 				}
-				this.renderCurrentSessionState();
 				this.showStatus("Resumed session in current cwd");
 				return result;
 			}
@@ -5072,8 +5068,20 @@ export class InteractiveMode {
 			this.ui.requestRender();
 		};
 
+		let chatRestoredBeforeSessionStart = false;
+		let reloadBoxDismissed = false;
+		const restoreChatBeforeSessionStart = () => {
+			if (chatRestoredBeforeSessionStart) {
+				return;
+			}
+			this.hideThinkingBlock = this.settingsManager.getHideThinkingBlock();
+			this.rebuildChatFromMessages();
+			chatRestoredBeforeSessionStart = true;
+		};
+
 		try {
-			await this.session.reload();
+			await this.session.reload({ beforeSessionStart: restoreChatBeforeSessionStart });
+			restoreChatBeforeSessionStart();
 			configureHttpDispatcher(this.settingsManager.getHttpIdleTimeoutMs());
 			this.keybindings.reload();
 			const activeHeader = this.customHeader ?? this.builtInHeader;
@@ -5081,7 +5089,6 @@ export class InteractiveMode {
 				activeHeader.setExpanded(this.toolOutputExpanded);
 			}
 			setRegisteredThemes(this.session.resourceLoader.getThemes().themes);
-			this.hideThinkingBlock = this.settingsManager.getHideThinkingBlock();
 			await this.themeController.applyFromSettings();
 			const editorPaddingX = this.settingsManager.getEditorPaddingX();
 			const autocompleteMaxVisible = this.settingsManager.getAutocompleteMaxVisible();
@@ -5096,8 +5103,6 @@ export class InteractiveMode {
 			this.setupAutocompleteProvider();
 			const runner = this.session.extensionRunner;
 			this.setupExtensionShortcuts(runner);
-			this.rebuildChatFromMessages();
-			dismissReloadBox(this.editor as Component);
 			this.showLoadedResources({
 				force: false,
 				showDiagnosticsWhenQuiet: true,
@@ -5112,8 +5117,12 @@ export class InteractiveMode {
 					? "Reloaded keybindings, extensions, skills, prompts, themes; saved project trust"
 					: "Reloaded keybindings, extensions, skills, prompts, themes",
 			);
+			dismissReloadBox(this.editor as Component);
+			reloadBoxDismissed = true;
 		} catch (error) {
-			dismissReloadBox(previousEditor as Component);
+			if (!reloadBoxDismissed) {
+				dismissReloadBox(previousEditor as Component);
+			}
 			this.showError(`Reload failed: ${error instanceof Error ? error.message : String(error)}`);
 		}
 	}
@@ -5187,7 +5196,6 @@ export class InteractiveMode {
 				this.showStatus("Import cancelled");
 				return;
 			}
-			this.renderCurrentSessionState();
 			this.showStatus(`Session imported from: ${inputPath}`);
 		} catch (error: unknown) {
 			if (error instanceof MissingSessionCwdError) {
@@ -5201,7 +5209,6 @@ export class InteractiveMode {
 					this.showStatus("Import cancelled");
 					return;
 				}
-				this.renderCurrentSessionState();
 				this.showStatus(`Session imported from: ${inputPath}`);
 				return;
 			}
@@ -5540,7 +5547,6 @@ export class InteractiveMode {
 			if (result.cancelled) {
 				return;
 			}
-			this.renderCurrentSessionState();
 			this.chatContainer.addChild(new Spacer(1));
 			this.chatContainer.addChild(new Text(`${theme.fg("accent", "✓ New session started")}`, 1, 1));
 			this.ui.requestRender();

packages/coding-agent/test/interactive-mode-clone-command.test.ts

@@ -41,7 +41,7 @@ describe("InteractiveMode /clone", () => {
 		await interactiveModePrototype.handleCloneCommand.call(context);
 
 		expect(fork).toHaveBeenCalledWith("leaf-123", { position: "at" });
-		expect(renderCurrentSessionState).toHaveBeenCalled();
+		expect(renderCurrentSessionState).not.toHaveBeenCalled();
 		expect(setText).toHaveBeenCalledWith("");
 		expect(showStatus).toHaveBeenCalledWith("Cloned to new session");
 		expect(showError).not.toHaveBeenCalled();

packages/coding-agent/test/suite/regressions/5943-session-start-notify.test.ts

@@ -0,0 +1,430 @@
+import { fauxAssistantMessage } from "@earendil-works/pi-ai";
+import { describe, expect, it, vi } from "vitest";
+import type { AgentSessionEvent } from "../../../src/core/agent-session.ts";
+import type { ExtensionUIContext } from "../../../src/core/extensions/index.ts";
+import { InteractiveMode } from "../../../src/modes/interactive/interactive-mode.ts";
+import { initTheme, type Theme, theme } from "../../../src/modes/interactive/theme/theme.ts";
+import { createHarness } from "../harness.ts";
+
+function createUiContext(
+	onNotify: (message: string, type: "info" | "warning" | "error" | undefined) => void,
+): ExtensionUIContext {
+	return {
+		select: async () => undefined,
+		confirm: async () => false,
+		input: async () => undefined,
+		notify: onNotify,
+		onTerminalInput: () => () => {},
+		setStatus: () => {},
+		setWorkingMessage: () => {},
+		setWorkingVisible: () => {},
+		setWorkingIndicator: () => {},
+		setHiddenThinkingLabel: () => {},
+		setWidget: () => {},
+		setFooter: () => {},
+		setHeader: () => {},
+		setTitle: () => {},
+		custom: async <T>() => undefined as T,
+		pasteToEditor: () => {},
+		setEditorText: () => {},
+		getEditorText: () => "",
+		editor: async () => undefined,
+		addAutocompleteProvider: () => {},
+		setEditorComponent: () => {},
+		getEditorComponent: () => undefined,
+		get theme() {
+			return theme;
+		},
+		getAllThemes: () => [],
+		getTheme: () => undefined,
+		setTheme: (_theme: string | Theme) => ({ success: false, error: "Theme switching not available in tests" }),
+		getToolsExpanded: () => false,
+		setToolsExpanded: () => {},
+	};
+}
+
+type RebindContext = {
+	unsubscribe?: () => void;
+	applyRuntimeSettings: () => void;
+	renderCurrentSessionState: () => void;
+	bindCurrentSessionExtensions: () => Promise<void>;
+	subscribeToAgent: () => void;
+	updateAvailableProviderCount: () => Promise<void>;
+	updateEditorBorderColor: () => void;
+	updateTerminalTitle: () => void;
+};
+
+type ReloadCommandContext = {
+	hideThinkingBlock: boolean;
+	session: {
+		isStreaming: boolean;
+		isCompacting: boolean;
+		reload: (options?: { beforeSessionStart?: () => void | Promise<void> }) => Promise<void>;
+		resourceLoader: { getThemes: () => { themes: [] } };
+		extensionRunner: unknown;
+		modelRegistry: { getError: () => string | undefined };
+	};
+	settingsManager: {
+		getHttpIdleTimeoutMs: () => number;
+		getHideThinkingBlock: () => boolean;
+		getEditorPaddingX: () => number;
+		getAutocompleteMaxVisible: () => number;
+		getShowHardwareCursor: () => boolean;
+		getClearOnShrink: () => boolean;
+	};
+	keybindings: { reload: () => void };
+	customHeader?: unknown;
+	builtInHeader?: unknown;
+	editorContainer: { clear: () => void; addChild: (component: unknown) => void };
+	ui: {
+		setFocus: (component: unknown) => void;
+		requestRender: (force?: boolean) => void;
+		setShowHardwareCursor: (enabled: boolean) => void;
+		setClearOnShrink: (enabled: boolean) => void;
+	};
+	editor: unknown;
+	defaultEditor: { setPaddingX: (padding: number) => void; setAutocompleteMaxVisible: (maxVisible: number) => void };
+	themeController: { applyFromSettings: () => Promise<void> };
+	resetExtensionUI: () => void;
+	rebuildChatFromMessages: () => void;
+	setupAutocompleteProvider: () => void;
+	setupExtensionShortcuts: (runner: unknown) => void;
+	showLoadedResources: (options: unknown) => void;
+	maybeSaveImplicitProjectTrustAfterReload: () => boolean;
+	showStatus: (message: string) => void;
+	showWarning: (message: string) => void;
+	showError: (message: string) => void;
+};
+
+type InteractiveModePrototype = {
+	rebindCurrentSession(this: RebindContext, options?: { renderBeforeBind?: boolean }): Promise<void>;
+	handleReloadCommand(this: ReloadCommandContext): Promise<void>;
+};
+
+const interactiveModePrototype = InteractiveMode.prototype as unknown as InteractiveModePrototype;
+
+type ReloadCommandContextOverrides = Omit<
+	Partial<ReloadCommandContext>,
+	"session" | "settingsManager" | "keybindings" | "editorContainer" | "ui" | "defaultEditor" | "themeController"
+> & {
+	session?: Partial<ReloadCommandContext["session"]>;
+	settingsManager?: Partial<ReloadCommandContext["settingsManager"]>;
+	keybindings?: Partial<ReloadCommandContext["keybindings"]>;
+	editorContainer?: Partial<ReloadCommandContext["editorContainer"]>;
+	ui?: Partial<ReloadCommandContext["ui"]>;
+	defaultEditor?: Partial<ReloadCommandContext["defaultEditor"]>;
+	themeController?: Partial<ReloadCommandContext["themeController"]>;
+};
+
+function createReloadCommandContext(overrides: ReloadCommandContextOverrides = {}): ReloadCommandContext {
+	const editor = overrides.editor ?? {};
+	return {
+		hideThinkingBlock: overrides.hideThinkingBlock ?? false,
+		session: {
+			isStreaming: false,
+			isCompacting: false,
+			reload: async (options) => {
+				await options?.beforeSessionStart?.();
+			},
+			resourceLoader: { getThemes: () => ({ themes: [] }) },
+			extensionRunner: {},
+			modelRegistry: { getError: () => undefined },
+			...overrides.session,
+		},
+		settingsManager: {
+			getHttpIdleTimeoutMs: () => 0,
+			getHideThinkingBlock: () => false,
+			getEditorPaddingX: () => 1,
+			getAutocompleteMaxVisible: () => 10,
+			getShowHardwareCursor: () => false,
+			getClearOnShrink: () => false,
+			...overrides.settingsManager,
+		},
+		keybindings: { reload: () => {}, ...overrides.keybindings },
+		editorContainer: { clear: () => {}, addChild: () => {}, ...overrides.editorContainer },
+		ui: {
+			setFocus: () => {},
+			requestRender: () => {},
+			setShowHardwareCursor: () => {},
+			setClearOnShrink: () => {},
+			...overrides.ui,
+		},
+		editor,
+		defaultEditor: { setPaddingX: () => {}, setAutocompleteMaxVisible: () => {}, ...overrides.defaultEditor },
+		themeController: { applyFromSettings: async () => {}, ...overrides.themeController },
+		customHeader: overrides.customHeader,
+		builtInHeader: overrides.builtInHeader,
+		resetExtensionUI: overrides.resetExtensionUI ?? (() => {}),
+		rebuildChatFromMessages: overrides.rebuildChatFromMessages ?? (() => {}),
+		setupAutocompleteProvider: overrides.setupAutocompleteProvider ?? (() => {}),
+		setupExtensionShortcuts: overrides.setupExtensionShortcuts ?? (() => {}),
+		showLoadedResources: overrides.showLoadedResources ?? (() => {}),
+		maybeSaveImplicitProjectTrustAfterReload: overrides.maybeSaveImplicitProjectTrustAfterReload ?? (() => false),
+		showStatus: overrides.showStatus ?? (() => {}),
+		showWarning: overrides.showWarning ?? (() => {}),
+		showError: overrides.showError ?? (() => {}),
+	};
+}
+
+type MessageEvent = Extract<AgentSessionEvent, { type: "message_start" | "message_end" }>;
+
+function getMessageText(event: MessageEvent): string {
+	const message = event.message;
+	if (!("content" in message)) {
+		return "";
+	}
+	const content = message.content;
+	if (typeof content === "string") {
+		return content;
+	}
+	return content
+		.filter((part): part is { type: "text"; text: string } => part.type === "text")
+		.map((part) => part.text)
+		.join("");
+}
+
+describe("regression #5943: session_start transient UI", () => {
+	it("renders replacement session state before session_start handlers can notify", async () => {
+		const events: string[] = [];
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_start", (_event, ctx) => {
+						ctx.ui.notify("Hello Error", "error");
+					});
+				},
+			],
+		});
+
+		try {
+			const context: RebindContext = {
+				applyRuntimeSettings: () => events.push("apply"),
+				renderCurrentSessionState: () => events.push("render"),
+				bindCurrentSessionExtensions: async () => {
+					events.push("bind");
+					await harness.session.bindExtensions({
+						uiContext: createUiContext((message) => events.push(`notify:${message}`)),
+						mode: "tui",
+					});
+				},
+				subscribeToAgent: () => events.push("subscribe"),
+				updateAvailableProviderCount: async () => {},
+				updateEditorBorderColor: () => {},
+				updateTerminalTitle: () => {},
+			};
+
+			await interactiveModePrototype.rebindCurrentSession.call(context, { renderBeforeBind: true });
+
+			expect(events).toEqual(["apply", "render", "subscribe", "bind", "notify:Hello Error"]);
+		} finally {
+			harness.cleanup();
+		}
+	});
+
+	it("subscribes before replacement session_start handlers send messages", async () => {
+		const events: string[] = [];
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_start", () => {
+						pi.sendMessage({
+							customType: "session-start",
+							content: "custom from start",
+							display: true,
+						});
+					});
+				},
+			],
+		});
+
+		try {
+			const context: RebindContext = {
+				applyRuntimeSettings: () => {},
+				renderCurrentSessionState: () => events.push("render"),
+				bindCurrentSessionExtensions: async () => {
+					events.push("bind");
+					await harness.session.bindExtensions({
+						uiContext: createUiContext(() => {}),
+						mode: "tui",
+					});
+				},
+				subscribeToAgent: () => {
+					events.push("subscribe");
+					harness.session.subscribe((event) => {
+						if (event.type !== "message_start" && event.type !== "message_end") {
+							return;
+						}
+						events.push(`${event.type}:${event.message.role}:${getMessageText(event)}`);
+					});
+				},
+				updateAvailableProviderCount: async () => {},
+				updateEditorBorderColor: () => {},
+				updateTerminalTitle: () => {},
+			};
+
+			await interactiveModePrototype.rebindCurrentSession.call(context, { renderBeforeBind: true });
+
+			expect(events).toEqual([
+				"render",
+				"subscribe",
+				"bind",
+				"message_start:custom:custom from start",
+				"message_end:custom:custom from start",
+			]);
+		} finally {
+			harness.cleanup();
+		}
+	});
+
+	it("subscribes before replacement session_start handlers send user messages", async () => {
+		const events: string[] = [];
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_start", () => {
+						pi.sendUserMessage("user from start");
+					});
+				},
+			],
+		});
+		harness.setResponses([fauxAssistantMessage("assistant from start")]);
+
+		try {
+			const context: RebindContext = {
+				applyRuntimeSettings: () => {},
+				renderCurrentSessionState: () => events.push("render"),
+				bindCurrentSessionExtensions: async () => {
+					events.push("bind");
+					await harness.session.bindExtensions({
+						uiContext: createUiContext(() => {}),
+						mode: "tui",
+					});
+				},
+				subscribeToAgent: () => {
+					events.push("subscribe");
+					harness.session.subscribe((event) => {
+						if (event.type !== "message_start" && event.type !== "message_end") {
+							return;
+						}
+						events.push(`${event.type}:${event.message.role}:${getMessageText(event)}`);
+					});
+				},
+				updateAvailableProviderCount: async () => {},
+				updateEditorBorderColor: () => {},
+				updateTerminalTitle: () => {},
+			};
+
+			await interactiveModePrototype.rebindCurrentSession.call(context, { renderBeforeBind: true });
+			await harness.session.agent.waitForIdle();
+
+			expect(events.slice(0, 3)).toEqual(["render", "subscribe", "bind"]);
+			expect(events).toContain("message_start:user:user from start");
+			expect(events).toContain("message_end:user:user from start");
+			expect(events).toContain("message_end:assistant:assistant from start");
+		} finally {
+			harness.cleanup();
+		}
+	});
+
+	it("runs the reload render hook before reload session_start handlers can notify", async () => {
+		const events: string[] = [];
+		const beforeSessionStart = vi.fn(() => {
+			events.push("render");
+		});
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_start", (event, ctx) => {
+						events.push(`start:${event.reason}`);
+						ctx.ui.notify(`notify:${event.reason}`, "error");
+					});
+				},
+			],
+		});
+
+		try {
+			await harness.session.bindExtensions({
+				uiContext: createUiContext((message) => events.push(message)),
+				mode: "tui",
+			});
+			expect(events).toEqual(["start:startup", "notify:startup"]);
+
+			events.length = 0;
+			await harness.session.reload({ beforeSessionStart });
+
+			expect(beforeSessionStart).toHaveBeenCalledTimes(1);
+			expect(events).toEqual(["render", "start:reload", "notify:reload"]);
+		} finally {
+			harness.cleanup();
+		}
+	});
+
+	it("refreshes hideThinkingBlock before rebuilding chat during reload", async () => {
+		initTheme("dark", false);
+		const events: string[] = [];
+		let context: ReloadCommandContext;
+		context = createReloadCommandContext({
+			settingsManager: { getHideThinkingBlock: () => true },
+			session: {
+				reload: async (options) => {
+					events.push("reload");
+					await options?.beforeSessionStart?.();
+					events.push(`start:${context.hideThinkingBlock}`);
+				},
+			},
+			rebuildChatFromMessages: () => {
+				events.push(`rebuild:${context.hideThinkingBlock}`);
+			},
+		});
+
+		await interactiveModePrototype.handleReloadCommand.call(context);
+
+		expect(context.hideThinkingBlock).toBe(true);
+		expect(events).toEqual(["reload", "rebuild:true", "start:true"]);
+	});
+
+	it("keeps the reload blocker focused until async reload completes", async () => {
+		initTheme("dark", false);
+		const editor = {};
+		let focused: unknown;
+		let chatRestored = false;
+		let markReloadWaiting!: () => void;
+		let finishReload!: () => void;
+		const reloadWaiting = new Promise<void>((resolve) => {
+			markReloadWaiting = resolve;
+		});
+		const reloadFinished = new Promise<void>((resolve) => {
+			finishReload = resolve;
+		});
+
+		const context = createReloadCommandContext({
+			editor,
+			session: {
+				reload: async (options) => {
+					await options?.beforeSessionStart?.();
+					markReloadWaiting();
+					await reloadFinished;
+				},
+			},
+			ui: {
+				setFocus: (component) => {
+					focused = component;
+				},
+			},
+			rebuildChatFromMessages: () => {
+				chatRestored = true;
+			},
+		});
+
+		const reloadPromise = interactiveModePrototype.handleReloadCommand.call(context);
+		await reloadWaiting;
+
+		expect(chatRestored).toBe(true);
+		expect(focused).not.toBe(editor);
+
+		finishReload();
+		await reloadPromise;
+
+		expect(focused).toBe(editor);
+	});
+});