pi-extensions/token-stats/index.ts

import type { ExtensionAPI, ExtensionContext, Theme } from "@mariozechner/pi-coding-agent";
import { safely } from "../shared/ctx.js";
import { formatTokens } from "../shared/format.js";
import { readExtensionBooleanSetting } from "../shared/pi-settings.js";
import {
	buildPiTokenStatsEntry,
	findDisplayableTokenStats,
	findLatestPiTokenStats,
	formatPiTokenRateStatus,
	normalizePiEntries,
	readPiTokenStatsEntry,
	summarizePiTokenStats,
} from "../shared/token-stats.js";

const EXTENSION_NAME = "token-stats";
const TOKEN_STATS_TURN_ENTRY_TYPE = "token-stats-turn";

function readTokenStatsEnabled(): boolean {
	return readExtensionBooleanSetting(EXTENSION_NAME, "enabled", true);
}

function renderFooter(
	pi: ExtensionAPI,
	ctx: ExtensionContext,
	theme: Theme,
	footerData: any,
	width: number,
): string[] {
	const dim = (text: string) => theme.fg("dim", text);
	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 [];
	};

	let totalIn = 0;
	let totalOut = 0;
	let totalCache = 0;
	let totalCost = 0;
	let compactions = 0;
	const entries = normalizePiEntries(ctx.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;
		}
	}

	const usage = ctx.getContextUsage();
	const curTokens = usage?.tokens;
	const maxTokens = usage?.contextWindow ?? ctx.model?.contextWindow ?? 0;
	const pct = usage?.percent ?? null;
	const latestTokenStats = findDisplayableTokenStats(entries);
	const tokenRateText =
		latestTokenStats === null
			? null
			: formatPiTokenRateStatus(summarizePiTokenStats(latestTokenStats));

	let pwd = ctx.sessionManager.getCwd();
	const home = process.env.HOME ?? process.env.USERPROFILE;
	if (home && pwd.startsWith(home)) pwd = `~${pwd.slice(home.length)}`;
	const branch = footerData.getGitBranch();
	if (branch) pwd = `${pwd} (${branch})`;
	const sessionName = ctx.sessionManager.getSessionName();
	if (sessionName) pwd = `${pwd} • ${sessionName}`;
	const pwdLine = dim(pwd);

	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)}`);

	let contextText: string;
	if (curTokens !== null && curTokens !== undefined && maxTokens > 0) {
		contextText = `${formatTokens(curTokens)}/${formatTokens(maxTokens)} (${
			pct === null ? "?" : pct.toFixed(0)
		}%)`;
	} else {
		contextText = `?/${formatTokens(maxTokens)}`;
	}
	let contextColored = dim(contextText);
	if (pct !== null && pct > 90) {
		contextColored = theme.fg("error", contextText);
	} else if (pct !== null && pct > 70) {
		contextColored = theme.fg("warning", contextText);
	}

	const dimmedLeft = leftParts.map(dim).join(" ");
	const dimmedLeftRaw = leftParts.join(" ");
	const tokenRateChunk = tokenRateText !== null ? ` ${dim(tokenRateText)}` : "";
	const tokenRateRaw = tokenRateText !== null ? ` ${tokenRateText}` : "";
	const compactionChunk = compactions > 0 ? ` ${dim(`C${compactions}`)}` : "";
	const compactionRaw = compactions > 0 ? ` C${compactions}` : "";
	const leftFull =
		(dimmedLeft.length > 0 ? `${dimmedLeft} ` : "")
		+ contextColored
		+ tokenRateChunk
		+ compactionChunk;
	const leftRawWidth =
		(dimmedLeftRaw.length > 0 ? dimmedLeftRaw.length + 1 : 0)
		+ contextText.length
		+ tokenRateRaw.length
		+ compactionRaw.length;

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

	const minSpace = 2;
	const statsLine =
		leftRawWidth + minSpace + rightRaw.length <= width
			? leftFull + " ".repeat(width - leftRawWidth - rightRaw.length) + rightFull
			: leftFull;

	const lines = [pwdLine, statsLine];
	const extStatuses = readExtensionStatuses();
	if (extStatuses.length > 0) {
		const statusLine = extStatuses
			.sort(([a], [b]) => a.localeCompare(b))
			.map(([, text]) => text)
			.join(" ");
		if (statusLine.length > 0) {
			lines.push(statusLine);
		}
	}

	return lines;
}

export default function (pi: ExtensionAPI) {
	let cachedCtx: ExtensionContext | null = null;
	let footerInstalled = false;
	let lastLoggedAssistantEntryId: string | null = null;

	const restoreLoggedAssistantEntryId = (ctx: ExtensionContext) => {
		for (const entry of normalizePiEntries(ctx.sessionManager.getBranch()).toReversed()) {
			const logged = readPiTokenStatsEntry(entry, TOKEN_STATS_TURN_ENTRY_TYPE);
			if (logged) {
				lastLoggedAssistantEntryId = logged.assistantEntryId;
				return;
			}
		}
		lastLoggedAssistantEntryId = null;
	};

	const appendTurnStatsEntry = (ctx: ExtensionContext) => {
		const entries = normalizePiEntries(ctx.sessionManager.getEntries());
		for (let index = entries.length - 1; index >= 0; index--) {
			const entry = entries[index];
			if (!entry || typeof entry !== "object") {
				continue;
			}
			const candidate = entry as {
				id?: string;
				type?: string;
				message?: {
					role?: string;
				};
			};
			if (candidate.type !== "message" || candidate.message?.role !== "assistant") {
				continue;
			}
			if (!candidate.id || candidate.id === lastLoggedAssistantEntryId) {
				return;
			}
			const stats = findLatestPiTokenStats([candidate]);
			if (!stats) {
				return;
			}
			pi.appendEntry(TOKEN_STATS_TURN_ENTRY_TYPE, buildPiTokenStatsEntry(candidate.id, stats));
			lastLoggedAssistantEntryId = candidate.id;
			return;
		}
	};

	const clearStatus = (ctx: ExtensionContext) => {
		if (!ctx.hasUI) {
			return;
		}
		safely(() => ctx.ui.setStatus("token-stats", undefined));
	};

	const applyFooter = (ctx: ExtensionContext) => {
		if (!ctx.hasUI) {
			return;
		}
		cachedCtx = ctx;
		clearStatus(ctx);
		if (!readTokenStatsEnabled()) {
			if (footerInstalled) {
				safely(() => ctx.ui.setFooter(undefined));
				footerInstalled = false;
			}
			clearStatus(ctx);
			return;
		}
		safely(() => {
			ctx.ui.setFooter((_tui, theme, footerData) => ({
				render(width: number): string[] {
					if (!cachedCtx) {
						return [];
					}
					try {
						return renderFooter(pi, cachedCtx, theme, footerData, width);
					} catch {
						return [];
					}
				},
				invalidate() {},
			}));
			footerInstalled = true;
		});
	};

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

	pi.on("session_tree", async (_event, ctx) => {
		applyFooter(ctx);
	});

	pi.on("turn_end", async (_event, ctx) => {
		appendTurnStatsEntry(ctx);
		applyFooter(ctx);
	});

	pi.on("session_shutdown", async () => {
		cachedCtx = null;
		footerInstalled = false;
		lastLoggedAssistantEntryId = null;
	});
}