@heyhuynhgiabuu/pi-pretty 0.6.25 → 0.6.26

package/src/index.ts

@@ -15,6 +15,7 @@ import type {
 	ExtensionAPI,
 	ExtensionCommandContext,
 	ExtensionContext,
+	MessageEndEvent,
 	MessageUpdateEvent,
 	SessionInfoChangedEvent,
 } from "@earendil-works/pi-coding-agent";
@@ -30,13 +31,17 @@ import { registerReadTool } from "./tools/read.js";
 import type { PiPrettyDeps, SdkTools } from "./types.js";
 import {
 	createThinkingLabelAnimator,
+	createThinkingTimer,
+	installPerRowThinkingLabels,
 	installWorkingIndicator,
+	type PerRowThinkingLabels,
 	resolveThinkingIndicatorSettings,
 	resolveWorkingIndicatorSettings,
-	type ThinkingLabelAnimator,
+	type ThinkingTimer,
 	thinkingBlockActive,
 	WORKING_INTERVAL_MS,
 	type WorkingIndicatorController,
+	workingTokens,
 } from "./working-indicator.js";
 
 // ---------------------------------------------------------------------------
@@ -160,13 +165,24 @@ export default async function piPrettyExtension(pi: ExtensionAPI, deps?: PiPrett
 	const workingSettings = resolveWorkingIndicatorSettings(config.workingIndicator);
 	const thinkingSettings = resolveThinkingIndicatorSettings(config.thinkingIndicator);
 	let workingController: WorkingIndicatorController | undefined;
-	let thinkingAnimator: ThinkingLabelAnimator | undefined;
+	let thinkingTimer: ThinkingTimer | undefined;
 	let thinkingInterval: ReturnType<typeof setInterval> | undefined;
 	let thinkingLastMessage: unknown;
-	let thinkingHiddenCache = true;
 	let thinkingHiddenCheckedAt = 0;
 	let workingSessionName: string | undefined;
 	let workingStreaming = false;
+	let workingStatsUpdatedAt = 0;
+	let perRowLabels: PerRowThinkingLabels | undefined;
+	/** Accumulated thinking time per message timestamp — later runs in the same
+	 * message resume from this total instead of rewinding to zero. */
+	const thinkingElapsedMs = new Map<number, number>();
+	const TOKEN_COUNT_FORMAT = new Intl.NumberFormat("en-US");
+
+	/** Message timestamps are the per-row identity the label patch keys on. */
+	const messageTimestamp = (message: unknown): number | undefined => {
+		const ts = (message as { timestamp?: unknown } | undefined)?.timestamp;
+		return typeof ts === "number" ? ts : undefined;
+	};
 
 	/** pi persists the thinking-visibility toggle as a top-level settings key. */
 	const thinkingHidden = (): boolean => {
@@ -182,6 +198,17 @@ export default async function piPrettyExtension(pi: ExtensionAPI, deps?: PiPrett
 		}
 	};
 
+	/**
+	 * True when the 500ms-cached hidden check is stale and settings now show
+	 * thinking revealed — the single re-check shared by the ticker and the
+	 * completed-phase message_update path (timestamp guard runs before the read).
+	 */
+	const thinkingRevealed = (): boolean => {
+		if (Date.now() - thinkingHiddenCheckedAt <= 500) return false;
+		thinkingHiddenCheckedAt = Date.now();
+		return !thinkingHidden();
+	};
+
 	const safeSessionName = (ctx: ExtensionContext): string | undefined => {
 		try {
 			return ctx.sessionManager?.getSessionName?.() ?? undefined;
@@ -218,6 +245,15 @@ export default async function piPrettyExtension(pi: ExtensionAPI, deps?: PiPrett
 	pi.on("session_start", async (_event: unknown, ctx: ExtensionContext) => {
 		if (ctx.mode === "tui") {
 			ctx.ui.setToolsExpanded(false);
+			// Per-row hidden-thinking labels: intercept the host's label fan-out so
+			// each row keeps its own duration. Falls back to the global-label
+			// behavior whenever the host class is missing or reshaped.
+			perRowLabels = thinkingSettings.enabled
+				? installPerRowThinkingLabels(
+						deps?.assistantMessageComponent ??
+							(hostSdk as { AssistantMessageComponent?: unknown }).AssistantMessageComponent,
+					)
+				: undefined;
 			workingSessionName = safeSessionName(ctx);
 			await installIndicator(ctx);
 		}
@@ -262,6 +298,8 @@ export default async function piPrettyExtension(pi: ExtensionAPI, deps?: PiPrett
 	pi.on("agent_start", async (_event: unknown, ctx: ExtensionContext) => {
 		if (ctx.mode !== "tui") return;
 		workingStreaming = true;
+		workingStatsUpdatedAt = 0;
+		workingController?.setStats(undefined);
 		workingController?.start();
 	});
 	const stopStreaming = async (_event: unknown, ctx: ExtensionContext): Promise<void> => {
@@ -273,69 +311,135 @@ export default async function piPrettyExtension(pi: ExtensionAPI, deps?: PiPrett
 	pi.on("agent_end", stopStreaming);
 	pi.on("agent_settled", stopStreaming);
 
+	const clearThinkingInterval = (): void => {
+		if (!thinkingInterval) return;
+		clearInterval(thinkingInterval);
+		thinkingInterval = undefined;
+	};
+
 	const stopThinkingShimmer = (): void => {
-		if (thinkingInterval) {
-			clearInterval(thinkingInterval);
-			thinkingInterval = undefined;
+		clearThinkingInterval();
+		perRowLabels?.clearActive();
+		try {
+			thinkingTimer?.restore();
+		} catch {
+			// Teardown must never throw — a host that explodes on restore is left
+			// with whatever label it has; the interval is already gone.
 		}
-		thinkingAnimator?.restore();
-		thinkingAnimator = undefined;
+		thinkingTimer = undefined;
 		thinkingLastMessage = undefined;
 	};
 
+	const completeThinkingShimmer = (): void => {
+		clearThinkingInterval();
+		const ts = messageTimestamp(thinkingLastMessage);
+		const frozenMs = thinkingTimer?.complete();
+		if (ts !== undefined && frozenMs !== undefined) {
+			thinkingElapsedMs.set(ts, frozenMs);
+			perRowLabels?.complete(ts, frozenMs);
+		}
+		// The completed row is no longer the animating one — drop the active mark
+		// so subsequent global writes freeze it at its own duration.
+		perRowLabels?.clearActive();
+		thinkingLastMessage = undefined;
+	};
+
+	/** Surface one warning without letting a throwing host escape the crash guard. */
+	const notifyThinkingFailure = (ctx: ExtensionContext, error: unknown): void => {
+		try {
+			ctx.ui?.notify?.(
+				`pi-pretty thinking indicator failed: ${error instanceof Error ? error.message : String(error)}`,
+				"warning",
+			);
+		} catch {
+			// A host that throws on notify too must not crash the timer path.
+		}
+	};
+
 	// Same 30fps cadence as the working row — but the interval lives ONLY while
 	// the streaming message is actively emitting a thinking block (each label
 	// change rebuilds chat children, so the moment text or tool calls stream the
-	// shimmer stops and pi's default label is restored).
+	// label freezes as `Thought for {duration}` until the run ends).
 	const startThinkingShimmer = (ctx: ExtensionContext): void => {
 		if (thinkingInterval || !thinkingSettings.enabled) return;
 		try {
-			thinkingAnimator = createThinkingLabelAnimator(ctx.ui, workingSettings, workingSessionName, thinkingSettings);
-			thinkingHiddenCache = thinkingHidden();
 			thinkingHiddenCheckedAt = Date.now();
-			if (!thinkingHiddenCache) {
-				stopThinkingShimmer();
-				return;
-			}
-			thinkingAnimator.tick();
+			if (!thinkingHidden()) return; // thinking visible: no label writes at all
+			const animator = createThinkingLabelAnimator(ctx.ui, workingSettings, workingSessionName, thinkingSettings);
+			const activeTs = messageTimestamp(thinkingLastMessage);
+			// Same message thinking again (interleaved runs): resume its total.
+			const resumeMs = activeTs !== undefined ? (thinkingElapsedMs.get(activeTs) ?? 0) : 0;
+			thinkingTimer = createThinkingTimer(animator, Date.now, resumeMs);
+			if (activeTs !== undefined) perRowLabels?.setActive(activeTs);
+			thinkingTimer.tick();
 			thinkingInterval = setInterval(() => {
-				if (!workingStreaming) {
-					stopThinkingShimmer();
-					return;
-				}
-				const now = Date.now();
-				if (now - thinkingHiddenCheckedAt > 500) {
-					thinkingHiddenCache = thinkingHidden();
-					thinkingHiddenCheckedAt = now;
-					if (!thinkingHiddenCache) {
+				try {
+					if (!workingStreaming) {
 						stopThinkingShimmer();
 						return;
 					}
-				}
-				if (!thinkingBlockActive(thinkingLastMessage)) {
-					// Thinking phase over: restore pi's default label and stop.
+					if (thinkingRevealed()) {
+						stopThinkingShimmer();
+						return;
+					}
+					if (!thinkingBlockActive(thinkingLastMessage)) {
+						// Thinking phase over: freeze its duration until the run ends.
+						completeThinkingShimmer();
+						return;
+					}
+					thinkingTimer?.tick();
+					// The label has no widget of its own — reuse the working widget's
+					// TUI handle so the frame lands even when deltas pause mid-thinking.
+					workingController?.requestRender();
+				} catch (error: unknown) {
+					// A throwing host call inside our timer would crash pi (uncaught in
+					// setInterval) — tear the label down and surface one warning instead.
 					stopThinkingShimmer();
-					return;
+					notifyThinkingFailure(ctx, error);
 				}
-				thinkingAnimator?.tick();
 			}, WORKING_INTERVAL_MS);
 		} catch (error: unknown) {
 			stopThinkingShimmer();
-			ctx.ui?.notify?.(
-				`pi-pretty thinking indicator failed: ${error instanceof Error ? error.message : String(error)}`,
-				"warning",
-			);
+			notifyThinkingFailure(ctx, error);
 		}
 	};
 
 	// Remember the streaming message; lazily start the shimmer on the first
-	// thinking delta (thinkingHidden() is re-checked at most every 500ms).
+	// thinking delta (thinkingHidden() is re-checked at most every 500ms). After
+	// the phase completes, deltas keep refreshing the visibility cache so a
+	// mid-run `hideThinkingBlock` toggle-off restores the default label.
 	pi.on("message_update", async (event: MessageUpdateEvent, ctx: ExtensionContext) => {
 		if (ctx.mode !== "tui" || !workingStreaming) return;
 		thinkingLastMessage = event.message;
-		if (thinkingInterval) return;
-		if (!thinkingBlockActive(event.message)) return;
-		startThinkingShimmer(ctx);
+		// Live token suffix on the working row (throttled to one update/second).
+		const now = Date.now();
+		if (now - workingStatsUpdatedAt >= 1000) {
+			workingStatsUpdatedAt = now;
+			const tokens = workingTokens(event.message);
+			workingController?.setStats(tokens > 0 ? ` (↓ ${TOKEN_COUNT_FORMAT.format(tokens)} tokens)` : undefined);
+		}
+		if (thinkingInterval) {
+			// Complete immediately on the first text/tool delta instead of waiting
+			// for the next animation frame.
+			if (!thinkingBlockActive(event.message)) completeThinkingShimmer();
+			return;
+		}
+		if (thinkingBlockActive(event.message)) {
+			// First thinking delta — or a later phase after tool calls — starts a
+			// fresh timer (the completed one is replaced wholesale).
+			startThinkingShimmer(ctx);
+			return;
+		}
+		if (thinkingTimer && thinkingRevealed()) stopThinkingShimmer();
+	});
+
+	// The completed wording belongs to the message that produced it. The label
+	// is global, so once that assistant message ends, restore pi's default —
+	// otherwise every later phase stamps its frozen duration onto every hidden
+	// thinking row in the transcript (the image-read "Thought for 5s" stack).
+	pi.on("message_end", async (event: MessageEndEvent, ctx: ExtensionContext) => {
+		if (ctx.mode !== "tui" || event.message.role !== "assistant") return;
+		if (thinkingInterval || thinkingTimer) stopThinkingShimmer();
 	});
 
 	pi.on("session_shutdown", async () => {
@@ -344,6 +448,9 @@ export default async function piPrettyExtension(pi: ExtensionAPI, deps?: PiPrett
 		workingController?.dispose();
 		workingController = undefined;
 		stopThinkingShimmer();
+		perRowLabels?.uninstall();
+		perRowLabels = undefined;
+		thinkingElapsedMs.clear();
 		workingStreaming = false;
 		// Intentionally keep the native FFF finder on session shutdown.
 		// Pi can emit shutdown/start during resume or session switching while the

package/src/tools/bash.ts

@@ -66,7 +66,7 @@ export function registerBashTool(
 						: rawCmd;
 			const commandLabel = theme.fg(ctx.isError ? "error" : "toolTitle", theme.bold(`$ ${cmd}`));
 			text.setText(
-				fillToolBackground(`\n${TOOL_RESULT_INDENT}${commandLabel}${t}\n`, undefined, ctx.expanded ? undefined : tw),
+				fillToolBackground(`\n${TOOL_RESULT_INDENT}${commandLabel}${t}`, undefined, ctx.expanded ? undefined : tw),
 			);
 			return text;
 		},

package/src/tools/find.ts

@@ -150,7 +150,7 @@ export function registerFindTool(
 			const pathPart = theme.fg("toolOutput", pathArg);
 			const limitPart = limit !== undefined && limit !== null ? theme.fg("dim", ` limit ${limit}`) : "";
 			const out = `${findLabel} ${patternPart}${inPart}${pathPart}${limitPart}`;
-			text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${out}\n`, ctx.isError ? BG_ERROR : undefined));
+			text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${out}`, ctx.isError ? BG_ERROR : undefined));
 			return text;
 		},
 
@@ -168,7 +168,7 @@ export function registerFindTool(
 					const noticeStr = d.notices?.length
 						? `\n${TOOL_RESULT_INDENT}${theme.fg("warning", `[${d.notices.join(". ")}]`)}`
 						: "";
-					text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${theme.fg("dim", "0 files")}${noticeStr}\n`));
+					text.setText(fillToolBackground(`${TOOL_RESULT_INDENT}${theme.fg("dim", "0 files")}${noticeStr}\n`));
 					return text;
 				}
 				if (!ctx.expanded) {
@@ -190,14 +190,14 @@ export function registerFindTool(
 				const duration = renderToolDuration(r);
 				text.setText(
 					fillToolBackground(
-						`\n${TOOL_RESULT_INDENT}${theme.fg("dim", `${d.matchCount} files`)}${duration ? `${FG_DIM}· ${duration}${RST}` : ""}\n${rendered}${noticeStr}\n`,
+						`${TOOL_RESULT_INDENT}${theme.fg("dim", `${d.matchCount} files`)}${duration ? `${FG_DIM}· ${duration}${RST}` : ""}\n${rendered}${noticeStr}\n`,
 					),
 				);
 				return text;
 			}
 			const fc = r.content?.[0] as TextContent | undefined;
 			text.setText(
-				fillToolBackground(`\n${TOOL_RESULT_INDENT}${theme.fg("dim", fc?.text?.slice(0, 120) ?? "0 files")}\n`),
+				fillToolBackground(`${TOOL_RESULT_INDENT}${theme.fg("dim", fc?.text?.slice(0, 120) ?? "0 files")}\n`),
 			);
 			return text;
 		},

package/src/tools/grep.ts

@@ -111,7 +111,7 @@ export function registerGrepTool(
 			if (limit !== undefined && limit !== null) out += theme.fg("dim", ` limit ${limit}`);
 			if (literal) out += theme.fg("dim", ` (literal)`);
 			if (caseInsensitive) out += theme.fg("dim", ` (case-insensitive)`);
-			text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${out}\n`, ctx.isError ? BG_ERROR : undefined));
+			text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${out}`, ctx.isError ? BG_ERROR : undefined));
 			return text;
 		},
 

package/src/tools/ls.ts

@@ -51,7 +51,7 @@ export function registerLsTool(
 			let out = theme.fg("toolTitle", theme.bold("ls"));
 			if (path) out += ` ${theme.fg("accent", path)}`;
 			if (limit !== undefined && limit !== null) out += theme.fg("toolOutput", ` (limit ${limit})`);
-			text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${out}\n`, ctx.isError ? BG_ERROR : undefined));
+			text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${out}`, ctx.isError ? BG_ERROR : undefined));
 			return text;
 		},
 
@@ -79,7 +79,7 @@ export function registerLsTool(
 					.join("\n");
 				text.setText(
 					fillToolBackground(
-						`\n${TOOL_RESULT_INDENT}${FG_DIM}${d.entryCount} entries${RST}${renderToolMetrics(result)}\n${rendered}\n`,
+						`${TOOL_RESULT_INDENT}${FG_DIM}${d.entryCount} entries${RST}${renderToolMetrics(result)}\n${rendered}\n`,
 					),
 				);
 				return text;

package/src/tools/read.ts

@@ -100,7 +100,7 @@ export function registerReadTool(
 
 			const path = String(args.path ?? "");
 			const label = theme.fg("error", theme.bold("→ read"));
-			text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${label} ${theme.fg("toolTitle", path)}\n`, BG_ERROR));
+			text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${label} ${theme.fg("toolTitle", path)}`, BG_ERROR));
 			return text;
 		},
 
@@ -137,7 +137,7 @@ export function registerReadTool(
 				if (!ctx.expanded) {
 					if (skillName) {
 						const header = renderSkillHeader(skillName, false, theme);
-						text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${header}\n`, BG_BASE));
+						text.setText(fillToolBackground(`\n${TOOL_RESULT_INDENT}${header}`, BG_BASE));
 						return text;
 					}
 					text.setText(

package/src/types.ts

@@ -187,4 +187,8 @@ export interface PiPrettyDeps {
 	sdk?: SdkTools;
 	TextComponent?: new (text?: string, x?: number, y?: number) => ComponentLike;
 	fffModule?: typeof import("@ff-labs/fff-node");
+	/** Host AssistantMessageComponent class for the per-row label patch; defaults
+	 * to the package export. Injectable so tests can verify against the exact
+	 * class instance they construct rows from. */
+	assistantMessageComponent?: unknown;
 }