pi-extensions/dark-mechanicus/footer.ts

/**
 * mechanicus-footer — replaces pi's default footer with a mechanicum-flavored
 * version. Changes from the default:
 *
 *   - Context usage rendered as "45k/131k (34%)" instead of "34%/131k" so you
 *     don't have to multiply the percentage in your head.
 *   - Compaction counter: appends "C<N>" when one or more compactions have
 *     occurred during the session. Walks sessionManager entries for the
 *     current count; no state tracking needed.
 *   - Everything else (↑in, ↓out, R cache, cost, model, thinking level,
 *     extension statuses on line 3+) is preserved in the same slots and
 *     formatting.
 *
 * The context percentage threshold colors still apply:
 *   > 90%  → error (rust)
 *   > 70%  → warning (gold)
 *   else   → dim (default)
 */

import type {
	ExtensionAPI,
	ExtensionContext,
	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;
	let lastTurnStartMs: number | null = null;
	let lastTurnDurationMs: number | null = null;

	const install = (ctx: ExtensionContext) => {
		if (!ctx.hasUI) return;
		cachedCtx = ctx;
		ctx.ui.setFooter((_tui, theme, footerData) => ({
			render(width: number): string[] {
				if (!cachedCtx) return [];
				try {
					return renderInner(cachedCtx, theme, footerData, width);
				} catch {
					// Any late render after the runner went stale: swallow and
					// return nothing rather than crashing node.
					return [];
				}
			},
			invalidate() {},
		}));
	};

	const renderInner = (
		c: ExtensionContext,
		theme: Theme,
		footerData: any,
		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;
				let totalOut = 0;
				let totalCache = 0;
				let totalCost = 0;
				let compactions = 0;
				let lastAssistantOutput = 0;
				const entries = normalizePiEntries(c.sessionManager.getEntries());
				for (const entry of entries) {
					if (entry.type === "compaction") compactions++;
					if (
						entry.type === "message" &&
						entry.message.role === "assistant"
					) {
						totalIn += entry.message.usage.input;
						totalOut += entry.message.usage.output;
						totalCache += entry.message.usage.cacheRead;
						totalCost += entry.message.usage.cost.total;
						lastAssistantOutput = entry.message.usage.output;
					}
				}

				// ── Context usage (current) ─────────────────────────────
				const usage = c.getContextUsage();
				const curTokens = usage?.tokens;
				const maxTokens =
					usage?.contextWindow ?? c.model?.contextWindow ?? 0;
				const pct = usage?.percent ?? null;

				// ── Line 1: cwd + git branch + session name ─────────────
				let pwd = c.sessionManager.getCwd();
				const home = process.env.HOME;
				if (home && pwd.startsWith(home)) pwd = `~${pwd.slice(home.length)}`;
				const branch = footerData.getGitBranch();
				if (branch) pwd = `${pwd} (${branch})`;
				const sessionName = c.sessionManager.getSessionName();
				if (sessionName) pwd = `${pwd} • ${sessionName}`;
				const pwdLine = dim(pwd);

				// ── Line 2: stats + right-aligned model ─────────────────
				const leftParts: string[] = [];
				if (totalIn) leftParts.push(`↑${formatTokens(totalIn)}`);
				if (totalOut) leftParts.push(`↓${formatTokens(totalOut)}`);
				if (totalCache) leftParts.push(`R${formatTokens(totalCache)}`);
				if (totalCost > 0) leftParts.push(`$${totalCost.toFixed(3)}`);

				// Generation speed (last turn): output tokens / duration
				if (
					lastTurnDurationMs !== null &&
					lastTurnDurationMs > 500 &&
					lastAssistantOutput > 0
				) {
					const tokPerSec =
						lastAssistantOutput / (lastTurnDurationMs / 1000);
					leftParts.push(`${tokPerSec.toFixed(1)}tok/s`);
				}

				// Context chunk — format: "45k/131k (34%)"
				let ctxRaw: string;
				let ctxColored: string;
				if (curTokens !== null && curTokens !== undefined && maxTokens > 0) {
					ctxRaw = `${formatTokens(curTokens)}/${formatTokens(maxTokens)} (${
						pct === null ? "?" : pct.toFixed(0)
					}%)`;
				} else {
					ctxRaw = `?/${formatTokens(maxTokens)}`;
				}
				if (pct !== null && pct > 90) ctxColored = theme.fg("error", ctxRaw);
				else if (pct !== null && pct > 70)
					ctxColored = theme.fg("warning", ctxRaw);
				else ctxColored = dim(ctxRaw);

				// Build left-side string. Non-context parts are dim; context keeps
				// its threshold color. We interleave raw lengths for padding math.
				const dimmedLeft = leftParts.map(dim).join(" ");
				const dimmedLeftRaw = leftParts.join(" ");
				const compactionChunk =
					compactions > 0 ? ` ${dim(`C${compactions}`)}` : "";
				const compactionRaw = compactions > 0 ? ` C${compactions}` : "";
				const leftFull =
					(dimmedLeft.length ? dimmedLeft + " " : "") +
					ctxColored +
					compactionChunk;
				const leftRawWidth =
					(dimmedLeftRaw.length ? dimmedLeftRaw.length + 1 : 0) +
					ctxRaw.length +
					compactionRaw.length;

				// Right side: model + thinking level
				const modelName = c.model?.id ?? "no-model";
				let rightRaw = modelName;
				if (c.model?.reasoning) {
					const level = pi.getThinkingLevel?.() ?? "off";
					rightRaw = `${modelName} • ${level === "off" ? "thinking off" : level}`;
				}
				const rightFull = dim(rightRaw);

				// Pad between
				const minSpace = 2;
				let statsLine: string;
				if (leftRawWidth + minSpace + rightRaw.length <= width) {
					const pad = " ".repeat(width - leftRawWidth - rightRaw.length);
					statsLine = leftFull + pad + rightFull;
				} else {
					statsLine = leftFull;
				}

				const lines = [pwdLine, statsLine];

				// ── Line 3+: extension statuses (alphabetical) ──────────
				const extStatuses = readExtensionStatuses();
				if (extStatuses.length > 0) {
					const sorted = extStatuses.sort(([a], [b]) =>
						a.localeCompare(b),
					);
					const statusLine = sorted.map(([, text]) => text).join(" ");
					lines.push(statusLine);
				}

				return lines;
	};

	pi.on("session_start", async (_event, ctx) => {
		install(ctx);
	});

	pi.on("turn_start", async () => {
		lastTurnStartMs = Date.now();
	});

	// Keep ctx fresh on every turn so token sums and context usage are current.
	pi.on("turn_end", async (_event, ctx) => {
		if (lastTurnStartMs !== null) {
			lastTurnDurationMs = Date.now() - lastTurnStartMs;
			lastTurnStartMs = null;
		}
		cachedCtx = ctx;
	});

	// On /quit or /reload pi marks the ExtensionRunner stale, but the TUI
	// may still fire a pending render timer. Accessing ctx.sessionManager
	// after that throws "This extension instance is stale…". Null the
	// captured ctx so render() short-circuits to [] and no late draw
	// touches the dead runner.
	pi.on("session_shutdown", async () => {
		cachedCtx = null;
	});

	pi.registerCommand("mechanicus-footer-off", {
		description: "Restore pi's default footer",
		handler: async (_args, ctx) => {
			ctx.ui.setFooter(undefined);
			ctx.ui.notify("Default footer restored.", "info");
		},
	});

	pi.registerCommand("mechanicus-footer-on", {
		description: "Re-install the mechanicum footer",
		handler: async (_args, ctx) => {
			install(ctx);
			ctx.ui.notify("Mechanicum footer installed.", "info");
		},
	});
}