Oh My Pi - Provider Quota Summary Extension

provider-quota-summary.ts

import type { ExtensionAPI, ExtensionContext } from "@oh-my-pi/pi-coding-agent";
import type { OAuthAccountSummary } from "@oh-my-pi/pi-ai";
import type { UsageReport } from "@oh-my-pi/pi-ai/usage";
import { resolveUsedFraction } from "@oh-my-pi/pi-ai/usage";
import { formatDuration } from "@oh-my-pi/pi-utils";

type Theme = ExtensionContext["ui"]["theme"];

interface MatchedAccountReport {
	account: OAuthAccountSummary;
	report?: UsageReport;
	fresh: boolean;
	/** Soonest credit expiry, or Infinity when there is none. Drives sorting and hasWarning. */
	earliestExpiry: number;
}

const STALE_THRESHOLD_MS = 10 * 60 * 1000;
const SEVEN_DAYS_MS = 7 * 24 * 60 * 60 * 1000;
interface ProviderInfo {
	name: string;
	icon: string;
}

const PROVIDERS: Record<string, ProviderInfo> = {
	"openai-codex": { name: "OpenAI", icon: "✳" },
	"openai": { name: "OpenAI", icon: "✳" },
	"anthropic": { name: "Anthropic", icon: "✦" },
	"claude": { name: "Anthropic", icon: "✦" },
	"google": { name: "Google", icon: "◈" },
	"google-gemini-cli": { name: "Google Gemini", icon: "◈" },
	"google-antigravity": { name: "Google Antigravity", icon: "◈" },
	"xai-oauth": { name: "xAI", icon: "𝕏" },
	"xai": { name: "xAI", icon: "𝕏" },
	"github": { name: "GitHub", icon: "🐙" },
	"github-copilot": { name: "GitHub Copilot", icon: "🐙" },
	"deepseek": { name: "DeepSeek", icon: "🌀" },
	"groq": { name: "Groq", icon: "⚡" },
	"ollama": { name: "Ollama", icon: "📦" },
	"mistral": { name: "Mistral", icon: "🌪" },
	"openrouter": { name: "OpenRouter", icon: "🔀" },
	"cursor": { name: "Cursor", icon: "▶" },
	"minimax-code": { name: "MiniMax", icon: "⚡" },
	"kimi": { name: "Kimi", icon: "🌙" },
	"kimi-code": { name: "Kimi", icon: "🌙" },
	"zai": { name: "ZAI", icon: "⚡" },
};
// Hoisted: constructing this per call is ~37x the cost of formatting with it,
// and it dominated the whole render (toLocaleString ≈ 77µs vs 2µs per date).
const CREDIT_DATE_FORMAT = new Intl.DateTimeFormat("tr-TR", { dateStyle: "short", timeStyle: "medium" });

function isFresh(report: UsageReport, now: number): boolean {
	const age = now - report.fetchedAt;
	return Boolean(report.fetchedAt && age >= 0 && age <= STALE_THRESHOLD_MS);
}

function maskEmail(rawEmail?: string): string {
	const email = rawEmail?.trim();
	if (!email?.includes("@")) return "[Bilinmeyen Hesap]";
	const [local, domain] = email.split("@");
	if (local.length <= 2) return `${local[0]}*@${domain}`;
	return `${local[0]}***${local[local.length - 1]}@${domain}`;
}

function metaString(report: UsageReport, key: "accountId" | "email" | "planType"): string | undefined {
	const value = report.metadata?.[key];
	return typeof value === "string" ? value.trim() || undefined : undefined;
}

function creditExpiry(credit: { status?: string; expiresAt?: string }, now: number): number | undefined {
	if (credit.status && credit.status !== "available") return undefined;
	const expiry = Date.parse(credit.expiresAt ?? "");
	return Number.isFinite(expiry) && expiry > now ? expiry : undefined;
}

function creditExpiries(report: UsageReport | undefined, now: number): number[] {
	return (report?.resetCredits?.credits ?? [])
		.map((credit) => creditExpiry(credit, now))
		.filter((expiry): expiry is number => expiry !== undefined)
		.sort((a, b) => a - b);
}

function matchOpenAIAccountsAndReports(
	accounts: OAuthAccountSummary[],
	reports: UsageReport[],
	now: number,
): MatchedAccountReport[] {
	const openAiReports = reports.filter((r) => r.provider === "openai-codex");
	const unclaimed = new Set(openAiReports);

	// Pass 1: match by accountId.
	const matched: MatchedAccountReport[] = accounts.map((account) => {
		const report = account.accountId
			? openAiReports.find((r) => unclaimed.has(r) && metaString(r, "accountId") === account.accountId)
			: undefined;
		if (report) unclaimed.delete(report);
		return { account, report, fresh: report ? isFresh(report, now) : false, earliestExpiry: Infinity };
	});

	// Pass 2: fall back to email, but only where it is unambiguous on both sides.
	// A duplicated email maps to null so the lookup below rejects it.
	const reportByEmail = new Map<string, UsageReport | null>();
	for (const report of unclaimed) {
		const email = metaString(report, "email")?.toLowerCase();
		if (email) reportByEmail.set(email, reportByEmail.has(email) ? null : report);
	}
	const accountEmailCounts = new Map<string, number>();
	for (const item of matched) {
		const email = item.report ? undefined : item.account.email?.trim().toLowerCase();
		if (email) accountEmailCounts.set(email, (accountEmailCounts.get(email) ?? 0) + 1);
	}
	for (const item of matched) {
		if (item.report) continue;
		const email = item.account.email?.trim().toLowerCase();
		const report = email && accountEmailCounts.get(email) === 1 ? reportByEmail.get(email) : null;
		if (report) {
			item.report = report;
			item.fresh = isFresh(report, now);
			unclaimed.delete(report);
		}
	}

	// Expiries come from fresh reports only, then sort by soonest, then by account position.
	for (const item of matched) {
		const expiries = item.fresh ? creditExpiries(item.report, now) : [];
		item.earliestExpiry = expiries.length > 0 ? expiries[0] : Infinity;
	}
	matched.sort((a, b) => {
		if (a.earliestExpiry !== b.earliestExpiry) return a.earliestExpiry - b.earliestExpiry;
		return (a.account.position ?? 0) - (b.account.position ?? 0);
	});

	return matched;
}

function getProviderInfo(provider: string): ProviderInfo {
	if (PROVIDERS[provider]) return PROVIDERS[provider];
	const lower = provider.toLowerCase();
	if (lower.includes("github") || lower.includes("copilot")) return { name: "GitHub Copilot", icon: "🐙" };
	if (lower.includes("antigravity")) return { name: "Google Antigravity", icon: "◈" };
	if (lower.includes("google") || lower.includes("gemini")) return { name: "Google Gemini", icon: "◈" };
	if (lower.includes("anthropic") || lower.includes("claude")) return { name: "Anthropic", icon: "✦" };
	if (lower.includes("openai") || lower.includes("codex")) return { name: "OpenAI", icon: "✳" };
	if (lower.includes("xai") || lower.includes("grok")) return { name: "xAI", icon: "𝕏" };
	const name = provider
		.split(/[-_]/)
		.map((w) => w.charAt(0).toUpperCase() + w.slice(1))
		.join(" ");
	return { name, icon: "•" };
}

function formatLimitLabel(limit: { id?: string; label?: string; window?: { label?: string } }): string {
	const rawLabel = (limit.label || "").trim();
	const windowLabel = (limit.window?.label || "").trim();
	const id = (limit.id || "").trim();

	// Check for "Usage (CounterName)" pattern from Antigravity / multi-counter providers
	const usageMatch = rawLabel.match(/^Usage \((.+)\)$/i);
	if (usageMatch) {
		const counter = usageMatch[1];
		if (windowLabel && windowLabel.toLowerCase() !== counter.toLowerCase() && windowLabel.toLowerCase() !== "default") {
			return `${counter} · ${windowLabel}`;
		}
		return counter;
	}

	// Check if ID has provider:counter:... structure like "google-antigravity:anthropic:default:daily"
	const idParts = id.split(":");
	if (idParts.length >= 4 && idParts[0] === "google-antigravity") {
		const counterRaw = idParts[1];
		const counter =
			counterRaw === "anthropic"
				? "Anthropic"
				: counterRaw === "google"
					? "Google"
					: counterRaw === "openai"
						? "OpenAI"
						: counterRaw.charAt(0).toUpperCase() + counterRaw.slice(1);
		if (windowLabel && windowLabel.toLowerCase() !== "default") {
			return `${counter} · ${windowLabel}`;
		}
		return `${counter} · ${rawLabel || "Limit"}`;
	}

	// If rawLabel and windowLabel are identical (e.g. "7 days" / "7 days" or "Weekly" / "Weekly")
	if (rawLabel && windowLabel && rawLabel.toLowerCase() === windowLabel.toLowerCase()) {
		return rawLabel;
	}

	// If rawLabel is generic like "Usage" or "Limit", prefer windowLabel
	if (rawLabel.toLowerCase() === "usage" || rawLabel.toLowerCase() === "limit" || !rawLabel) {
		return (windowLabel && windowLabel.toLowerCase() !== "default" ? windowLabel : rawLabel) || "Limit";
	}

	// If windowLabel is generic or missing, use rawLabel
	if (!windowLabel || windowLabel.toLowerCase() === "default") {
		return rawLabel;
	}

	// If rawLabel already contains windowLabel (e.g. "Grok Build (Weekly)" with window "Weekly")
	if (rawLabel.toLowerCase().includes(windowLabel.toLowerCase())) {
		return rawLabel;
	}

	// Otherwise, if both are distinct informative labels
	return `${rawLabel} · ${windowLabel}`;
}

function formatQuotaAccountLine(
	branch: string,
	accountLabel: string,
	report: UsageReport | undefined,
	fresh: boolean,
	now: number,
	titleColor: "warning" | "accent",
	separator: string,
	theme: Theme,
): string {
	const planType = report ? metaString(report, "planType") : undefined;
	const planStr = planType ? ` ${theme.fg("dim", `(${planType})`)}` : "";
	const head = `${theme.fg("dim", branch)} ${theme.fg("accent", accountLabel)}${planStr}`;
	if (!report) return `     ${head}${separator}${theme.fg("muted", "Kota bilgisi mevcut değil")}`;
	if (!fresh) return `     ${head}${separator}Kota bilgisi: ${theme.fg("warning", "veri eski")}`;

	const parts = (report.limits ?? []).map((limit) => {
		const used = resolveUsedFraction(limit);
		const resetsAt = limit.window?.resetsAt;
		const resetRemainingMs = typeof resetsAt === "number" && resetsAt > now ? resetsAt - now : undefined;
		const durationMs = limit.window?.durationMs;
		const percentRemaining =
			typeof used === "number" && Number.isFinite(used)
				? Math.max(0, Math.min(1, 1 - used)) * 100
				: undefined;
		const reset = resetRemainingMs === undefined ? undefined : formatDuration(resetRemainingMs);
		const resetRemainingFraction =
			resetRemainingMs !== undefined &&
			typeof durationMs === "number" &&
			Number.isFinite(durationMs) &&
			durationMs > 0
				? Math.min(1, resetRemainingMs / durationMs)
				: undefined;
		const value = percentRemaining === undefined ? "belirsiz" : styleGauge(percentRemaining, theme);
		const resetText = reset
			? ` (Sıfırlama: ${styleResetDuration(reset, resetRemainingFraction, theme)})`
			: "";
		const limitLabel = formatLimitLabel(limit);
		return `${limitLabel}: ${value}${resetText}`;
	});

	if (report.provider === "openai-codex" && report.resetCredits) {
		const availableCount = report.resetCredits.availableCount ?? 0;
		const credits = [`↻${availableCount}`];
		const expiries = creditExpiries(report, now);
		for (const expiry of expiries) {
			credits.push(`${CREDIT_DATE_FORMAT.format(expiry)} (${formatDuration(expiry - now)})`);
		}
		const unknownCount = availableCount - expiries.length;
		if (unknownCount > 0) credits.push(`+${unknownCount} tarih bilinmiyor`);
		parts.push(theme.fg(titleColor, credits.join(" │ ")));
	}

	return parts.length > 0 ? `     ${head}${separator}${parts.join(separator)}` : `     ${head}`;
}

interface AccountEntry {
	label: string;
	report?: UsageReport;
	fresh: boolean;
}

function formatQuotaSummary(
	accounts: OAuthAccountSummary[],
	reports: UsageReport[],
	now: number,
	theme: Theme,
): string | undefined {
	const matchedOpenAI = matchOpenAIAccountsAndReports(accounts, reports, now);
	const claimed = new Set(matchedOpenAI.map((item) => item.report));
	const additional = reports
		.filter((report) => !claimed.has(report))
		.sort((a, b) => a.provider.localeCompare(b.provider));
	const hasWarning = matchedOpenAI.some((item) => item.earliestExpiry <= now + SEVEN_DAYS_MS);
	const titleColor = hasWarning ? "warning" : "accent";
	const separator = theme.fg("dim", " │ ");
	const title = `${theme.fg(titleColor, "◆")} ${theme.bold(theme.fg(titleColor, "Provider Kota Özeti"))}`;
	const groupMap = new Map<string, AccountEntry[]>();

	if (matchedOpenAI.length > 0) {
		groupMap.set(
			"openai-codex",
			matchedOpenAI.map((item) => ({
				label: maskEmail(item.account.email),
				report: item.report,
				fresh: item.fresh,
			})),
		);
	}

	for (const report of additional) {
		const email = metaString(report, "email");
		const list = groupMap.get(report.provider) ?? [];
		const label = email ? maskEmail(email) : `[Hesap ${list.length + 1}]`;
		list.push({ label, report, fresh: isFresh(report, now) });
		groupMap.set(report.provider, list);
	}

	if (groupMap.size === 0) return undefined;

	const lines = [title];
	for (const [provider, accounts] of groupMap) {
		const pInfo = getProviderInfo(provider);
		lines.push(`  ${theme.fg(titleColor, pInfo.icon)} ${theme.bold(theme.fg(titleColor, pInfo.name))}`);
		for (let i = 0; i < accounts.length; i++) {
			const acc = accounts[i];
			const isLast = i === accounts.length - 1;
			const branch = isLast ? "└─" : "├─";
			lines.push(
				formatQuotaAccountLine(
					branch,
					acc.label,
					acc.report,
					acc.fresh,
					now,
					titleColor,
					separator,
					theme,
				),
			);
		}
	}

	return lines.length > 1 ? lines.join("\n") : undefined;
}

function styleGauge(percentRemaining: number, theme: Theme): string {
	const color = percentRemaining >= 50 ? "success" : percentRemaining >= 20 ? "warning" : "error";
	const filled = Math.round(percentRemaining / 20);
	const bar = `${theme.fg(color, "■".repeat(filled))}${theme.fg("dim", "□".repeat(5 - filled))}`;
	return `${theme.fg(color, `${percentRemaining.toFixed(0)}%`)} ${bar}`;
}

function styleResetDuration(duration: string, remainingFraction: number | undefined, theme: Theme): string {
	if (remainingFraction === undefined) return duration;
	const color = remainingFraction > 2 / 3 ? "success" : remainingFraction > 1 / 3 ? "warning" : "error";
	return theme.fg(color, duration);
}

export default function quotaSummaryExtension(pi: ExtensionAPI) {
	let pendingTimer: ReturnType<ExtensionContext["setTimeout"]> | undefined;
	let activeController: AbortController | null = null;

	function cancelInFlight(ctx: ExtensionContext) {
		if (pendingTimer !== undefined) ctx.clearTimer(pendingTimer);
		pendingTimer = undefined;
		activeController?.abort();
		activeController = null;
	}

	// Deferred via ctx.setTimeout rather than a detached promise: an unmanaged async
	// throw becomes an uncaughtException and the postmortem handler kills the session.
	function scheduleFetchAndShow(ctx: ExtensionContext) {
		if (!ctx.hasUI || !process.stdout.isTTY) return;

		cancelInFlight(ctx);
		pendingTimer = ctx.setTimeout(() => {
			pendingTimer = undefined;
			return fetchAndShow(ctx);
		}, 0);
	}

	async function fetchAndShow(ctx: ExtensionContext) {
		const controller = new AbortController();
		activeController = controller;

		try {
			const reports = await ctx.modelRegistry.authStorage.fetchUsageReports({
				baseUrlResolver: (provider) => ctx.modelRegistry.getProviderBaseUrl(provider),
				// Only controller.signal marks a cancellation; a timeout still reports the failure.
				signal: AbortSignal.any([controller.signal, AbortSignal.timeout(10_000)]),
			});
			if (controller.signal.aborted) return;

			const accounts = ctx.modelRegistry.authStorage.listOAuthAccounts("openai-codex");
			if (accounts.length === 0 && !reports?.length) return;

			const summary = formatQuotaSummary(accounts, reports ?? [], Date.now(), ctx.ui.theme);
			if (summary) {
				ctx.ui.notify(summary, "info");
			}
		} catch {
			if (controller.signal.aborted) return;
			ctx.ui.notify("Quota bilgileri alınamadı. Ayrıntılar için /usage komutunu kullanın.", "warning");
		} finally {
			if (activeController === controller) activeController = null;
		}
	}

	pi.on("session_start", (_event, ctx) => {
		scheduleFetchAndShow(ctx);
	});

	pi.on("session_switch", (_event, ctx) => {
		scheduleFetchAndShow(ctx);
	});

	pi.on("session_before_switch", (_event, ctx) => {
		cancelInFlight(ctx);
	});

	pi.on("session_shutdown", (_event, ctx) => {
		cancelInFlight(ctx);
	});
}