fix(tui): harden overlay focus restoration

2026-06-18 15:54:04 +08:00 · 2026-05-30 12:48:05 -07:00
parent 82f29ea442
commit 91a2f86600
8 changed files with 542 additions and 176 deletions
@@ -2371,7 +2371,7 @@ const result = await ctx.ui.custom<string | null>(
 );
 ```

-For advanced positioning (anchors, margins, percentages, responsive visibility), pass `overlayOptions`. Use `onHandle` to control visibility programmatically:
+For advanced positioning (anchors, margins, percentages, responsive visibility), pass `overlayOptions`. Use `onHandle` to control focus or visibility programmatically:

 ```typescript
 const result = await ctx.ui.custom<string | null>(
@@ -2379,12 +2379,19 @@ const result = await ctx.ui.custom<string | null>(
  {
    overlay: true,
    overlayOptions: { anchor: "top-right", width: "50%", margin: 2 },
-    onHandle: (handle) => { /* handle.setHidden(true/false) */ }
+    onHandle: (handle) => {
+      handle.focus(); // focus this overlay and bring it to the visual front
+      // handle.unfocus({ target: editorComponent }); // release input to a specific component
+      // handle.setHidden(true/false); // toggle visibility
+      // handle.hide(); // permanently remove
+    }
  }
 );
 ```

-See [tui.md](tui.md) for the full `OverlayOptions` API and [overlay-qa-tests.ts](../examples/extensions/overlay-qa-tests.ts) for examples.
+A focused visible overlay can reclaim input after temporary non-overlay custom UI closes. If you intentionally want another component to keep input while the overlay stays visible, call `handle.unfocus({ target })`. Passing `{ target: null }` releases the overlay without focusing another component.
+
+See [tui.md](tui.md) for the full `OverlayOptions` and `OverlayHandle` API and [overlay-qa-tests.ts](../examples/extensions/overlay-qa-tests.ts) for examples.

 ### Custom Editor

@@ -145,8 +145,11 @@ const result = await ctx.ui.custom<string | null>(
      // Responsive: hide on narrow terminals
      visible: (termWidth, termHeight) => termWidth >= 80,
    },
-    // Get handle for programmatic visibility control
+    // Get handle for programmatic focus and visibility control
    onHandle: (handle) => {
+      // handle.focus() - focus this overlay and bring it to the visual front
+      // handle.unfocus() - release input to normal fallback
+      // handle.unfocus({ target }) - release input to a specific component or null
      // handle.setHidden(true/false) - toggle visibility
      // handle.hide() - permanently remove
    },
@@ -154,6 +157,12 @@ const result = await ctx.ui.custom<string | null>(
 );
 ```

+### Overlay Focus
+
+A focused visible overlay keeps input ownership across temporary non-overlay UI. If an overlay opens another `ctx.ui.custom()` component without `{ overlay: true }`, that replacement UI receives input while it is active; when it closes, the focused overlay can reclaim input.
+
+Use `handle.unfocus()` when a visible overlay should stop owning input and let TUI fall back to another visible capturing overlay or the previous focus target. Use `handle.unfocus({ target })` when a specific component should receive input while the overlay stays visible. Passing `{ target: null }` intentionally leaves no focused component until focus is set again.
+
 ### Overlay Lifecycle

 Overlay components are disposed when closed. Don't reuse references - create fresh instances:
@@ -21,7 +21,7 @@

 import type { ExtensionAPI, ExtensionCommandContext, Theme } from "@earendil-works/pi-coding-agent";
 import type { Component, OverlayAnchor, OverlayHandle, OverlayOptions, TUI } from "@earendil-works/pi-tui";
-import { CURSOR_MARKER, matchesKey, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
+import { Input, matchesKey, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
 import { spawn } from "child_process";

 // Global handle for toggle demo (in real code, use a more elegant pattern)
@@ -1031,77 +1031,66 @@ class TimerPanel extends BaseOverlay {

 // === Focus cycling demo ===

+type FocusPanelColor = "error" | "success" | "accent";
+type FocusPanelConfig = { label: string; color: FocusPanelColor; options: OverlayOptions };
+type FocusPanelEntry = { panel: FocusPanel; handle: OverlayHandle };
+
+const FOCUS_PANEL_CONFIGS = [
+	{ label: "Alpha", color: "error", options: { row: 2, col: 4, width: 34 } },
+	{ label: "Beta", color: "success", options: { row: 5, col: 28, width: 34 } },
+	{ label: "Gamma", color: "accent", options: { row: 8, col: 52, width: 34 } },
+] satisfies FocusPanelConfig[];
+
 class FocusDemoController extends BaseOverlay {
-	private tui: TUI;
-	private panels: FocusPanel[] = [];
-	private handles: OverlayHandle[] = [];
-	private done: () => void;
+	private readonly tui: TUI;
+	private entries: FocusPanelEntry[] = [];
+	private readonly done: () => void;
+	private closed = false;

 	constructor(tui: TUI, theme: Theme, done: () => void) {
 		super(theme);
 		this.tui = tui;
 		this.done = done;
-		const colors = ["error", "success", "accent"] as const;
-		const labels = ["Alpha", "Beta", "Gamma"];
-		const positions = [
-			{ row: 2, col: 4 },
-			{ row: 5, col: 28 },
-			{ row: 8, col: 52 },
-		];

-		for (let i = 0; i < 3; i++) {
-			const panel = new FocusPanel(theme, labels[i]!, colors[i]!, this);
-			const handle = this.tui.showOverlay(panel, {
-				nonCapturing: true,
-				...positions[i]!,
-				width: 34,
-			});
-			panel.handle = handle;
-			this.panels.push(panel);
-			this.handles.push(handle);
+		for (const config of FOCUS_PANEL_CONFIGS) {
+			const panel = new FocusPanel({ theme, config, controller: this });
+			const handle = this.tui.showOverlay(panel, { nonCapturing: true, ...config.options });
+			this.entries.push({ panel, handle });
 		}

 		this.focusFirstOpenPanel();
 	}

 	focusNext(current: FocusPanel, direction: 1 | -1 = 1): void {
-		const open = this.openPanelIndexes();
-		if (open.length === 0) {
-			this.close();
-			return;
-		}
-
-		const currentIndex = this.panels.indexOf(current);
-		const currentOpenPosition = open.indexOf(currentIndex);
-		const nextOpenPosition =
-			currentOpenPosition === -1 ? 0 : (currentOpenPosition + direction + open.length) % open.length;
-		this.handles[open[nextOpenPosition]!]!.focus();
-		this.tui.requestRender();
+		const openEntries = this.openEntries();
+		const currentOpenPosition = openEntries.findIndex((entry) => entry.panel === current);
+		if (currentOpenPosition === -1) throw new Error(`Panel ${current.label} is not open`);
+		const nextOpenPosition = (currentOpenPosition + direction + openEntries.length) % openEntries.length;
+		this.focusEntryAt(openEntries, nextOpenPosition);
 	}

 	dismiss(panel: FocusPanel): void {
-		const index = this.panels.indexOf(panel);
-		if (index === -1 || panel.closed) return;
+		const openEntries = this.openEntries();
+		const currentOpenPosition = openEntries.findIndex((candidate) => candidate.panel === panel);
+		if (currentOpenPosition === -1) return;
+		const entry = openEntries[currentOpenPosition];
+		if (!entry) throw new Error(`Invalid focus panel index ${currentOpenPosition}`);
+		const remainingEntries = openEntries.filter((candidate) => candidate.panel !== panel);

-		panel.closed = true;
-		this.handles[index]?.hide();
-		if (this.openPanelIndexes().length === 0) {
+		entry.panel.closed = true;
+		entry.handle.hide();
+		if (remainingEntries.length === 0) {
 			this.close();
 			return;
 		}

-		this.focusNext(panel);
+		this.focusEntryAt(remainingEntries, currentOpenPosition % remainingEntries.length);
 	}

 	close(): void {
-		for (let i = 0; i < this.handles.length; i++) {
-			if (!this.panels[i]?.closed) {
-				this.panels[i]!.closed = true;
-				this.handles[i]!.hide();
-			}
-		}
-		this.handles = [];
-		this.panels = [];
+		if (this.closed) return;
+		this.closed = true;
+		this.hidePanels();
 		this.done();
 	}

@@ -1115,7 +1104,7 @@ class FocusDemoController extends BaseOverlay {

 	render(width: number): string[] {
 		const th = this.theme;
-		const focused = this.panels.find((panel) => panel.handle?.isFocused())?.label ?? "Controller";
+		const focused = this.entries.find((entry) => entry.handle.isFocused())?.panel.label ?? "Controller";
 		return this.box(
 			[
 				"",
@@ -1139,36 +1128,62 @@ class FocusDemoController extends BaseOverlay {
 	}

 	override dispose(): void {
-		for (const handle of this.handles) handle.hide();
+		if (this.closed) return;
+		this.closed = true;
+		this.hidePanels();
 	}

 	private focusFirstOpenPanel(): void {
-		const firstOpen = this.openPanelIndexes()[0];
-		if (firstOpen !== undefined) {
-			this.handles[firstOpen]!.focus();
+		const firstOpen = this.openEntries()[0];
+		if (firstOpen) {
+			firstOpen.handle.focus();
 			this.tui.requestRender();
 		}
 	}

-	private openPanelIndexes(): number[] {
-		return this.panels.flatMap((panel, index) => (panel.closed ? [] : [index]));
+	private focusEntryAt(entries: FocusPanelEntry[], index: number): void {
+		const entry = entries[index];
+		if (!entry) throw new Error(`Invalid focus panel index ${index}`);
+		entry.handle.focus();
+		this.tui.requestRender();
+	}
+
+	private hidePanels(): void {
+		for (const entry of this.entries) {
+			if (!entry.panel.closed) {
+				entry.panel.closed = true;
+				entry.handle.hide();
+			}
+		}
+		this.entries = [];
+	}
+
+	private openEntries(): FocusPanelEntry[] {
+		return this.entries.filter((entry) => !entry.panel.closed);
 	}
 }

 class FocusPanel extends BaseOverlay {
-	handle: OverlayHandle | null = null;
+	focused = false;
 	closed = false;
 	readonly label: string;
-	private color: "error" | "success" | "accent";
-	private controller: FocusDemoController;
-	private text = "";
-	private cursor = 0;
+	private readonly color: FocusPanelColor;
+	private readonly controller: FocusDemoController;
+	private readonly input = new Input();
 	private inputs: string[] = [];

-	constructor(theme: Theme, label: string, color: "error" | "success" | "accent", controller: FocusDemoController) {
+	constructor({
+		theme,
+		config,
+		controller,
+	}: {
+		theme: Theme;
+		config: FocusPanelConfig;
+		controller: FocusDemoController;
+	}) {
 		super(theme);
-		this.label = label;
-		this.color = color;
+		this.label = config.label;
+		this.color = config.color;
 		this.controller = controller;
 	}

@@ -1181,52 +1196,47 @@ class FocusPanel extends BaseOverlay {
 			this.controller.dismiss(this);
 		} else if (matchesKey(data, "ctrl+c")) {
 			this.controller.close();
-		} else if (matchesKey(data, "backspace")) {
-			if (this.cursor > 0) {
-				this.text = this.text.slice(0, this.cursor - 1) + this.text.slice(this.cursor);
-				this.cursor--;
-			}
-			this.inputs.push("Backspace");
-		} else if (matchesKey(data, "left")) {
-			this.cursor = Math.max(0, this.cursor - 1);
-			this.inputs.push("←");
-		} else if (matchesKey(data, "right")) {
-			this.cursor = Math.min(this.text.length, this.cursor + 1);
-			this.inputs.push("→");
 		} else if (matchesKey(data, "return")) {
 			this.inputs.push("Enter");
 		} else if (matchesKey(data, "up")) {
 			this.inputs.push("↑");
 		} else if (matchesKey(data, "down")) {
 			this.inputs.push("↓");
-		} else if (data.length === 1 && data.charCodeAt(0) >= 32) {
-			this.text = this.text.slice(0, this.cursor) + data + this.text.slice(this.cursor);
-			this.cursor++;
-			this.inputs.push(JSON.stringify(data));
+		} else if (matchesKey(data, "left")) {
+			this.input.handleInput(data);
+			this.inputs.push("←");
+		} else if (matchesKey(data, "right")) {
+			this.input.handleInput(data);
+			this.inputs.push("→");
+		} else if (matchesKey(data, "backspace")) {
+			this.input.handleInput(data);
+			this.inputs.push("Backspace");
 		} else {
+			this.input.handleInput(data);
 			this.inputs.push(JSON.stringify(data));
 		}
 	}

 	render(width: number): string[] {
 		const th = this.theme;
-		const focused = this.handle?.isFocused() ?? false;
 		const innerW = Math.max(1, width - 2);
-		const border = (c: string) => th.fg(focused ? this.color : "dim", c);
+		const border = (c: string) => th.fg(this.focused ? this.color : "dim", c);
 		const padLine = (s: string) => truncateToWidth(s, innerW, "...", true);
 		const recent = this.inputs.length === 0 ? "(none)" : this.inputs.slice(-6).join(" ");
 		const lines: string[] = [];

+		this.input.focused = this.focused;
+		const [inputLine = ""] = this.input.render(Math.max(1, innerW - 8));
 		lines.push(border(`╭${"─".repeat(innerW)}╮`));
 		lines.push(
 			border("│") +
 				padLine(
-					` ${th.fg(this.color, this.label)} ${focused ? th.fg("success", "FOCUSED") : th.fg("dim", "visible")}`,
+					` ${th.fg(this.color, this.label)} ${this.focused ? th.fg("success", "FOCUSED") : th.fg("dim", "visible")}`,
 				) +
 				border("│"),
 		);
 		lines.push(border("│") + padLine("") + border("│"));
-		lines.push(border("│") + padLine(` Input: ${this.renderInput(focused)}`) + border("│"));
+		lines.push(border("│") + padLine(` Input: ${inputLine}`) + border("│"));
 		lines.push(border("│") + padLine(` Keys: ${recent}`) + border("│"));
 		lines.push(border("│") + padLine(th.fg("dim", " Tab/Shift+Tab focus")) + border("│"));
 		lines.push(border("│") + padLine(th.fg("dim", " Esc/Ctrl+D dismiss")) + border("│"));
@@ -1234,15 +1244,6 @@ class FocusPanel extends BaseOverlay {

 		return lines;
 	}
-
-	private renderInput(focused: boolean): string {
-		if (!focused) return this.text || this.theme.fg("dim", "(empty)");
-
-		const before = this.text.slice(0, this.cursor);
-		const cursorChar = this.cursor < this.text.length ? this.text[this.cursor] : " ";
-		const after = this.text.slice(this.cursor + 1);
-		return `${before}${CURSOR_MARKER}\x1b[7m${cursorChar}\x1b[27m${after}`;
-	}
 }

 // === Streaming input panel test (/overlay-streaming) ===
@@ -1,7 +1,9 @@
 import { homedir } from "node:os";
 import * as path from "node:path";
-import { type AutocompleteProvider, CombinedAutocompleteProvider, Container } from "@earendil-works/pi-tui";
+import { type AutocompleteProvider, CombinedAutocompleteProvider } from "@earendil-works/pi-tui";
 import { beforeAll, describe, expect, test, vi } from "vitest";
+import { type Component, Container, type Focusable, TUI } from "../../tui/src/tui.ts";
+import { VirtualTerminal } from "../../tui/test/virtual-terminal.ts";
 import type { AutocompleteProviderFactory } from "../src/core/extensions/types.ts";
 import type { SourceInfo } from "../src/core/source-info.ts";
 import { InteractiveMode } from "../src/modes/interactive/interactive-mode.ts";
@@ -17,6 +19,41 @@ function renderAll(container: Container, width = 120): string {
 	return container.children.flatMap((child) => child.render(width)).join("\n");
 }

+class TestFocusableComponent implements Component, Focusable {
+	focused = false;
+	inputs: string[] = [];
+	private readonly label: string;
+	private text = "";
+
+	constructor(label: string) {
+		this.label = label;
+	}
+
+	handleInput(data: string): void {
+		this.inputs.push(data);
+	}
+
+	getText(): string {
+		return this.text;
+	}
+
+	setText(text: string): void {
+		this.text = text;
+	}
+
+	render(): string[] {
+		return [this.label];
+	}
+
+	invalidate(): void {}
+}
+
+async function flushTui(tui: TUI, terminal: VirtualTerminal): Promise<void> {
+	tui.requestRender(true);
+	await Promise.resolve();
+	await terminal.waitForRender();
+}
+
 function normalizeRenderedOutput(container: Container, width = 220): string {
 	return renderAll(container, width)
 		.replace(/\u001b\[[0-9;]*m/g, "")
@@ -148,6 +185,78 @@ describe("InteractiveMode.createExtensionUIContext setTheme", () => {
 	});
 });

+describe("InteractiveMode.showExtensionCustom", () => {
+	beforeAll(() => {
+		initTheme("dark");
+	});
+
+	test("overlay custom UI reclaims input after non-overlay custom UI closes", async () => {
+		const terminal = new VirtualTerminal(80, 24);
+		const ui = new TUI(terminal);
+		const editorContainer = new Container();
+		const editor = new TestFocusableComponent("EDITOR");
+		const palette = new TestFocusableComponent("PALETTE");
+		const overlay = new TestFocusableComponent("OVERLAY");
+		const replacement = new TestFocusableComponent("REPLACEMENT");
+		let closeOverlay: (value: string) => void = () => {
+			throw new Error("closeOverlay was not initialized");
+		};
+		let closeReplacement: (value: string) => void = () => {
+			throw new Error("closeReplacement was not initialized");
+		};
+		const fakeThis = {
+			editor,
+			editorContainer,
+			keybindings: {},
+			ui,
+		};
+		const showExtensionCustom = <T>(
+			factory: (tui: TUI, theme: unknown, keybindings: unknown, done: (result: T) => void) => Component,
+			options?: { overlay?: boolean },
+		): Promise<T> =>
+			(InteractiveMode as any).prototype.showExtensionCustom.call(fakeThis, factory, options) as Promise<T>;
+
+		editorContainer.addChild(editor);
+		ui.addChild(editorContainer);
+		ui.addChild(palette);
+		ui.setFocus(palette);
+		ui.start();
+		try {
+			const overlayPromise = showExtensionCustom<string>(
+				(_tui, _theme, _keybindings, done) => {
+					closeOverlay = done;
+					return overlay;
+				},
+				{ overlay: true },
+			);
+			await flushTui(ui, terminal);
+			expect(overlay.focused).toBe(true);
+
+			const replacementPromise = showExtensionCustom<string>((_tui, _theme, _keybindings, done) => {
+				closeReplacement = done;
+				return replacement;
+			});
+			await flushTui(ui, terminal);
+			expect(replacement.focused).toBe(true);
+
+			closeReplacement("done");
+			await replacementPromise;
+			await flushTui(ui, terminal);
+			terminal.sendInput("x");
+			await flushTui(ui, terminal);
+
+			expect(overlay.inputs).toEqual(["x"]);
+			expect(editor.inputs).toEqual([]);
+			expect(overlay.focused).toBe(true);
+
+			closeOverlay("closed");
+			await overlayPromise;
+		} finally {
+			ui.stop();
+		}
+	});
+});
+
 describe("InteractiveMode.createExtensionUIContext addAutocompleteProvider", () => {
 	test("stores wrapper factories and rebuilds autocomplete immediately", () => {
 		const wrapper: AutocompleteProviderFactory = (current) => current;
@@ -2,14 +2,6 @@

 ## [Unreleased]

-### Added
-
- Added `OverlayHandle.unfocus({ target })` for explicitly releasing overlay focus to a chosen component while overlays remain visible.
-
-### 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
@@ -116,11 +116,17 @@ handle.setHidden(true);     // Temporarily hide (can show again)
 handle.setHidden(false);    // Show again after hiding
 handle.isHidden();          // Check if temporarily hidden
 handle.focus();             // Focus and bring to visual front
-handle.unfocus();           // Release focus to the next visible overlay or previous target
+handle.unfocus();           // Release focus to normal fallback
 handle.unfocus({ target: baseComponent }); // Release this overlay to a specific component
-handle.unfocus({ target: null });   // Release this overlay without focusing another component
+handle.unfocus({ target: null });   // Release this overlay and leave focus empty
 handle.isFocused();         // Check if overlay has focus

+handle.unfocus();
+// Overlay loses focus; TUI falls back to another visible capturing overlay or the previous focus target.
+
+handle.unfocus({ target: null });
+// Overlay loses focus; no component receives input until focus is set again.
+
 // A focused visible overlay reclaims keyboard input after temporary replacement UI
 // releases focus. If you want a specific component to receive input while overlays remain
 // visible, call handle.unfocus({ target: component }).
@@ -194,25 +194,32 @@ export interface OverlayHandle {
 	isHidden(): boolean;
 	/** Focus this overlay and bring it to the visual front */
 	focus(): void;
-	/** Release focus to the next visible overlay or the previous target, or to an explicit target when provided */
+	/** Release focus to the next visible capturing overlay or previous target, or to an explicit target when provided */
 	unfocus(options?: OverlayUnfocusOptions): void;
 	/** Check if this overlay currently has focus */
 	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;
+	options?: OverlayOptions;
 	preFocus: Component | null;
 	hidden: boolean;
 	focusOrder: number;
-	restoreFocus: OverlayRestoreFocusState;
 };

+type OverlayBlockedFocusResume = { status: "restore-overlay" } | { status: "focus-target"; target: Component | null };
+type EligibleOverlayFocusRestoreState = { status: "eligible"; overlay: OverlayStackEntry };
+type BlockedOverlayFocusRestoreState = {
+	status: "blocked";
+	overlay: OverlayStackEntry;
+	blockedBy: Component;
+	resume: OverlayBlockedFocusResume;
+};
+type ActiveOverlayFocusRestoreState = EligibleOverlayFocusRestoreState | BlockedOverlayFocusRestoreState;
+type OverlayFocusRestoreState = { status: "inactive" } | ActiveOverlayFocusRestoreState;
+type OverlayFocusRestorePolicy = "clear" | "preserve";
+
 /**
 * Container - a component that contains other components
 */
@@ -282,6 +289,7 @@ export class TUI extends Container {
 	// Overlay stack for modal components rendered on top of base content
 	private focusOrderCounter = 0;
 	private overlayStack: OverlayStackEntry[] = [];
+	private overlayFocusRestore: OverlayFocusRestoreState = { status: "inactive" };

 	constructor(terminal: Terminal, showHardwareCursor?: boolean) {
 		super();
@@ -322,36 +330,62 @@ export class TUI extends Container {
 	}

 	setFocus(component: Component | null): void {
+		this.setFocusInternal({ component, overlayFocusRestore: "clear" });
+	}
+
+	private setFocusInternal({
+		component,
+		overlayFocusRestore,
+	}: {
+		component: Component | null;
+		overlayFocusRestore: OverlayFocusRestorePolicy;
+	}): 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;
+		const restoreState = this.getVisibleOverlayFocusRestore();
 		if (nextFocus && !nextFocusIsOverlay) {
-			const restoreOverlay = this.getRestoreFocusOverlay();
-			if (
-				restoreOverlay?.restoreFocus.status === "blocked" &&
-				restoreOverlay.restoreFocus.blockedBy === previousFocus
-			) {
-				nextFocus = restoreOverlay.component;
+			if (restoreState.status === "blocked" && restoreState.blockedBy === previousFocus) {
+				if (restoreState.resume.status === "focus-target" || !this.isComponentMounted(restoreState.blockedBy)) {
+					nextFocus = this.resolveBlockedOverlayFocusResume(restoreState);
+				} else {
+					this.overlayFocusRestore = {
+						status: "blocked",
+						overlay: restoreState.overlay,
+						blockedBy: nextFocus,
+						resume: restoreState.resume,
+					};
+				}
 			} else if (
 				previousFocusedOverlay &&
-				previousFocusedOverlay.restoreFocus.status !== "inactive" &&
+				restoreState.status !== "inactive" &&
+				restoreState.overlay === previousFocusedOverlay &&
 				!this.isOverlayFocusAncestor(previousFocusedOverlay, nextFocus)
 			) {
-				previousFocusedOverlay.restoreFocus = { status: "blocked", blockedBy: nextFocus };
+				this.overlayFocusRestore = {
+					status: "blocked",
+					overlay: previousFocusedOverlay,
+					blockedBy: nextFocus,
+					resume: { status: "restore-overlay" },
+				};
+			}
+		} else if (nextFocus === null) {
+			if (restoreState.status === "blocked" && restoreState.blockedBy === previousFocus) {
+				nextFocus = this.resolveBlockedOverlayFocusResume(restoreState);
+			} else if (overlayFocusRestore === "clear") {
+				this.clearOverlayFocusRestore();
 			}
 		}

-		// Clear focused flag on old component
 		if (isFocusable(this.focusedComponent)) {
 			this.focusedComponent.focused = false;
 		}

 		this.focusedComponent = nextFocus;

-		// Set focused flag on new component
 		if (isFocusable(nextFocus)) {
 			nextFocus.focused = true;
 		}
@@ -360,26 +394,33 @@ export class TUI extends Container {
 			? this.overlayStack.find((entry) => entry.component === nextFocus && this.isOverlayVisible(entry))
 			: undefined;
 		if (focusedOverlay) {
-			this.markOverlayRestoreFocusEligible(focusedOverlay);
+			this.overlayFocusRestore = { status: "eligible", overlay: focusedOverlay };
 		}
 	}

-	private markOverlayRestoreFocusEligible(entry: OverlayStackEntry): void {
-		for (const overlay of this.overlayStack) {
-			this.clearOverlayRestoreFocus(overlay);
+	private clearOverlayFocusRestore(): void {
+		this.overlayFocusRestore = { status: "inactive" };
+	}
+
+	private clearOverlayFocusRestoreFor(overlay: OverlayStackEntry): void {
+		if (this.overlayFocusRestore.status !== "inactive" && this.overlayFocusRestore.overlay === overlay) {
+			this.clearOverlayFocusRestore();
 		}
-		entry.restoreFocus = { status: "eligible" };
 	}

-	private clearOverlayRestoreFocus(entry: OverlayStackEntry): void {
-		entry.restoreFocus = { status: "inactive" };
+	private resolveBlockedOverlayFocusResume(restoreState: BlockedOverlayFocusRestoreState): Component | null {
+		if (restoreState.resume.status === "restore-overlay") return restoreState.overlay.component;
+		this.clearOverlayFocusRestore();
+		return restoreState.resume.target;
 	}

-	private getRestoreFocusOverlay(): RestorableOverlayStackEntry | undefined {
-		return this.overlayStack.find(
-			(overlay): overlay is RestorableOverlayStackEntry =>
-				overlay.restoreFocus.status !== "inactive" && this.isOverlayVisible(overlay),
-		);
+	private getVisibleOverlayFocusRestore(): OverlayFocusRestoreState {
+		const restoreState = this.overlayFocusRestore;
+		if (restoreState.status === "inactive") return restoreState;
+		if (!this.overlayStack.includes(restoreState.overlay) || !this.isOverlayVisible(restoreState.overlay)) {
+			return { status: "inactive" };
+		}
+		return restoreState;
 	}

 	private isOverlayFocusAncestor(entry: OverlayStackEntry, component: Component): boolean {
@@ -393,6 +434,24 @@ export class TUI extends Container {
 		return false;
 	}

+	private retargetOverlayPreFocus(removed: OverlayStackEntry): void {
+		for (const overlay of this.overlayStack) {
+			if (overlay !== removed && overlay.preFocus === removed.component) {
+				overlay.preFocus = removed.preFocus;
+			}
+		}
+	}
+
+	private isComponentMounted(component: Component): boolean {
+		return this.children.some((child) => this.containsComponent(child, component));
+	}
+
+	private containsComponent(root: Component, target: Component): boolean {
+		if (root === target) return true;
+		if (!(root instanceof Container)) return false;
+		return root.children.some((child) => this.containsComponent(child, target));
+	}
+
 	/**
 	 * Show an overlay component with configurable positioning and sizing.
 	 * Returns a handle to control the overlay's visibility.
@@ -400,11 +459,10 @@ export class TUI extends Container {
 	showOverlay(component: Component, options?: OverlayOptions): OverlayHandle {
 		const entry: OverlayStackEntry = {
 			component,
-			options,
+			...(options === undefined ? {} : { options }),
 			preFocus: this.focusedComponent,
 			hidden: false,
 			focusOrder: ++this.focusOrderCounter,
-			restoreFocus: { status: "inactive" },
 		};
 		this.overlayStack.push(entry);
 		// Only focus if overlay is actually visible
@@ -419,7 +477,8 @@ export class TUI extends Container {
 			hide: () => {
 				const index = this.overlayStack.indexOf(entry);
 				if (index !== -1) {
-					this.clearOverlayRestoreFocus(entry);
+					this.clearOverlayFocusRestoreFor(entry);
+					this.retargetOverlayPreFocus(entry);
 					this.overlayStack.splice(index, 1);
 					// Restore focus if this overlay had focus
 					if (this.focusedComponent === component) {
@@ -435,7 +494,7 @@ export class TUI extends Container {
 				entry.hidden = hidden;
 				// Update focus when hiding/showing
 				if (hidden) {
-					this.clearOverlayRestoreFocus(entry);
+					this.clearOverlayFocusRestoreFor(entry);
 					// If this overlay had focus, move focus to next visible or preFocus
 					if (this.focusedComponent === component) {
 						const topVisible = this.getTopmostVisibleOverlay();
@@ -459,17 +518,29 @@ export class TUI extends Container {
 			},
 			unfocus: (unfocusOptions) => {
 				const isFocused = this.focusedComponent === component;
-				const hasPendingRestore = entry.restoreFocus.status !== "inactive";
-				// Nothing to release: we neither hold focus nor have a pending reclaim.
+				const restoreState = this.overlayFocusRestore;
+				const hasPendingRestore = restoreState.status !== "inactive" && restoreState.overlay === entry;
 				if (!isFocused && !hasPendingRestore) return;
-				// True when this overlay is waiting to reclaim focus from the component that
-				// currently holds it; computed before clearing the (about-to-be-reset) state.
-				const blockedByFocused =
-					entry.restoreFocus.status === "blocked" && this.focusedComponent === entry.restoreFocus.blockedBy;
-				this.clearOverlayRestoreFocus(entry);
-				// Move focus only if we currently hold it, or the caller named an explicit
-				// target and we're not mid-reclaim from the focused component.
-				if (isFocused || (unfocusOptions && !blockedByFocused)) {
+				if (
+					restoreState.status === "blocked" &&
+					restoreState.overlay === entry &&
+					this.focusedComponent === restoreState.blockedBy
+				) {
+					if (unfocusOptions) {
+						this.overlayFocusRestore = {
+							status: "blocked",
+							overlay: entry,
+							blockedBy: restoreState.blockedBy,
+							resume: { status: "focus-target", target: unfocusOptions.target },
+						};
+					} else {
+						this.clearOverlayFocusRestore();
+					}
+					this.requestRender();
+					return;
+				}
+				this.clearOverlayFocusRestoreFor(entry);
+				if (isFocused || unfocusOptions) {
 					const topVisible = this.getTopmostVisibleOverlay();
 					const fallbackTarget = topVisible && topVisible !== entry ? topVisible.component : entry.preFocus;
 					this.setFocus(unfocusOptions ? unfocusOptions.target : fallbackTarget);
@@ -482,9 +553,11 @@ export class TUI extends Container {

 	/** Hide the topmost overlay and restore previous focus. */
 	hideOverlay(): void {
-		const overlay = this.overlayStack.pop();
+		const overlay = this.overlayStack[this.overlayStack.length - 1];
 		if (!overlay) return;
-		this.clearOverlayRestoreFocus(overlay);
+		this.clearOverlayFocusRestoreFor(overlay);
+		this.retargetOverlayPreFocus(overlay);
+		this.overlayStack.pop();
 		if (this.focusedComponent === overlay.component) {
 			// Find topmost visible overlay, or fall back to preFocus
 			const topVisible = this.getTopmostVisibleOverlay();
@@ -666,20 +739,22 @@ export class TUI extends Container {
 			if (topVisible) {
 				this.setFocus(topVisible.component);
 			} else {
-				// No visible overlays, restore to preFocus
-				this.setFocus(focusedOverlay.preFocus);
+				this.setFocusInternal({ component: focusedOverlay.preFocus, overlayFocusRestore: "preserve" });
 			}
 		}

 		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);
+			const restoreState = this.getVisibleOverlayFocusRestore();
+			if (restoreState.status === "eligible") {
+				this.setFocus(restoreState.overlay.component);
+			} else if (restoreState.status === "blocked" && restoreState.blockedBy !== this.focusedComponent) {
+				if (restoreState.resume.status === "restore-overlay") {
+					this.setFocus(restoreState.overlay.component);
+				} else {
+					this.clearOverlayFocusRestore();
+					this.setFocus(restoreState.resume.target);
+				}
 			}
 		}

@@ -1,7 +1,7 @@
 import assert from "node:assert";
 import { describe, it } from "node:test";
 import type { Component, Focusable } from "../src/tui.ts";
-import { TUI } from "../src/tui.ts";
+import { Container, TUI } from "../src/tui.ts";
 import { VirtualTerminal } from "./virtual-terminal.ts";

 class StaticOverlay implements Component {
@@ -222,6 +222,35 @@ describe("TUI overlay non-capturing", () => {
 			}
 		});

+		it("removed focused child overlay does not become parent overlay fallback", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const child = new FocusableOverlay(["CHILD"]);
+			const parent = new FocusableOverlay(["PARENT"]);
+			tui.addChild(new EmptyContent());
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				const childHandle = tui.showOverlay(child, { nonCapturing: true });
+				childHandle.focus();
+				const parentHandle = tui.showOverlay(parent);
+				assert.strictEqual(parent.focused, true);
+
+				childHandle.hide();
+				parentHandle.hide();
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+
+				assert.deepStrictEqual(editor.inputs, ["x"]);
+				assert.deepStrictEqual(child.inputs, []);
+				assert.deepStrictEqual(parent.inputs, []);
+				assert.strictEqual(editor.focused, true);
+			} finally {
+				tui.stop();
+			}
+		});
+
 		it("microtask-deferred sub-overlay pattern (showExtensionCustom simulation) restores focus", async () => {
 			const terminal = new VirtualTerminal(80, 24);
 			const tui = new TUI(terminal);
@@ -234,18 +263,18 @@ describe("TUI overlay non-capturing", () => {
 			try {
 				// Simulate showExtensionCustom: factory creates timer synchronously,
 				// then .then() pushes controller as a microtask
-				let timerHandle: ReturnType<typeof tui.showOverlay>;
+				let timerHandle: ReturnType<typeof tui.showOverlay> | null = null;
 				let doneFn: () => void = () => {
 					throw new Error("doneFn was not initialized");
 				};

 				const overlayPromise = new Promise<void>((resolve) => {
 					doneFn = () => {
+						if (!timerHandle) throw new Error("timerHandle was not initialized");
 						timerHandle.hide();
 						tui.hideOverlay();
 						resolve();
 					};
-					// Factory runs synchronously: creates timer sub-overlay
 					timerHandle = tui.showOverlay(timer, { nonCapturing: true });
 					// .then() runs as microtask — same as showExtensionCustom
 					Promise.resolve(controller).then((c) => {
@@ -253,8 +282,7 @@ describe("TUI overlay non-capturing", () => {
 					});
 				});

-				// Wait for .then() microtask and renders to settle
-				await new Promise<void>((r) => setTimeout(r, 50));
+				await Promise.resolve();
 				await renderAndFlush(tui, terminal);

 				assert.strictEqual(controller.focused, true);
@@ -391,20 +419,110 @@ describe("TUI overlay non-capturing", () => {
 			}
 		});

+		it("blocked replacement can move focus internally before overlay restore", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const base = new Container();
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const firstReplacement = new FocusableOverlay(["FIRST"]);
+			const secondReplacement = new FocusableOverlay(["SECOND"]);
+			const overlay = new FocusableOverlay(["OVERLAY"]);
+			overlay.handleInput = (data: string) => {
+				overlay.inputs.push(data);
+				if (data === "b") tui.setFocus(firstReplacement);
+			};
+			firstReplacement.handleInput = (data: string) => {
+				firstReplacement.inputs.push(data);
+				if (data === "n") tui.setFocus(secondReplacement);
+			};
+			secondReplacement.handleInput = (data: string) => {
+				secondReplacement.inputs.push(data);
+				if (data === "\r") {
+					base.clear();
+					base.addChild(editor);
+					tui.setFocus(editor);
+				}
+			};
+			base.addChild(editor);
+			base.addChild(firstReplacement);
+			base.addChild(secondReplacement);
+			tui.addChild(base);
+			tui.setFocus(editor);
+			tui.start();
+			try {
+				tui.showOverlay(overlay);
+				terminal.sendInput("b");
+				await renderAndFlush(tui, terminal);
+				terminal.sendInput("n");
+				await renderAndFlush(tui, terminal);
+				terminal.sendInput("2");
+				terminal.sendInput("\r");
+				await renderAndFlush(tui, terminal);
+
+				assert.deepStrictEqual(overlay.inputs, ["b"]);
+				assert.deepStrictEqual(firstReplacement.inputs, ["n"]);
+				assert.deepStrictEqual(secondReplacement.inputs, ["2", "\r"]);
+				assert.strictEqual(overlay.focused, true);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("removed replacement restores overlay even when overlay preFocus differs from next focus", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const base = new Container();
+			const editor = new FocusableOverlay(["EDITOR"]);
+			const palette = new FocusableOverlay(["PALETTE"]);
+			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") {
+					base.clear();
+					base.addChild(editor);
+					tui.setFocus(editor);
+				}
+			};
+			base.addChild(editor);
+			base.addChild(palette);
+			base.addChild(replacement);
+			tui.addChild(base);
+			tui.setFocus(palette);
+			tui.start();
+			try {
+				tui.showOverlay(overlay);
+				terminal.sendInput("b");
+				await renderAndFlush(tui, terminal);
+				terminal.sendInput("\r");
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+
+				assert.deepStrictEqual(overlay.inputs, ["b", "x"]);
+				assert.deepStrictEqual(replacement.inputs, ["\r"]);
+				assert.deepStrictEqual(editor.inputs, []);
+				assert.strictEqual(overlay.focused, true);
+			} finally {
+				tui.stop();
+			}
+		});
+
 		it("unfocus target releases a blocked overlay while replacement remains focused", async () => {
 			const terminal = new VirtualTerminal(80, 24);
 			const tui = new TUI(terminal);
-			const editor = new FocusableOverlay(["EDITOR"]);
+			const fallback = new FocusableOverlay(["FALLBACK"]);
+			const target = new FocusableOverlay(["TARGET"]);
 			const replacement = new FocusableOverlay(["REPLACEMENT"]);
 			const overlay = new FocusableOverlay(["OVERLAY"]);
 			replacement.handleInput = (data: string) => {
 				replacement.inputs.push(data);
-				if (data === "\r") {
-					tui.setFocus(editor);
-				}
+				if (data === "\r") tui.setFocus(fallback);
 			};
 			tui.addChild(new EmptyContent());
-			tui.setFocus(editor);
 			tui.start();
 			try {
 				const overlayHandle = tui.showOverlay(overlay);
@@ -412,20 +530,21 @@ describe("TUI overlay non-capturing", () => {
 					overlay.inputs.push(data);
 					if (data === "b") {
 						tui.setFocus(replacement);
-						overlayHandle.unfocus({ target: editor });
+						overlayHandle.unfocus({ target });
 					}
 				};
+
 				terminal.sendInput("b");
 				await renderAndFlush(tui, terminal);
 				assert.strictEqual(replacement.focused, true);
-
 				terminal.sendInput("\r");
 				terminal.sendInput("x");
 				await renderAndFlush(tui, terminal);
-				assert.deepStrictEqual(replacement.inputs, ["\r"]);
+
 				assert.deepStrictEqual(overlay.inputs, ["b"]);
-				assert.deepStrictEqual(editor.inputs, ["x"]);
-				assert.strictEqual(editor.focused, true);
+				assert.deepStrictEqual(replacement.inputs, ["\r"]);
+				assert.deepStrictEqual(fallback.inputs, []);
+				assert.deepStrictEqual(target.inputs, ["x"]);
 			} finally {
 				tui.stop();
 			}
@@ -541,6 +660,54 @@ describe("TUI overlay non-capturing", () => {
 			}
 		});

+		it("setFocus(null) explicitly clears visible overlay restore", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const overlay = new FocusableOverlay(["OVERLAY"]);
+			tui.addChild(new EmptyContent());
+			tui.start();
+			try {
+				tui.showOverlay(overlay);
+				tui.setFocus(null);
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(overlay.inputs, []);
+				assert.strictEqual(overlay.focused, false);
+			} finally {
+				tui.stop();
+			}
+		});
+
+		it("blocked replacement setFocus(null) resumes the visible overlay", async () => {
+			const terminal = new VirtualTerminal(80, 24);
+			const tui = new TUI(terminal);
+			const replacement = new FocusableOverlay(["REPLACEMENT"]);
+			const overlay = new FocusableOverlay(["OVERLAY"]);
+			replacement.handleInput = (data: string) => {
+				replacement.inputs.push(data);
+				if (data === "\r") tui.setFocus(null);
+			};
+			overlay.handleInput = (data: string) => {
+				overlay.inputs.push(data);
+				if (data === "b") tui.setFocus(replacement);
+			};
+			tui.addChild(new EmptyContent());
+			tui.start();
+			try {
+				tui.showOverlay(overlay);
+				terminal.sendInput("b");
+				await renderAndFlush(tui, terminal);
+				terminal.sendInput("\r");
+				terminal.sendInput("x");
+				await renderAndFlush(tui, terminal);
+				assert.deepStrictEqual(replacement.inputs, ["\r"]);
+				assert.deepStrictEqual(overlay.inputs, ["b", "x"]);
+				assert.strictEqual(overlay.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);