@heyhuynhgiabuu/pi-pretty 0.6.25 → 0.6.27

package/src/working-indicator.ts

@@ -1,19 +1,17 @@
 /**
- * omp-style shimmer working indicator (widget takeover).
+ * omp-style shimmer working indicator.
  *
  * Ports the shimmer sweep from oh-my-pi's `modes/theme/shimmer.ts`: the sweep
  * is discretized into one pre-rendered frame per band position, cycled at
  * 1000/30 ms so the band travels 30 cells/second — omp's exact speed.
  *
- * Why a widget: the host `Loader` extends `Text` with a hardcoded paddingX=1,
- * and `setWorkingIndicator({ frames })` cannot change host component layout.
- * To render the row flush-left we hide the host loader
- * (`setWorkingVisible(false)`) and install our own zero-padding component via
- * `setWidget(..., "aboveEditor")`, animating it with a 33ms interval.
+ * On hosts that expose Pi's embedded working-status API, frames are delegated
+ * to the host so the row is rendered in the editor's top border. Older hosts
+ * without that API use the zero-padding widget fallback above the editor.
  *
- * Frame layout: `<dim spinner> <shimmer text> <dim interrupt hint>`, one full
- * sweep per phrase — `texts` rotates through phrases chapter by chapter with
- * a continuous spinner phase.
+ * Frame layout: `<dim spinner> <shimmer text>`, one full sweep per phrase —
+ * `texts` rotates through phrases chapter by chapter with a continuous spinner
+ * phase. The extension adds no interrupt hint or token-count suffix.
  */
 
 import { truncateToWidth } from "@earendil-works/pi-tui";
@@ -61,6 +59,7 @@ export interface WorkingIndicatorSettings {
 	mode: WorkingIndicatorMode;
 	palette: WorkingIndicatorPalette;
 	bold: boolean;
+	/** @deprecated Retained for config compatibility; no interrupt hint is rendered. */
 	hint: boolean;
 	/** Tint mid/high tiers (and dim the spinner) with a per-session accent color. */
 	sessionAccent: boolean;
@@ -261,7 +260,7 @@ export interface WorkingFrameOptions {
 	spinnerColor?: string;
 	/** Wrap every frame in italic (pi's hidden-thinking label is italic). */
 	italic?: boolean;
-	/** Pre-colored suffix (e.g. the dim interrupt hint) appended to every frame. */
+	/** @deprecated Retained for call-site compatibility; suffixes are not rendered. */
 	hint?: string;
 	intervalMs?: number;
 }
@@ -277,7 +276,6 @@ export function buildWorkingFrames(
 ): { frames: string[]; intervalMs: number } {
 	const intervalMs = options.intervalMs ?? WORKING_INTERVAL_MS;
 	const bold = options.bold ?? true;
-	const suffix = options.hint ?? "";
 	const withSpinner = options.spinner ?? true;
 	const spinnerColor = options.spinnerColor ?? options.ansi.low;
 	const highOpen = bold ? `${BOLD_OPEN}${options.ansi.high}` : options.ansi.high;
@@ -331,7 +329,7 @@ export function buildWorkingFrames(
 					text,
 					ranges,
 					Array.from({ length: total }, () => "mid" as Tier),
-				)}${suffix}`,
+				)}`,
 			);
 			return { frames, intervalMs };
 		}
@@ -344,7 +342,7 @@ export function buildWorkingFrames(
 						text,
 						ranges,
 						Array.from({ length: total }, () => "high" as Tier),
-					)}${suffix}`,
+					)}`,
 				);
 				continue;
 			}
@@ -353,19 +351,35 @@ export function buildWorkingFrames(
 				const goingRight = sweep < range;
 				const head = goingRight ? sweep : cycleCells - sweep;
 				const tiers = Array.from({ length: total }, (_, i) => tierFor(kittIntensity(i, head, goingRight)));
-				frames.push(`${spinnerPrefix(spinnerIndex++)}${paintFrame(text, ranges, tiers)}${suffix}`);
+				frames.push(`${spinnerPrefix(spinnerIndex++)}${paintFrame(text, ranges, tiers)}`);
 			}
 		} else {
 			const period = total + CLASSIC_PADDING * 2;
 			for (let pos = 0; pos < period; pos++) {
 				const tiers = Array.from({ length: total }, (_, i) => tierFor(classicIntensity(i, pos)));
-				frames.push(`${spinnerPrefix(spinnerIndex++)}${paintFrame(text, ranges, tiers)}${suffix}`);
+				frames.push(`${spinnerPrefix(spinnerIndex++)}${paintFrame(text, ranges, tiers)}`);
 			}
 		}
 	}
 	return { frames, intervalMs };
 }
 
+/** @deprecated The working row no longer displays token stats. */
+export function workingTokens(message: unknown): number {
+	const usage = (message as { usage?: { output?: unknown } } | undefined)?.usage;
+	const output = Number(usage?.output);
+	if (Number.isFinite(output) && output > 0) return Math.round(output);
+	const content = (message as { content?: Array<{ type?: string; text?: unknown; thinking?: unknown }> } | undefined)
+		?.content;
+	if (!Array.isArray(content)) return 0;
+	let chars = 0;
+	for (const block of content) {
+		if (block?.type === "text" && typeof block.text === "string") chars += block.text.length;
+		else if (block?.type === "thinking" && typeof block.thinking === "string") chars += block.thinking.length;
+	}
+	return Math.round(chars / 4);
+}
+
 // ─── Widget ──────────────────────────────────────────────────────────────────
 
 /** Slice of the pi-tui TUI instance the widget needs. */
@@ -375,8 +389,8 @@ export interface WidgetTuiLike {
 
 /**
  * Zero-padding working-row component. The host renders whatever `render()`
- * returns with no margins, so the row sits flush-left — the thing the frames
- * API could not do.
+ * returns with no margins, so the legacy widget row sits flush-left without
+ * relying on the host loader's padding.
  */
 export class WorkingWidget {
 	#frames: string[] = [];
@@ -397,6 +411,16 @@ export class WorkingWidget {
 		this.#tui = tui;
 	}
 
+	/** Ask the host for a render — lets non-widget consumers (the thinking label
+	 * ticker) reuse this widget's TUI handle instead of relying on natural
+	 * streaming renders. */
+	requestRender(): void {
+		this.#tui?.requestRender();
+	}
+
+	/** @deprecated Token stats are owned by Pi and are no longer rendered here. */
+	setStats(_text: string | undefined): void {}
+
 	start(): void {
 		if (this.#disposed || this.#started) return;
 		this.#started = true;
@@ -462,16 +486,161 @@ export interface ThinkingUiLike {
 	setHiddenThinkingLabel(label?: string): void;
 }
 
+// ─── Per-row hidden-thinking labels ──────────────────────────────────────────
+
+/** Controller for the host-component label interception. */
+export interface PerRowThinkingLabels {
+	/** Mark the message currently streaming a thinking block (its row animates). */
+	setActive(timestamp: number): void;
+	/** Freeze a message's completed thinking duration in ms (its row shows
+	 * `Thought for X` derived from this). */
+	complete(timestamp: number, durationMs: number): void;
+	/** Drop the active mark (run ended); completed rows keep their durations. */
+	clearActive(): void;
+	/** Restore the original prototype setter. */
+	uninstall(): void;
+}
+
+const PER_ROW_PATCH = Symbol.for("pi-pretty.per-row-thinking-labels");
+
+interface RowInternals {
+	lastMessage?: { timestamp?: unknown };
+}
+
+/**
+ * Intercept the host's per-row label fan-out method
+ * (`AssistantMessageComponent.prototype.setHiddenThinkingLabel`): each row
+ * substitutes the incoming global label with its own — the streaming row
+ * animates, completed rows stay frozen at their own `Thought for X`, unknown
+ * rows fall back to pi's default. Any per-call failure passes the incoming
+ * label through unchanged, so a host internals change degrades to pi-pretty's
+ * global-label behavior instead of breaking. Durations live per session.
+ */
+export function installPerRowThinkingLabels(componentClass: unknown): PerRowThinkingLabels | undefined {
+	const proto = (componentClass as { prototype?: { setHiddenThinkingLabel?: unknown } } | undefined)?.prototype;
+	const original = proto?.setHiddenThinkingLabel;
+	if (typeof componentClass !== "function" || !proto || typeof original !== "function") return undefined;
+	const existing = (original as { [PER_ROW_PATCH]?: PerRowThinkingLabels })[PER_ROW_PATCH];
+	// Reuse the installed patch but start from clean state: a session_start
+	// without an intervening shutdown (defensive) must not resurrect stale rows.
+	if (existing) {
+		existing.clearActive();
+		return existing;
+	}
+
+	const state: { activeTs?: number; completed: Map<number, number> } = { completed: new Map() };
+	const patched = function (this: unknown, label: string | undefined): void {
+		const own = ownLabel(this, label);
+		(original as (this: unknown, label: string | undefined) => void).call(this, own);
+	};
+	const ownLabel = (component: unknown, incoming: string | undefined): string | undefined => {
+		try {
+			const ts = (component as RowInternals | undefined)?.lastMessage?.timestamp;
+			if (typeof ts !== "number") return incoming;
+			if (ts === state.activeTs) return incoming;
+			const ms = state.completed.get(ts);
+			return ms === undefined ? THINKING_LABEL : `${THOUGHT_LABEL} ${formatThinkingDuration(ms)}`;
+		} catch {
+			return incoming;
+		}
+	};
+	const controller: PerRowThinkingLabels = {
+		setActive(timestamp: number): void {
+			state.activeTs = timestamp;
+		},
+		complete(timestamp: number, durationMs: number): void {
+			state.completed.set(timestamp, durationMs);
+		},
+		clearActive(): void {
+			state.activeTs = undefined;
+		},
+		uninstall(): void {
+			if (proto.setHiddenThinkingLabel === patched) {
+				proto.setHiddenThinkingLabel = original as { setHiddenThinkingLabel?: unknown }["setHiddenThinkingLabel"];
+			}
+			state.completed.clear();
+			state.activeTs = undefined;
+		},
+	};
+	(patched as { [PER_ROW_PATCH]?: PerRowThinkingLabels })[PER_ROW_PATCH] = controller;
+	proto.setHiddenThinkingLabel = patched as NonNullable<{ setHiddenThinkingLabel?: unknown }["setHiddenThinkingLabel"]>;
+	return controller;
+}
+
 export interface ThinkingLabelAnimator {
-	/** Apply the next shimmer frame to the hidden-thinking label. */
-	tick(): void;
+	/** Apply the next shimmer frame, rebuilding frames when the label changes. */
+	tick(label?: string): void;
+	/** Apply a non-animated label (the host supplies its normal thinking style). */
+	show(label: string): void;
 	/** Restore pi's default static label. */
 	restore(): void;
-	/** Pre-rendered label frames (exposed for diagnostics). */
+	/** Current pre-rendered label frames (exposed for diagnostics). */
 	readonly frames: readonly string[];
 }
 
-const THINKING_LABEL = "Thinking...";
+export interface ThinkingTimer {
+	/** Render the active label with the current elapsed duration. */
+	tick(): void;
+	/** Freeze the elapsed duration; returns the total in milliseconds, or
+	 * undefined when the timer was already terminal. */
+	complete(): number | undefined;
+	/** Restore pi's default label and make the timer terminal. */
+	restore(): void;
+}
+
+const THINKING_LABEL = " Thinking...";
+const THOUGHT_LABEL = " Thought for";
+const padThinkingLabel = (label: string): string => (label.startsWith(" ") ? label : ` ${label}`);
+
+/** Format sub-second elapsed time as milliseconds, then use compact whole seconds. */
+export function formatThinkingDuration(elapsedMs: number): string {
+	const wholeMs = Math.max(0, Math.floor(elapsedMs));
+	if (wholeMs < 1000) return `${wholeMs}ms`;
+	const totalSeconds = Math.floor(wholeMs / 1000);
+	const seconds = totalSeconds % 60;
+	const totalMinutes = Math.floor(totalSeconds / 60);
+	const minutes = totalMinutes % 60;
+	const hours = Math.floor(totalMinutes / 60);
+	if (hours > 0) return `${hours}h ${String(minutes).padStart(2, "0")}m ${String(seconds).padStart(2, "0")}s`;
+	if (totalMinutes > 0) return `${totalMinutes}m ${String(seconds).padStart(2, "0")}s`;
+	return `${seconds}s`;
+}
+
+/**
+ * Couple elapsed-time wording to the label animator through an injected clock.
+ * The timer starts on the first thinking delta, freezes once, and is terminal
+ * after completion or restore. `startElapsedMs` resumes a later run in the
+ * same message from its accumulated thinking time (per-message semantics:
+ * the host renders one label line per thinking run, all sharing the message
+ * label, so a fresh-zero restart would visibly rewind earlier lines).
+ */
+export function createThinkingTimer(
+	animator: ThinkingLabelAnimator,
+	now: () => number = Date.now,
+	startElapsedMs = 0,
+): ThinkingTimer {
+	const startedAt = now() - startElapsedMs;
+	let terminal = false;
+	const elapsed = (): number => now() - startedAt;
+	const duration = (): string => formatThinkingDuration(elapsed());
+	return {
+		tick(): void {
+			if (terminal) return;
+			animator.tick(`${THINKING_LABEL} ${duration()}`);
+		},
+		complete(): number | undefined {
+			if (terminal) return undefined;
+			terminal = true;
+			const ms = elapsed();
+			animator.show(`${THOUGHT_LABEL} ${formatThinkingDuration(ms)}`);
+			return ms;
+		},
+		restore(): void {
+			terminal = true;
+			animator.restore();
+		},
+	};
+}
 
 /**
  * True while a streaming assistant message is currently emitting a thinking
@@ -489,9 +658,9 @@ export function thinkingBlockActive(message: unknown): boolean {
  * Animate pi's hidden-thinking label ("Thinking...") with the same shimmer as
  * the working row. Pi renders the label as static italic `thinkingText` text;
  * the only extension lever is `setHiddenThinkingLabel(label)`, which rebuilds
- * chat children — the streaming component already rebuilds per delta, so the
- * animator is driven at a throttled cadence (100ms) from index.ts while the
- * agent streams and thinking blocks are hidden.
+ * chat children, so index.ts drives it at 30fps only while the streaming
+ * message's last block is thinking. The elapsed label is rebuilt when its
+ * displayed duration changes; the shimmer frame still advances every tick.
  *
  * Tiers: low = theme `thinkingText` (pi's own label look), mid/high = session
  * accent when enabled; every frame is italic like pi's label; no spinner.
@@ -504,6 +673,7 @@ export function createThinkingLabelAnimator(
 ): ThinkingLabelAnimator {
 	const noop: ThinkingLabelAnimator = {
 		tick() {},
+		show() {},
 		restore() {},
 		frames: [],
 	};
@@ -519,25 +689,42 @@ export function createThinkingLabelAnimator(
 		const accent = hexToAnsiFg(sessionAccentHex(sessionName));
 		if (accent) ansi = { ...ansi, mid: accent, high: accent };
 	}
-	const { frames } = buildWorkingFrames([THINKING_LABEL], {
-		mode: workSettings.mode,
-		ansi,
-		bold: workSettings.bold,
-		spinner: false,
-		italic: true,
-	});
-	if (frames.length === 0) return noop;
+	const buildLabelFrames = (label: string): string[] => {
+		const padded = padThinkingLabel(label);
+		const text = padded.trimStart();
+		const leftPadding = padded.slice(0, padded.length - text.length);
+		return buildWorkingFrames([text], {
+			mode: workSettings.mode,
+			ansi,
+			bold: workSettings.bold,
+			spinner: false,
+			italic: true,
+		}).frames.map((frame) => `${leftPadding}${frame}`);
+	};
+
+	let currentLabel = THINKING_LABEL;
+	let frames = buildLabelFrames(currentLabel);
 
 	let index = 0;
-	let applied = false;
+	let appliedLabel: string | undefined;
 	return {
-		frames,
-		tick(): void {
-			// A single frame (static mode) never changes — skip redundant rebuilds.
-			if (applied && frames.length === 1) return;
+		get frames(): readonly string[] {
+			return frames;
+		},
+		tick(label = THINKING_LABEL): void {
+			const paddedLabel = padThinkingLabel(label);
+			if (paddedLabel !== currentLabel) {
+				frames = buildLabelFrames(paddedLabel);
+				currentLabel = paddedLabel;
+			}
+			// A static frame changes only when its elapsed-time label changes.
+			if (appliedLabel === currentLabel && frames.length === 1) return;
 			ui.setHiddenThinkingLabel(frames[index % frames.length]);
 			index++;
-			applied = true;
+			appliedLabel = currentLabel;
+		},
+		show(label: string): void {
+			ui.setHiddenThinkingLabel(padThinkingLabel(label));
 		},
 		restore(): void {
 			// undefined → pi falls back to its default static label.
@@ -548,15 +735,24 @@ export function createThinkingLabelAnimator(
 
 // ─── Installation ────────────────────────────────────────────────────────────
 
-/** Structural slice of pi's Theme needed for the hint, palette, and thinking label. */
+/** Structural slice of pi's Theme needed for the palette and thinking label. */
 export interface WorkingThemeLike {
-	/** Narrowed to the only color the hint needs; pi's ThemeColor satisfies it. */
+	/** Retained for compatibility with older theme adapters. */
 	fg?: (name: "dim", text: string) => string;
 	getFgAnsi?: (name: "dim" | "muted" | "accent" | "thinkingText") => string;
 }
 
+export interface WorkingStatusOptionsLike {
+	frames?: string[];
+	intervalMs?: number;
+}
+
 export interface WorkingUiLike {
 	theme?: WorkingThemeLike;
+	/** Hosts with Pi's embedded working-status API render this in the editor border. */
+	setWorkingIndicator?: (options?: WorkingStatusOptionsLike) => void;
+	/** Clear the host message appended after a custom native frame. */
+	setWorkingMessage?: (message?: string) => void;
 	setWorkingVisible(visible: boolean): void;
 	setWidget(
 		key: string,
@@ -565,6 +761,7 @@ export interface WorkingUiLike {
 	): void;
 }
 
+/** @deprecated Retained for compatibility with the former hint-resolution seam. */
 export interface KeybindingsLike {
 	getKeys(name: string): string[];
 }
@@ -574,60 +771,26 @@ export interface WorkingIndicatorController {
 	start(): void;
 	/** Streaming ended — hide the row and pause the animation. */
 	stop(): void;
-	/** Remove the widget and stop the animation. */
+	/** Remove the indicator and stop the animation. */
 	dispose(): void;
+	/** @deprecated Token stats are owned by Pi and are no longer rendered here. */
+	setStats(text: string | undefined): void;
+	/** Request a host render for consumers that drive another animated label. */
+	requestRender(): void;
 	/** Pre-rendered frames (exposed for diagnostics). */
 	readonly frames: readonly string[];
 }
 
 const WIDGET_KEY = "pi-pretty-working";
 
-async function loadHostKeybindings(): Promise<KeybindingsLike | undefined> {
-	try {
-		const tui = (await import("@earendil-works/pi-tui")) as typeof import("@earendil-works/pi-tui");
-		return tui.getKeybindings?.();
-	} catch {
-		return undefined;
-	}
-}
-
-function formatKeyText(keys: string[]): string {
-	const darwin = process.platform === "darwin";
-	return keys
-		.join("/")
-		.split("/")
-		.map((key) =>
-			key
-				.split("+")
-				.map((part) => (darwin && part.toLowerCase() === "alt" ? "option" : part))
-				.join("+"),
-		)
-		.join("/");
-}
-
-async function resolveHint(
-	ui: WorkingUiLike,
-	deps?: { getKeybindings?: () => KeybindingsLike | undefined | Promise<KeybindingsLike | undefined> },
-): Promise<string | undefined> {
-	try {
-		const loader = deps?.getKeybindings ?? loadHostKeybindings;
-		const keybindings = await loader();
-		const keys = keybindings?.getKeys?.("app.interrupt");
-		if (!keys?.length) return undefined;
-		const text = ` (${formatKeyText(keys)} to interrupt)`;
-		return ui.theme?.fg ? ui.theme.fg("dim", text) : text;
-	} catch {
-		return undefined;
-	}
-}
-
 /**
- * Take over pi's working row with our own flush-left shimmer widget.
+ * Install the custom working shimmer and return its lifecycle controller.
  *
- * Installs a zero-padding component above the editor and hides the host
- * loader; returns a controller the host lifecycle drives (`start` on
- * agent_start, `stop` on agent_end). No-op (host defaults untouched) when
- * disabled or the phrase list is empty.
+ * Current Pi hosts receive the pre-rendered frames through the embedded
+ * working-status API, which places them in the editor's top border. Older
+ * hosts use a zero-padding widget above the editor. The host lifecycle drives
+ * the controller (`start` on agent_start, `stop` on agent_end); disabled or
+ * empty phrase settings leave host defaults untouched.
  */
 export async function installWorkingIndicator(
 	ui: WorkingUiLike,
@@ -639,11 +802,14 @@ export async function installWorkingIndicator(
 		start() {},
 		stop() {},
 		dispose() {},
+		setStats() {},
+		requestRender() {},
 		frames: [],
 	};
 	if (!settings.enabled) return noopController;
 	const texts = settings.texts.map((t) => t.trim()).filter((t) => t.length > 0);
 	if (texts.length === 0) return noopController;
+	void deps; // Retained for compatibility; interrupt hints are no longer resolved here.
 
 	const ansiResolved = resolvePaletteAnsi(settings.palette, ui.theme);
 	let ansi = ansiResolved;
@@ -656,20 +822,73 @@ export async function installWorkingIndicator(
 			spinnerColor = hexToAnsiFg(dimAccentHex(accentHex));
 		}
 	}
-	const hint = settings.hint ? await resolveHint(ui, deps) : undefined;
 	const { frames, intervalMs } = buildWorkingFrames(texts, {
 		mode: settings.mode,
 		ansi,
 		bold: settings.bold,
 		spinnerColor,
-		hint,
 	});
 	if (frames.length === 0) return noopController;
 
+	const setWorkingIndicator = ui.setWorkingIndicator;
+	const setWorkingMessage = ui.setWorkingMessage;
+	if (typeof setWorkingIndicator === "function" && typeof setWorkingMessage === "function") {
+		let started = false;
+		let disposed = false;
+		const updateNativeMessage = (): void => {
+			// Keep Pi's built-in message hidden because the custom frame already
+			// supplies the working label; no hint or token suffix is added here.
+			setWorkingMessage.call(ui, "");
+		};
+
+		try {
+			setWorkingIndicator.call(ui, { frames, intervalMs });
+			// Keep the native status row hidden until the extension lifecycle starts it.
+			// An empty message is intentional: all visible content lives in our frames.
+			updateNativeMessage();
+			ui.setWorkingVisible(false);
+		} catch (error: unknown) {
+			try {
+				setWorkingIndicator.call(ui);
+				setWorkingMessage.call(ui);
+			} catch {
+				// Best-effort rollback; the caller still restores host visibility.
+			}
+			throw error;
+		}
+		return {
+			start: () => {
+				if (disposed || started) return;
+				started = true;
+				updateNativeMessage();
+				ui.setWorkingVisible(true);
+			},
+			stop: () => {
+				if (!started) return;
+				started = false;
+				ui.setWorkingVisible(false);
+			},
+			dispose: () => {
+				if (disposed) return;
+				if (started) {
+					started = false;
+					ui.setWorkingVisible(false);
+				}
+				disposed = true;
+				setWorkingIndicator.call(ui);
+				setWorkingMessage.call(ui);
+			},
+			setStats: (_text: string | undefined) => {},
+			requestRender: () => updateNativeMessage(),
+			frames,
+		};
+	}
+
 	const widget = new WorkingWidget();
 	widget.setFrames(frames, intervalMs);
 	// Widget first, visibility second: if setWidget throws (older host without
-	// the API) the host loader was never hidden and pi's default remains.
+	// the native working-status API) the host loader was never hidden and pi's
+	// default remains.
 	ui.setWidget(
 		WIDGET_KEY,
 		(tui) => {
@@ -686,6 +905,8 @@ export async function installWorkingIndicator(
 			widget.dispose();
 			ui.setWidget(WIDGET_KEY, undefined);
 		},
+		setStats: (_text: string | undefined) => {},
+		requestRender: () => widget.requestRender(),
 		frames,
 	};
 }