Refactor extension structure

README.md

@@ -13,7 +13,8 @@ Warhammer 40k "Dark Mechanicum" flavoring on top of pi's interactive TUI.
 | Extension | What it does |
 |---|---|
 | [`ai-server/`](ai-server/) | Remote llama.cpp provider over mTLS. Dynamic model discovery. Admin slash commands (load / unload / ctx / preset / restart / refresh). Custom SSE stream implementation with tool calls, reasoning, cache token reporting. See [ai-server/README.md](ai-server/README.md) for the full setup. |
-| [`dark-mechanicus/`](dark-mechanicus/) | TUI customization bundle for the dark-mechanicus theme — loaded as one extension via `index.ts`. Includes: `indicator.ts` (working indicator: `⚙ <quote> · <elapsed>`, pulsing cog, 45-quote pool), `banner.ts` (cog-and-skull header art), `footer.ts` (custom context+compaction+tok/s footer), `status-line.ts` (third footer line with rotating flavor text), `session-names.ts` (auto `<adj>-<noun> · <NNN>` session names + tab title), `thinking-label.ts` (`Cogitating...` for folded thinking blocks), `markdown-body-color.ts` (forces lavender body text). |
+| [`token-stats.ts`](token-stats.ts) + [`shared/token-stats.ts`](shared/token-stats.ts) | Footer owner for context-window + token-rate display. Tracks prefill and generation speed, including reasoning/thinking tokens, and reads `tokenStats.enabled` from `~/.pi/agent/settings.json`. |
+| [`dark-mechanicus/`](dark-mechanicus/) | TUI customization bundle for the dark-mechanicus theme — loaded as one extension via `index.ts`. Includes: `indicator.ts` (working indicator: `⚙ <quote> · <elapsed>`, pulsing cog, 45-quote pool), `banner.ts` (cog-and-skull header art), `status-line.ts` (third footer line with rotating flavor text), `session-names.ts` (auto `<adj>-<noun> · <NNN>` session names + tab title), `thinking-label.ts` (`Cogitating...` for folded thinking blocks), `markdown-body-color.ts` (forces lavender body text). Display toggles now come from `darkMechanicus` settings instead of slash commands. |
 | [`local-llama.ts`](local-llama.ts) | Tiny provider registration for a local llama-server (e.g. `127.0.0.1:8088`). |
 
 ## Theme
@@ -36,7 +37,7 @@ Activate with `/settings` inside pi (or set `"theme": "dark-mechanicus"` in
 
 ## Tests
 
-Twenty-eight tests total, no external dependencies. Runs with Node 22+'s
+Thirty-two tests total, no external dependencies. Runs with Node 22+'s
 built-in test runner + `--experimental-strip-types`:
 
 ```bash
@@ -48,6 +49,7 @@ node --experimental-strip-types --test tests/*.test.ts
 | `tests/messages.test.ts` | 13 unit tests over `ai-server/messages.ts` — pi Context → OpenAI payload conversion (system prompts, user/assistant/tool-result roles, tool calls, image-only messages). |
 | `tests/router-utils.test.ts` | 9 unit tests over `ai-server/router-utils.ts` — `extractCtxSize`, `isShardArtefact` (the filter that hides GGUF multi-shard phantoms from the model picker). |
 | `tests/integration.test.ts` | 6 live-endpoint tests: `/health`, `/models`, model-entry shape, mTLS enforcement, publicly-trusted cert (Let's Encrypt contract), chat completion usage shape including `prompt_tokens_details.cached_tokens`. Auto-skip if the server is unreachable. |
+| `tests/token-stats.test.ts` | 4 unit tests over `shared/token-stats.ts` — timing metadata parsing and rate calculation, including thinking-token-aware generation speed. |
 
 Stream-parsing edge cases (SSE framing, tool-call splits across chunks,
 reasoning deltas, abort mid-stream) remain deferred — they need a mock
@@ -80,11 +82,11 @@ pi-extensions/
 │   ├── admin.ts                router HTTP client + SSH helpers
 │   ├── router-utils.ts         pure helpers (test-friendly)
 │   └── README.md               full mTLS + systemd + Caddy setup notes
+├── token-stats.ts              footer owner for context + token rate display
 ├── dark-mechanicus/            theme TUI bundle (one extension, multi-file)
 │   ├── index.ts                entry — sequences each module's registrar
 │   ├── indicator.ts            working indicator (cog + quote + timer)
 │   ├── banner.ts               TUI header art
-│   ├── footer.ts               custom footer layout
 │   ├── status-line.ts          rotating flavor status line
 │   ├── session-names.ts        auto-name generator + tab title
 │   ├── thinking-label.ts       "Cogitating..." for folded thinking blocks

ai-server/stream.ts

@@ -12,6 +12,7 @@ import {
 	createAssistantMessageEventStream,
 	parseStreamingJson,
 } from "@mariozechner/pi-ai";
+import { captureFirstOutput, finalizePiTokenStats, type PiTokenStats } from "../shared/token-stats.js";
 import {
 	AI_SERVER_CHAT_PATH,
 	AI_SERVER_URL,
@@ -56,6 +57,9 @@ export function streamAiServer(
 	const stream = createAssistantMessageEventStream();
 
 	(async () => {
+		const tokenStats: PiTokenStats = {
+			requestStartMs: Date.now(),
+		};
 		const output: AssistantMessage = {
 			role: "assistant",
 			content: [],
@@ -76,6 +80,7 @@ export function streamAiServer(
 
 		let currentBlock: CurrentBlock | null = null;
 		let streamEnded = false;
+		let thinkingTokens: number | undefined;
 
 		const endWithError = (reason: "error" | "aborted", message: string) => {
 			if (streamEnded) return;
@@ -249,6 +254,14 @@ export function streamAiServer(
 						if (data.id && !output.responseId) output.responseId = data.id;
 
 						if (data.usage) {
+							const reportedThinkingTokens =
+								data.usage.completion_tokens_details?.reasoning_tokens;
+							if (
+								typeof reportedThinkingTokens === "number"
+								&& Number.isFinite(reportedThinkingTokens)
+							) {
+								thinkingTokens = reportedThinkingTokens;
+							}
 							output.usage.input =
 								data.usage.prompt_tokens ?? output.usage.input;
 							output.usage.output =
@@ -278,6 +291,7 @@ export function streamAiServer(
 						const reasoning: string | undefined =
 							delta.reasoning_content ?? delta.reasoning;
 						if (reasoning) {
+							captureFirstOutput(tokenStats, Date.now());
 							if (!currentBlock || currentBlock.kind !== "thinking") {
 								finishCurrentBlock();
 								currentBlock = { kind: "thinking", thinking: "" };
@@ -301,6 +315,7 @@ export function streamAiServer(
 
 						// ── Text ──
 						if (typeof delta.content === "string" && delta.content.length > 0) {
+							captureFirstOutput(tokenStats, Date.now());
 							if (!currentBlock || currentBlock.kind !== "text") {
 								finishCurrentBlock();
 								currentBlock = { kind: "text", text: "" };
@@ -325,6 +340,7 @@ export function streamAiServer(
 						// ── Tool calls ──
 						if (Array.isArray(delta.tool_calls)) {
 							for (const tc of delta.tool_calls) {
+								captureFirstOutput(tokenStats, Date.now());
 								const tcId: string | undefined = tc.id;
 								const tcName: string | undefined = tc.function?.name;
 								const tcArgs: string | undefined = tc.function?.arguments;
@@ -402,6 +418,15 @@ export function streamAiServer(
 
 					finishCurrentBlock();
 					calculateCost(model, output.usage);
+					(output as AssistantMessage & { piTokenStats?: PiTokenStats }).piTokenStats = finalizePiTokenStats(
+						tokenStats,
+						{
+							input: output.usage.input,
+							output: output.usage.output,
+							thinking: thinkingTokens,
+						},
+						Date.now(),
+					);
 
 					if (options?.signal?.aborted) {
 						endWithError("aborted", "Request aborted");

dark-mechanicus/banner.ts

@@ -1,4 +1,5 @@
 import type { ExtensionAPI, Theme } from "@mariozechner/pi-coding-agent";
+import { readExtensionEnumSetting } from "../shared/pi-settings.js";
 
 // ─── Compact cog-and-skull art (17 lines × ~38 cols) ─────────────────────
 
@@ -187,12 +188,17 @@ function renderFull(theme: Theme, width: number): string[] {
 // ─── Extension entry ─────────────────────────────────────────────────────
 
 type Mode = "compact" | "full" | "off";
-let currentMode: Mode = "compact";
+const EXTENSION_NAME = "dark-mechanicus";
+const BANNER_MODES: readonly Mode[] = ["compact", "full", "off"];
+
+function readBannerMode(): Mode {
+	return readExtensionEnumSetting(EXTENSION_NAME, "bannerMode", BANNER_MODES, "full");
+}
 
 export default function (pi: ExtensionAPI) {
-	const apply = (ctx: any, mode: Mode) => {
+	const apply = (ctx: any) => {
 		if (!ctx.hasUI) return;
-		currentMode = mode;
+		const mode = readBannerMode();
 		if (mode === "off") {
 			ctx.ui.setHeader(undefined);
 			return;
@@ -207,41 +213,10 @@ export default function (pi: ExtensionAPI) {
 	};
 
 	pi.on("session_start", async (_event, ctx) => {
-		apply(ctx, "full");
+		apply(ctx);
 	});
 
-	pi.registerCommand("mechanicus-banner-compact", {
+	pi.on("turn_end", async (_event, ctx) => {
-		description: "Compact text-only mechanicus banner (default)",
+		apply(ctx);
-		handler: async (_args, ctx) => {
-			apply(ctx, "compact");
-			ctx.ui.notify("Compact banner active.", "info");
-		},
-	});
-
-	pi.registerCommand("mechanicus-banner-full", {
-		description: "Full banner with the large cog-and-skull art (44 lines tall)",
-		handler: async (_args, ctx) => {
-			apply(ctx, "full");
-			ctx.ui.notify(
-				`Full banner active (${ART_HEIGHT} lines). /mechanicus-banner-compact to revert.`,
-				"info",
-			);
-		},
-	});
-
-	pi.registerCommand("mechanicus-banner-off", {
-		description: "Restore pi's default header",
-		handler: async (_args, ctx) => {
-			apply(ctx, "off");
-			ctx.ui.notify("Default header restored.", "info");
-		},
-	});
-
-	pi.registerCommand("mechanicus-banner-on", {
-		description: "Re-install the current mechanicus banner mode",
-		handler: async (_args, ctx) => {
-			apply(ctx, currentMode === "off" ? "compact" : currentMode);
-			ctx.ui.notify(`Banner active (${currentMode === "off" ? "compact" : currentMode}).`, "info");
-		},
 	});
 }

dark-mechanicus/footer.ts

@@ -23,6 +23,7 @@ import type {
 	Theme,
 } from "@mariozechner/pi-coding-agent";
 import { formatTokens } from "../shared/format.js";
+import { normalizePiEntries } from "../shared/token-stats.js";
 
 export default function (pi: ExtensionAPI) {
 	let cachedCtx: ExtensionContext | null = null;
@@ -54,6 +55,29 @@ export default function (pi: ExtensionAPI) {
 		width: number,
 	): string[] => {
 		const dim = (s: string) => theme.fg("dim", s);
+		const readExtensionStatuses = (): Array<[string, string]> => {
+			const statuses = footerData?.getExtensionStatuses?.();
+			if (!statuses) {
+				return [];
+			}
+			if (Array.isArray(statuses)) {
+				return statuses as Array<[string, string]>;
+			}
+			if (typeof statuses === "object") {
+				const candidate = statuses as {
+					[Symbol.iterator]?: unknown;
+				};
+				if (typeof candidate[Symbol.iterator] === "function") {
+					try {
+						return Array.from(candidate as Iterable<[string, string]>);
+					} catch {
+						return [];
+					}
+				}
+				return Object.entries(statuses as Record<string, string>);
+			}
+			return [];
+		};
 
 				// ── Walk session entries for token totals + compactions ─
 				let totalIn = 0;
@@ -62,7 +86,7 @@ export default function (pi: ExtensionAPI) {
 				let totalCost = 0;
 				let compactions = 0;
 				let lastAssistantOutput = 0;
-				const entries = c.sessionManager.getEntries();
+				const entries = normalizePiEntries(c.sessionManager.getEntries());
 				for (const entry of entries) {
 					if (entry.type === "compaction") compactions++;
 					if (
@@ -165,9 +189,9 @@ export default function (pi: ExtensionAPI) {
 				const lines = [pwdLine, statsLine];
 
 				// ── Line 3+: extension statuses (alphabetical) ──────────
-				const extStatuses = footerData.getExtensionStatuses();
+				const extStatuses = readExtensionStatuses();
-				if (extStatuses.size > 0) {
+				if (extStatuses.length > 0) {
-					const sorted = Array.from(extStatuses.entries()).sort(([a], [b]) =>
+					const sorted = extStatuses.sort(([a], [b]) =>
 						a.localeCompare(b),
 					);
 					const statusLine = sorted.map(([, text]) => text).join(" ");

dark-mechanicus/index.ts

@@ -1,18 +1,15 @@
 /**
  * dark-mechanicus — bundle of TUI customization extensions for the
  * "dark-mechanicus" theme. Loaded as a single pi extension because the
- * pieces are designed to compose (banner + status line + footer + indicator
+ * pieces are designed to compose (banner + status line + indicator
  * + thinking label all share visual language with themes/dark-mechanicus.json).
  *
- * Each module is self-contained — registers its own commands and event
+ * Display toggles are configured via ~/.pi/agent/settings.json under
- * listeners. This index only sequences their initialization. Disabling any
+ * extensions["dark-mechanicus"]. This index only sequences initialization.
- * single module is done via its own `/<name>-off` command, not by editing
- * this file.
  */
 
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 import banner from "./banner.js";
-import footer from "./footer.js";
 import indicator from "./indicator.js";
 import markdownBodyColor from "./markdown-body-color.js";
 import sessionNames from "./session-names.js";
@@ -21,7 +18,6 @@ import thinkingLabel from "./thinking-label.js";
 
 export default function (pi: ExtensionAPI) {
 	banner(pi);
-	footer(pi);
 	indicator(pi);
 	markdownBodyColor(pi);
 	sessionNames(pi);

dark-mechanicus/indicator.ts

@@ -14,14 +14,14 @@
  * Fisher-Yates shuffle on every turn_start so the order differs per
  * assistant response.
  *
- * Commands:
+ * Controlled via ~/.pi/agent/settings.json under
- *   /dark-mechanicus-indicator-off   Restore pi's default spinner
+ * extensions["dark-mechanicus"].indicatorEnabled.
- *   /dark-mechanicus-indicator-on    Re-enable the mechanicum indicator
  */
 
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 import { safely } from "../shared/ctx.js";
 import { formatElapsed } from "../shared/format.js";
+import { readExtensionBooleanSetting } from "../shared/pi-settings.js";
 
 // ─── Quote pools (15 per category) ───────────────────────────────────────
 
@@ -118,10 +118,13 @@ function buildFrames(quote: string, elapsedMs: number): string[] {
 
 // ─── Extension ───────────────────────────────────────────────────────────
 
-type Mode = "on" | "off";
+const EXTENSION_NAME = "dark-mechanicus";
+
+function readIndicatorEnabled(): boolean {
+	return readExtensionBooleanSetting(EXTENSION_NAME, "indicatorEnabled", true);
+}
 
 export default function (pi: ExtensionAPI) {
-	let mode: Mode = "on";
 	let shuffled = shuffle(ALL_QUOTES);
 	let startedAt: number | null = null;
 	let tickerHandle: NodeJS.Timeout | null = null;
@@ -146,7 +149,7 @@ export default function (pi: ExtensionAPI) {
 	};
 
 	const applyFrame = (ctx: any) => {
-		if (!ctx.hasUI || mode === "off" || startedAt === null) return;
+		if (!ctx.hasUI || !readIndicatorEnabled() || startedAt === null) return;
 		// Guard against stale-ctx: the setInterval ticker can fire between
 		// session_shutdown being queued and our turn_end/session_shutdown
 		// handlers running. Any ctx.ui.* access on a stale runner throws.
@@ -181,12 +184,18 @@ export default function (pi: ExtensionAPI) {
 
 	// ── Seed at session start so pi's default spinner never appears ──
 	pi.on("session_start", async (_event, ctx) => {
+		if (!readIndicatorEnabled()) return;
 		suppressWorkingMessage(ctx);
 	});
 
 	// ── Turn lifecycle ───────────────────────────────────────────────
 	pi.on("turn_start", async (_event, ctx) => {
-		if (mode === "off") return;
+		if (!readIndicatorEnabled()) {
+			stopTicker();
+			startedAt = null;
+			safely(() => ctx.ui.setWorkingIndicator(undefined));
+			return;
+		}
 		suppressWorkingMessage(ctx);
 		startedAt = Date.now();
 		shuffled = shuffle(ALL_QUOTES);
@@ -204,26 +213,4 @@ export default function (pi: ExtensionAPI) {
 		stopTicker();
 		startedAt = null;
 	});
-
-	// ── Commands ─────────────────────────────────────────────────────
-	pi.registerCommand("dark-mechanicus-indicator-off", {
-		description: "Restore pi's default spinner and skip cogitation summaries",
-		handler: async (_args, ctx) => {
-			mode = "off";
-			stopTicker();
-			ctx.ui.setWorkingIndicator(undefined);
-			ctx.ui.notify("Default spinner restored.", "info");
-		},
-	});
-
-	pi.registerCommand("dark-mechanicus-indicator-on", {
-		description: "Re-enable the pulsing cog indicator + cogitation summary",
-		handler: async (_args, ctx) => {
-			mode = "on";
-			ctx.ui.notify(
-				"Dark mechanicum indicator re-armed. Takes effect on the next turn.",
-				"info",
-			);
-		},
-	});
 }

dark-mechanicus/status-line.ts

@@ -12,6 +12,7 @@
  */
 
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { readExtensionBooleanSetting } from "../shared/pi-settings.js";
 
 const STATUSES = [
 	"HERETEK FORGE · ACTIVE",
@@ -35,9 +36,19 @@ function pick<T>(arr: readonly T[]): T {
 	return arr[Math.floor(Math.random() * arr.length)]!;
 }
 
+const EXTENSION_NAME = "dark-mechanicus";
+
+function readStatusLineEnabled(): boolean {
+	return readExtensionBooleanSetting(EXTENSION_NAME, "statusLineEnabled", true);
+}
+
 export default function (pi: ExtensionAPI) {
 	const applyStatus = (ctx: any, body?: string) => {
 		if (!ctx.hasUI) return;
+		if (!readStatusLineEnabled()) {
+			ctx.ui.setStatus("mechanicus", undefined);
+			return;
+		}
 		const text = body ?? pick(STATUSES);
 		const theme = ctx.ui.theme;
 		const line = `${theme.fg("warning", "⚙")} ${theme.fg("muted", text)}`;
@@ -57,16 +68,12 @@ export default function (pi: ExtensionAPI) {
 	pi.registerCommand("mechanicus-status-cycle", {
 		description: "Force a new mechanicum status line immediately",
 		handler: async (_args, ctx) => {
+			if (!readStatusLineEnabled()) {
+				ctx.ui.notify("mechanicus status line is disabled in settings.", "info");
+				return;
+			}
 			applyStatus(ctx);
 			ctx.ui.notify("Status re-transmuted.", "info");
 		},
 	});
-
-	pi.registerCommand("mechanicus-status-off", {
-		description: "Remove the mechanicum status line from the footer",
-		handler: async (_args, ctx) => {
-			ctx.ui.setStatus("mechanicus", undefined);
-			ctx.ui.notify("Status line cleared.", "info");
-		},
-	});
 }

scripts/install-client.sh

@@ -15,7 +15,8 @@
 #   certs            mTLS client cert + root CA from the Caddy host
 #   ssh              SSH key-auth to the ai-server host (admin commands)
 #   ai-server        the pi extension (mTLS llama.cpp router provider)
-#   dark-mechanicus  theme TUI bundle (banner, footer, indicator, etc.)
+#   dark-mechanicus  theme TUI bundle (banner, indicator, status line, etc.)
+#   token-stats      footer owner for context + token rate display
 #   themes           JSON theme palette files (themes/*.json)
 #   local-llama      tiny stub provider for a local llama-server
 #   ai-complete      shell CLI for direct llama-server access (installed to
@@ -63,7 +64,7 @@ SKIP_VERIFY=0
 LEGACY_NO_CERTS=0
 LEGACY_NO_SSH=0
 
-ALL_COMPONENTS=(certs ssh ai-server dark-mechanicus themes local-llama ai-complete shared)
+ALL_COMPONENTS=(certs ssh ai-server dark-mechanicus token-stats themes local-llama ai-complete shared)
 
 usage() { sed -n '2,/^$/p' "$0" | sed 's/^#\{0,1\} \{0,1\}//'; exit 0; }
 
@@ -110,10 +111,12 @@ if [[ $LEGACY_NO_SSH -eq 1 ]]; then
 	COMPONENTS=$(echo "$COMPONENTS" | tr ',' '\n' | grep -vx 'ssh' | paste -sd, -)
 fi
 
-# Auto-add 'shared' if dark-mechanicus is selected (it imports from ../shared/)
+# Auto-add 'shared' if dark-mechanicus or token-stats is selected (they import from ../shared/ or ./shared/)
-if [[ ",$COMPONENTS," == *",dark-mechanicus,"* && ",$COMPONENTS," != *",shared,"* ]]; then
+if [[ ",$COMPONENTS," == *",dark-mechanicus,"* || ",$COMPONENTS," == *",token-stats,"* ]]; then
+	if [[ ",$COMPONENTS," != *",shared,"* ]]; then
 	COMPONENTS="$COMPONENTS,shared"
 	fi
+fi
 
 # Normalize and dedupe
 COMPONENTS=$(echo "$COMPONENTS" | tr ',' '\n' | sed '/^$/d' | awk '!seen[$0]++' | paste -sd, -)
@@ -182,6 +185,17 @@ if want dark-mechanicus; then
 	else
 		echo "    (dark-mechanicus/ not in repo, skipping)"
 	fi
+
+	if want token-stats; then
+		if [[ -f token-stats.ts ]]; then
+			echo; echo "==> [token-stats] syncing footer extension"
+			if [[ -f "$PI_DIR/extensions/token-stats.ts" && $UPDATE_ONLY -eq 1 ]]; then
+				echo "    keep existing token-stats.ts (--update-only)"
+			else
+				rsync "${RSYNC_FLAGS[@]}" token-stats.ts "$PI_DIR/extensions/"
+			fi
+		fi
+	fi
 fi
 
 if want shared; then

shared/pi-settings.ts

@@ -0,0 +1,59 @@
+import { existsSync, readFileSync } from "node:fs";
+import { join } from "node:path";
+
+const HOME = process.env.HOME ?? process.env.USERPROFILE ?? "";
+const PI_SETTINGS_PATH = join(HOME, ".pi", "agent", "settings.json");
+
+function isPlainObject(value: unknown): value is Record<string, unknown> {
+	return !!value && typeof value === "object" && !Array.isArray(value);
+}
+
+function toCamelCase(value: string): string {
+	return value.replace(/-([a-z])/g, (_match, letter: string) => letter.toUpperCase());
+}
+
+function readPiSettings(): Record<string, unknown> | null {
+	if (!existsSync(PI_SETTINGS_PATH)) {
+		return null;
+	}
+	try {
+		return JSON.parse(readFileSync(PI_SETTINGS_PATH, "utf8")) as Record<string, unknown>;
+	} catch {
+		return null;
+	}
+}
+
+export function readExtensionSettings(extensionName: string): Record<string, unknown> | null {
+	const settings = readPiSettings();
+	if (!settings) {
+		return null;
+	}
+	const config = settings[extensionName] ?? settings[toCamelCase(extensionName)];
+	return isPlainObject(config) ? config : null;
+}
+
+export function readExtensionBooleanSetting(
+	extensionName: string,
+	key: string,
+	defaultValue: boolean,
+): boolean {
+	const config = readExtensionSettings(extensionName);
+	if (!config || typeof config[key] !== "boolean") {
+		return defaultValue;
+	}
+	return config[key] as boolean;
+}
+
+export function readExtensionEnumSetting<T extends string>(
+	extensionName: string,
+	key: string,
+	allowedValues: readonly T[],
+	defaultValue: T,
+): T {
+	const config = readExtensionSettings(extensionName);
+	const value = config?.[key];
+	if (typeof value !== "string") {
+		return defaultValue;
+	}
+	return (allowedValues as readonly string[]).includes(value) ? (value as T) : defaultValue;
+}

shared/token-stats.ts

@@ -0,0 +1,153 @@
+export interface PiTokenStats {
+	requestStartMs: number;
+	firstOutputMs?: number;
+	responseEndMs?: number;
+	inputTokens?: number;
+	outputTokens?: number;
+	thinkingTokens?: number;
+}
+
+export interface PiTokenRateSummary {
+	processingTokensPerSecond: number | null;
+	generationTokensPerSecond: number | null;
+}
+
+function readNumber(value: unknown): number | null {
+	if (typeof value !== "number" || !Number.isFinite(value)) {
+		return null;
+	}
+	return value;
+}
+
+export function captureFirstOutput(stats: PiTokenStats, nowMs: number): void {
+	if (stats.firstOutputMs === undefined && Number.isFinite(nowMs)) {
+		stats.firstOutputMs = nowMs;
+	}
+}
+
+export function finalizePiTokenStats(
+	stats: PiTokenStats,
+	usage: { input: number; output: number; thinking?: number },
+	nowMs: number,
+): PiTokenStats {
+	return {
+		requestStartMs: stats.requestStartMs,
+		firstOutputMs: stats.firstOutputMs,
+		responseEndMs: nowMs,
+		inputTokens: usage.input,
+		outputTokens: usage.output,
+		thinkingTokens: usage.thinking,
+	};
+}
+
+export function readPiTokenStats(message: unknown): PiTokenStats | null {
+	if (!message || typeof message !== "object") {
+		return null;
+	}
+	const candidate = (message as { piTokenStats?: unknown }).piTokenStats;
+	if (!candidate || typeof candidate !== "object") {
+		return null;
+	}
+	const stats = candidate as Record<string, unknown>;
+	const requestStartMs = readNumber(stats.requestStartMs);
+	if (requestStartMs === null) {
+		return null;
+	}
+	const firstOutputMs = readNumber(stats.firstOutputMs) ?? undefined;
+	const responseEndMs = readNumber(stats.responseEndMs) ?? undefined;
+	const inputTokens = readNumber(stats.inputTokens) ?? undefined;
+	const outputTokens = readNumber(stats.outputTokens) ?? undefined;
+	const thinkingTokens = readNumber(stats.thinkingTokens) ?? undefined;
+	return {
+		requestStartMs,
+		firstOutputMs,
+		responseEndMs,
+		inputTokens,
+		outputTokens,
+		thinkingTokens,
+	};
+}
+
+export function normalizePiEntries(entries: unknown): readonly unknown[] {
+	if (Array.isArray(entries)) {
+		return entries;
+	}
+	if (entries && typeof entries === "object") {
+		const candidate = entries as {
+			entries?: unknown;
+			[Symbol.iterator]?: unknown;
+		};
+		if (Array.isArray(candidate.entries)) {
+			return candidate.entries;
+		}
+		if (typeof candidate[Symbol.iterator] === "function") {
+			try {
+				return Array.from(candidate as Iterable<unknown>);
+			} catch {
+				return [];
+			}
+		}
+	}
+	return [];
+}
+
+export function findLatestPiTokenStats(entries: unknown): PiTokenStats | null {
+	const normalizedEntries = normalizePiEntries(entries);
+	for (let index = normalizedEntries.length - 1; index >= 0; index--) {
+		const entry = normalizedEntries[index];
+		if (!entry || typeof entry !== "object") {
+			continue;
+		}
+		const candidate = entry as {
+			type?: string;
+			message?: {
+				role?: string;
+			};
+		};
+		if (candidate.type !== "message") {
+			continue;
+		}
+		if (candidate.message?.role !== "assistant") {
+			continue;
+		}
+		const stats = readPiTokenStats(candidate.message);
+		if (stats) {
+			return stats;
+		}
+		break;
+	}
+	return null;
+}
+
+function computeRate(tokens: number | undefined, startMs: number | undefined, endMs: number | undefined): number | null {
+	if (
+		tokens === undefined
+		|| startMs === undefined
+		|| endMs === undefined
+		|| tokens <= 0
+		|| endMs <= startMs
+	) {
+		return null;
+	}
+	return tokens / ((endMs - startMs) / 1000);
+}
+
+export function summarizePiTokenStats(stats: PiTokenStats): PiTokenRateSummary {
+	const processingEndMs = stats.firstOutputMs ?? stats.responseEndMs;
+	const generatedTokens = (stats.outputTokens ?? 0) + (stats.thinkingTokens ?? 0);
+	return {
+		processingTokensPerSecond: computeRate(stats.inputTokens, stats.requestStartMs, processingEndMs),
+		generationTokensPerSecond: computeRate(generatedTokens, stats.firstOutputMs, stats.responseEndMs),
+	};
+}
+
+export function formatPiTokenRateStatus(summary: PiTokenRateSummary): string | null {
+	const parts: string[] = [];
+	if (summary.processingTokensPerSecond !== null) {
+		parts.push(`P ${summary.processingTokensPerSecond.toFixed(1)}tok/s`);
+	}
+	if (summary.generationTokensPerSecond !== null) {
+		parts.push(`G ${summary.generationTokensPerSecond.toFixed(1)}tok/s`);
+	}
+	return parts.length > 0 ? parts.join("  ") : null;
+}

tests/token-stats.test.ts

@@ -0,0 +1,111 @@
+import assert from "node:assert/strict";
+import { test } from "node:test";
+import {
+	findLatestPiTokenStats,
+	formatPiTokenRateStatus,
+	normalizePiEntries,
+	readPiTokenStats,
+	summarizePiTokenStats,
+} from "../shared/token-stats.ts";
+
+test("summarizePiTokenStats reports separate processing and generation rates", () => {
+	const summary = summarizePiTokenStats({
+		requestStartMs: 1000,
+		firstOutputMs: 3000,
+		responseEndMs: 7000,
+		inputTokens: 800,
+		outputTokens: 200,
+		thinkingTokens: 300,
+	});
+	assert.equal(summary.processingTokensPerSecond, 400);
+	assert.equal(summary.generationTokensPerSecond, 125);
+	assert.equal(formatPiTokenRateStatus(summary), "P 400.0tok/s  G 125.0tok/s");
+});
+
+test("summarizePiTokenStats falls back to response end when no output token arrived", () => {
+	const summary = summarizePiTokenStats({
+		requestStartMs: 1000,
+		responseEndMs: 5000,
+		inputTokens: 200,
+		outputTokens: 0,
+	});
+	assert.equal(summary.processingTokensPerSecond, 50);
+	assert.equal(summary.generationTokensPerSecond, null);
+	assert.equal(formatPiTokenRateStatus(summary), "P 50.0tok/s");
+});
+
+test("readPiTokenStats rejects missing timing metadata", () => {
+	assert.equal(readPiTokenStats({ role: "assistant" }), null);
+	assert.deepEqual(
+		readPiTokenStats({
+			piTokenStats: {
+				requestStartMs: 10,
+				firstOutputMs: 20,
+				responseEndMs: 30,
+				inputTokens: 40,
+				outputTokens: 50,
+				thinkingTokens: 60,
+			},
+		}),
+		{
+			requestStartMs: 10,
+			firstOutputMs: 20,
+			responseEndMs: 30,
+			inputTokens: 40,
+			outputTokens: 50,
+			thinkingTokens: 60,
+		},
+	);
+});
+
+test("findLatestPiTokenStats ignores older assistant messages after the latest assistant turn", () => {
+	assert.deepEqual(
+		findLatestPiTokenStats([
+			{
+				type: "message",
+				message: {
+					role: "assistant",
+					piTokenStats: {
+						requestStartMs: 1,
+						responseEndMs: 2,
+						outputTokens: 3,
+					},
+				},
+			},
+			{
+				type: "message",
+				message: {
+					role: "user",
+				},
+			},
+			{
+				type: "message",
+				message: {
+					role: "assistant",
+					piTokenStats: {
+						requestStartMs: 10,
+						firstOutputMs: 20,
+						responseEndMs: 30,
+						inputTokens: 40,
+						outputTokens: 50,
+						thinkingTokens: 60,
+					},
+				},
+			},
+		]),
+		{
+			requestStartMs: 10,
+			firstOutputMs: 20,
+			responseEndMs: 30,
+			inputTokens: 40,
+			outputTokens: 50,
+			thinkingTokens: 60,
+		},
+	);
+});
+
+test("findLatestPiTokenStats treats non-array entries as empty", () => {
+	assert.equal(findLatestPiTokenStats(undefined), null);
+	assert.equal(findLatestPiTokenStats({}), null);
+	assert.deepEqual(normalizePiEntries(undefined), []);
+});