@kata-sh/cli 0.1.0 → 0.1.2

package/src/resources/extensions/kata/metrics.ts

@@ -0,0 +1,397 @@
+/**
+ * Kata Metrics — Token & Cost Tracking
+ *
+ * Accumulates per-unit usage data across auto-mode sessions.
+ * Data is extracted from session entries before each context wipe,
+ * written to .kata/metrics.json, and surfaced in the dashboard.
+ *
+ * Data flow:
+ *   1. Before newSession() wipes context, snapshotUnitMetrics() scans
+ *      session entries for AssistantMessage usage data
+ *   2. The unit record is appended to the in-memory ledger and flushed to disk
+ *   3. The dashboard overlay and progress widget read from the in-memory ledger
+ *   4. On crash recovery or fresh start, the ledger is loaded from disk
+ */
+
+import { readFileSync, writeFileSync, mkdirSync } from "node:fs";
+import { join } from "node:path";
+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
+import { kataRoot } from "./paths.js";
+
+// ─── Types ────────────────────────────────────────────────────────────────────
+
+export interface TokenCounts {
+  input: number;
+  output: number;
+  cacheRead: number;
+  cacheWrite: number;
+  total: number;
+}
+
+export interface UnitMetrics {
+  type: string; // e.g. "research-milestone", "execute-task"
+  id: string; // e.g. "M001/S01/T01"
+  model: string; // model ID used
+  startedAt: number; // ms timestamp
+  finishedAt: number; // ms timestamp
+  tokens: TokenCounts;
+  cost: number; // total USD cost
+  toolCalls: number;
+  assistantMessages: number;
+  userMessages: number;
+}
+
+export interface MetricsLedger {
+  version: 1;
+  projectStartedAt: number;
+  units: UnitMetrics[];
+}
+
+// ─── Phase classification ─────────────────────────────────────────────────────
+
+export type MetricsPhase =
+  | "research"
+  | "planning"
+  | "execution"
+  | "completion"
+  | "reassessment";
+
+export function classifyUnitPhase(unitType: string): MetricsPhase {
+  switch (unitType) {
+    case "research-milestone":
+    case "research-slice":
+      return "research";
+    case "plan-milestone":
+    case "plan-slice":
+      return "planning";
+    case "execute-task":
+      return "execution";
+    case "complete-slice":
+      return "completion";
+    case "reassess-roadmap":
+      return "reassessment";
+    default:
+      return "execution";
+  }
+}
+
+// ─── In-memory state ──────────────────────────────────────────────────────────
+
+let ledger: MetricsLedger | null = null;
+let basePath: string = "";
+
+// ─── Public API ───────────────────────────────────────────────────────────────
+
+/**
+ * Initialize the metrics system for a given project.
+ * Loads existing ledger from disk if present.
+ */
+export function initMetrics(base: string): void {
+  basePath = base;
+  ledger = loadLedger(base);
+}
+
+/**
+ * Reset in-memory state. Called when auto-mode stops.
+ */
+export function resetMetrics(): void {
+  ledger = null;
+  basePath = "";
+}
+
+/**
+ * Snapshot usage metrics from the current session before it's wiped.
+ * Scans session entries for AssistantMessage usage data.
+ */
+export function snapshotUnitMetrics(
+  ctx: ExtensionContext,
+  unitType: string,
+  unitId: string,
+  startedAt: number,
+  model: string,
+): UnitMetrics | null {
+  if (!ledger) return null;
+
+  const entries = ctx.sessionManager.getEntries();
+  if (!entries || entries.length === 0) return null;
+
+  const tokens: TokenCounts = {
+    input: 0,
+    output: 0,
+    cacheRead: 0,
+    cacheWrite: 0,
+    total: 0,
+  };
+  let cost = 0;
+  let toolCalls = 0;
+  let assistantMessages = 0;
+  let userMessages = 0;
+
+  for (const entry of entries) {
+    if (entry.type !== "message") continue;
+    const msg = (entry as any).message;
+    if (!msg) continue;
+
+    if (msg.role === "assistant") {
+      assistantMessages++;
+      if (msg.usage) {
+        tokens.input += msg.usage.input ?? 0;
+        tokens.output += msg.usage.output ?? 0;
+        tokens.cacheRead += msg.usage.cacheRead ?? 0;
+        tokens.cacheWrite += msg.usage.cacheWrite ?? 0;
+        tokens.total += msg.usage.totalTokens ?? 0;
+        if (msg.usage.cost) {
+          cost += msg.usage.cost.total ?? 0;
+        }
+      }
+      // Count tool calls in this message
+      if (msg.content && Array.isArray(msg.content)) {
+        for (const block of msg.content) {
+          if (block.type === "tool_call") toolCalls++;
+        }
+      }
+    } else if (msg.role === "user") {
+      userMessages++;
+    }
+  }
+
+  const unit: UnitMetrics = {
+    type: unitType,
+    id: unitId,
+    model,
+    startedAt,
+    finishedAt: Date.now(),
+    tokens,
+    cost,
+    toolCalls,
+    assistantMessages,
+    userMessages,
+  };
+
+  ledger.units.push(unit);
+  saveLedger(basePath, ledger);
+
+  return unit;
+}
+
+/**
+ * Get the current ledger (read-only).
+ */
+export function getLedger(): MetricsLedger | null {
+  return ledger;
+}
+
+// ─── Aggregation helpers ──────────────────────────────────────────────────────
+
+export interface PhaseAggregate {
+  phase: MetricsPhase;
+  units: number;
+  tokens: TokenCounts;
+  cost: number;
+  duration: number; // ms
+}
+
+export interface SliceAggregate {
+  sliceId: string;
+  units: number;
+  tokens: TokenCounts;
+  cost: number;
+  duration: number;
+}
+
+export interface ModelAggregate {
+  model: string;
+  units: number;
+  tokens: TokenCounts;
+  cost: number;
+}
+
+export interface ProjectTotals {
+  units: number;
+  tokens: TokenCounts;
+  cost: number;
+  duration: number;
+  toolCalls: number;
+  assistantMessages: number;
+  userMessages: number;
+}
+
+function emptyTokens(): TokenCounts {
+  return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 };
+}
+
+function addTokens(a: TokenCounts, b: TokenCounts): TokenCounts {
+  return {
+    input: a.input + b.input,
+    output: a.output + b.output,
+    cacheRead: a.cacheRead + b.cacheRead,
+    cacheWrite: a.cacheWrite + b.cacheWrite,
+    total: a.total + b.total,
+  };
+}
+
+export function aggregateByPhase(units: UnitMetrics[]): PhaseAggregate[] {
+  const map = new Map<MetricsPhase, PhaseAggregate>();
+  for (const u of units) {
+    const phase = classifyUnitPhase(u.type);
+    let agg = map.get(phase);
+    if (!agg) {
+      agg = { phase, units: 0, tokens: emptyTokens(), cost: 0, duration: 0 };
+      map.set(phase, agg);
+    }
+    agg.units++;
+    agg.tokens = addTokens(agg.tokens, u.tokens);
+    agg.cost += u.cost;
+    agg.duration += u.finishedAt - u.startedAt;
+  }
+  // Return in a stable order
+  const order: MetricsPhase[] = [
+    "research",
+    "planning",
+    "execution",
+    "completion",
+    "reassessment",
+  ];
+  return order.map((p) => map.get(p)).filter((a): a is PhaseAggregate => !!a);
+}
+
+export function aggregateBySlice(units: UnitMetrics[]): SliceAggregate[] {
+  const map = new Map<string, SliceAggregate>();
+  for (const u of units) {
+    const parts = u.id.split("/");
+    // Slice ID is parts[0]/parts[1] if it exists, else parts[0]
+    const sliceId = parts.length >= 2 ? `${parts[0]}/${parts[1]}` : parts[0];
+    let agg = map.get(sliceId);
+    if (!agg) {
+      agg = { sliceId, units: 0, tokens: emptyTokens(), cost: 0, duration: 0 };
+      map.set(sliceId, agg);
+    }
+    agg.units++;
+    agg.tokens = addTokens(agg.tokens, u.tokens);
+    agg.cost += u.cost;
+    agg.duration += u.finishedAt - u.startedAt;
+  }
+  return Array.from(map.values()).sort((a, b) =>
+    a.sliceId.localeCompare(b.sliceId),
+  );
+}
+
+export function aggregateByModel(units: UnitMetrics[]): ModelAggregate[] {
+  const map = new Map<string, ModelAggregate>();
+  for (const u of units) {
+    let agg = map.get(u.model);
+    if (!agg) {
+      agg = { model: u.model, units: 0, tokens: emptyTokens(), cost: 0 };
+      map.set(u.model, agg);
+    }
+    agg.units++;
+    agg.tokens = addTokens(agg.tokens, u.tokens);
+    agg.cost += u.cost;
+  }
+  return Array.from(map.values()).sort((a, b) => b.cost - a.cost);
+}
+
+export function getProjectTotals(units: UnitMetrics[]): ProjectTotals {
+  const totals: ProjectTotals = {
+    units: units.length,
+    tokens: emptyTokens(),
+    cost: 0,
+    duration: 0,
+    toolCalls: 0,
+    assistantMessages: 0,
+    userMessages: 0,
+  };
+  for (const u of units) {
+    totals.tokens = addTokens(totals.tokens, u.tokens);
+    totals.cost += u.cost;
+    totals.duration += u.finishedAt - u.startedAt;
+    totals.toolCalls += u.toolCalls;
+    totals.assistantMessages += u.assistantMessages;
+    totals.userMessages += u.userMessages;
+  }
+  return totals;
+}
+
+// ─── Formatting helpers ───────────────────────────────────────────────────────
+
+export function formatCost(cost: number): string {
+  if (cost < 0.01) return `$${cost.toFixed(4)}`;
+  if (cost < 1) return `$${cost.toFixed(3)}`;
+  return `$${cost.toFixed(2)}`;
+}
+
+/**
+ * Compute a projected remaining cost based on completed slice averages.
+ *
+ * Filters to slice-level entries (sliceId contains "/") to exclude bare milestone
+ * aggregates from the average. Returns [] when fewer than 2 slice-level entries
+ * exist (insufficient data for a reliable projection).
+ *
+ * If `budgetCeiling` is provided and `totalCost >= budgetCeiling`, a warning line
+ * is appended to the result.
+ */
+export function formatCostProjection(
+  completedSlices: SliceAggregate[],
+  remainingCount: number,
+  budgetCeiling?: number,
+): string[] {
+  const sliceLevel = completedSlices.filter((s) => s.sliceId.includes("/"));
+  if (sliceLevel.length < 2) return [];
+
+  const totalCost = sliceLevel.reduce((sum, s) => sum + s.cost, 0);
+  const avgCost = totalCost / sliceLevel.length;
+  const projected = avgCost * remainingCount;
+
+  const projLine = `Projected remaining: ${formatCost(projected)} (${formatCost(avgCost)}/slice avg × ${remainingCount} remaining)`;
+  const result: string[] = [projLine];
+
+  if (budgetCeiling !== undefined && totalCost >= budgetCeiling) {
+    result.push(
+      `Budget ceiling ${formatCost(budgetCeiling)} reached (spent ${formatCost(totalCost)})`,
+    );
+  }
+
+  return result;
+}
+
+export function formatTokenCount(count: number): string {
+  if (count < 1000) return `${count}`;
+  if (count < 1_000_000) return `${(count / 1000).toFixed(1)}k`;
+  return `${(count / 1_000_000).toFixed(2)}M`;
+}
+
+// ─── Disk I/O ─────────────────────────────────────────────────────────────────
+
+function metricsPath(base: string): string {
+  return join(kataRoot(base), "metrics.json");
+}
+
+function loadLedger(base: string): MetricsLedger {
+  try {
+    const raw = readFileSync(metricsPath(base), "utf-8");
+    const parsed = JSON.parse(raw);
+    if (parsed.version === 1 && Array.isArray(parsed.units)) {
+      return parsed as MetricsLedger;
+    }
+  } catch {
+    // File doesn't exist or is corrupt — start fresh
+  }
+  return {
+    version: 1,
+    projectStartedAt: Date.now(),
+    units: [],
+  };
+}
+
+function saveLedger(base: string, data: MetricsLedger): void {
+  try {
+    mkdirSync(kataRoot(base), { recursive: true });
+    writeFileSync(
+      metricsPath(base),
+      JSON.stringify(data, null, 2) + "\n",
+      "utf-8",
+    );
+  } catch {
+    // Don't let metrics failures break auto-mode
+  }
+}