Pi Agent Cost Tracker

status-cost.ts

/**
 * Status Cost — daily cost tracking with auto-reset on reset day.
 * Persisted to ~/.pi/cost.json
 * Commands: /cost, /cost-reset
 */
import { readFileSync, unlinkSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

const COST_FILE = join(homedir(), ".pi", "cost.json");

interface CostHistory {
  dailyCost: Record<string, number>;
  resetDay: number;
}

const formatCurrency = (n: number) => `$${n.toFixed(n < 0.01 ? 4 : n < 1 ? 3 : 2)}`;
const formatDate = (date: Date) =>
  `${date.getFullYear()}-${String(date.getMonth() + 1).padStart(2, "0")}-${String(date.getDate()).padStart(2, "0")}`;
const monthlyCost = (history: CostHistory) => Object.values(history.dailyCost).reduce((a, b) => a + b, 0);
const today = () => formatDate(new Date());

const load = (): CostHistory => {
  try {
    const history = JSON.parse(readFileSync(COST_FILE, "utf-8"));
    const dailyCost = history.dailyCost ?? {};
    const resetDay = history.resetDay ?? 1;
    const now = new Date();
    const cutoff = formatDate(
      now.getDate() >= resetDay
        ? new Date(now.getFullYear(), now.getMonth(), resetDay)
        : new Date(now.getFullYear(), now.getMonth() - 1, resetDay),
    );
    for (const key of Object.keys(dailyCost)) {
      if (key < cutoff) delete dailyCost[key];
    }
    return { dailyCost, resetDay };
  } catch {
    return { dailyCost: {}, resetDay: 1 };
  }
};
const save = (history: CostHistory) => writeFileSync(COST_FILE, JSON.stringify(history));

export default function (pi: ExtensionAPI) {
  const status = (ctx: any, history: CostHistory) =>
    ctx.ui.setStatus("cost", `💰 ${formatCurrency(monthlyCost(history))}`);

  pi.on("session_start", async (_e, ctx) => {
    const history = load();
    status(ctx, history);
  });

  pi.on("message_end", async (event, ctx) => {
    if (event.message.role !== "assistant") return;
    const cost = event.message.usage?.cost?.total;
    if (!cost) return;
    const history = load();
    const day = today();
    history.dailyCost[day] = (history.dailyCost[day] ?? 0) + cost;
    save(history);
    status(ctx, history);
  });

  pi.registerCommand("cost", {
    description: "Show monthly and daily cost",
    handler: async (_args, ctx) => {
      const history = load();
      const lines = [
        `Monthly Total: ${formatCurrency(monthlyCost(history))} | Today: ${formatCurrency(history.dailyCost[today()] ?? 0)}`,
        `Reset day: ${history.resetDay}`,
      ];
      const days = Object.entries(history.dailyCost)
        .sort(([a], [b]) => b.localeCompare(a))
        .slice(0, 7);
      if (days.length) lines.push(days.map(([k, c]) => `${k}: ${formatCurrency(c as number)}`).join(" | "));
      ctx.ui.notify(lines.join("\n"), "info");
    },
  });

  pi.registerCommand("cost-reset", {
    description: "Reset cost or change reset day",
    handler: async (_args, ctx) => {
      const action = await ctx.ui.select("Cost reset", ["Change reset day", "Delete all records"]);
      if (!action) return;
      const history = load();
      if (action === "Change reset day") {
        const input = await ctx.ui.input(
          "Reset day",
          `Current: ${history.resetDay}. Enter 1-28 (e.g. 15 = every 15th):`,
        );
        const day = parseInt(input ?? "", 10);
        if (!day || day < 1 || day > 28) return ctx.ui.notify("Day must be 1-28", "error");
        history.resetDay = day;
        save(history);
        status(ctx, history);
        ctx.ui.notify(`Reset day set to ${day}`, "info");
      } else {
        const confirm = await ctx.ui.input(
          "Delete all records",
          'Type "delete" to confirm. This will delete all cost history and reset settings.',
        );
        if (confirm?.toLowerCase() !== "delete") return;
        try {
          unlinkSync(COST_FILE);
        } catch {
          /* file may not exist */
        }
        status(ctx, { dailyCost: {}, resetDay: 1 });
        ctx.ui.notify("All cost records deleted", "info");
      }
    },
  });
}