@aaroncarry/pi-usage 0.1.0

package/src/credentials.ts

@@ -0,0 +1,217 @@
+/**
+ * Credential resolution for pi-usage.
+ *
+ * Prefers pi's model registry (`ctx.modelRegistry.getProviderAuth`), which
+ * refreshes OAuth tokens before they expire and persists them to auth.json.
+ * The auth.json fallback covers providers the registry does not know; it
+ * mirrors pi's config value interpolation rules ($ENV / ${ENV}, $$ and $!
+ * escapes, !command).
+ */
+
+import { execSync } from "node:child_process";
+import { readFileSync } from "node:fs";
+import { homedir } from "node:os";
+import { join } from "node:path";
+
+/** Location of pi's agent directory (mirrors pi's config.ts rules). */
+export function getAgentDir(): string {
+	return process.env.PI_CODING_AGENT_DIR || join(homedir(), ".pi", "agent");
+}
+
+export interface StoredApiKeyCredential {
+	type: "api_key";
+	key?: string;
+	env?: Record<string, string>;
+}
+
+export interface StoredOAuthCredential {
+	type: "oauth";
+	access?: string;
+	refresh?: string;
+	expires?: number;
+}
+
+export type StoredCredential = StoredApiKeyCredential | StoredOAuthCredential;
+
+/** Every provider id that has a recognizable credential in auth.json. */
+export function readStoredCredentialIds(agentDir: string): string[] {
+	let raw: unknown;
+	try {
+		raw = JSON.parse(readFileSync(join(agentDir, "auth.json"), "utf8"));
+	} catch {
+		return [];
+	}
+	if (typeof raw !== "object" || raw === null || Array.isArray(raw)) return [];
+	return Object.entries(raw as Record<string, unknown>)
+		.filter(([, value]) => typeof value === "object" && value !== null && !Array.isArray(value))
+		.map(([id]) => id);
+}
+
+export function readStoredCredential(agentDir: string, providerId: string): StoredCredential | undefined {
+	let raw: unknown;
+	try {
+		raw = JSON.parse(readFileSync(join(agentDir, "auth.json"), "utf8"));
+	} catch {
+		return undefined;
+	}
+	if (typeof raw !== "object" || raw === null || Array.isArray(raw)) return undefined;
+	const entries = raw as Record<string, unknown>;
+	const credential = entries[providerId];
+	if (typeof credential !== "object" || credential === null) return undefined;
+	const record = credential as Record<string, unknown>;
+	if (record.type === "api_key") {
+		return {
+			type: "api_key",
+			key: typeof record.key === "string" ? record.key : undefined,
+			env: readEnvRecord(record.env),
+		};
+	}
+	if (record.type === "oauth") {
+		return {
+			type: "oauth",
+			access: typeof record.access === "string" ? record.access : undefined,
+			refresh: typeof record.refresh === "string" ? record.refresh : undefined,
+			expires: typeof record.expires === "number" ? record.expires : undefined,
+		};
+	}
+	return undefined;
+}
+
+/** Subset of pi-ai's AuthResult that pi-usage relies on. */
+export interface RegistryAuth {
+	auth: { apiKey?: string; headers?: Record<string, string> };
+}
+
+export type ProviderAuthLookup = (providerId: string) => Promise<RegistryAuth | undefined>;
+
+const BEARER_RE = /^Bearer\s+(.+)$/iu;
+
+/** Extract the bearer token/api key from a registry auth result. */
+export function tokenFromRegistryAuth(auth: RegistryAuth | undefined): string | undefined {
+	if (!auth) return undefined;
+	if (auth.auth.apiKey) return auth.auth.apiKey;
+	const authorization = Object.entries(auth.auth.headers ?? {}).find(
+		([name]) => name.toLowerCase() === "authorization",
+	)?.[1];
+	return typeof authorization === "string" ? BEARER_RE.exec(authorization)?.[1] : undefined;
+}
+
+/** OAuth tokens this close to expiry are considered stale in the auth.json fallback. */
+const OAUTH_FALLBACK_MIN_VALIDITY_MS = 60_000;
+
+/**
+ * Resolve a usable token for a provider. Throws with an actionable message
+ * when no credential is configured or the fallback OAuth token is expired;
+ * the service turns the message into the account's `error`.
+ */
+export async function resolveProviderToken(
+	providerId: string,
+	lookup: ProviderAuthLookup | undefined,
+	agentDir: string,
+): Promise<string> {
+	try {
+		const fromRegistry = tokenFromRegistryAuth(await lookup?.(providerId));
+		if (fromRegistry) return fromRegistry;
+	} catch {
+		// Registry lookup failed (unknown provider, refresh error); fall through.
+	}
+	const credential = readStoredCredential(agentDir, providerId);
+	if (!credential) {
+		throw new Error(`No credential configured for "${providerId}" (run /login or pi auth)`);
+	}
+	if (credential.type === "api_key") {
+		const key = credential.key ? resolveConfigValue(credential.key, credential.env) : undefined;
+		if (key) return key;
+		throw new Error(`Could not resolve API key for "${providerId}"`);
+	}
+	if (!credential.access) {
+		throw new Error(`OAuth credential for "${providerId}" has no access token`);
+	}
+	if (credential.expires !== undefined && Date.now() > credential.expires - OAUTH_FALLBACK_MIN_VALIDITY_MS) {
+		throw new Error(`OAuth token for "${providerId}" is expired; run /login or "pi auth check" to refresh`);
+	}
+	return credential.access;
+}
+
+/**
+ * Resolve a config value following pi's resolve-config-value rules:
+ * "!command" runs through the shell, "$ENV"/"${ENV}" interpolate, "$$"/"$!"
+ * escape literals, anything else is literal.
+ */
+export function resolveConfigValue(config: string, env?: Record<string, string>): string | undefined {
+	if (config.startsWith("!")) return runCommand(config.slice(1));
+	return interpolateTemplate(config, env);
+}
+
+const ENV_NAME_RE = /^[A-Za-z_][A-Za-z0-9_]*$/;
+const ENV_NAME_PREFIX_RE = /^[A-Za-z_][A-Za-z0-9_]*/;
+
+function interpolateTemplate(config: string, env?: Record<string, string>): string | undefined {
+	let out = "";
+	let index = 0;
+	while (index < config.length) {
+		const dollar = config.indexOf("$", index);
+		if (dollar < 0) {
+			out += config.slice(index);
+			break;
+		}
+		out += config.slice(index, dollar);
+		const next = config[dollar + 1];
+		if (next === "$" || next === "!") {
+			out += next;
+			index = dollar + 2;
+			continue;
+		}
+		if (next === "{") {
+			const end = config.indexOf("}", dollar + 2);
+			if (end < 0) {
+				out += "$";
+				index = dollar + 1;
+				continue;
+			}
+			const name = config.slice(dollar + 2, end);
+			const value = ENV_NAME_RE.test(name) ? envValue(name, env) : undefined;
+			if (value === undefined) return undefined;
+			out += value;
+			index = end + 1;
+			continue;
+		}
+		const match = ENV_NAME_PREFIX_RE.exec(config.slice(dollar + 1));
+		if (match?.[0]) {
+			const value = envValue(match[0], env);
+			if (value === undefined) return undefined;
+			out += value;
+			index = dollar + 1 + match[0].length;
+			continue;
+		}
+		out += "$";
+		index = dollar + 1;
+	}
+	return out;
+}
+
+function envValue(name: string, env?: Record<string, string>): string | undefined {
+	return env?.[name] || process.env[name] || undefined;
+}
+
+function runCommand(command: string): string | undefined {
+	try {
+		const output = execSync(command, {
+			encoding: "utf8",
+			timeout: 10_000,
+			stdio: ["ignore", "pipe", "ignore"],
+		});
+		return output.trim() || undefined;
+	} catch {
+		return undefined;
+	}
+}
+
+function readEnvRecord(value: unknown): Record<string, string> | undefined {
+	if (typeof value !== "object" || value === null || Array.isArray(value)) return undefined;
+	const out: Record<string, string> = {};
+	for (const [name, envValueEntry] of Object.entries(value)) {
+		if (typeof envValueEntry === "string") out[name] = envValueEntry;
+	}
+	return Object.keys(out).length > 0 ? out : undefined;
+}

package/src/format.ts

@@ -0,0 +1,35 @@
+/** Shared display formatting for panel and status line. */
+
+import type { MoneyBalance } from "./types.ts";
+
+const CURRENCY_SYMBOLS: Record<string, string> = { CNY: "¥", USD: "$", EUR: "€" };
+
+export function formatMoney(balance: MoneyBalance): string {
+	const symbol = CURRENCY_SYMBOLS[balance.currency] ?? (balance.currency ? `${balance.currency} ` : "");
+	return `${symbol}${balance.amount.toFixed(2)}`;
+}
+
+export function formatBar(percent: number, cells = 10): string {
+	const clamped = Math.max(0, Math.min(100, percent));
+	const filled = Math.round((clamped / 100) * cells);
+	return "█".repeat(filled) + "░".repeat(cells - filled);
+}
+
+/** Human-readable reset countdown: "now", "16m", "5h 12m", "6d 18h". */
+export function formatResetDuration(resetsAt: number, now = Date.now()): string {
+	if (resetsAt - now <= 0) return "now";
+	const totalMinutes = Math.max(1, Math.ceil((resetsAt - now) / 60_000));
+	const days = Math.floor(totalMinutes / 1_440);
+	const hours = Math.floor(totalMinutes / 60) % 24;
+	const minutes = totalMinutes % 60;
+	if (days > 0) return hours > 0 ? `${days}d ${hours}h` : `${days}d`;
+	if (hours > 0) return minutes > 0 ? `${hours}h ${minutes}m` : `${hours}h`;
+	return `${minutes}m`;
+}
+
+/** " · resets in 5h 12m" / " · resets now" / "" when unknown. */
+export function formatResetSuffix(resetsAt: number | undefined, now = Date.now()): string {
+	if (resetsAt === undefined) return "";
+	const duration = formatResetDuration(resetsAt, now);
+	return duration === "now" ? " · resets now" : ` · resets in ${duration}`;
+}

package/src/index.ts

@@ -0,0 +1,168 @@
+/**
+ * pi-usage extension entry point.
+ *
+ * Registers the /usage command (prints an inline usage card into the session),
+ * keeps a footer status line updated with the active account's quota plus
+ * session consumption, and runs a session-scoped background refresh.
+ */
+
+import type { AutocompleteItem } from "@earendil-works/pi-tui";
+import type { ExtensionAPI, ExtensionContext, ExtensionUIContext } from "@earendil-works/pi-coding-agent";
+import { isRecord, loadBalanceConfig, saveStatusMode, type StatusMode } from "./config.ts";
+import { getAgentDir, tokenFromRegistryAuth } from "./credentials.ts";
+import { sumSessionUsage, type SessionUsageTotals } from "./session-usage.ts";
+import { BalanceService } from "./service.ts";
+import { buildUsageCard, type UsageCardData } from "./ui/card.ts";
+import { formatStatusLine } from "./ui/statusline.ts";
+import type { CredentialResolver } from "./types.ts";
+
+const USAGE_ENTRY_TYPE = "usage-report";
+/** Bound /usage fetches so a hanging provider API cannot stall the command. */
+const USAGE_FETCH_TIMEOUT_MS = 15_000;
+
+const STATUS_MODES = ["active", "all", "off"] as const;
+
+function isStatusMode(value: string): value is StatusMode {
+	return (STATUS_MODES as readonly string[]).includes(value);
+}
+
+export default function (pi: ExtensionAPI) {
+	let service: BalanceService | undefined;
+	let ui: ExtensionUIContext | undefined;
+	let activeProviderId: string | undefined;
+	let consumption: SessionUsageTotals | undefined;
+	/** `/usage <mode>` override for this session; wins over flag and usage.json. */
+	let statusOverride: StatusMode | undefined;
+	/** `--usage-status` CLI flag; wins over usage.json. */
+	let flagStatus: StatusMode | undefined;
+
+	pi.registerFlag("usage-status", {
+		description: "Footer status line mode: active | all | off (overrides usage.json)",
+		type: "string",
+	});
+
+	/**
+	 * Registry-backed credential resolution: the token comes from
+	 * getProviderAuth (which refreshes OAuth tokens before expiry), the base
+	 * URL from the registered provider — the auto-detect adapter needs it to
+	 * probe relay billing endpoints. Throwing sends the caller to the
+	 * auth.json fallback.
+	 */
+	function credentialResolverFor(ctx: ExtensionContext): CredentialResolver {
+		return async (providerId) => {
+			const baseUrl = ctx.modelRegistry.getProvider(providerId)?.baseUrl;
+			const auth = await ctx.modelRegistry.getProviderAuth(providerId).catch(() => undefined);
+			if (!auth) {
+				throw new Error(`Provider "${providerId}" not found in model registry`);
+			}
+			const headers: Record<string, string> = {};
+			for (const [name, value] of Object.entries(auth.auth.headers ?? {})) {
+				if (typeof value === "string") headers[name] = value;
+			}
+			const token = tokenFromRegistryAuth({ auth: { apiKey: auth.auth.apiKey, headers } });
+			if (!token) throw new Error(`No usable credential for "${providerId}"`);
+			return { token, baseUrl };
+		};
+	}
+
+	function updateStatus(): void {
+		if (!service || !ui) return;
+		const mode = statusOverride ?? flagStatus ?? service.config.status ?? "active";
+		const text = formatStatusLine({
+			balances: service.getAll(),
+			mode,
+			activeProviderId,
+			theme: ui.theme,
+			consumption,
+		});
+		ui.setStatus("usage", text);
+	}
+
+	function ensureService(ctx: ExtensionContext): BalanceService {
+		if (service) return service;
+		const agentDir = getAgentDir();
+		const next = new BalanceService({ agentDir, config: loadBalanceConfig(agentDir) });
+		next.setCredentialResolver(credentialResolverFor(ctx));
+		service = next;
+		return next;
+	}
+
+	pi.registerEntryRenderer(USAGE_ENTRY_TYPE, (entry, _options, theme) => {
+		if (!isRecord(entry.data) || !Array.isArray(entry.data.balances)) {
+			return buildUsageCard({ balances: [], generatedAt: Date.now() }, theme);
+		}
+		const data: UsageCardData = {
+			balances: entry.data.balances as UsageCardData["balances"],
+			activeProviderId: typeof entry.data.activeProviderId === "string" ? entry.data.activeProviderId : undefined,
+			generatedAt: typeof entry.data.generatedAt === "number" ? entry.data.generatedAt : Date.now(),
+		};
+		return buildUsageCard(data, theme);
+	});
+
+	pi.registerCommand("usage", {
+		description: "Show usage card, or set the footer status line mode",
+		getArgumentCompletions: (prefix: string): AutocompleteItem[] | null => {
+			const matches = STATUS_MODES.filter((mode) => mode.startsWith(prefix.toLowerCase()));
+			return matches.length > 0 ? matches.map((mode) => ({ value: mode, label: mode })) : null;
+		},
+		handler: async (args, ctx) => {
+			const argument = args?.trim().toLowerCase();
+			if (!argument) {
+				const current = ensureService(ctx);
+				await current.refreshAll({ signal: AbortSignal.timeout(USAGE_FETCH_TIMEOUT_MS) });
+				pi.appendEntry(USAGE_ENTRY_TYPE, {
+					balances: current.getAll(),
+					activeProviderId,
+					generatedAt: Date.now(),
+				});
+				return;
+			}
+			if (!isStatusMode(argument)) {
+				ctx.ui.notify(`/usage: unknown mode "${argument}" — use active, all, or off`, "warning");
+				return;
+			}
+			statusOverride = argument;
+			saveStatusMode(getAgentDir(), argument);
+			updateStatus();
+			ctx.ui.notify(`Footer status line: ${argument} (saved to usage.json)`, "info");
+		},
+	});
+
+	pi.on("session_start", async (_event, ctx) => {
+		// Session replacement tears the old service down via session_shutdown;
+		// build a fresh one bound to the new session context.
+		const cliFlag = pi.getFlag("usage-status");
+		flagStatus = typeof cliFlag === "string" && isStatusMode(cliFlag) ? cliFlag : undefined;
+		statusOverride = undefined;
+		const agentDir = getAgentDir();
+		const next = new BalanceService({ agentDir, config: loadBalanceConfig(agentDir) });
+		next.setCredentialResolver(credentialResolverFor(ctx));
+		next.onChange(updateStatus);
+		service?.stop();
+		service = next;
+		ui = ctx.ui;
+		activeProviderId = ctx.model?.provider;
+		consumption = sumSessionUsage(ctx.sessionManager.getEntries());
+		next.start();
+		updateStatus();
+	});
+
+	// Session tokens/cost are local data: refresh immediately after every turn
+	// instead of waiting for the provider quota polling cycle.
+	pi.on("turn_end", async (_event, ctx) => {
+		consumption = sumSessionUsage(ctx.sessionManager.getEntries());
+		updateStatus();
+	});
+
+	pi.on("model_select", async (event) => {
+		activeProviderId = event.model.provider;
+		updateStatus();
+	});
+
+	pi.on("session_shutdown", async () => {
+		service?.stop();
+		service = undefined;
+		ui = undefined;
+		consumption = undefined;
+	});
+}

package/src/parse.ts

@@ -0,0 +1,21 @@
+/** Shared response-field parsing helpers for provider adapters. */
+
+/** Accept numbers or numeric strings ("0.00" style). */
+export function toNumber(value: unknown): number | undefined {
+	if (typeof value === "number" && Number.isFinite(value)) return value;
+	if (typeof value === "string" && value.trim() !== "" && Number.isFinite(Number(value))) return Number(value);
+	return undefined;
+}
+
+/**
+ * Normalize timestamps across providers: epoch seconds, epoch milliseconds,
+ * or ISO-8601 strings. Returns epoch milliseconds, or undefined.
+ */
+export function parseTimestamp(value: unknown): number | undefined {
+	if (typeof value === "number" && Number.isFinite(value)) return value < 1e12 ? value * 1000 : value;
+	if (typeof value === "string" && value.trim() !== "") {
+		const parsed = Date.parse(value);
+		if (!Number.isNaN(parsed)) return parsed;
+	}
+	return undefined;
+}

package/src/providers/anthropic.ts

@@ -0,0 +1,76 @@
+/**
+ * Anthropic Claude (subscription OAuth) adapter.
+ *
+ * Query method borrowed from CodexBar: the Claude Code OAuth access token
+ * stored by pi in auth.json is accepted by the OAuth usage endpoint with the
+ * matching beta header.
+ */
+
+import { parseTimestamp, toNumber } from "../parse.ts";
+import type { AccountBalance, ProviderAdapter, ProviderFetchArgs, UsageWindow } from "../types.ts";
+
+const ENDPOINT = "https://api.anthropic.com/api/oauth/usage";
+const BETA_HEADER = "oauth-2025-04-20";
+
+interface AnthropicUsageWindow {
+	utilization?: unknown;
+	resets_at?: unknown;
+}
+
+interface AnthropicUsageResponse {
+	five_hour?: unknown;
+	seven_day?: unknown;
+	extra_usage?: {
+		is_enabled?: unknown;
+		monthly_limit?: unknown;
+		used_credits?: unknown;
+	} | null;
+}
+
+function parseWindow(label: string, raw: unknown): UsageWindow | undefined {
+	if (typeof raw !== "object" || raw === null) return undefined;
+	const window = raw as AnthropicUsageWindow;
+	const usedPercent = toNumber(window.utilization);
+	if (usedPercent === undefined) return undefined;
+	return { label, usedPercent, resetsAt: parseTimestamp(window.resets_at) };
+}
+
+export const anthropicAdapter: ProviderAdapter = {
+	id: "anthropic",
+	label: "Claude",
+	async fetch({ token, signal, fetchImpl }: ProviderFetchArgs): Promise<AccountBalance> {
+		const response = await fetchImpl(ENDPOINT, {
+			headers: {
+				Authorization: `Bearer ${token}`,
+				"anthropic-beta": BETA_HEADER,
+				Accept: "application/json",
+			},
+			signal,
+		});
+		if (!response.ok) {
+			throw new Error(`Claude usage API returned HTTP ${response.status}`);
+		}
+		const body: unknown = await response.json();
+		if (typeof body !== "object" || body === null) {
+			throw new Error("Claude usage API returned an unexpected response");
+		}
+		const usage = body as AnthropicUsageResponse;
+		const windows: UsageWindow[] = [];
+		const fiveHour = parseWindow("5h", usage.five_hour);
+		if (fiveHour) windows.push(fiveHour);
+		const sevenDay = parseWindow("weekly", usage.seven_day);
+		if (sevenDay) windows.push(sevenDay);
+		if (windows.length === 0) {
+			throw new Error("Claude usage API returned no usage windows");
+		}
+		const notes: string[] = [];
+		const extra = usage.extra_usage;
+		const monthlyLimit = toNumber(extra?.monthly_limit);
+		const usedCredits = toNumber(extra?.used_credits);
+		if (extra?.is_enabled === true && monthlyLimit !== undefined) {
+			const used = usedCredits !== undefined ? ` ${usedCredits}` : "";
+			notes.push(`extra usage:${used} of ${monthlyLimit}`);
+		}
+		return { providerId: "anthropic", label: "Claude", windows, notes, fetchedAt: Date.now() };
+	},
+};