pi-zentui 0.15.0 → 0.15.1

package/extensions/zentui/editor-transfer.ts

@@ -0,0 +1,73 @@
+export type EditorComponentFactory = (...args: never[]) => unknown;
+
+export type EditorTransferUi<Factory = EditorComponentFactory> = {
+	getEditorText?: () => unknown;
+	setEditorText?: (text: string) => void;
+	setEditorComponent?: (factory: Factory | undefined) => void;
+	getEditorComponent?: () => Factory | undefined;
+};
+
+export type EditorTransferFailureReason =
+	| "unsupported-transfer-api"
+	| "editor-factory-snapshot-failed"
+	| "editor-text-snapshot-failed"
+	| "editor-text-preparation-failed"
+	| "editor-replacement-failed-with-rollback"
+	| "editor-replacement-rollback-failed";
+
+export type EditorTransferResult =
+	| { ok: true }
+	| { ok: false; reason: EditorTransferFailureReason };
+
+/**
+ * Prepare the active editor's public expanded text before Pi replaces its factory.
+ * This intentionally transfers prompt contents only; editor-private state is never inspected.
+ */
+export function replaceEditorComponentWithExpandedText<Factory>(
+	ui: EditorTransferUi<Factory>,
+	factory: Factory | undefined,
+): EditorTransferResult {
+	if (
+		typeof ui.getEditorText !== "function" ||
+		typeof ui.setEditorText !== "function" ||
+		typeof ui.setEditorComponent !== "function" ||
+		typeof ui.getEditorComponent !== "function"
+	) {
+		return { ok: false, reason: "unsupported-transfer-api" };
+	}
+
+	let previousFactory: Factory | undefined;
+	try {
+		previousFactory = ui.getEditorComponent();
+	} catch {
+		return { ok: false, reason: "editor-factory-snapshot-failed" };
+	}
+
+	let expandedText: unknown;
+	try {
+		expandedText = ui.getEditorText();
+	} catch {
+		return { ok: false, reason: "editor-text-snapshot-failed" };
+	}
+	if (typeof expandedText !== "string") {
+		return { ok: false, reason: "editor-text-snapshot-failed" };
+	}
+
+	try {
+		ui.setEditorText(expandedText);
+	} catch {
+		return { ok: false, reason: "editor-text-preparation-failed" };
+	}
+
+	try {
+		ui.setEditorComponent(factory);
+	} catch {
+		try {
+			ui.setEditorComponent(previousFactory);
+			return { ok: false, reason: "editor-replacement-failed-with-rollback" };
+		} catch {
+			return { ok: false, reason: "editor-replacement-rollback-failed" };
+		}
+	}
+	return { ok: true };
+}

package/extensions/zentui/format.ts

@@ -1,5 +1,4 @@
 import { homedir, hostname, userInfo } from "node:os";
-import type { AssistantMessage } from "@earendil-works/pi-ai";
 import type { ExtensionContext, Theme } from "@earendil-works/pi-coding-agent";
 import { truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
 import type {
@@ -81,23 +80,40 @@ export type UsageTotals = {
 
 export type ContextColorTier = "normal" | "warning" | "error";
 
+type SessionUsage = {
+	input?: unknown;
+	output?: unknown;
+	cacheRead?: unknown;
+	cacheWrite?: unknown;
+	cost?: unknown;
+};
+
 type SessionEntry = {
 	type?: string;
 	id?: string | number;
 	timestamp?: string | number;
+	usage?: SessionUsage;
 	message?: {
 		role?: string;
-		usage?: AssistantMessage["usage"];
+		usage?: SessionUsage;
 	};
 };
 
+type SelectedUsage = {
+	usage: SessionUsage | undefined;
+	location: "message" | "entry";
+	isAssistant: boolean;
+};
+
 type UsageCacheEntry = {
 	key: string;
 	totals: UsageTotals;
 };
 
+const MAX_USAGE_TOTAL = Number.MAX_VALUE;
+
 let usageTotalsCache: UsageCacheEntry | undefined;
-let usageTotalsComputeCount = 0;
+let usageTotalsAggregationPassCount = 0;
 
 export function formatCount(value: number): string {
 	if (value < 1000) return value.toString();
@@ -130,26 +146,78 @@ function calculateCacheHitRate(
 	cacheWrite: number,
 ): number | undefined {
 	const promptTokens = input + cacheRead + cacheWrite;
-	return promptTokens > 0 ? (cacheRead / promptTokens) * 100 : undefined;
+	if (promptTokens === 0) return undefined;
+	if (Number.isFinite(promptTokens)) return (cacheRead / promptTokens) * 100;
+
+	const scale = Math.max(input, cacheRead, cacheWrite);
+	const scaledPromptTokens = input / scale + cacheRead / scale + cacheWrite / scale;
+	return (cacheRead / scale / scaledPromptTokens) * 100;
+}
+
+function normalizeUsageNumber(value: unknown): number {
+	return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : 0;
+}
+
+function usageCostTotal(usage: SessionUsage | undefined): number {
+	if (typeof usage?.cost !== "object" || usage.cost === null) return 0;
+	return normalizeUsageNumber((usage.cost as { total?: unknown }).total);
+}
+
+function addUsageTotal(total: number, value: number): number {
+	const sum = total + value;
+	return Number.isFinite(sum) ? sum : MAX_USAGE_TOTAL;
+}
+
+function usageForEntry(entry: SessionEntry): SelectedUsage | undefined {
+	if (entry.type === "message") {
+		const role = entry.message?.role;
+		if (role !== "assistant" && role !== "toolResult") return undefined;
+		return {
+			usage: entry.message?.usage,
+			location: "message",
+			isAssistant: role === "assistant",
+		};
+	}
+	if (entry.type === "compaction" || entry.type === "branch_summary") {
+		return { usage: entry.usage, location: "entry", isAssistant: false };
+	}
+	return undefined;
+}
+
+function normalizedUsage(usage: SessionUsage | undefined) {
+	return {
+		input: normalizeUsageNumber(usage?.input),
+		output: normalizeUsageNumber(usage?.output),
+		cacheRead: normalizeUsageNumber(usage?.cacheRead),
+		cacheWrite: normalizeUsageNumber(usage?.cacheWrite),
+		cost: usageCostTotal(usage),
+	};
 }
 
-function entryIdentity(entry: SessionEntry | undefined): string {
-	if (!entry) return "";
-	const usage = entry.message?.usage;
-	const usageKey = usage
-		? `${usage.input ?? 0}:${usage.output ?? 0}:${usage.cacheRead ?? 0}:${usage.cacheWrite ?? 0}:${usage.cost?.total ?? 0}`
-		: "";
-	return `${entry.id ?? ""}|${entry.timestamp ?? ""}|${entry.type ?? ""}|${entry.message?.role ?? ""}|${usageKey}`;
+function entryIdentity(entry: SessionEntry): string {
+	const selected = usageForEntry(entry);
+	if (!selected) return "unsupported";
+	const usage = normalizedUsage(selected.usage);
+	return JSON.stringify([
+		entry.id ?? null,
+		entry.timestamp ?? null,
+		entry.type ?? null,
+		entry.message?.role ?? null,
+		selected.location,
+		usage.input,
+		usage.output,
+		usage.cacheRead,
+		usage.cacheWrite,
+		usage.cost,
+	]);
 }
 
 function buildUsageFingerprint(entries: readonly SessionEntry[]): string {
-	const first = entries[0];
-	const last = entries[entries.length - 1];
-	return `${entries.length}\0${entryIdentity(first)}\0${entryIdentity(last)}`;
+	return entries.map(entryIdentity).join("\0");
 }
 
 function computeUsageTotals(entries: readonly SessionEntry[]): UsageTotals {
-	usageTotalsComputeCount += 1;
+	usageTotalsAggregationPassCount += 1;
 	let input = 0;
 	let output = 0;
 	let cacheRead = 0;
@@ -158,18 +226,18 @@ function computeUsageTotals(entries: readonly SessionEntry[]): UsageTotals {
 	let cost = 0;
 
 	for (const entry of entries) {
-		if (entry.type !== "message" || entry.message?.role !== "assistant") continue;
-		const usage = entry.message.usage;
-		const entryInput = usage?.input ?? 0;
-		const entryCacheRead = usage?.cacheRead ?? 0;
-		const entryCacheWrite = usage?.cacheWrite ?? 0;
-
-		input += entryInput;
-		output += usage?.output ?? 0;
-		cacheRead += entryCacheRead;
-		cacheWrite += entryCacheWrite;
-		cost += usage?.cost?.total ?? 0;
-		latestCacheHitRate = calculateCacheHitRate(entryInput, entryCacheRead, entryCacheWrite);
+		const selected = usageForEntry(entry);
+		if (!selected) continue;
+		const usage = normalizedUsage(selected.usage);
+
+		input = addUsageTotal(input, usage.input);
+		output = addUsageTotal(output, usage.output);
+		cacheRead = addUsageTotal(cacheRead, usage.cacheRead);
+		cacheWrite = addUsageTotal(cacheWrite, usage.cacheWrite);
+		cost = addUsageTotal(cost, usage.cost);
+		if (selected.isAssistant) {
+			latestCacheHitRate = calculateCacheHitRate(usage.input, usage.cacheRead, usage.cacheWrite);
+		}
 	}
 
 	return Object.freeze({ input, output, cacheRead, cacheWrite, latestCacheHitRate, cost });
@@ -179,23 +247,26 @@ export function invalidateUsageTotalsCache(): void {
 	usageTotalsCache = undefined;
 }
 
-/** Test helper: number of full usage scans performed since process start / last reset. */
-export function __usageTotalsComputeCount(): number {
-	return usageTotalsComputeCount;
+/** Test helper: counts aggregation passes only; every cache lookup still fingerprints all entries. */
+export function __usageTotalsAggregationPassCount(): number {
+	return usageTotalsAggregationPassCount;
 }
 
-/** Test helper: reset memoization counters/cache. */
+/** Test helper: reset the aggregation-pass counter and cached totals. */
 export function __resetUsageTotalsCacheForTests(): void {
 	usageTotalsCache = undefined;
-	usageTotalsComputeCount = 0;
+	usageTotalsAggregationPassCount = 0;
 }
 
 export function getUsageTotals(ctx: ExtensionContext): UsageTotals {
 	const sessionManager = ctx.sessionManager as {
-		getEntries?: () => SessionEntry[];
-		getBranch: () => SessionEntry[];
+		getEntries?: () => readonly SessionEntry[];
+		getBranch: () => readonly SessionEntry[];
 	};
-	const entries = sessionManager.getEntries?.() ?? sessionManager.getBranch();
+	const entries =
+		typeof sessionManager.getEntries === "function"
+			? sessionManager.getEntries()
+			: sessionManager.getBranch();
 	const key = buildUsageFingerprint(entries);
 	if (usageTotalsCache?.key === key) return usageTotalsCache.totals;
 

package/extensions/zentui/index.ts

@@ -42,6 +42,10 @@ import {
 	saveUiFeaturesPatch,
 	type UiFeaturesConfig,
 } from "./config";
+import {
+	type EditorTransferFailureReason,
+	replaceEditorComponentWithExpandedText,
+} from "./editor-transfer";
 import {
 	disposeFixedEditor,
 	installFixedEditorProbe,
@@ -80,10 +84,30 @@ type ZentuiEditorFactory = EditorFactory & {
 
 type ApplyUiResult = {
 	editorBlocked: boolean;
+	editorReason?: string;
 };
 
+type EditorChangeResult = { ok: true } | { ok: false; reason: string };
+
 type EditorInstallMode = "none" | "standalone" | "wrapper";
 
+function editorTransferFailureMessage(reason: EditorTransferFailureReason): string {
+	switch (reason) {
+		case "unsupported-transfer-api":
+			return "this Pi version cannot safely transfer expanded editor text; reload Pi to apply this change";
+		case "editor-factory-snapshot-failed":
+			return "the current editor factory could not be read safely; reload Pi to apply this change";
+		case "editor-text-snapshot-failed":
+			return "expanded editor text could not be read safely; reload Pi to apply this change";
+		case "editor-text-preparation-failed":
+			return "expanded editor text could not be prepared safely; reload Pi to apply this change";
+		case "editor-replacement-failed-with-rollback":
+			return "the editor replacement failed; the previous factory was reapplied, but editor instance identity is not guaranteed";
+		case "editor-replacement-rollback-failed":
+			return "the editor replacement and previous-factory rollback both failed; reload Pi before editing";
+	}
+}
+
 function isZentuiEditorFactory(factory: EditorFactory | undefined): boolean {
 	return Boolean((factory as ZentuiEditorFactory | undefined)?.[ZENTUI_EDITOR_FACTORY]);
 }
@@ -266,15 +290,28 @@ export default function (pi: ExtensionAPI) {
 		if (footerInstalled) scheduleProjectRefresh(ctx, { force: true });
 	};
 
-	const installPrototypePatches = () => {
-		if (prototypePatchesInstalled) return;
-		const cleanupSelectorBorderStyle = installSelectorBorderStyle(getActiveTheme, getCurrentConfig);
-		const cleanupUserMessageStyle = installUserMessageStyle(getActiveTheme, getCurrentConfig);
-		cleanupPrototypePatches = () => {
-			cleanupSelectorBorderStyle();
-			cleanupUserMessageStyle();
-		};
-		prototypePatchesInstalled = true;
+	const installPrototypePatches = (): boolean => {
+		if (prototypePatchesInstalled) return true;
+		let cleanupSelectorBorderStyle: (() => void) | undefined;
+		try {
+			cleanupSelectorBorderStyle = installSelectorBorderStyle(getActiveTheme, getCurrentConfig);
+			const cleanupUserMessageStyle = installUserMessageStyle(getActiveTheme, getCurrentConfig);
+			cleanupPrototypePatches = () => {
+				cleanupSelectorBorderStyle?.();
+				cleanupUserMessageStyle();
+			};
+			prototypePatchesInstalled = true;
+			return true;
+		} catch {
+			try {
+				cleanupSelectorBorderStyle?.();
+			} catch {
+				// Best effort: each installer is locally transactional as well.
+			}
+			cleanupPrototypePatches = () => {};
+			prototypePatchesInstalled = false;
+			return false;
+		}
 	};
 
 	const uninstallPrototypePatches = () => {
@@ -283,6 +320,43 @@ export default function (pi: ExtensionAPI) {
 		prototypePatchesInstalled = false;
 	};
 
+	const clearEditorOwnership = () => {
+		wrappedEditorFactory = undefined;
+		installedEditorFactory = undefined;
+		editorInstallMode = "none";
+		editorInstalled = false;
+	};
+
+	const trackZentuiEditorFactory = (factory: EditorFactory) => {
+		const baseFactory = getZentuiEditorBaseFactory(factory);
+		wrappedEditorFactory = baseFactory;
+		installedEditorFactory = factory;
+		editorInstallMode = baseFactory ? "wrapper" : "standalone";
+		editorInstalled = true;
+	};
+
+	const observeEditorFactory = (
+		ctx: ExtensionContext,
+	): { known: true; factory: EditorFactory | undefined } | { known: false } => {
+		try {
+			return { known: true, factory: ctx.ui.getEditorComponent() };
+		} catch {
+			return { known: false };
+		}
+	};
+
+	const reconcileObservedEditorOwnership = (ctx: ExtensionContext) => {
+		const observed = observeEditorFactory(ctx);
+		if (!observed.known) return observed;
+		if (observed.factory && isZentuiEditorFactory(observed.factory)) {
+			trackZentuiEditorFactory(observed.factory);
+		} else {
+			uninstallPrototypePatches();
+			clearEditorOwnership();
+		}
+		return observed;
+	};
+
 	const makeEditorFactory = (ctx: ExtensionContext): ZentuiEditorFactory => {
 		const sessionTheme = ctx.ui.theme;
 		const factory = ((tui: TUI, theme: EditorTheme, keybindings: KeybindingsManager) =>
@@ -329,53 +403,73 @@ export default function (pi: ExtensionAPI) {
 		return factory;
 	};
 
-	const installEditor = (ctx: ExtensionContext): boolean => {
+	const replaceEditor = (
+		ctx: ExtensionContext,
+		factory: EditorFactory | undefined,
+	): EditorChangeResult => {
+		const result = replaceEditorComponentWithExpandedText(ctx.ui, factory);
+		return result.ok ? result : { ok: false, reason: editorTransferFailureMessage(result.reason) };
+	};
+
+	const installEditor = (ctx: ExtensionContext): EditorChangeResult => {
 		const currentFactory = ctx.ui.getEditorComponent();
 		if (currentFactory && currentFactory === installedEditorFactory) {
 			editorInstalled = true;
-			return true;
+			return { ok: true };
 		}
 
-		installPrototypePatches();
 		const currentZentuiBaseFactory = getZentuiEditorBaseFactory(currentFactory);
-		if (currentFactory && isZentuiEditorFactory(currentFactory)) {
-			wrappedEditorFactory = currentZentuiBaseFactory;
-			const nextFactory = currentZentuiBaseFactory
-				? makeWrappedEditorFactory(ctx, currentZentuiBaseFactory)
-				: makeEditorFactory(ctx);
-			ctx.ui.setEditorComponent(nextFactory);
-			installedEditorFactory = nextFactory;
-			editorInstallMode = currentZentuiBaseFactory ? "wrapper" : "standalone";
-		} else if (currentFactory) {
-			wrappedEditorFactory = currentFactory;
-			const nextFactory = makeWrappedEditorFactory(ctx, currentFactory);
-			ctx.ui.setEditorComponent(nextFactory);
-			installedEditorFactory = nextFactory;
-			editorInstallMode = "wrapper";
-		} else {
-			wrappedEditorFactory = undefined;
-			const nextFactory = makeEditorFactory(ctx);
-			ctx.ui.setEditorComponent(nextFactory);
-			installedEditorFactory = nextFactory;
-			editorInstallMode = "standalone";
+		const baseFactory =
+			currentZentuiBaseFactory ??
+			(currentFactory && !isZentuiEditorFactory(currentFactory) ? currentFactory : undefined);
+		const nextFactory = baseFactory
+			? makeWrappedEditorFactory(ctx, baseFactory)
+			: makeEditorFactory(ctx);
+		const replacement = replaceEditor(ctx, nextFactory);
+		if (!replacement.ok) return replacement;
+		if (!installPrototypePatches()) {
+			const rollback = replaceEditor(ctx, currentFactory);
+			reconcileObservedEditorOwnership(ctx);
+			return {
+				ok: false,
+				reason: rollback.ok
+					? "editor chrome patch installation failed; the previous factory was reapplied, but editor instance identity is not guaranteed"
+					: "editor chrome patch installation failed and the previous factory could not be restored; reload Pi before editing",
+			};
 		}
-		editorInstalled = true;
-		return true;
+
+		trackZentuiEditorFactory(nextFactory);
+		return { ok: true };
 	};
 
-	const uninstallEditor = (ctx: ExtensionContext): boolean => {
-		const currentFactory = ctx.ui.getEditorComponent();
-		if (currentFactory && !isZentuiEditorFactory(currentFactory)) return false;
+	const uninstallEditor = (ctx: ExtensionContext): EditorChangeResult => {
+		const observed = observeEditorFactory(ctx);
+		if (!observed.known) {
+			return {
+				ok: false,
+				reason:
+					"the current editor factory could not be observed safely; reload Pi to apply this change",
+			};
+		}
+		const currentFactory = observed.factory;
+		if (!currentFactory || !isZentuiEditorFactory(currentFactory)) {
+			uninstallPrototypePatches();
+			clearEditorOwnership();
+			return { ok: true };
+		}
 
-		uninstallPrototypePatches();
-		ctx.ui.setEditorComponent(
-			editorInstallMode === "wrapper" && wrappedEditorFactory ? wrappedEditorFactory : undefined,
+		const replacement = replaceEditor(
+			ctx,
+			getZentuiEditorBaseFactory(currentFactory) ??
+				(editorInstallMode === "wrapper" && wrappedEditorFactory
+					? wrappedEditorFactory
+					: undefined),
 		);
-		wrappedEditorFactory = undefined;
-		installedEditorFactory = undefined;
-		editorInstallMode = "none";
-		editorInstalled = false;
-		return true;
+		if (!replacement.ok) return replacement;
+
+		uninstallPrototypePatches();
+		clearEditorOwnership();
+		return { ok: true };
 	};
 
 	const installStatusLine = (ctx: ExtensionContext) => {
@@ -413,12 +507,17 @@ export default function (pi: ExtensionAPI) {
 		const result: ApplyUiResult = { editorBlocked: false };
 		if (!isTuiContext(ctx)) return result;
 		activeTheme = ctx.ui.theme;
+		let editorChange: EditorChangeResult | undefined;
 		if (currentConfig.features.editor) {
 			const currentFactory = ctx.ui.getEditorComponent();
 			const editorMissingOrReplaced = !editorInstalled || !isZentuiEditorFactory(currentFactory);
-			if (editorMissingOrReplaced) result.editorBlocked = !installEditor(ctx);
+			if (editorMissingOrReplaced) editorChange = installEditor(ctx);
 		} else if (editorInstalled || prototypePatchesInstalled) {
-			result.editorBlocked = !uninstallEditor(ctx);
+			editorChange = uninstallEditor(ctx);
+		}
+		if (editorChange && !editorChange.ok) {
+			result.editorBlocked = true;
+			result.editorReason = editorChange.reason;
 		}
 
 		if (currentConfig.features.statusLine) {
@@ -450,11 +549,16 @@ export default function (pi: ExtensionAPI) {
 	const scheduleEditorReconciliation = (ctx: ExtensionContext) => {
 		sessionLifecycle.defer(() => {
 			if (!isTuiContext(ctx) || !currentConfig.features.editor) return;
-			const currentFactory = ctx.ui.getEditorComponent();
-			if (currentFactory && currentFactory !== installedEditorFactory) {
-				applyConfiguredUi(ctx);
+			const observed = observeEditorFactory(ctx);
+			if (!observed.known || observed.factory === installedEditorFactory) return;
+			if (!observed.factory || !isZentuiEditorFactory(observed.factory)) {
+				uninstallPrototypePatches();
+				clearEditorOwnership();
 				refresh();
+				return;
 			}
+			trackZentuiEditorFactory(observed.factory);
+			refresh();
 		});
 	};
 
@@ -464,28 +568,29 @@ export default function (pi: ExtensionAPI) {
 		try {
 			disposeFixedEditor(ctx);
 			if (isTuiContext(ctx)) removeFixedEditorProbe(ctx);
-			uninstallPrototypePatches();
 			stopSessionTimer();
 			stopProjectRefresh();
 			requestFooterRender = undefined;
 			getActiveExtensionStatuses = () => new Map();
 			if (isTuiContext(ctx)) {
 				ctx.ui.setFooter(undefined);
-				const currentFactory = ctx.ui.getEditorComponent();
-				if (!currentFactory || isZentuiEditorFactory(currentFactory)) {
-					ctx.ui.setEditorComponent(
-						getZentuiEditorBaseFactory(currentFactory) ??
-							(editorInstallMode === "wrapper" && wrappedEditorFactory
-								? wrappedEditorFactory
-								: undefined),
-					);
+				const before = observeEditorFactory(ctx);
+				if (before.known) {
+					const currentFactory = before.factory;
+					if (currentFactory && isZentuiEditorFactory(currentFactory)) {
+						replaceEditor(
+							ctx,
+							getZentuiEditorBaseFactory(currentFactory) ??
+								(editorInstallMode === "wrapper" && wrappedEditorFactory
+									? wrappedEditorFactory
+									: undefined),
+						);
+					}
+					reconcileObservedEditorOwnership(ctx);
 				}
 			}
-			wrappedEditorFactory = undefined;
-			installedEditorFactory = undefined;
-			editorInstallMode = "none";
+			uninstallPrototypePatches();
 			footerInstalled = false;
-			editorInstalled = false;
 			activeTheme = undefined;
 		} finally {
 			requestFooterRender = undefined;
@@ -520,9 +625,7 @@ export default function (pi: ExtensionAPI) {
 			const result = applyConfiguredUi(ctx);
 			return {
 				applied: !(patch.editor !== undefined && result.editorBlocked),
-				reason: result.editorBlocked
-					? "another extension is currently managing the editor; reload Pi to apply this change"
-					: undefined,
+				reason: result.editorBlocked ? result.editorReason : undefined,
 			};
 		},
 		setFooterSegments(patch: Partial<FooterSegmentsConfig>, ctx: ExtensionContext) {

package/extensions/zentui/prototype-patch-registry.ts

@@ -79,6 +79,7 @@ export function installPrototypePatch(
 	if (!(record && record.method === method && target[method] === record.wrapper)) {
 		const predecessor = target[method];
 		if (typeof predecessor !== "function") {
+			if (registry.size === 0) delete target[ZENTUI_PROTOTYPE_PATCH_REGISTRY];
 			throw new TypeError(`Cannot patch ${method}: predecessor is not a function`);
 		}
 		const nextRecord: PatchRecord = {
@@ -98,7 +99,13 @@ export function installPrototypePatch(
 		nextRecord.wrapper = wrapper;
 		record = nextRecord;
 		registry.set(adapter, record);
-		target[method] = wrapper;
+		try {
+			target[method] = wrapper;
+		} catch (error) {
+			registry.delete(adapter);
+			if (registry.size === 0) delete target[ZENTUI_PROTOTYPE_PATCH_REGISTRY];
+			throw error;
+		}
 	}
 
 	const token = Symbol(adapter);

package/extensions/zentui/selector-border.ts

@@ -65,11 +65,17 @@ export function installSelectorBorderStyle(
 		getTheme,
 		getConfig,
 	);
-	const cleanupSettings = patchSelectorBorderStyle(
-		SettingsSelectorComponent.prototype as unknown as PatchableSelectorPrototype,
-		getTheme,
-		getConfig,
-	);
+	let cleanupSettings: Cleanup;
+	try {
+		cleanupSettings = patchSelectorBorderStyle(
+			SettingsSelectorComponent.prototype as unknown as PatchableSelectorPrototype,
+			getTheme,
+			getConfig,
+		);
+	} catch (error) {
+		cleanupModel();
+		throw error;
+	}
 	return () => {
 		cleanupModel();
 		cleanupSettings();

package/extensions/zentui/ui.ts

@@ -17,7 +17,7 @@ import {
 	safeThemeFg,
 } from "./style";
 
-const SPLIT_POLISHED_FRAME: unique symbol = Symbol.for("pi-zentui.polished-frame");
+const LEGACY_SPLIT_POLISHED_FRAME = Symbol.for("pi-zentui.polished-frame");
 
 type ViewportCounts = {
 	above?: string;
@@ -30,7 +30,12 @@ type PolishedFrameSplit = {
 	viewport: ViewportCounts;
 };
 
-const POLISHED_FRAME_SPLITS = new WeakMap<string[], PolishedFrameSplit>();
+type PolishedFrameProvenance = {
+	rows: readonly string[];
+	split: PolishedFrameSplit;
+};
+
+const POLISHED_FRAME_SPLITS = new WeakMap<string[], PolishedFrameProvenance>();
 
 type AutocompleteEditorInternals = {
 	autocompleteList?: Pick<Component, "render">;
@@ -58,7 +63,6 @@ type WrappedEditor = EditorComponent &
 		setAutocompleteProvider?: (provider: AutocompleteProvider) => void;
 		setPaddingX?: (padding: number) => void;
 		setAutocompleteMaxVisible?: (maxVisible: number) => void;
-		[SPLIT_POLISHED_FRAME]?: (lines: string[]) => PolishedFrameSplit | undefined;
 	};
 
 type EditorMeta = {
@@ -78,9 +82,58 @@ type PolishedFrameOptions = {
 	modelMeta: EditorMeta;
 	thinkingLevel: string | undefined;
 	rightStatus?: string;
-	splitBaseFrame?: (lines: string[]) => PolishedFrameSplit | undefined;
+	ownedFrame?: PolishedFrameSplit;
+	trustedBaseFrame?: boolean;
 };
 
+type PolishedFrameResult = { lines: string[]; decorated: boolean };
+
+type AutocompleteCount = { known: true; count: number } | { known: false };
+
+function isStringArray(value: unknown): value is string[] {
+	return Array.isArray(value) && value.every((line) => typeof line === "string");
+}
+
+function isViewportCounts(value: unknown): value is ViewportCounts {
+	if (!value || typeof value !== "object") return false;
+	const counts = value as Record<string, unknown>;
+	return [counts.above, counts.below].every(
+		(count) => count === undefined || (typeof count === "string" && /^[1-9]\d*$/.test(count)),
+	);
+}
+
+function isPolishedFrameSplit(value: unknown, baseLineCount: number): value is PolishedFrameSplit {
+	if (!value || typeof value !== "object") return false;
+	const split = value as Record<string, unknown>;
+	return (
+		isStringArray(split.editorLines) &&
+		isStringArray(split.trailingLines) &&
+		split.trailingLines.length <= baseLineCount &&
+		isViewportCounts(split.viewport)
+	);
+}
+
+function readAutocompleteCount(
+	source: AutocompleteEditorInternals,
+	width: number,
+	baseLineCount: number,
+): AutocompleteCount {
+	try {
+		const showing = source.isShowingAutocomplete;
+		if (typeof showing !== "function") return { known: true, count: 0 };
+		if (!showing.call(source)) return { known: true, count: 0 };
+		const list = source.autocompleteList;
+		if (!list || typeof list.render !== "function") return { known: false };
+		const rendered = list.render(width);
+		if (!isStringArray(rendered) || rendered.length <= 0 || rendered.length >= baseLineCount) {
+			return { known: false };
+		}
+		return { known: true, count: rendered.length };
+	} catch {
+		return { known: false };
+	}
+}
+
 function clampRenderedLines(lines: string[], width: number): string[] {
 	const maxWidth = Math.max(0, width);
 	return lines.map((line) => truncateToWidth(line, maxWidth, ""));
@@ -264,42 +317,49 @@ function renderPolishedFrame({
 	modelMeta,
 	thinkingLevel,
 	rightStatus,
-	splitBaseFrame,
-}: PolishedFrameOptions): string[] {
-	if (width <= 2) return clampRenderedLines(baseRendered, width);
+	ownedFrame,
+	trustedBaseFrame = false,
+}: PolishedFrameOptions): PolishedFrameResult {
+	if (width <= 2) return { lines: clampRenderedLines(baseRendered, width), decorated: false };
 
 	const reset = "\x1b[0m";
 	const colorSource = config.colorSources.editor;
 	const { prompt, promptWidth, rail, railWidth } = getEditorChromeWidths(config, uiTheme, reset);
 	const innerWidth = Math.max(0, width - railWidth);
 	const copyFriendlyContinuation = " ".repeat(promptWidth);
-	const isShowingAutocomplete =
-		typeof autocompleteSource.isShowingAutocomplete === "function"
-			? Boolean(autocompleteSource.isShowingAutocomplete())
-			: false;
-
-	if (baseRendered.length < 2) return clampRenderedLines(baseRendered, width);
-
-	const ownedFrame = splitBaseFrame?.(baseRendered);
-	const { autocompleteList } = autocompleteSource;
-	const autocompleteCount =
-		!ownedFrame && isShowingAutocomplete && typeof autocompleteList?.render === "function"
-			? autocompleteList.render(innerWidth).length
-			: 0;
+
+	if (baseRendered.length < 2) {
+		return { lines: clampRenderedLines(baseRendered, width), decorated: false };
+	}
+	if (ownedFrame && !isPolishedFrameSplit(ownedFrame, baseRendered.length)) {
+		return { lines: clampRenderedLines(baseRendered, width), decorated: false };
+	}
+
+	const autocomplete = ownedFrame
+		? { known: true, count: ownedFrame.trailingLines.length }
+		: readAutocompleteCount(autocompleteSource, innerWidth, baseRendered.length);
+	if (!autocomplete.known) {
+		return { lines: clampRenderedLines(baseRendered, width), decorated: false };
+	}
 	const editorFrame =
-		!ownedFrame && autocompleteCount > 0 && autocompleteCount < baseRendered.length
-			? baseRendered.slice(0, -autocompleteCount)
+		!ownedFrame && autocomplete.count > 0
+			? baseRendered.slice(0, -autocomplete.count)
 			: baseRendered;
 	const autocompleteLines = ownedFrame
 		? ownedFrame.trailingLines
-		: autocompleteCount > 0 && autocompleteCount < baseRendered.length
-			? baseRendered.slice(-autocompleteCount)
+		: autocomplete.count > 0
+			? baseRendered.slice(-autocomplete.count)
 			: [];
-	if (editorFrame.length < 2) return clampRenderedLines(baseRendered, width);
+	if (editorFrame.length < 2) {
+		return { lines: clampRenderedLines(baseRendered, width), decorated: false };
+	}
 
-	const editorLines = ownedFrame?.editorLines ?? editorFrame.slice(1, -1);
 	const parsedTop = parseEditorBorder(editorFrame[0] ?? "", "above");
 	const parsedBottom = parseEditorBorder(editorFrame.at(-1) ?? "", "below");
+	if (!ownedFrame && !trustedBaseFrame && (!parsedTop || !parsedBottom)) {
+		return { lines: clampRenderedLines(baseRendered, width), decorated: false };
+	}
+	const editorLines = ownedFrame?.editorLines ?? editorFrame.slice(1, -1);
 	const viewport = ownedFrame?.viewport ?? {
 		above: parsedTop?.count,
 		below: parsedBottom?.count,
@@ -365,11 +425,14 @@ function renderPolishedFrame({
 
 	const clamped = clampRenderedLines(renderedLines, width);
 	POLISHED_FRAME_SPLITS.set(clamped, {
-		editorLines,
-		trailingLines: autocompleteLines.length > 0 ? clamped.slice(-autocompleteLines.length) : [],
-		viewport,
+		rows: Object.freeze([...clamped]),
+		split: {
+			editorLines,
+			trailingLines: autocompleteLines.length > 0 ? clamped.slice(-autocompleteLines.length) : [],
+			viewport,
+		},
 	});
-	return clamped;
+	return { lines: clamped, decorated: true };
 }
 
 export class PolishedEditor extends CustomEditor {
@@ -404,34 +467,53 @@ export class PolishedEditor extends CustomEditor {
 		const { railWidth } = getEditorChromeWidths(config, this.uiTheme, "\x1b[0m");
 		const innerWidth = Math.max(0, width - railWidth);
 		const rendered = super.render(innerWidth);
-		const modelMeta = this.getModelMeta();
-		const result = renderPolishedFrame({
-			width,
-			baseRendered: rendered,
-			autocompleteSource: this as unknown as AutocompleteEditorInternals,
-			uiTheme: this.uiTheme,
-			config,
-			modelMeta,
-			thinkingLevel: this.getThinkingLevel(),
-		});
-		return result;
-	}
-
-	[SPLIT_POLISHED_FRAME](lines: string[]): PolishedFrameSplit | undefined {
-		return (
-			POLISHED_FRAME_SPLITS.get(lines) ?? splitPolishedFrame(lines, this.getConfig(), this.uiTheme)
-		);
+		try {
+			return renderPolishedFrame({
+				width,
+				baseRendered: rendered,
+				autocompleteSource: this as unknown as AutocompleteEditorInternals,
+				uiTheme: this.uiTheme,
+				config,
+				modelMeta: this.getModelMeta(),
+				thinkingLevel: this.getThinkingLevel(),
+				trustedBaseFrame: true,
+			}).lines;
+		} catch {
+			return clampRenderedLines(rendered, width);
+		}
 	}
 }
 
 export class WrappedPolishedEditor implements EditorComponent {
+	declare readonly addToHistory?: (text: string) => void;
+	declare readonly insertTextAtCursor?: (text: string) => void;
+	declare readonly setAutocompleteProvider?: (provider: AutocompleteProvider) => void;
+	declare readonly setPaddingX?: (padding: number) => void;
+	declare readonly setAutocompleteMaxVisible?: (maxVisible: number) => void;
+
 	constructor(
 		private readonly base: WrappedEditor,
 		private readonly uiTheme: Theme,
 		private readonly getConfig: () => PolishedTuiConfig,
 		private readonly getModelMeta: () => EditorMeta,
 		private readonly getThinkingLevel: () => string | undefined,
-	) {}
+	) {
+		if (typeof base.addToHistory === "function") {
+			this.addToHistory = (text) => base.addToHistory?.(text);
+		}
+		if (typeof base.insertTextAtCursor === "function") {
+			this.insertTextAtCursor = (text) => base.insertTextAtCursor?.(text);
+		}
+		if (typeof base.setAutocompleteProvider === "function") {
+			this.setAutocompleteProvider = (provider) => base.setAutocompleteProvider?.(provider);
+		}
+		if (typeof base.setPaddingX === "function") {
+			this.setPaddingX = (padding) => base.setPaddingX?.(padding);
+		}
+		if (typeof base.setAutocompleteMaxVisible === "function") {
+			this.setAutocompleteMaxVisible = (maxVisible) => base.setAutocompleteMaxVisible?.(maxVisible);
+		}
+	}
 
 	get focused(): boolean {
 		return Boolean(this.base.focused);
@@ -516,26 +598,42 @@ export class WrappedPolishedEditor implements EditorComponent {
 		const config = this.getConfig();
 		const { railWidth } = getEditorChromeWidths(config, this.uiTheme, "\x1b[0m");
 		const innerWidth = Math.max(0, width - railWidth);
-		const rendered = this.base.render(innerWidth);
-		const vimStatus = readVimStatus(this.base, this.uiTheme);
-		const modelMeta = this.getModelMeta();
-		return renderPolishedFrame({
-			width,
-			baseRendered: rendered,
-			autocompleteSource: this.base,
-			uiTheme: this.uiTheme,
-			config,
-			modelMeta,
-			thinkingLevel: this.getThinkingLevel(),
-			rightStatus: vimStatus,
-			splitBaseFrame: this.base[SPLIT_POLISHED_FRAME]?.bind(this.base),
-		});
-	}
-
-	[SPLIT_POLISHED_FRAME](lines: string[]): PolishedFrameSplit | undefined {
-		return (
-			POLISHED_FRAME_SPLITS.get(lines) ?? splitPolishedFrame(lines, this.getConfig(), this.uiTheme)
-		);
+		let rendered: string[];
+		try {
+			rendered = this.base.render(innerWidth);
+		} catch {
+			return clampRenderedLines(this.base.render(width), width);
+		}
+		let result: PolishedFrameResult | undefined;
+		try {
+			const provenance = POLISHED_FRAME_SPLITS.get(rendered);
+			const provenanceMatches = Boolean(
+				provenance &&
+					provenance.rows.length === rendered.length &&
+					provenance.rows.every((line, index) => line === rendered[index]),
+			);
+			const ownedFrame = provenanceMatches ? provenance?.split : undefined;
+			const hasMutatedProvenance = Boolean(provenance && !provenanceMatches);
+			const hasUntrustedLegacySplitter = LEGACY_SPLIT_POLISHED_FRAME in this.base;
+			const hasUnprovenPolishedFrame =
+				!ownedFrame && Boolean(splitPolishedFrame(rendered, config, this.uiTheme));
+			if (!hasMutatedProvenance && !hasUntrustedLegacySplitter && !hasUnprovenPolishedFrame) {
+				result = renderPolishedFrame({
+					width,
+					baseRendered: rendered,
+					autocompleteSource: this.base,
+					uiTheme: this.uiTheme,
+					config,
+					modelMeta: this.getModelMeta(),
+					thinkingLevel: this.getThinkingLevel(),
+					rightStatus: readVimStatus(this.base, this.uiTheme),
+					ownedFrame,
+				});
+			}
+		} catch {
+			// Decoration is optional; re-render the base at the caller's width below.
+		}
+		return result?.decorated ? result.lines : clampRenderedLines(this.base.render(width), width);
 	}
 
 	invalidate(): void {
@@ -554,30 +652,10 @@ export class WrappedPolishedEditor implements EditorComponent {
 		this.base.setText(text);
 	}
 
-	addToHistory(text: string): void {
-		this.base.addToHistory?.(text);
-	}
-
-	insertTextAtCursor(text: string): void {
-		this.base.insertTextAtCursor?.(text);
-	}
-
 	getExpandedText(): string {
 		return this.base.getExpandedText?.() ?? this.base.getText();
 	}
 
-	setAutocompleteProvider(provider: AutocompleteProvider): void {
-		this.base.setAutocompleteProvider?.(provider);
-	}
-
-	setPaddingX(padding: number): void {
-		this.base.setPaddingX?.(padding);
-	}
-
-	setAutocompleteMaxVisible(maxVisible: number): void {
-		this.base.setAutocompleteMaxVisible?.(maxVisible);
-	}
-
 	getLines(): string[] {
 		return this.base.getLines?.() ?? this.base.getText().split("\n");
 	}

package/extensions/zentui/user-message.ts

@@ -225,24 +225,30 @@ export function installUserMessageStyle(
 			return Reflect.apply(predecessor, receiver, args);
 		},
 	);
-	const cleanupRender = installPrototypePatch(
-		prototype,
-		"render",
-		"user-message-render",
-		({ predecessor, receiver, args }) => {
-			const width = args[0];
-			if (typeof width !== "number") return Reflect.apply(predecessor, receiver, args);
-			const lines = renderZentuiUserMessage(
-				receiver as PatchableUserMessagePrototype,
-				width,
-				getTheme(),
-				getConfig(),
-			);
-			if (!lines) return Reflect.apply(predecessor, receiver, args);
-			if (lines.length === 0) return lines;
-			return withPromptZoneMarkers(lines);
-		},
-	);
+	let cleanupRender: Cleanup;
+	try {
+		cleanupRender = installPrototypePatch(
+			prototype,
+			"render",
+			"user-message-render",
+			({ predecessor, receiver, args }) => {
+				const width = args[0];
+				if (typeof width !== "number") return Reflect.apply(predecessor, receiver, args);
+				const lines = renderZentuiUserMessage(
+					receiver as PatchableUserMessagePrototype,
+					width,
+					getTheme(),
+					getConfig(),
+				);
+				if (!lines) return Reflect.apply(predecessor, receiver, args);
+				if (lines.length === 0) return lines;
+				return withPromptZoneMarkers(lines);
+			},
+		);
+	} catch (error) {
+		cleanupInvalidate();
+		throw error;
+	}
 	let cleaned = false;
 	return () => {
 		if (cleaned) return;

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "pi-zentui",
-	"version": "0.15.0",
+	"version": "0.15.1",
 	"description": "A Starship-inspired statusline and Opencode-style TUI for Pi.",
 	"type": "module",
 	"license": "MIT",
@@ -50,9 +50,9 @@
 	},
 	"devDependencies": {
 		"@biomejs/biome": "^2.5.0",
-		"@earendil-works/pi-ai": "^0.80.6",
-		"@earendil-works/pi-coding-agent": "^0.80.6",
-		"@earendil-works/pi-tui": "^0.80.6",
+		"@earendil-works/pi-ai": "^0.82.1",
+		"@earendil-works/pi-coding-agent": "^0.82.1",
+		"@earendil-works/pi-tui": "^0.82.1",
 		"@types/node": "^26.0.0",
 		"typescript": "^7.0.2",
 		"vitest": "^4.1.8"