Pi Dev Colorful Statusline Extension

statusline.ts

/**
 * Colorful Statusline Extension
 *
 * Replaces the built-in footer with a colorful, labeled version that mirrors
 * the default layout:
 *   line 1:  path + git branch + session name
 *   line 2:  token stats + context %   ...right-aligned...   model + thinking
 *   line 3:  extension statuses (from ctx.ui.setStatus)
 *
 * Enhancements over the default:
 *   - a distinct theme color for every field
 *   - human-readable labels ("In", "Out", "Cache", "Hit", "Cost", "Ctx" ...)
 *
 * Usage:
 *   pi -e ./statusline.ts            # auto-enables on startup
 *   /statusline                      # toggle at runtime
 *
 * The footer is stateless: render() recomputes everything from the live theme
 * proxy and current ctx on every frame, so theme switches apply immediately.
 */

import type { AssistantMessage } from "@earendil-works/pi-ai";
import type {
        ExtensionAPI,
        ExtensionCommandContext,
        ExtensionContext,
        ReadonlyFooterDataProvider,
        Theme,
        ThemeColor,
} from "@earendil-works/pi-coding-agent";
import type { TUI } from "@earendil-works/pi-tui";
import { truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
import { isAbsolute, relative, resolve, sep } from "node:path";

const GAP = "   "; // visual gap between stat groups on line 2

/** Compact token formatter, matching the built-in footer. */
function formatTokens(n: number): string {
        if (n < 1000) return `${n}`;
        if (n < 10000) return `${(n / 1000).toFixed(1)}k`;
        if (n < 1_000_000) return `${Math.round(n / 1000)}k`;
        if (n < 10_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
        return `${Math.round(n / 1_000_000)}M`;
}

/** Replace $HOME with ~ for display, matching the built-in footer. */
function formatCwd(cwd: string, home: string | undefined): string {
        if (!home) return cwd;
        const r = relative(resolve(home), resolve(cwd));
        const inside = r === "" || (!r.startsWith(`..${sep}`) && r !== ".." && !isAbsolute(r));
        if (!inside) return cwd;
        return r === "" ? "~" : `~${sep}${r}`;
}

/** Collapse whitespace/control chars so a status can't break the line. */
function sanitize(text: string): string {
        return text.replace(/[\r\n\t]/g, " ").replace(/ +/g, " ").trim();
}

/** One labeled, colored stat group: "label value". */
function stat(theme: Theme, label: string, value: string, color: ThemeColor): string {
        return `${theme.fg("muted", label)} ${theme.fg(color, value)}`;
}

interface FooterComponent {
        dispose?(): void;
        invalidate(): void;
        render(width: number): string[];
}

/**
 * Build the footer component. `theme` is the live proxy captured from the
 * factory, so colors track the active theme without any caching.
 */
function makeFooter(
        theme: Theme,
        footerData: ReadonlyFooterDataProvider,
        ctx: ExtensionContext,
        api: ExtensionAPI,
): FooterComponent {
        // Re-render when the git branch changes (branch comes from footerData).
        const unsub = footerData.onBranchChange(() => {});

        return {
                dispose: unsub,
                invalidate() {
                        /* stateless — nothing to cache-clear */
                },
                render(width: number): string[] {
                        // --- Aggregate token usage across the whole session -----------------
                        let input = 0;
                        let output = 0;
                        let cacheRead = 0;
                        let cacheWrite = 0;
                        let cost = 0;
                        let hitRate: number | undefined;

                        for (const e of ctx.sessionManager.getEntries()) {
                                if (e.type === "message" && e.message.role === "assistant") {
                                        const u = (e.message as AssistantMessage).usage;
                                        input += u.input;
                                        output += u.output;
                                        cacheRead += u.cacheRead;
                                        cacheWrite += u.cacheWrite;
                                        cost += u.cost.total;
                                        const prompt = u.input + u.cacheRead + u.cacheWrite;
                                        hitRate = prompt > 0 ? (u.cacheRead / prompt) * 100 : hitRate;
                                }
                        }

                        // --- Line 1: path + git branch + session name ------------------------
                        const pwd = formatCwd(
                                ctx.sessionManager.getCwd(),
                                process.env.HOME ?? process.env.USERPROFILE,
                        );
                        const seg1: string[] = [theme.fg("dim", `📁 ${pwd}`)];

                        const branch = footerData.getGitBranch();
                        if (branch) seg1.push(theme.fg("success", `🌿 ${branch}`));

                        const sessionName = ctx.sessionManager.getSessionName();
                        if (sessionName) seg1.push(theme.fg("accent", `💬 ${sessionName}`));

                        const line1 = truncateToWidth(seg1.join(" "), width, theme.fg("dim", "…"));

                        // --- Line 2: stats (left) + model (right) ----------------------------
                        const groups: string[] = [];
                        if (input) groups.push(stat(theme, "📥 In", formatTokens(input), "accent"));
                        if (output) groups.push(stat(theme, "📤 Out", formatTokens(output), "text"));
                        if (cacheRead) groups.push(stat(theme, "💾 Cache", formatTokens(cacheRead), "mdCode"));
                        if (cacheWrite) groups.push(stat(theme, "⤴️ W", formatTokens(cacheWrite), "warning"));
                        if ((cacheRead > 0 || cacheWrite > 0) && hitRate !== undefined) {
                                groups.push(stat(theme, "🎯 Hit", `${hitRate.toFixed(1)}%`, "success"));
                        }
                        if (cost > 0) {
                                groups.push(stat(theme, "💰 Cost", `$${cost.toFixed(3)}`, cost > 1 ? "warning" : "success"));
                        }

                        // Context %, colored by threshold (matches built-in semantics).
                        const usage = ctx.getContextUsage();
                        const window = usage?.contextWindow ?? ctx.model?.contextWindow ?? 0;
                        const pct = usage?.percent ?? 0;
                        const known = usage?.percent != null;
                        const ctxText = known
                                ? `${pct.toFixed(1)}%/${formatTokens(window)}`
                                : `?/${formatTokens(window)}`;
                        const ctxColor: ThemeColor = pct > 90 ? "error" : pct > 70 ? "warning" : "accent";
                        groups.push(stat(theme, "📊 Ctx", ctxText, ctxColor));

                        let left = groups.join(theme.fg("dim", GAP));
                        if (visibleWidth(left) > width) {
                                left = truncateToWidth(left, width, theme.fg("dim", "…"));
                        }

                        // Right side: provider + model + thinking level.
                        const model = ctx.model;
                        const rightPieces = [theme.fg("accent", `🤖 ${model?.id ?? "no-model"}`)];
                        if (model?.reasoning) {
                                const lvl = api.getThinkingLevel();
                                rightPieces.push(
                                        theme.fg("muted", "•"),
                                        theme.fg("thinkingText", lvl === "off" ? "🧠 think off" : `🧠 think ${lvl}`),
                                );
                        }
                        let right = rightPieces.join(" ");
                        if (footerData.getAvailableProviderCount() > 1 && model) {
                                const withProvider = `${theme.fg("muted", `(${model.provider})`)} ${right}`;
                                if (visibleWidth(left) + 2 + visibleWidth(withProvider) <= width) {
                                        right = withProvider;
                                }
                        }

                        const minPad = 2;
                        const fits = visibleWidth(left) + minPad + visibleWidth(right) <= width;
                        let line2: string;
                        if (fits) {
                                const pad = " ".repeat(width - visibleWidth(left) - visibleWidth(right));
                                line2 = left + pad + right;
                        } else {
                                const avail = width - visibleWidth(left) - minPad;
                                if (avail > 0) {
                                        const t = truncateToWidth(right, avail, "");
                                        line2 = left + " ".repeat(Math.max(0, width - visibleWidth(left) - visibleWidth(t))) + t;
                                } else {
                                        line2 = left;
                                }
                        }

                        const lines = [line1, line2];

                        // --- Line 3: extension statuses (preserve built-in behavior) --------
                        const statuses = footerData.getExtensionStatuses();
                        if (statuses.size > 0) {
                                const text = Array.from(statuses.values())
                                        .sort()
                                        .map(sanitize)
                                        .join(" ");
                                lines.push(truncateToWidth(text, width, theme.fg("dim", "…")));
                        }

                        return lines;
                },
        };
}

export default function (api: ExtensionAPI) {
        let enabled = false;

        const enable = (ctx: ExtensionContext) => {
                ctx.ui.setFooter((_tui: TUI, theme: Theme, footerData: ReadonlyFooterDataProvider) =>
                        makeFooter(theme, footerData, ctx, api),
                );
                enabled = true;
        };
        const disable = (ctx: ExtensionContext) => {
                ctx.ui.setFooter(undefined);
                enabled = false;
        };

        // Auto-enable whenever a session starts (startup / reload / new / resume / fork).
        api.on("session_start", (_event, ctx) => {
                if (!enabled) enable(ctx);
        });

        // Manual toggle.
        api.registerCommand("statusline", {
                description: "Toggle the colorful statusline footer",
                handler: async (_args, ctx: ExtensionCommandContext) => {
                        if (enabled) {
                                disable(ctx);
                                ctx.ui.notify("Statusline: default footer restored", "info");
                        } else {
                                enable(ctx);
                                ctx.ui.notify("Statusline: colorful footer enabled", "info");
                        }
                },
        });

        // CLI flag:  pi --statusline   (forces on at startup)
        api.registerFlag("statusline", {
                description: "Enable the colorful statusline footer",
                type: "boolean",
                default: false,
        });
}