@latentminds/pi-quotas 0.2.6 → 0.3.1

package/src/extensions/command-tokens/tokens-display.ts

@@ -0,0 +1,357 @@
+import type { Theme } from "@mariozechner/pi-coding-agent";
+import { DynamicBorder } from "@mariozechner/pi-coding-agent";
+import type { Component } from "@mariozechner/pi-tui";
+import { Loader, matchesKey, truncateToWidth } from "@mariozechner/pi-tui";
+import pkg from "../../../package.json" with { type: "json" };
+import type {
+  AggregateResult,
+} from "../../lib/session-tokens.js";
+import {
+  formatCost,
+  formatNumber,
+  formatTokenSummary,
+} from "../../lib/session-tokens.js";
+
+type TokensState =
+  | { type: "loading" }
+  | { type: "loaded"; result: AggregateResult };
+
+type TabId = "overview" | "models" | "sessions";
+
+const TABS: Array<{ id: TabId; label: string }> = [
+  { id: "overview", label: "Overview" },
+  { id: "models", label: "Models" },
+  { id: "sessions", label: "Sessions" },
+];
+
+function renderProgressBar(
+  value: number,
+  max: number,
+  width: number,
+  theme: Theme,
+): string {
+  if (max <= 0 || value <= 0) return theme.fg("dim", "░".repeat(width));
+  const ratio = Math.min(1, value / max);
+  const filled = Math.round(ratio * width);
+  const parts: string[] = [];
+  for (let i = 0; i < width; i++) {
+    parts.push(i < filled ? theme.fg("accent", "█") : theme.fg("dim", "░"));
+  }
+  return parts.join("");
+}
+
+export class TokensComponent implements Component {
+  private state: TokensState = { type: "loading" };
+  private loader: Loader | null = null;
+  private activeTab: TabId = "overview";
+  private scrollOffset = 0;
+
+  constructor(
+    private theme: Theme,
+    private tui: any,
+    private onClose: () => void,
+    private onRefetch: () => void,
+    private cwd?: string,
+  ) {
+    this.startLoader();
+  }
+
+  private startLoader(): void {
+    this.loader = new Loader(
+      this.tui,
+      (s: string) => this.theme.fg("accent", s),
+      (s: string) => this.theme.fg("muted", s),
+      "Scanning sessions...",
+    );
+  }
+
+  destroy(): void {
+    this.loader?.stop();
+    this.loader = null;
+  }
+
+  setState(state: TokensState): void {
+    if (state.type === "loading") {
+      this.loader?.stop();
+      this.startLoader();
+      this.scrollOffset = 0;
+    } else if (this.state.type === "loading") {
+      this.loader?.stop();
+      this.loader = null;
+    }
+    this.state = state;
+  }
+
+  handleInput(data: string): boolean {
+    if (matchesKey(data, "escape") || data === "q") {
+      this.onClose();
+      return true;
+    }
+    if (data === "r") {
+      this.onRefetch();
+      return true;
+    }
+    if (data === "\t" || data === "right") {
+      const idx = TABS.findIndex((t) => t.id === this.activeTab);
+      this.activeTab = TABS[(idx + 1) % TABS.length].id;
+      this.scrollOffset = 0;
+      this.tui.requestRender();
+      return true;
+    }
+    if (data === "left") {
+      const idx = TABS.findIndex((t) => t.id === this.activeTab);
+      this.activeTab = TABS[(idx - 1 + TABS.length) % TABS.length].id;
+      this.scrollOffset = 0;
+      this.tui.requestRender();
+      return true;
+    }
+    if (data === "1") {
+      this.activeTab = "overview";
+      this.scrollOffset = 0;
+      this.tui.requestRender();
+      return true;
+    }
+    if (data === "2") {
+      this.activeTab = "models";
+      this.scrollOffset = 0;
+      this.tui.requestRender();
+      return true;
+    }
+    if (data === "3") {
+      this.activeTab = "sessions";
+      this.scrollOffset = 0;
+      this.tui.requestRender();
+      return true;
+    }
+    if (matchesKey(data, "down") || data === "j") {
+      this.scrollOffset++;
+      this.tui.requestRender();
+      return true;
+    }
+    if (matchesKey(data, "up") || data === "k") {
+      this.scrollOffset = Math.max(0, this.scrollOffset - 1);
+      this.tui.requestRender();
+      return true;
+    }
+    return false;
+  }
+
+  render(width: number): string[] {
+    const lines: string[] = [];
+    const border = new DynamicBorder((s: string) => this.theme.fg("border", s));
+    lines.push(...border.render(width));
+    lines.push(
+      truncateToWidth(
+        ` ${this.theme.fg("accent", this.theme.bold("Token Usage"))}`,
+        width,
+      ),
+    );
+
+    // Tab bar
+    const tabParts = TABS.map((tab) => {
+      if (tab.id === this.activeTab) {
+        return this.theme.fg("accent", this.theme.bold(`[${tab.label}]`));
+      }
+      return this.theme.fg("dim", ` ${tab.label} `);
+    });
+    lines.push(truncateToWidth(`  ${tabParts.join("  ")}`, width));
+    lines.push("");
+
+    if (this.state.type === "loading") {
+      lines.push(
+        ...(this.loader
+          ? this.loader.render(width)
+          : [this.theme.fg("muted", "  Scanning sessions...")]),
+      );
+    } else {
+      const contentLines = this.renderTab(width);
+      // Apply scroll
+      const maxVisible = Math.max(1, 20); // approximate visible lines
+      const visible = contentLines.slice(
+        this.scrollOffset,
+        this.scrollOffset + maxVisible,
+      );
+      lines.push(...visible);
+    }
+
+    lines.push("");
+    lines.push(
+      this.theme.fg(
+        "dim",
+        `  pi-quotas v${pkg.version}  ·  1/2/3 switch tab  r refresh  q/Esc close`,
+      ),
+    );
+    lines.push(...border.render(width));
+    return lines;
+  }
+
+  private renderTab(width: number): string[] {
+    if (this.state.type !== "loaded") return [];
+    const { result } = this.state;
+
+    switch (this.activeTab) {
+      case "overview":
+        return this.renderOverview(result, width);
+      case "models":
+        return this.renderModels(result, width);
+      case "sessions":
+        return this.renderSessions(result, width);
+    }
+  }
+
+  private renderOverview(result: AggregateResult, width: number): string[] {
+    const lines: string[] = [];
+    const { totals } = result;
+
+    lines.push(
+      truncateToWidth(
+        `  ${this.theme.fg("accent", this.theme.bold("Summary"))}`,
+        width,
+      ),
+    );
+    lines.push("");
+
+    const statRows = [
+      ["Sessions", String(result.sessionCount)],
+      ["Messages", String(result.messageCount)],
+      ["Input tokens", formatNumber(totals.input)],
+      ["Output tokens", formatNumber(totals.output)],
+      ["Cache read", formatNumber(totals.cacheRead)],
+      ["Cache write", formatNumber(totals.cacheWrite)],
+      ["Total tokens", formatNumber(totals.totalTokens)],
+      ["Total cost", formatCost(totals.costTotal)],
+    ];
+
+    const labelWidth = Math.max(...statRows.map(([l]) => l.length));
+    for (const [label, value] of statRows) {
+      lines.push(
+        truncateToWidth(
+          `  ${this.theme.fg("dim", label.padEnd(labelWidth))}  ${this.theme.fg("accent", value)}`,
+          width,
+        ),
+      );
+    }
+
+    // Provider breakdown
+    if (result.byProvider.length > 0) {
+      lines.push("");
+      lines.push(
+        truncateToWidth(
+          `  ${this.theme.fg("accent", this.theme.bold("By Provider"))}`,
+          width,
+        ),
+      );
+      lines.push("");
+      const maxCost = Math.max(
+        ...result.byProvider.map((p) => p.tokens.costTotal),
+      );
+      for (const provider of result.byProvider) {
+        const barWidth = Math.min(20, Math.max(8, width - 40));
+        const bar = renderProgressBar(
+          provider.tokens.costTotal,
+          maxCost,
+          barWidth,
+          this.theme,
+        );
+        const cost = formatCost(provider.tokens.costTotal);
+        const msgs = `${provider.messageCount} msgs`;
+        lines.push(
+          truncateToWidth(
+            `  ${this.theme.fg("accent", provider.provider.padEnd(20))} ${bar} ${this.theme.fg("accent", cost)} ${this.theme.fg("dim", msgs)}`,
+            width,
+          ),
+        );
+      }
+    }
+
+    return lines;
+  }
+
+  private renderModels(result: AggregateResult, width: number): string[] {
+    const lines: string[] = [];
+
+    if (result.byModel.length === 0) {
+      lines.push(this.theme.fg("dim", "  No assistant messages found"));
+      return lines;
+    }
+
+    const maxCost = Math.max(...result.byModel.map((m) => m.tokens.costTotal));
+
+    for (const model of result.byModel) {
+      const barWidth = Math.min(16, Math.max(8, width - 50));
+      const bar = renderProgressBar(
+        model.tokens.costTotal,
+        maxCost,
+        barWidth,
+        this.theme,
+      );
+      const cost = formatCost(model.tokens.costTotal);
+      const label =
+        model.provider === model.model
+          ? model.model
+          : `${model.provider}/${model.model}`;
+      const shortLabel = label.length > 30 ? label.slice(0, 27) + "..." : label;
+
+      lines.push(
+        truncateToWidth(`  ${this.theme.fg("accent", shortLabel)}`, width),
+      );
+      lines.push(
+        truncateToWidth(
+          `    ${bar} ${this.theme.fg("accent", cost)}  ${this.theme.fg("dim", formatTokenSummary(model.tokens))}`,
+          width,
+        ),
+      );
+    }
+
+    return lines;
+  }
+
+  private renderSessions(result: AggregateResult, width: number): string[] {
+    const lines: string[] = [];
+
+    if (result.bySession.length === 0) {
+      lines.push(this.theme.fg("dim", "  No sessions found"));
+      return lines;
+    }
+
+    const maxCost = Math.max(
+      ...result.bySession.map((s) => s.tokens.costTotal),
+    );
+
+    for (const session of result.bySession) {
+      const dateStr = session.created.toISOString().slice(0, 10);
+      const name = session.name ?? session.sessionId.slice(0, 8);
+      const cost = formatCost(session.tokens.costTotal);
+      const msgs = `${session.messageCount} msgs`;
+
+      const barWidth = Math.min(12, Math.max(6, width - 50));
+      const bar = renderProgressBar(
+        session.tokens.costTotal,
+        maxCost,
+        barWidth,
+        this.theme,
+      );
+
+      lines.push(
+        truncateToWidth(
+          `  ${this.theme.fg("dim", dateStr)} ${this.theme.fg("accent", name)} ${bar} ${this.theme.fg("accent", cost)} ${this.theme.fg("dim", msgs)}`,
+          width,
+        ),
+      );
+    }
+
+    if (result.sessionCount > result.bySession.length) {
+      lines.push("");
+      lines.push(
+        this.theme.fg(
+          "dim",
+          `  Showing ${result.bySession.length} of ${result.sessionCount} sessions`,
+        ),
+      );
+    }
+
+    return lines;
+  }
+
+  invalidate(): void {}
+}

package/src/extensions/quota-warnings/index.ts

@@ -1,4 +1,7 @@
-import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+} from "@mariozechner/pi-coding-agent";
 import {
   QUOTAS_CONFIG_UPDATED_EVENT,
   QUOTAS_EXTENSIONS_REGISTER_EVENT,
@@ -6,7 +9,11 @@ import {
   type QuotasConfigUpdatedPayload,
   configLoader,
 } from "../../config.js";
-import { fetchProviderQuotas, isSupportedProvider } from "../../lib/quotas.js";
+import {
+  fetchProviderQuotas,
+  isSupportedProvider,
+  PROVIDER_LABELS,
+} from "../../lib/quotas.js";
 import {
   assessWindow,
   formatTimeRemaining,
@@ -24,9 +31,11 @@ function shouldNotify(key: string, severity: RiskSeverity): boolean {
   const current = alertState.get(key);
   if (!current) return true;
   const order: RiskSeverity[] = ["none", "warning", "high", "critical"];
-  if (order.indexOf(severity) > order.indexOf(current.lastSeverity)) return true;
+  if (order.indexOf(severity) > order.indexOf(current.lastSeverity))
+    return true;
   if (severity === "high" || severity === "critical") return true;
-  if (severity === "warning") return Date.now() - current.lastNotifiedAt >= COOLDOWN_MS;
+  if (severity === "warning")
+    return Date.now() - current.lastNotifiedAt >= COOLDOWN_MS;
   return false;
 }
 
@@ -50,7 +59,10 @@ export default async function (pi: ExtensionAPI) {
     if (onlyNew && now - lastFetchAt < MIN_FETCH_INTERVAL_MS) return;
     lastFetchAt = now;
 
-    const result = await fetchProviderQuotas(ctx.modelRegistry.authStorage, provider);
+    const result = await fetchProviderQuotas(
+      ctx.modelRegistry.authStorage,
+      provider,
+    );
     if (!result.success) return;
 
     const risky = result.data.windows
@@ -59,19 +71,23 @@ export default async function (pi: ExtensionAPI) {
     if (risky.length === 0) return;
 
     const toNotify = onlyNew
-      ? risky.filter((entry) => shouldNotify(`${provider}:${entry.window.label}`, entry.assessment.severity))
+      ? risky.filter((entry) =>
+        shouldNotify(
+          `${provider}:${entry.window.label}`,
+          entry.assessment.severity,
+        ),
+      )
       : risky;
     if (toNotify.length === 0) return;
 
     for (const entry of toNotify) {
-      markNotified(`${provider}:${entry.window.label}`, entry.assessment.severity);
+      markNotified(
+        `${provider}:${entry.window.label}`,
+        entry.assessment.severity,
+      );
     }
 
-    const providerName = provider === "openai-codex"
-      ? "Codex"
-      : provider === "github-copilot"
-        ? "GitHub Copilot"
-        : "Anthropic";
+    const providerName = PROVIDER_LABELS[provider];
 
     const lines = toNotify.map(({ window, assessment }) => {
       const projected = Math.round(assessment.projectedPercent);
@@ -79,7 +95,11 @@ export default async function (pi: ExtensionAPI) {
       return `- ${window.label}: ${used}% used, projected ${projected}% (${assessment.severity}), resets in ${formatTimeRemaining(window.resetsAt)}`;
     });
 
-    const level = toNotify.some((entry) => entry.assessment.severity === "critical" || entry.assessment.severity === "high")
+    const level = toNotify.some(
+      (entry) =>
+        entry.assessment.severity === "critical" ||
+        entry.assessment.severity === "high",
+    )
       ? "error"
       : "warning";
     ctx.ui.notify(`${providerName} quota warning:\n${lines.join("\n")}`, level);
@@ -128,7 +148,9 @@ export default async function (pi: ExtensionAPI) {
 
   pi.events.on(QUOTAS_EXTENSIONS_REQUEST_EVENT, () => {
     if (configLoader.getConfig().quotaWarnings) {
-      pi.events.emit(QUOTAS_EXTENSIONS_REGISTER_EVENT, { feature: "quotaWarnings" });
+      pi.events.emit(QUOTAS_EXTENSIONS_REGISTER_EVENT, {
+        feature: "quotaWarnings",
+      });
     }
   });
 }

package/src/extensions/token-status/index.ts

@@ -0,0 +1,241 @@
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+  ThemeColor,
+} from "@mariozechner/pi-coding-agent";
+import {
+  QUOTAS_CONFIG_UPDATED_EVENT,
+  QUOTAS_EXTENSIONS_REGISTER_EVENT,
+  QUOTAS_EXTENSIONS_REQUEST_EVENT,
+  type QuotasConfigUpdatedPayload,
+  configLoader,
+} from "../../config.js";
+import { aggregateAllSessions, formatCost } from "../../lib/session-tokens.js";
+
+const EXTENSION_ID = "pi-quotas-token-status";
+const REFRESH_INTERVAL_MS = 60_000;
+
+/** Go tier limits (approximate, from docs) */
+const GO_LIMITS = {
+  rolling5h: 12, // $12 per 5 hours
+  weekly: 30, // $30 per week
+  monthly: 60, // $60 per month
+} as const;
+
+interface RollingWindowCosts {
+  rolling5h: number;
+  weekly: number;
+  monthly: number;
+}
+
+function computeWindowTimestamps(now: Date): {
+  since5h: Date;
+  sinceWeekly: Date;
+  sinceMonthly: Date;
+} {
+  return {
+    since5h: new Date(now.getTime() - 5 * 60 * 60 * 1000),
+    sinceWeekly: new Date(now.getTime() - 7 * 24 * 60 * 60 * 1000),
+    sinceMonthly: new Date(now.getTime() - 30 * 24 * 60 * 60 * 1000),
+  };
+}
+
+async function computeRollingCosts(cwd?: string): Promise<RollingWindowCosts> {
+  const now = new Date();
+  const { since5h, sinceWeekly, sinceMonthly } = computeWindowTimestamps(now);
+
+  // Run all three aggregations in parallel
+  const [result5h, resultWeekly, resultMonthly] = await Promise.all([
+    aggregateAllSessions({ cwd, since: since5h }),
+    aggregateAllSessions({ cwd, since: sinceWeekly }),
+    aggregateAllSessions({ cwd, since: sinceMonthly }),
+  ]);
+
+  return {
+    rolling5h: result5h.totals.costTotal,
+    weekly: resultWeekly.totals.costTotal,
+    monthly: resultMonthly.totals.costTotal,
+  };
+}
+
+function formatWindowForStatus(
+  theme: ExtensionContext["ui"]["theme"],
+  label: string,
+  cost: number,
+  limit: number,
+): string {
+  const percent = limit > 0 ? Math.round((cost / limit) * 100) : 0;
+  const costStr = formatCost(cost);
+
+  // Color by usage level
+  let color: ThemeColor;
+  if (percent >= 90) {
+    color = "error";
+  } else if (percent >= 70) {
+    color = "warning";
+  } else {
+    color = "dim";
+  }
+
+  return `${theme.fg(color, `${label}:`)}${theme.fg("accent", costStr)}`;
+}
+
+function formatTokenStatus(
+  theme: ExtensionContext["ui"]["theme"],
+  costs: RollingWindowCosts,
+): string {
+  const parts = [
+    formatWindowForStatus(theme, "5h", costs.rolling5h, GO_LIMITS.rolling5h),
+    formatWindowForStatus(theme, "wk", costs.weekly, GO_LIMITS.weekly),
+    formatWindowForStatus(theme, "mo", costs.monthly, GO_LIMITS.monthly),
+  ];
+  return parts.join(" · ");
+}
+
+function isGoProvider(provider: string | undefined): boolean {
+  if (!provider) return false;
+  return provider === "opencode-go" || provider.startsWith("opencode-go/");
+}
+
+function createTokenStatusRefresher() {
+  let refreshTimer: ReturnType<typeof setInterval> | undefined;
+  let activeContext: ExtensionContext | undefined;
+  let lastCosts: RollingWindowCosts | undefined;
+  let inFlight = false;
+  let queued = false;
+
+  async function update(ctx: ExtensionContext): Promise<void> {
+    if (!ctx.hasUI) return;
+    if (inFlight) {
+      queued = true;
+      return;
+    }
+    inFlight = true;
+    try {
+      const costs = await computeRollingCosts(ctx.cwd);
+      if (!ctx.hasUI) return;
+      lastCosts = costs;
+      const status = formatTokenStatus(ctx.ui.theme, costs);
+      ctx.ui.setStatus(EXTENSION_ID, status);
+    } catch {
+      ctx.ui.setStatus(
+        EXTENSION_ID,
+        ctx.ui.theme.fg("warning", "token tracking unavailable"),
+      );
+    } finally {
+      inFlight = false;
+      if (queued) {
+        queued = false;
+        void update(ctx);
+      }
+    }
+  }
+
+  return {
+    async refreshFor(ctx: ExtensionContext): Promise<void> {
+      activeContext = ctx;
+      if (!isGoProvider(ctx.model?.provider)) {
+        ctx.ui.setStatus(EXTENSION_ID, undefined);
+        return;
+      }
+      await update(ctx);
+    },
+    start(): void {
+      if (refreshTimer) clearInterval(refreshTimer);
+      refreshTimer = setInterval(() => {
+        if (activeContext) void update(activeContext);
+      }, REFRESH_INTERVAL_MS);
+      refreshTimer.unref?.();
+    },
+    stop(ctx?: ExtensionContext): void {
+      if (refreshTimer) clearInterval(refreshTimer);
+      refreshTimer = undefined;
+      activeContext = undefined;
+      lastCosts = undefined;
+      ctx?.ui.setStatus(EXTENSION_ID, undefined);
+    },
+    renderLast(ctx: ExtensionContext): boolean {
+      if (!lastCosts || !ctx.hasUI) return false;
+      ctx.ui.setStatus(
+        EXTENSION_ID,
+        formatTokenStatus(ctx.ui.theme, lastCosts),
+      );
+      return true;
+    },
+  };
+}
+
+export default async function (pi: ExtensionAPI) {
+  await configLoader.load();
+  const refresher = createTokenStatusRefresher();
+  let enabled = configLoader.getConfig().tokenStatus;
+  let currentContext: ExtensionContext | undefined;
+
+  function scheduleRefresh(ctx: ExtensionContext): void {
+    void refresher.refreshFor(ctx).catch(() => {
+      if (ctx.hasUI)
+        ctx.ui.setStatus(
+          EXTENSION_ID,
+          ctx.ui.theme.fg("warning", "token tracking unavailable"),
+        );
+    });
+  }
+
+  pi.events.on(QUOTAS_CONFIG_UPDATED_EVENT, (data: unknown) => {
+    const config = (data as QuotasConfigUpdatedPayload).config;
+    enabled = config.tokenStatus;
+    if (!enabled) {
+      refresher.stop(currentContext);
+      return;
+    }
+    if (currentContext) {
+      refresher.start();
+      scheduleRefresh(currentContext);
+    }
+  });
+
+  pi.on("session_start", (_event, ctx) => {
+    currentContext = ctx;
+    if (!enabled) return;
+    if (!isGoProvider(ctx.model?.provider)) {
+      ctx.ui.setStatus(EXTENSION_ID, undefined);
+      return;
+    }
+    refresher.start();
+    scheduleRefresh(ctx);
+  });
+
+  pi.on("turn_end", (_event, ctx) => {
+    currentContext = ctx;
+    if (!enabled) return;
+    if (!isGoProvider(ctx.model?.provider)) return;
+    scheduleRefresh(ctx);
+  });
+
+  pi.on("model_select", (_event, ctx) => {
+    currentContext = ctx;
+    if (!enabled) {
+      refresher.stop(ctx);
+      return;
+    }
+    if (!isGoProvider(ctx.model?.provider)) {
+      refresher.stop(ctx);
+      return;
+    }
+    refresher.start();
+    scheduleRefresh(ctx);
+  });
+
+  pi.on("session_shutdown", async (_event, ctx) => {
+    currentContext = undefined;
+    refresher.stop(ctx);
+  });
+
+  pi.events.on(QUOTAS_EXTENSIONS_REQUEST_EVENT, () => {
+    if (configLoader.getConfig().tokenStatus) {
+      pi.events.emit(QUOTAS_EXTENSIONS_REGISTER_EVENT, {
+        feature: "tokenStatus",
+      });
+    }
+  });
+}