fix(tui): keep focused overlays interactive after UI

Preserve a focused visible overlay as the input owner across extension custom UI replacement, while letting the replacement receive and close its own input before focus is restored. Fixes #5129.

packages/tui/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Fixed
+
+- Fixed focused visible overlays losing input to base components after focus restoration behind the overlay ([#5129](https://github.com/earendil-works/pi/issues/5129)).
+
 ## [0.78.0] - 2026-05-29
 
 ### Fixed

packages/tui/README.md

@@ -119,6 +119,10 @@ handle.focus();             // Focus and bring to visual front
 handle.unfocus();           // Release focus to previous target
 handle.isFocused();         // Check if overlay has focus
 
+// A focused visible overlay keeps keyboard input until hidden or explicitly unfocused.
+// If you want a base component to receive input while the overlay remains visible,
+// call handle.unfocus() before focusing the base component.
+
 // Hide topmost overlay
 tui.hideOverlay();
 

packages/tui/src/tui.ts

@@ -194,6 +194,19 @@ export interface OverlayHandle {
 	isFocused(): boolean;
 }
 
+type ActiveOverlayRestoreFocusState = { status: "eligible" } | { status: "blocked"; blockedBy: Component };
+type OverlayRestoreFocusState = { status: "inactive" } | ActiveOverlayRestoreFocusState;
+type RestorableOverlayStackEntry = OverlayStackEntry & { restoreFocus: ActiveOverlayRestoreFocusState };
+
+type OverlayStackEntry = {
+	component: Component;
+	options: OverlayOptions | undefined;
+	preFocus: Component | null;
+	hidden: boolean;
+	focusOrder: number;
+	restoreFocus: OverlayRestoreFocusState;
+};
+
 /**
  * Container - a component that contains other components
  */
@@ -262,13 +275,7 @@ export class TUI extends Container {
 
 	// Overlay stack for modal components rendered on top of base content
 	private focusOrderCounter = 0;
-	private overlayStack: {
+	private overlayStack: OverlayStackEntry[] = [];
-		component: Component;
-		options?: OverlayOptions;
-		preFocus: Component | null;
-		hidden: boolean;
-		focusOrder: number;
-	}[] = [];
 
 	constructor(terminal: Terminal, showHardwareCursor?: boolean) {
 		super();
@@ -309,17 +316,76 @@ export class TUI extends Container {
 	}
 
 	setFocus(component: Component | null): void {
+		const previousFocus = this.focusedComponent;
+		let nextFocus = component;
+		const previousFocusedOverlay = previousFocus
+			? this.overlayStack.find((entry) => entry.component === previousFocus && this.isOverlayVisible(entry))
+			: undefined;
+		const nextFocusIsOverlay = nextFocus ? this.overlayStack.some((entry) => entry.component === nextFocus) : false;
+		if (nextFocus && !nextFocusIsOverlay) {
+			const restoreOverlay = this.getRestoreFocusOverlay();
+			if (
+				restoreOverlay?.restoreFocus.status === "blocked" &&
+				restoreOverlay.restoreFocus.blockedBy === previousFocus
+			) {
+				restoreOverlay.restoreFocus = { status: "eligible" };
+				nextFocus = restoreOverlay.component;
+			} else if (
+				previousFocusedOverlay &&
+				previousFocusedOverlay.restoreFocus.status !== "inactive" &&
+				!this.isOverlayFocusAncestor(previousFocusedOverlay, nextFocus)
+			) {
+				previousFocusedOverlay.restoreFocus = { status: "blocked", blockedBy: nextFocus };
+			}
+		}
+
 		// Clear focused flag on old component
 		if (isFocusable(this.focusedComponent)) {
 			this.focusedComponent.focused = false;
 		}
 
-		this.focusedComponent = component;
+		this.focusedComponent = nextFocus;
 
 		// Set focused flag on new component
-		if (isFocusable(component)) {
+		if (isFocusable(nextFocus)) {
-			component.focused = true;
+			nextFocus.focused = true;
 		}
+
+		const focusedOverlay = nextFocus
+			? this.overlayStack.find((entry) => entry.component === nextFocus && this.isOverlayVisible(entry))
+			: undefined;
+		if (focusedOverlay) {
+			this.markOverlayRestoreFocusEligible(focusedOverlay);
+		}
+	}
+
+	private markOverlayRestoreFocusEligible(entry: OverlayStackEntry): void {
+		for (const overlay of this.overlayStack) {
+			this.clearOverlayRestoreFocus(overlay);
+		}
+		entry.restoreFocus = { status: "eligible" };
+	}
+
+	private clearOverlayRestoreFocus(entry: OverlayStackEntry): void {
+		entry.restoreFocus = { status: "inactive" };
+	}
+
+	private getRestoreFocusOverlay(): RestorableOverlayStackEntry | undefined {
+		return this.overlayStack.find(
+			(overlay): overlay is RestorableOverlayStackEntry =>
+				overlay.restoreFocus.status !== "inactive" && this.isOverlayVisible(overlay),
+		);
+	}
+
+	private isOverlayFocusAncestor(entry: OverlayStackEntry, component: Component): boolean {
+		const visited = new Set<Component>();
+		let current = entry.preFocus;
+		while (current && !visited.has(current)) {
+			visited.add(current);
+			if (current === component) return true;
+			current = this.overlayStack.find((overlay) => overlay.component === current)?.preFocus ?? null;
+		}
+		return false;
 	}
 
 	/**
@@ -327,12 +393,13 @@ export class TUI extends Container {
 	 * Returns a handle to control the overlay's visibility.
 	 */
 	showOverlay(component: Component, options?: OverlayOptions): OverlayHandle {
-		const entry = {
+		const entry: OverlayStackEntry = {
 			component,
 			options,
 			preFocus: this.focusedComponent,
 			hidden: false,
 			focusOrder: ++this.focusOrderCounter,
+			restoreFocus: { status: "inactive" },
 		};
 		this.overlayStack.push(entry);
 		// Only focus if overlay is actually visible
@@ -347,6 +414,7 @@ export class TUI extends Container {
 			hide: () => {
 				const index = this.overlayStack.indexOf(entry);
 				if (index !== -1) {
+					this.clearOverlayRestoreFocus(entry);
 					this.overlayStack.splice(index, 1);
 					// Restore focus if this overlay had focus
 					if (this.focusedComponent === component) {
@@ -362,6 +430,7 @@ export class TUI extends Container {
 				entry.hidden = hidden;
 				// Update focus when hiding/showing
 				if (hidden) {
+					this.clearOverlayRestoreFocus(entry);
 					// If this overlay had focus, move focus to next visible or preFocus
 					if (this.focusedComponent === component) {
 						const topVisible = this.getTopmostVisibleOverlay();
@@ -379,14 +448,13 @@ export class TUI extends Container {
 			isHidden: () => entry.hidden,
 			focus: () => {
 				if (!this.overlayStack.includes(entry) || !this.isOverlayVisible(entry)) return;
-				if (this.focusedComponent !== component) {
-					this.setFocus(component);
-				}
 				entry.focusOrder = ++this.focusOrderCounter;
+				this.setFocus(component);
 				this.requestRender();
 			},
 			unfocus: () => {
 				if (this.focusedComponent !== component) return;
+				this.clearOverlayRestoreFocus(entry);
 				const topVisible = this.getTopmostVisibleOverlay();
 				this.setFocus(topVisible && topVisible !== entry ? topVisible.component : entry.preFocus);
 				this.requestRender();
@@ -399,6 +467,7 @@ export class TUI extends Container {
 	hideOverlay(): void {
 		const overlay = this.overlayStack.pop();
 		if (!overlay) return;
+		this.clearOverlayRestoreFocus(overlay);
 		if (this.focusedComponent === overlay.component) {
 			// Find topmost visible overlay, or fall back to preFocus
 			const topVisible = this.getTopmostVisibleOverlay();
@@ -414,7 +483,7 @@ export class TUI extends Container {
 	}
 
 	/** Check if an overlay entry is currently visible */
-	private isOverlayVisible(entry: (typeof this.overlayStack)[number]): boolean {
+	private isOverlayVisible(entry: OverlayStackEntry): boolean {
 		if (entry.hidden) return false;
 		if (entry.options?.visible) {
 			return entry.options.visible(this.terminal.columns, this.terminal.rows);
@@ -423,7 +492,7 @@ export class TUI extends Container {
 	}
 
 	/** Find the topmost visible capturing overlay, if any */
-	private getTopmostVisibleOverlay(): (typeof this.overlayStack)[number] | undefined {
+	private getTopmostVisibleOverlay(): OverlayStackEntry | undefined {
 		for (let i = this.overlayStack.length - 1; i >= 0; i--) {
 			if (this.overlayStack[i].options?.nonCapturing) continue;
 			if (this.isOverlayVisible(this.overlayStack[i])) {
@@ -584,6 +653,18 @@ export class TUI extends Container {
 			}
 		}
 
+		const focusIsOverlay = this.overlayStack.some((o) => o.component === this.focusedComponent);
+		if (!focusIsOverlay) {
+			const overlayToRestore = this.getRestoreFocusOverlay();
+			if (
+				overlayToRestore &&
+				(overlayToRestore.restoreFocus.status === "eligible" ||
+					overlayToRestore.restoreFocus.blockedBy !== this.focusedComponent)
+			) {
+				this.setFocus(overlayToRestore.component);
+			}
+		}
+
 		// Pass input to focused component (including Ctrl+C)
 		// The focused component can decide how to handle Ctrl+C
 		if (this.focusedComponent?.handleInput) {

packages/tui/test/overlay-non-capturing.test.ts

@@ -235,7 +235,9 @@ describe("TUI overlay non-capturing", () => {
 				// Simulate showExtensionCustom: factory creates timer synchronously,
 				// then .then() pushes controller as a microtask
 				let timerHandle: ReturnType<typeof tui.showOverlay>;
-				let doneFn: () => void;
+				let doneFn: () => void = () => {
+					throw new Error("doneFn was not initialized");
+				};
 
 				const overlayPromise = new Promise<void>((resolve) => {
 					doneFn = () => {
@@ -259,7 +261,7 @@ describe("TUI overlay non-capturing", () => {
 				assert.strictEqual(editor.focused, false);
 
 				// Simulate Esc: cleanup + close (from inside handleInput)
-				doneFn!();
+				doneFn();
 				// Now await the promise (simulating showExtensionCustom resolving)
 				await overlayPromise;
 				await renderAndFlush(tui, terminal);
@@ -305,6 +307,294 @@ describe("TUI overlay non-capturing", () => {
 			}
 		});
 
+		it("active base focus replacement receives close input before overlay restore", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const replacement = new FocusableOverlay(["REPLACEMENT"]);
+			const overlay = new FocusableOverlay(["OVERLAY"]);
+			overlay.handleInput = (data: string) => {
+				overlay.inputs.push(data);
+				if (data === "b") {
+					tui.setFocus(replacement);
+				}
+			};
+			replacement.handleInput = (data: string) => {
+				replacement.inputs.push(data);
+				if (data === "\r") {
+					tui.setFocus(editor);
+				}
+			};
+			tui.addChild(new EmptyContent());
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				tui.showOverlay(overlay);
+				assert.strictEqual(overlay.focused, true);
+				terminal.sendInput("b");
+				await renderAndFlush(tui, terminal);
+				assert.strictEqual(replacement.focused, true);
+
+				terminal.sendInput("\r");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(replacement.inputs, ["\r"]);
+				assert.deepStrictEqual(overlay.inputs, ["b"]);
+				assert.strictEqual(overlay.focused, true);
+
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(overlay.inputs, ["b", "x"]);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("active replacement still receives input when it is another overlay preFocus", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const replacement = new FocusableOverlay(["REPLACEMENT"]);
+			const passive = new FocusableOverlay(["PASSIVE"]);
+			const overlay = new FocusableOverlay(["OVERLAY"]);
+			overlay.handleInput = (data: string) => {
+				overlay.inputs.push(data);
+				if (data === "b") {
+					tui.setFocus(replacement);
+				}
+			};
+			replacement.handleInput = (data: string) => {
+				replacement.inputs.push(data);
+				if (data === "\r") {
+					tui.setFocus(editor);
+				}
+			};
+			tui.addChild(new EmptyContent());
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				tui.setFocus(replacement);
+				tui.showOverlay(passive, { nonCapturing: true });
+				tui.setFocus(editor);
+				tui.showOverlay(overlay);
+				terminal.sendInput("b");
+				await renderAndFlush(tui, terminal);
+				assert.strictEqual(replacement.focused, true);
+
+				terminal.sendInput("1");
+				terminal.sendInput("\r");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(replacement.inputs, ["1", "\r"]);
+				assert.deepStrictEqual(overlay.inputs, ["b"]);
+				assert.strictEqual(overlay.focused, true);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("handleInput restores focus to a visible focused overlay after base focus steal", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const replacement = new FocusableOverlay(["REPLACEMENT"]);
+			const overlay = new FocusableOverlay(["OVERLAY"]);
+			tui.addChild(new EmptyContent());
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				tui.showOverlay(overlay);
+				assert.strictEqual(overlay.focused, true);
+				tui.setFocus(replacement);
+				tui.setFocus(editor);
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(overlay.inputs, ["x"]);
+				assert.deepStrictEqual(editor.inputs, []);
+				assert.strictEqual(overlay.focused, true);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("handleInput restores focus to explicitly focused raw sub-overlay after base focus steal", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const controller = new FocusableOverlay(["CONTROLLER"]);
+			const subOverlay = new FocusableOverlay(["SUB"]);
+			tui.addChild(new EmptyContent());
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				tui.showOverlay(controller);
+				const subHandle = tui.showOverlay(subOverlay, { nonCapturing: true });
+				subHandle.focus();
+				tui.setFocus(editor);
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(subOverlay.inputs, ["x"]);
+				assert.deepStrictEqual(controller.inputs, []);
+				assert.deepStrictEqual(editor.inputs, []);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("passive non-capturing overlay does not regain input after base focus", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const passive = new FocusableOverlay(["PASSIVE"]);
+			tui.addChild(new EmptyContent());
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				tui.showOverlay(passive, { nonCapturing: true });
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(editor.inputs, ["x"]);
+				assert.deepStrictEqual(passive.inputs, []);
+				assert.strictEqual(editor.focused, true);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("explicitly focused non-capturing overlay regains input after base focus steal", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const overlay = new FocusableOverlay(["NC"]);
+			tui.addChild(new EmptyContent());
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				const handle = tui.showOverlay(overlay, { nonCapturing: true });
+				handle.focus();
+				tui.setFocus(editor);
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(overlay.inputs, ["x"]);
+				assert.deepStrictEqual(editor.inputs, []);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("unfocus() prevents visible overlay from regaining input", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const overlay = new FocusableOverlay(["OVERLAY"]);
+			tui.addChild(new EmptyContent());
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				const handle = tui.showOverlay(overlay);
+				handle.unfocus();
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(editor.inputs, ["x"]);
+				assert.deepStrictEqual(overlay.inputs, []);
+				assert.strictEqual(editor.focused, true);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("temporarily invisible focused overlay falls back without losing restore eligibility", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const overlay = new FocusableOverlay(["OVERLAY"]);
+			let visible = true;
+			tui.addChild(new EmptyContent());
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				tui.showOverlay(overlay, { visible: () => visible });
+				tui.setFocus(editor);
+				visible = false;
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(editor.inputs, ["x"]);
+				assert.deepStrictEqual(overlay.inputs, []);
+				visible = true;
+				terminal.sendInput("y");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(editor.inputs, ["x"]);
+				assert.deepStrictEqual(overlay.inputs, ["y"]);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("temporarily invisible focused overlay with null preFocus restores when visible again", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const overlay = new FocusableOverlay(["OVERLAY"]);
+			let visible = true;
+			tui.addChild(new EmptyContent());
+			tui.start();
+			try {
+				tui.showOverlay(overlay, { visible: () => visible });
+				visible = false;
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(overlay.inputs, []);
+				visible = true;
+				terminal.sendInput("y");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(overlay.inputs, ["y"]);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("cyclic overlay preFocus ancestry does not hang focus changes", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const overlay = new FocusableOverlay(["OVERLAY"]);
+			tui.addChild(new EmptyContent());
+			tui.setFocus(overlay);
+			tui.start();
+			try {
+				const handle = tui.showOverlay(overlay, { nonCapturing: true });
+				handle.focus();
+				tui.setFocus(editor);
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(editor.inputs, ["x"]);
+				assert.deepStrictEqual(overlay.inputs, []);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("handleInput restores the focus-order top overlay after base focus steal", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const lower = new FocusableOverlay(["LOWER"]);
+			const upper = new FocusableOverlay(["UPPER"]);
+			tui.addChild(new EmptyContent());
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				const lowerHandle = tui.showOverlay(lower);
+				tui.showOverlay(upper);
+				lowerHandle.focus();
+				tui.setFocus(editor);
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(lower.inputs, ["x"]);
+				assert.deepStrictEqual(upper.inputs, []);
+				assert.deepStrictEqual(editor.inputs, []);
+			} finally {
+				tui.stop();
+			}
+		});
+
 		it("hideOverlay() does not reassign focus when topmost overlay is non-capturing", async () => {
 			const terminal = new VirtualTerminal(80, 24);
 			const tui = new TUI(terminal);