@heyhuynhgiabuu/pi-pretty 0.6.25 → 0.6.26

package/src/working-indicator.ts

@@ -366,6 +366,22 @@ export function buildWorkingFrames(
 	return { frames, intervalMs };
 }
 
+/** Estimate streamed output tokens: provider `usage.output` when exposed, else visible chars ÷ 4. */
+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. */
@@ -386,6 +402,7 @@ export class WorkingWidget {
 	#tui: WidgetTuiLike | undefined;
 	#started = false;
 	#disposed = false;
+	#stats: string | undefined;
 
 	setFrames(frames: string[], intervalMs: number): void {
 		this.#frames = frames;
@@ -397,6 +414,19 @@ 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();
+	}
+
+	/** Dynamic right-side status segment (e.g. live token count), styled dim at
+	 * render time so theme changes don't bake stale ANSI into it. */
+	setStats(text: string | undefined): void {
+		this.#stats = text || undefined;
+	}
+
 	start(): void {
 		if (this.#disposed || this.#started) return;
 		this.#started = true;
@@ -421,7 +451,8 @@ export class WorkingWidget {
 	render(width: number): string[] {
 		if (!this.#started || this.#frames.length === 0) return [];
 		const frame = this.#frames[this.#index % this.#frames.length] ?? "";
-		return [truncateToWidth(frame, Math.max(1, width))];
+		const suffix = this.#stats ? `${FG_DIM}${this.#stats}${RESET_FG}` : "";
+		return [truncateToWidth(frame + suffix, Math.max(1, width))];
 	}
 
 	invalidate(): void {
@@ -462,17 +493,158 @@ 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 Xs` 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 Xs`, 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 for ${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[];
 }
 
+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...";
 
+/** Format elapsed milliseconds as whole seconds with compact padded units. */
+export function formatThinkingDuration(elapsedMs: number): string {
+	const totalSeconds = Math.floor(elapsedMs / 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 for ${formatThinkingDuration(ms)}`);
+			return ms;
+		},
+		restore(): void {
+			terminal = true;
+			animator.restore();
+		},
+	};
+}
+
 /**
  * True while a streaming assistant message is currently emitting a thinking
  * block (its last content block). Used to gate label ticks to the thinking
@@ -489,9 +661,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 only when its
+ * whole-second text 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 +676,7 @@ export function createThinkingLabelAnimator(
 ): ThinkingLabelAnimator {
 	const noop: ThinkingLabelAnimator = {
 		tick() {},
+		show() {},
 		restore() {},
 		frames: [],
 	};
@@ -519,25 +692,37 @@ 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[] =>
+		buildWorkingFrames([label], {
+			mode: workSettings.mode,
+			ansi,
+			bold: workSettings.bold,
+			spinner: false,
+			italic: true,
+		}).frames;
+
+	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 {
+			if (label !== currentLabel) {
+				frames = buildLabelFrames(label);
+				currentLabel = label;
+			}
+			// 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(label);
 		},
 		restore(): void {
 			// undefined → pi falls back to its default static label.
@@ -576,6 +761,10 @@ export interface WorkingIndicatorController {
 	stop(): void;
 	/** Remove the widget and stop the animation. */
 	dispose(): void;
+	/** Update the dim right-side status segment (live token count). */
+	setStats(text: string | undefined): void;
+	/** Request a host render via the widget's TUI handle. */
+	requestRender(): void;
 	/** Pre-rendered frames (exposed for diagnostics). */
 	readonly frames: readonly string[];
 }
@@ -639,6 +828,8 @@ export async function installWorkingIndicator(
 		start() {},
 		stop() {},
 		dispose() {},
+		setStats() {},
+		requestRender() {},
 		frames: [],
 	};
 	if (!settings.enabled) return noopController;
@@ -686,6 +877,8 @@ export async function installWorkingIndicator(
 			widget.dispose();
 			ui.setWidget(WIDGET_KEY, undefined);
 		},
+		setStats: (text: string | undefined) => widget.setStats(text),
+		requestRender: () => widget.requestRender(),
 		frames,
 	};
 }