Commit cd11bc97bbd7

Vincent Demeester <vincent@sbr.pm>
2026-06-08 23:19:16
feat(pi): add cost-tracker agent extension
1 parent b727c3e
Changed files (1)
dots
pi
agent
dots/pi/agent/extensions/cost-tracker.ts
@@ -0,0 +1,122 @@
+/**
+ * Cost Tracker Extension
+ *
+ * Shows session cost and token usage in the footer status bar,
+ * similar to GitHub Copilot's premium request multiplier display.
+ */
+
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+
+// Premium request multipliers (cost relative to Sonnet baseline)
+const MULTIPLIERS: Record<string, number> = {
+	// Anthropic
+	"claude-sonnet-4-20250514": 1,
+	"claude-sonnet-4-5-20250514": 1,
+	"claude-sonnet-4-5-20241022": 1,
+	"claude-haiku-3-5-20241022": 0.33,
+	"claude-opus-4-20250514": 3,
+	"claude-opus-4-5-20250120": 3,
+	// OpenAI
+	"gpt-4.1": 0,
+	"gpt-4o": 1,
+	"gpt-4o-mini": 0,
+	"o3": 3,
+	"o3-mini": 0.33,
+	"o4-mini": 0.33,
+	// Google
+	"gemini-2.5-pro": 1,
+	"gemini-2.5-flash": 0.33,
+	"gemini-2.0-flash": 0,
+};
+
+function getMultiplier(modelId: string): string {
+	for (const [pattern, mult] of Object.entries(MULTIPLIERS)) {
+		if (modelId.includes(pattern)) {
+			return `${mult}x`;
+		}
+	}
+	return "?x";
+}
+
+function formatCost(cost: number): string {
+	if (cost < 0.01) return `$${cost.toFixed(4)}`;
+	if (cost < 1) return `$${cost.toFixed(3)}`;
+	return `$${cost.toFixed(2)}`;
+}
+
+function formatTokens(n: number): string {
+	if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
+	if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k`;
+	return `${n}`;
+}
+
+export default function (pi: ExtensionAPI) {
+	let totalCost = 0;
+	let totalInput = 0;
+	let totalOutput = 0;
+	let totalCacheRead = 0;
+	let currentModel = "";
+
+	function updateStatus(ctx: { ui: any }) {
+		const theme = ctx.ui.theme;
+		const mult = getMultiplier(currentModel);
+		const parts = [
+			theme.fg("dim", "cost:"),
+			theme.fg("accent", formatCost(totalCost)),
+			theme.fg("dim", `[${mult}]`),
+			theme.fg("dim", `in:${formatTokens(totalInput)}`),
+			theme.fg("dim", `out:${formatTokens(totalOutput)}`),
+		];
+		if (totalCacheRead > 0) {
+			parts.push(theme.fg("dim", `cache:${formatTokens(totalCacheRead)}`));
+		}
+		ctx.ui.setStatus("cost-tracker", parts.join(" "));
+	}
+
+	pi.on("session_start", async (_event, ctx) => {
+		// Restore from existing session entries
+		for (const entry of ctx.sessionManager.getBranch()) {
+			if (entry.type === "message" && entry.message.role === "assistant") {
+				const msg = entry.message as any;
+				if (msg.usage) {
+					totalInput += msg.usage.input || 0;
+					totalOutput += msg.usage.output || 0;
+					totalCacheRead += msg.usage.cacheRead || 0;
+					if (msg.usage.cost) {
+						totalCost += msg.usage.cost.total || 0;
+					}
+				}
+				if (msg.model) currentModel = msg.model;
+			}
+		}
+		updateStatus(ctx);
+	});
+
+	pi.on("model_select", async (event, ctx) => {
+		currentModel = (event as any).modelId || "";
+		updateStatus(ctx);
+	});
+
+	pi.on("turn_end", async (_event, ctx) => {
+		// Recalculate from session to stay accurate
+		totalCost = 0;
+		totalInput = 0;
+		totalOutput = 0;
+		totalCacheRead = 0;
+		for (const entry of ctx.sessionManager.getBranch()) {
+			if (entry.type === "message" && entry.message.role === "assistant") {
+				const msg = entry.message as any;
+				if (msg.usage) {
+					totalInput += msg.usage.input || 0;
+					totalOutput += msg.usage.output || 0;
+					totalCacheRead += msg.usage.cacheRead || 0;
+					if (msg.usage.cost) {
+						totalCost += msg.usage.cost.total || 0;
+					}
+				}
+				if (msg.model) currentModel = msg.model;
+			}
+		}
+		updateStatus(ctx);
+	});
+}