feat(tui): add overlay compositing for ctx.ui.custom() (#558)

Adds overlay rendering capability to the TUI, enabling floating modal
components that render on top of existing content without clearing the screen.

- Add showOverlay(), hideOverlay(), hasOverlay() methods to TUI
- Implement ANSI-aware line compositing via extractSegments()
- Support overlay stack (multiple overlays, later on top)
- Add { overlay: true } option to ctx.ui.custom()
- Add overlay-test.ts example extension

Also fixes pre-existing bug where bash tool output cached visual lines
at fixed terminal width, causing crashes on terminal resize.

Co-authored-by: Nico Bailon <nico.bailon@gmail.com>

packages/coding-agent/CHANGELOG.md

@@ -16,6 +16,7 @@
 - Built-in renderers used automatically for tool overrides without custom `renderCall`/`renderResult`
 - `ssh.ts` example: remote tool execution via `--ssh user@host:/path`
 - Wayland clipboard support for `/copy` command using wl-copy with xclip/xsel fallback ([#570](https://github.com/badlogic/pi-mono/pull/570) by [@OgulcanCelik](https://github.com/OgulcanCelik))
+- `ctx.ui.custom()` now accepts `{ overlay: true }` option for floating modal components that composite over existing content without clearing the screen ([#558](https://github.com/badlogic/pi-mono/pull/558) by [@nicobailon](https://github.com/nicobailon))
 
 ### Fixed
 

packages/coding-agent/examples/extensions/overlay-test.ts

@@ -0,0 +1,145 @@
+/**
+ * Overlay Test - validates overlay compositing with inline text inputs
+ *
+ * Usage: pi --extension ./examples/extensions/overlay-test.ts
+ *
+ * Run /overlay-test to show a floating overlay with:
+ * - Inline text inputs within menu items
+ * - Edge case tests (wide chars, styled text, emoji)
+ */
+
+import type { ExtensionAPI, ExtensionCommandContext, Theme } from "@mariozechner/pi-coding-agent";
+import { matchesKey, visibleWidth } from "@mariozechner/pi-tui";
+
+export default function (pi: ExtensionAPI) {
+	pi.registerCommand("overlay-test", {
+		description: "Test overlay rendering with edge cases",
+		handler: async (_args: string, ctx: ExtensionCommandContext) => {
+			const result = await ctx.ui.custom<{ action: string; query?: string } | undefined>(
+				(_tui, theme, _keybindings, done) => new OverlayTestComponent(theme, done),
+				{ overlay: true },
+			);
+
+			if (result) {
+				const msg = result.query ? `${result.action}: "${result.query}"` : result.action;
+				ctx.ui.notify(msg, "info");
+			}
+		},
+	});
+}
+
+class OverlayTestComponent {
+	readonly width = 70;
+
+	private selected = 0;
+	private items = [
+		{ label: "Search", hasInput: true, text: "", cursor: 0 },
+		{ label: "Run", hasInput: true, text: "", cursor: 0 },
+		{ label: "Settings", hasInput: false, text: "", cursor: 0 },
+		{ label: "Cancel", hasInput: false, text: "", cursor: 0 },
+	];
+
+	constructor(
+		private theme: Theme,
+		private done: (result: { action: string; query?: string } | undefined) => void,
+	) {}
+
+	handleInput(data: string): void {
+		if (matchesKey(data, "escape")) {
+			this.done(undefined);
+			return;
+		}
+
+		const current = this.items[this.selected]!;
+
+		if (matchesKey(data, "return")) {
+			this.done({ action: current.label, query: current.hasInput ? current.text : undefined });
+			return;
+		}
+
+		if (matchesKey(data, "up")) {
+			this.selected = Math.max(0, this.selected - 1);
+		} else if (matchesKey(data, "down")) {
+			this.selected = Math.min(this.items.length - 1, this.selected + 1);
+		} else if (current.hasInput) {
+			if (matchesKey(data, "backspace")) {
+				if (current.cursor > 0) {
+					current.text = current.text.slice(0, current.cursor - 1) + current.text.slice(current.cursor);
+					current.cursor--;
+				}
+			} else if (matchesKey(data, "left")) {
+				current.cursor = Math.max(0, current.cursor - 1);
+			} else if (matchesKey(data, "right")) {
+				current.cursor = Math.min(current.text.length, current.cursor + 1);
+			} else if (data.length === 1 && data.charCodeAt(0) >= 32) {
+				current.text = current.text.slice(0, current.cursor) + data + current.text.slice(current.cursor);
+				current.cursor++;
+			}
+		}
+	}
+
+	render(_width: number): string[] {
+		const w = this.width;
+		const th = this.theme;
+		const innerW = w - 2;
+		const lines: string[] = [];
+
+		const pad = (s: string, len: number) => {
+			const vis = visibleWidth(s);
+			return s + " ".repeat(Math.max(0, len - vis));
+		};
+
+		const row = (content: string) => th.fg("border", "│") + pad(content, innerW) + th.fg("border", "│");
+
+		lines.push(th.fg("border", `╭${"─".repeat(innerW)}╮`));
+		lines.push(row(` ${th.fg("accent", "🧪 Overlay Test")}`));
+		lines.push(row(""));
+
+		// Edge cases - full width lines to test compositing at boundaries
+		lines.push(row(` ${th.fg("dim", "─── Edge Cases (borders should align) ───")}`));
+		lines.push(row(` Wide: ${th.fg("warning", "中文日本語한글テスト漢字繁體简体ひらがなカタカナ가나다라마바")}`));
+		lines.push(
+			row(
+				` Styled: ${th.fg("error", "RED")} ${th.fg("success", "GREEN")} ${th.fg("warning", "YELLOW")} ${th.fg("accent", "ACCENT")} ${th.fg("dim", "DIM")} ${th.fg("error", "more")} ${th.fg("success", "colors")}`,
+			),
+		);
+		lines.push(row(" Emoji: 👨‍👩‍👧‍👦 🇯🇵 🚀 💻 🎉 🔥 😀 🎯 🌟 💡 🎨 🔧 📦 🏆 🌈 🎪 🎭 🎬 🎮 🎲"));
+		lines.push(row(""));
+
+		// Menu with inline inputs
+		lines.push(row(` ${th.fg("dim", "─── Actions ───")}`));
+
+		for (let i = 0; i < this.items.length; i++) {
+			const item = this.items[i]!;
+			const isSelected = i === this.selected;
+			const prefix = isSelected ? " ▶ " : "   ";
+
+			let content: string;
+			if (item.hasInput) {
+				const label = isSelected ? th.fg("accent", `${item.label}:`) : th.fg("text", `${item.label}:`);
+
+				let inputDisplay = item.text;
+				if (isSelected) {
+					const before = inputDisplay.slice(0, item.cursor);
+					const cursorChar = item.cursor < inputDisplay.length ? inputDisplay[item.cursor] : " ";
+					const after = inputDisplay.slice(item.cursor + 1);
+					inputDisplay = `${before}\x1b[7m${cursorChar}\x1b[27m${after}`;
+				}
+				content = `${prefix + label} ${inputDisplay}`;
+			} else {
+				content = prefix + (isSelected ? th.fg("accent", item.label) : th.fg("text", item.label));
+			}
+
+			lines.push(row(content));
+		}
+
+		lines.push(row(""));
+		lines.push(row(` ${th.fg("dim", "↑↓ navigate • type to input • Enter select • Esc cancel")}`));
+		lines.push(th.fg("border", `╰${"─".repeat(innerW)}╯`));
+
+		return lines;
+	}
+
+	invalidate(): void {}
+	dispose(): void {}
+}

packages/coding-agent/src/core/extensions/types.ts

@@ -99,6 +99,7 @@ export interface ExtensionUIContext {
 			keybindings: KeybindingsManager,
 			done: (result: T) => void,
 		) => (Component & { dispose?(): void }) | Promise<Component & { dispose?(): void }>,
+		options?: { overlay?: boolean },
 	): Promise<T>;
 
 	/** Set the text in the core input editor. */

packages/coding-agent/src/modes/interactive/components/tool-execution.ts

@@ -355,24 +355,29 @@ export class ToolExecutionComponent extends Container {
 					// Show all lines when expanded
 					this.contentBox.addChild(new Text(`\n${styledOutput}`, 0, 0));
 				} else {
-					// Use visual line truncation when collapsed
+					// Use visual line truncation when collapsed with width-aware caching
-					// Box has paddingX=1, so content width = terminal.columns - 2
+					const textContent = `\n${styledOutput}`;
-					const { visualLines, skippedCount } = truncateToVisualLines(
+					let cachedWidth: number | undefined;
-						`\n${styledOutput}`,
+					let cachedLines: string[] | undefined;
-						BASH_PREVIEW_LINES,
+					let cachedSkipped: number | undefined;
-						this.ui.terminal.columns - 2,
-					);
 
-					if (skippedCount > 0) {
-						this.contentBox.addChild(
-							new Text(theme.fg("toolOutput", `\n... (${skippedCount} earlier lines)`), 0, 0),
-						);
-					}
-
-					// Add pre-rendered visual lines as a raw component
 					this.contentBox.addChild({
-						render: () => visualLines,
+						render: (width: number) => {
-						invalidate: () => {},
+							if (cachedLines === undefined || cachedWidth !== width) {
+								const result = truncateToVisualLines(textContent, BASH_PREVIEW_LINES, width);
+								cachedLines = result.visualLines;
+								cachedSkipped = result.skippedCount;
+								cachedWidth = width;
+							}
+							return cachedSkipped && cachedSkipped > 0
+								? ["", theme.fg("toolOutput", `... (${cachedSkipped} earlier lines)`), ...cachedLines]
+								: cachedLines;
+						},
+						invalidate: () => {
+							cachedWidth = undefined;
+							cachedLines = undefined;
+							cachedSkipped = undefined;
+						},
 					});
 				}
 			}

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

@@ -932,7 +932,7 @@ export class InteractiveMode {
 			setFooter: (factory) => this.setExtensionFooter(factory),
 			setHeader: (factory) => this.setExtensionHeader(factory),
 			setTitle: (title) => this.ui.terminal.setTitle(title),
-			custom: (factory) => this.showExtensionCustom(factory),
+			custom: (factory, options) => this.showExtensionCustom(factory, options),
 			setEditorText: (text) => this.editor.setText(text),
 			getEditorText: () => this.editor.getText(),
 			editor: (title, prefill) => this.showExtensionEditor(title, prefill),
@@ -1188,9 +1188,7 @@ export class InteractiveMode {
 		}
 	}
 
-	/**
+	/** Show a custom component with keyboard focus. Overlay mode renders on top of existing content. */
-	 * Show a custom component with keyboard focus.
-	 */
 	private async showExtensionCustom<T>(
 		factory: (
 			tui: TUI,
@@ -1198,28 +1196,55 @@ export class InteractiveMode {
 			keybindings: KeybindingsManager,
 			done: (result: T) => void,
 		) => (Component & { dispose?(): void }) | Promise<Component & { dispose?(): void }>,
+		options?: { overlay?: boolean },
 	): Promise<T> {
 		const savedText = this.editor.getText();
+		const isOverlay = options?.overlay ?? false;
 
-		return new Promise((resolve) => {
+		const restoreEditor = () => {
-			let component: Component & { dispose?(): void };
-
-			const close = (result: T) => {
-				component.dispose?.();
 			this.editorContainer.clear();
 			this.editorContainer.addChild(this.editor);
 			this.editor.setText(savedText);
 			this.ui.setFocus(this.editor);
 			this.ui.requestRender();
-				resolve(result);
 		};
 
-			Promise.resolve(factory(this.ui, theme, this.keybindings, close)).then((c) => {
+		return new Promise((resolve, reject) => {
+			let component: Component & { dispose?(): void };
+			let closed = false;
+
+			const close = (result: T) => {
+				if (closed) return;
+				closed = true;
+				if (isOverlay) this.ui.hideOverlay();
+				else restoreEditor();
+				// Note: both branches above already call requestRender
+				resolve(result);
+				try {
+					component?.dispose?.();
+				} catch {
+					/* ignore dispose errors */
+				}
+			};
+
+			Promise.resolve(factory(this.ui, theme, this.keybindings, close))
+				.then((c) => {
+					if (closed) return;
 					component = c;
+					if (isOverlay) {
+						const w = (component as { width?: number }).width;
+						this.ui.showOverlay(component, w ? { width: w } : undefined);
+					} else {
 						this.editorContainer.clear();
 						this.editorContainer.addChild(component);
 						this.ui.setFocus(component);
 						this.ui.requestRender();
+					}
+				})
+				.catch((err) => {
+					if (closed) return;
+					if (!isOverlay) restoreEditor();
+					reject(err);
 				});
 		});
 	}

packages/tui/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- Overlay compositing for `ctx.ui.custom()` with `{ overlay: true }` option ([#558](https://github.com/badlogic/pi-mono/pull/558) by [@nicobailon](https://github.com/nicobailon))
+
 ## [0.38.0] - 2026-01-08
 
 ### Added

packages/tui/src/tui.ts

@@ -8,7 +8,7 @@ import * as path from "node:path";
 import { isKeyRelease, matchesKey } from "./keys.js";
 import type { Terminal } from "./terminal.js";
 import { getCapabilities, setCellDimensions } from "./terminal-image.js";
-import { visibleWidth } from "./utils.js";
+import { extractSegments, sliceByColumn, sliceWithWidth, visibleWidth } from "./utils.js";
 
 /**
  * Component interface - all components must implement this
@@ -93,6 +93,13 @@ export class TUI extends Container {
 	private inputBuffer = ""; // Buffer for parsing terminal responses
 	private cellSizeQueryPending = false;
 
+	// Overlay stack for modal components rendered on top of base content
+	private overlayStack: {
+		component: Component;
+		options?: { row?: number; col?: number; width?: number };
+		preFocus: Component | null;
+	}[] = [];
+
 	constructor(terminal: Terminal) {
 		super();
 		this.terminal = terminal;
@@ -102,6 +109,32 @@ export class TUI extends Container {
 		this.focusedComponent = component;
 	}
 
+	/** Show an overlay component centered (or at specified position). */
+	showOverlay(component: Component, options?: { row?: number; col?: number; width?: number }): void {
+		this.overlayStack.push({ component, options, preFocus: this.focusedComponent });
+		this.setFocus(component);
+		this.terminal.hideCursor();
+		this.requestRender();
+	}
+
+	/** Hide the topmost overlay and restore previous focus. */
+	hideOverlay(): void {
+		const overlay = this.overlayStack.pop();
+		if (!overlay) return;
+		this.setFocus(overlay.preFocus);
+		if (this.overlayStack.length === 0) this.terminal.hideCursor();
+		this.requestRender();
+	}
+
+	hasOverlay(): boolean {
+		return this.overlayStack.length > 0;
+	}
+
+	override invalidate(): void {
+		super.invalidate();
+		for (const overlay of this.overlayStack) overlay.component.invalidate?.();
+	}
+
 	start(): void {
 		this.terminal.start(
 			(data) => this.handleInput(data),
@@ -215,12 +248,88 @@ export class TUI extends Container {
 		return line.includes("\x1b_G") || line.includes("\x1b]1337;File=");
 	}
 
+	/** Composite all overlays into content lines (in stack order, later = on top). */
+	private compositeOverlays(lines: string[], termWidth: number, termHeight: number): string[] {
+		if (this.overlayStack.length === 0) return lines;
+		const result = [...lines];
+		const viewportStart = Math.max(0, result.length - termHeight);
+
+		for (const { component, options } of this.overlayStack) {
+			const w =
+				options?.width !== undefined
+					? Math.max(1, Math.min(options.width, termWidth - 4))
+					: Math.max(1, Math.min(80, termWidth - 4));
+			const overlayLines = component.render(w);
+			const h = overlayLines.length;
+
+			const row = Math.max(0, Math.min(options?.row ?? Math.floor((termHeight - h) / 2), termHeight - h));
+			const col = Math.max(0, Math.min(options?.col ?? Math.floor((termWidth - w) / 2), termWidth - w));
+
+			for (let i = 0; i < h; i++) {
+				const idx = viewportStart + row + i;
+				if (idx >= 0 && idx < result.length) {
+					result[idx] = this.compositeLineAt(result[idx], overlayLines[i], col, w, termWidth);
+				}
+			}
+		}
+		return result;
+	}
+
+	private static readonly SEGMENT_RESET = "\x1b[0m\x1b]8;;\x07";
+
+	/** Splice overlay content into a base line at a specific column. Single-pass optimized. */
+	private compositeLineAt(
+		baseLine: string,
+		overlayLine: string,
+		startCol: number,
+		overlayWidth: number,
+		totalWidth: number,
+	): string {
+		if (this.containsImage(baseLine)) return baseLine;
+
+		// Single pass through baseLine extracts both before and after segments
+		const afterStart = startCol + overlayWidth;
+		const base = extractSegments(baseLine, startCol, afterStart, totalWidth - afterStart, true);
+
+		// Extract overlay with width tracking
+		const overlay = sliceWithWidth(overlayLine, 0, overlayWidth);
+
+		// Pad segments to target widths
+		const beforePad = Math.max(0, startCol - base.beforeWidth);
+		const overlayPad = Math.max(0, overlayWidth - overlay.width);
+		const actualBeforeWidth = Math.max(startCol, base.beforeWidth);
+		const actualOverlayWidth = Math.max(overlayWidth, overlay.width);
+		const afterTarget = Math.max(0, totalWidth - actualBeforeWidth - actualOverlayWidth);
+		const afterPad = Math.max(0, afterTarget - base.afterWidth);
+
+		// Compose result - widths are tracked so no final visibleWidth check needed
+		const r = TUI.SEGMENT_RESET;
+		const result =
+			base.before +
+			" ".repeat(beforePad) +
+			r +
+			overlay.text +
+			" ".repeat(overlayPad) +
+			r +
+			base.after +
+			" ".repeat(afterPad);
+
+		// Only truncate if wide char at after boundary caused overflow (rare)
+		const resultWidth = actualBeforeWidth + actualOverlayWidth + Math.max(afterTarget, base.afterWidth);
+		return resultWidth <= totalWidth ? result : sliceByColumn(result, 0, totalWidth, true);
+	}
+
 	private doRender(): void {
 		const width = this.terminal.columns;
 		const height = this.terminal.rows;
 
 		// Render all components to get new lines
-		const newLines = this.render(width);
+		let newLines = this.render(width);
+
+		// Composite overlays into the rendered lines (before differential compare)
+		if (this.overlayStack.length > 0) {
+			newLines = this.compositeOverlays(newLines, width, height);
+		}
 
 		// Width changed - need full re-render
 		const widthChanged = this.previousWidth !== 0 && this.previousWidth !== width;

packages/tui/src/utils.ts

@@ -135,21 +135,29 @@ export function visibleWidth(str: string): number {
 /**
  * Extract ANSI escape sequences from a string at the given position.
  */
-function extractAnsiCode(str: string, pos: number): { code: string; length: number } | null {
+export function extractAnsiCode(str: string, pos: number): { code: string; length: number } | null {
-	if (pos >= str.length || str[pos] !== "\x1b" || str[pos + 1] !== "[") {
+	if (pos >= str.length || str[pos] !== "\x1b") return null;
+
+	const next = str[pos + 1];
+
+	// CSI sequence: ESC [ ... m/G/K/H/J
+	if (next === "[") {
+		let j = pos + 2;
+		while (j < str.length && !/[mGKHJ]/.test(str[j]!)) j++;
+		if (j < str.length) return { code: str.substring(pos, j + 1), length: j + 1 - pos };
 		return null;
 	}
 
+	// OSC sequence: ESC ] ... BEL or ESC ] ... ST (ESC \)
+	// Used for hyperlinks (OSC 8), window titles, etc.
+	if (next === "]") {
 		let j = pos + 2;
-	while (j < str.length && str[j] && !/[mGKHJ]/.test(str[j]!)) {
+		while (j < str.length) {
+			if (str[j] === "\x07") return { code: str.substring(pos, j + 1), length: j + 1 - pos };
+			if (str[j] === "\x1b" && str[j + 1] === "\\") return { code: str.substring(pos, j + 2), length: j + 2 - pos };
 			j++;
 		}
-
+		return null;
-	if (j < str.length) {
-		return {
-			code: str.substring(pos, j + 1),
-			length: j + 1 - pos,
-		};
 	}
 
 	return null;
@@ -308,6 +316,11 @@ class AnsiCodeTracker {
 		this.bgColor = null;
 	}
 
+	/** Clear all state for reuse. */
+	clear(): void {
+		this.reset();
+	}
+
 	getActiveCodes(): string {
 		const codes: string[] = [];
 		if (this.bold) codes.push("1");
@@ -711,3 +724,140 @@ export function truncateToWidth(text: string, maxWidth: number, ellipsis: string
 	// Add reset code before ellipsis to prevent styling leaking into it
 	return `${result}\x1b[0m${ellipsis}`;
 }
+
+/**
+ * Extract a range of visible columns from a line. Handles ANSI codes and wide chars.
+ * @param strict - If true, exclude wide chars at boundary that would extend past the range
+ */
+export function sliceByColumn(line: string, startCol: number, length: number, strict = false): string {
+	return sliceWithWidth(line, startCol, length, strict).text;
+}
+
+/** Like sliceByColumn but also returns the actual visible width of the result. */
+export function sliceWithWidth(
+	line: string,
+	startCol: number,
+	length: number,
+	strict = false,
+): { text: string; width: number } {
+	if (length <= 0) return { text: "", width: 0 };
+	const endCol = startCol + length;
+	let result = "",
+		resultWidth = 0,
+		currentCol = 0,
+		i = 0,
+		pendingAnsi = "";
+
+	while (i < line.length) {
+		const ansi = extractAnsiCode(line, i);
+		if (ansi) {
+			if (currentCol >= startCol && currentCol < endCol) result += ansi.code;
+			else if (currentCol < startCol) pendingAnsi += ansi.code;
+			i += ansi.length;
+			continue;
+		}
+
+		let textEnd = i;
+		while (textEnd < line.length && !extractAnsiCode(line, textEnd)) textEnd++;
+
+		for (const { segment } of segmenter.segment(line.slice(i, textEnd))) {
+			const w = graphemeWidth(segment);
+			const inRange = currentCol >= startCol && currentCol < endCol;
+			const fits = !strict || currentCol + w <= endCol;
+			if (inRange && fits) {
+				if (pendingAnsi) {
+					result += pendingAnsi;
+					pendingAnsi = "";
+				}
+				result += segment;
+				resultWidth += w;
+			}
+			currentCol += w;
+			if (currentCol >= endCol) break;
+		}
+		i = textEnd;
+		if (currentCol >= endCol) break;
+	}
+	return { text: result, width: resultWidth };
+}
+
+// Pooled tracker instance for extractSegments (avoids allocation per call)
+const pooledStyleTracker = new AnsiCodeTracker();
+
+/**
+ * Extract "before" and "after" segments from a line in a single pass.
+ * Used for overlay compositing where we need content before and after the overlay region.
+ * Preserves styling from before the overlay that should affect content after it.
+ */
+export function extractSegments(
+	line: string,
+	beforeEnd: number,
+	afterStart: number,
+	afterLen: number,
+	strictAfter = false,
+): { before: string; beforeWidth: number; after: string; afterWidth: number } {
+	let before = "",
+		beforeWidth = 0,
+		after = "",
+		afterWidth = 0;
+	let currentCol = 0,
+		i = 0;
+	let pendingAnsiBefore = "";
+	let afterStarted = false;
+	const afterEnd = afterStart + afterLen;
+
+	// Track styling state so "after" inherits styling from before the overlay
+	pooledStyleTracker.clear();
+
+	while (i < line.length) {
+		const ansi = extractAnsiCode(line, i);
+		if (ansi) {
+			// Track all SGR codes to know styling state at afterStart
+			pooledStyleTracker.process(ansi.code);
+			// Include ANSI codes in their respective segments
+			if (currentCol < beforeEnd) {
+				pendingAnsiBefore += ansi.code;
+			} else if (currentCol >= afterStart && currentCol < afterEnd && afterStarted) {
+				// Only include after we've started "after" (styling already prepended)
+				after += ansi.code;
+			}
+			i += ansi.length;
+			continue;
+		}
+
+		let textEnd = i;
+		while (textEnd < line.length && !extractAnsiCode(line, textEnd)) textEnd++;
+
+		for (const { segment } of segmenter.segment(line.slice(i, textEnd))) {
+			const w = graphemeWidth(segment);
+
+			if (currentCol < beforeEnd) {
+				if (pendingAnsiBefore) {
+					before += pendingAnsiBefore;
+					pendingAnsiBefore = "";
+				}
+				before += segment;
+				beforeWidth += w;
+			} else if (currentCol >= afterStart && currentCol < afterEnd) {
+				const fits = !strictAfter || currentCol + w <= afterEnd;
+				if (fits) {
+					// On first "after" grapheme, prepend inherited styling from before overlay
+					if (!afterStarted) {
+						after += pooledStyleTracker.getActiveCodes();
+						afterStarted = true;
+					}
+					after += segment;
+					afterWidth += w;
+				}
+			}
+
+			currentCol += w;
+			// Early exit: done with "before" only, or done with both segments
+			if (afterLen <= 0 ? currentCol >= beforeEnd : currentCol >= afterEnd) break;
+		}
+		i = textEnd;
+		if (afterLen <= 0 ? currentCol >= beforeEnd : currentCol >= afterEnd) break;
+	}
+
+	return { before, beforeWidth, after, afterWidth };
+}