1/**
  2 * Cost Tracker Extension
  3 *
  4 * Shows session cost and token usage in the footer status bar,
  5 * similar to GitHub Copilot's premium request multiplier display.
  6 */
  7
  8import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
  9
 10// Premium request multipliers (cost relative to Sonnet baseline)
 11const MULTIPLIERS: Record<string, number> = {
 12	// Anthropic
 13	"claude-sonnet-4-20250514": 1,
 14	"claude-sonnet-4-5-20250514": 1,
 15	"claude-sonnet-4-5-20241022": 1,
 16	"claude-haiku-3-5-20241022": 0.33,
 17	"claude-opus-4-20250514": 3,
 18	"claude-opus-4-5-20250120": 3,
 19	// OpenAI
 20	"gpt-4.1": 0,
 21	"gpt-4o": 1,
 22	"gpt-4o-mini": 0,
 23	"o3": 3,
 24	"o3-mini": 0.33,
 25	"o4-mini": 0.33,
 26	// Google
 27	"gemini-2.5-pro": 1,
 28	"gemini-2.5-flash": 0.33,
 29	"gemini-2.0-flash": 0,
 30};
 31
 32function getMultiplier(modelId: string): string {
 33	for (const [pattern, mult] of Object.entries(MULTIPLIERS)) {
 34		if (modelId.includes(pattern)) {
 35			return `${mult}x`;
 36		}
 37	}
 38	return "?x";
 39}
 40
 41function formatCost(cost: number): string {
 42	if (cost < 0.01) return `$${cost.toFixed(4)}`;
 43	if (cost < 1) return `$${cost.toFixed(3)}`;
 44	return `$${cost.toFixed(2)}`;
 45}
 46
 47function formatTokens(n: number): string {
 48	if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
 49	if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k`;
 50	return `${n}`;
 51}
 52
 53export default function (pi: ExtensionAPI) {
 54	let totalCost = 0;
 55	let totalInput = 0;
 56	let totalOutput = 0;
 57	let totalCacheRead = 0;
 58	let currentModel = "";
 59
 60	function updateStatus(ctx: { ui: any }) {
 61		const theme = ctx.ui.theme;
 62		const mult = getMultiplier(currentModel);
 63		const parts = [
 64			theme.fg("dim", "cost:"),
 65			theme.fg("accent", formatCost(totalCost)),
 66			theme.fg("dim", `[${mult}]`),
 67			theme.fg("dim", `in:${formatTokens(totalInput)}`),
 68			theme.fg("dim", `out:${formatTokens(totalOutput)}`),
 69		];
 70		if (totalCacheRead > 0) {
 71			parts.push(theme.fg("dim", `cache:${formatTokens(totalCacheRead)}`));
 72		}
 73		ctx.ui.setStatus("cost-tracker", parts.join(" "));
 74	}
 75
 76	pi.on("session_start", async (_event, ctx) => {
 77		// Restore from existing session entries
 78		for (const entry of ctx.sessionManager.getBranch()) {
 79			if (entry.type === "message" && entry.message.role === "assistant") {
 80				const msg = entry.message as any;
 81				if (msg.usage) {
 82					totalInput += msg.usage.input || 0;
 83					totalOutput += msg.usage.output || 0;
 84					totalCacheRead += msg.usage.cacheRead || 0;
 85					if (msg.usage.cost) {
 86						totalCost += msg.usage.cost.total || 0;
 87					}
 88				}
 89				if (msg.model) currentModel = msg.model;
 90			}
 91		}
 92		updateStatus(ctx);
 93	});
 94
 95	pi.on("model_select", async (event, ctx) => {
 96		currentModel = (event as any).modelId || "";
 97		updateStatus(ctx);
 98	});
 99
100	pi.on("turn_end", async (_event, ctx) => {
101		// Recalculate from session to stay accurate
102		totalCost = 0;
103		totalInput = 0;
104		totalOutput = 0;
105		totalCacheRead = 0;
106		for (const entry of ctx.sessionManager.getBranch()) {
107			if (entry.type === "message" && entry.message.role === "assistant") {
108				const msg = entry.message as any;
109				if (msg.usage) {
110					totalInput += msg.usage.input || 0;
111					totalOutput += msg.usage.output || 0;
112					totalCacheRead += msg.usage.cacheRead || 0;
113					if (msg.usage.cost) {
114						totalCost += msg.usage.cost.total || 0;
115					}
116				}
117				if (msg.model) currentModel = msg.model;
118			}
119		}
120		updateStatus(ctx);
121	});
122}