Agent-JAE/default-extensions/cost-tracker.ts

import type { ExtensionAPI, ExtensionContext } from "@jaeswift/jae-coding-agent";

interface CostState {
  inputTokens: number;
  outputTokens: number;
  cacheReadTokens: number;
  cacheWriteTokens: number;
  totalCost: number;
  budgetLimit: number | null;
  turns: number;
}

const COST_PER_1K_INPUT = 0.003;
const COST_PER_1K_OUTPUT = 0.015;
const COST_PER_1K_CACHE_READ = 0.0003;
const COST_PER_1K_CACHE_WRITE = 0.00375;

export default function (pi: ExtensionAPI) {
  const state: CostState = {
    inputTokens: 0,
    outputTokens: 0,
    cacheReadTokens: 0,
    cacheWriteTokens: 0,
    totalCost: 0,
    budgetLimit: null,
    turns: 0,
  };

  function recalcCost() {
    state.totalCost =
      (state.inputTokens / 1000) * COST_PER_1K_INPUT +
      (state.outputTokens / 1000) * COST_PER_1K_OUTPUT +
      (state.cacheReadTokens / 1000) * COST_PER_1K_CACHE_READ +
      (state.cacheWriteTokens / 1000) * COST_PER_1K_CACHE_WRITE;
  }

  function updateStatus(ctx: ExtensionContext) {
    const cost = state.totalCost.toFixed(4);
    const budget = state.budgetLimit ? ` / $${state.budgetLimit.toFixed(2)}` : "";
    const tokens = ((state.inputTokens + state.outputTokens) / 1000).toFixed(1);
    ctx.ui.setStatus("cost-tracker", `\u{1F4B0} $${cost}${budget} | ${tokens}k tok | ${state.turns} turns`);
  }

  pi.on("turn_start", async (_event, ctx) => {
    state.turns++;
    updateStatus(ctx);
  });

  pi.on("agent_end", async (event, ctx) => {
    for (const message of event.messages) {
      if (message && typeof message === "object" && "usage" in message) {
        const u = (message as any).usage;
        state.inputTokens += u.input || 0;
        state.outputTokens += u.output || 0;
        state.cacheReadTokens += u.cacheRead || 0;
        state.cacheWriteTokens += u.cacheWrite || 0;
      }
    }
    recalcCost();
    updateStatus(ctx);

    if (state.budgetLimit && state.totalCost >= state.budgetLimit) {
      ctx.ui.notify(
        `\u26A0\uFE0F Budget limit of $${state.budgetLimit} reached! Total: $${state.totalCost.toFixed(4)}`,
        "warning",
      );
    }
  });

  pi.registerCommand("cost", {
    description: "Show session cost breakdown",
    handler: async (_args, ctx) => {
      ctx.ui.notify(
        `\u{1F4CA} Cost Report:\n` +
          `  Input:       ${state.inputTokens.toLocaleString()} tokens ($${((state.inputTokens / 1000) * COST_PER_1K_INPUT).toFixed(4)})\n` +
          `  Output:      ${state.outputTokens.toLocaleString()} tokens ($${((state.outputTokens / 1000) * COST_PER_1K_OUTPUT).toFixed(4)})\n` +
          `  Cache Read:  ${state.cacheReadTokens.toLocaleString()} tokens\n` +
          `  Cache Write: ${state.cacheWriteTokens.toLocaleString()} tokens\n` +
          `  Total Cost:  $${state.totalCost.toFixed(4)}\n` +
          `  Turns:       ${state.turns}`,
        "info",
      );
    },
  });

  pi.registerCommand("budget", {
    description: "Set session budget limit: /budget <amount>",
    handler: async (args, ctx) => {
      const amount = parseFloat(args || "");
      if (isNaN(amount) || amount <= 0) {
        ctx.ui.notify(
          state.budgetLimit
            ? `Current budget: $${state.budgetLimit.toFixed(2)} | Spent: $${state.totalCost.toFixed(4)}`
            : "No budget set. Usage: /budget 5.00",
          "info",
        );
      } else {
        state.budgetLimit = amount;
        ctx.ui.notify(`\u{1F4B0} Budget set to $${amount.toFixed(2)}`, "info");
        updateStatus(ctx);
      }
    },
  });
}