@narumitw/pi-usage 0.24.0

package/src/usage.ts

@@ -0,0 +1,676 @@
+import {
+	BorderedLoader,
+	type ExtensionAPI,
+	type ExtensionCommandContext,
+	type ExtensionContext,
+} from "@earendil-works/pi-coding-agent";
+import {
+	awaitWithDeadline,
+	errorMessage,
+	runWithConcurrency,
+	sanitizeDisplayText,
+	UsageCache,
+} from "./core.js";
+import { formatProviderStates, formatUsageStatusline } from "./format.js";
+import {
+	adapterForProvider,
+	isStaleExtensionContextError,
+	providerIsConfigured,
+	queryProviderUsage,
+	resolveUsageAuth,
+	SUPPORTED_ADAPTERS,
+} from "./query.js";
+import type {
+	PiModel,
+	ProviderUsageState,
+	ResolvedUsageAuth,
+	UsageDisplayState,
+	UsageProviderAdapter,
+} from "./types.js";
+
+const CACHE_TTL_MS = 5 * 60 * 1000;
+const DEFAULT_TIMEOUT_MS = 15_000;
+const ALL_PROVIDER_CONCURRENCY = 2;
+const FAILURE_BACKOFF_MS = 30_000;
+const MAX_ACCOUNT_STATES = 32;
+const STATUS_KEY = "usage";
+
+const REFRESH_CURRENT = "Refresh current usage";
+const VIEW_ANOTHER = "View another configured provider…";
+const VIEW_ALL = "View all configured providers…";
+const CLOSE = "Close";
+const MENU_ACTIONS = [REFRESH_CURRENT, VIEW_ANOTHER, VIEW_ALL, CLOSE];
+
+type QueryOutcome = {
+	state: ProviderUsageState;
+	fingerprint?: string;
+	authState?: "unavailable";
+};
+
+type StableCurrent = {
+	outcome: QueryOutcome;
+	model: PiModel | undefined;
+};
+
+type LoaderResult<T> = { ok: true; value: T } | { ok: false; error: unknown };
+
+export default function usageExtension(pi: ExtensionAPI) {
+	const cache = new UsageCache(CACHE_TTL_MS);
+	const failureBackoff = new Map<string, { until: number; message: string }>();
+	const latestQueries = new Map<string, number>();
+	const activeControllers = new Set<AbortController>();
+	let querySequence = 0;
+	let activeCurrentIdentity: string | undefined;
+	let sessionActive = false;
+	let statusGeneration = 0;
+	let statusRefreshTimer: ReturnType<typeof setTimeout> | undefined;
+	let statusController: AbortController | undefined;
+
+	const clearStatusTimer = () => {
+		if (statusRefreshTimer) clearTimeout(statusRefreshTimer);
+		statusRefreshTimer = undefined;
+	};
+
+	const safeSetStatus = (ctx: ExtensionContext, value: string | undefined): boolean => {
+		try {
+			ctx.ui.setStatus(STATUS_KEY, value);
+			return true;
+		} catch (error) {
+			if (isStaleExtensionContextError(error)) return false;
+			throw error;
+		}
+	};
+
+	const clearStatus = (ctx: ExtensionContext) => {
+		statusGeneration += 1;
+		statusController?.abort();
+		statusController = undefined;
+		clearStatusTimer();
+		safeSetStatus(ctx, undefined);
+	};
+
+	const scheduleStatusRefresh = (ctx: ExtensionContext, model: PiModel) => {
+		clearStatusTimer();
+		const generation = statusGeneration;
+		statusRefreshTimer = setTimeout(() => {
+			statusRefreshTimer = undefined;
+			if (!sessionActive || generation !== statusGeneration) return;
+			startStatusRefresh(ctx, model, true);
+		}, CACHE_TTL_MS);
+		statusRefreshTimer.unref?.();
+	};
+
+	const publishStatus = (
+		ctx: ExtensionContext,
+		outcome: QueryOutcome,
+		model: PiModel,
+		shouldSchedule: boolean,
+	) => {
+		if (outcome.state.status === "unsupported") {
+			clearStatusTimer();
+			safeSetStatus(ctx, undefined);
+			return;
+		}
+		if (outcome.state.status !== "ready") {
+			if (
+				safeSetStatus(
+					ctx,
+					outcome.state.status === "auth-unavailable" ? "auth unavailable" : "usage error",
+				)
+			) {
+				if (shouldSchedule && sessionActive) scheduleStatusRefresh(ctx, model);
+			}
+			return;
+		}
+		const value = formatUsageStatusline(outcome.state.report, model);
+		if (!safeSetStatus(ctx, value)) return;
+		if (shouldSchedule && sessionActive) scheduleStatusRefresh(ctx, model);
+	};
+
+	const transitionCurrentIdentity = (nextIdentity: string, providerId: string) => {
+		if (!activeCurrentIdentity || activeCurrentIdentity === nextIdentity) {
+			activeCurrentIdentity = nextIdentity;
+			return;
+		}
+		const previousProviderId = activeCurrentIdentity.split(":", 1)[0] ?? "";
+		for (const id of new Set([previousProviderId, providerId])) {
+			if (!id) continue;
+			cache.clearProvider(id);
+			for (const key of failureBackoff.keys()) {
+				if (key.startsWith(`${id}:`)) failureBackoff.delete(key);
+			}
+			for (const key of latestQueries.keys()) {
+				if (key.startsWith(`${id}:`)) latestQueries.delete(key);
+			}
+		}
+		activeCurrentIdentity = nextIdentity;
+	};
+
+	const queryAdapterState = async (
+		ctx: ExtensionContext,
+		adapter: UsageProviderAdapter,
+		displayState: UsageDisplayState,
+		force: boolean,
+		signal: AbortSignal,
+	): Promise<QueryOutcome> => {
+		const startedAt = Date.now();
+		let auth: ResolvedUsageAuth | undefined;
+		try {
+			auth = await awaitWithDeadline(
+				resolveUsageAuth(ctx, adapter),
+				signal,
+				DEFAULT_TIMEOUT_MS,
+				`resolving ${adapter.displayName} runtime auth`,
+			);
+		} catch (error) {
+			if (isStaleExtensionContextError(error) || isAbortError(error)) throw error;
+			if (displayState === "current") {
+				transitionCurrentIdentity(`${adapter.id}:auth-error`, adapter.id);
+			}
+			return {
+				state: {
+					providerId: adapter.id,
+					providerName: adapter.displayName,
+					displayState,
+					status: isTimeoutError(error) ? "query-failed" : "auth-unavailable",
+					message: errorMessage(error),
+				},
+			};
+		}
+		if (!auth) {
+			if (displayState === "current") {
+				transitionCurrentIdentity(`${adapter.id}:unavailable`, adapter.id);
+			}
+			return {
+				state: {
+					providerId: adapter.id,
+					providerName: adapter.displayName,
+					displayState,
+					status: "auth-unavailable",
+					message: `No runtime credential is configured for ${adapter.displayName}.`,
+				},
+				authState: "unavailable",
+			};
+		}
+		if (displayState === "current") {
+			transitionCurrentIdentity(`${adapter.id}:${auth.fingerprint}`, adapter.id);
+		}
+
+		const cached = !force ? cache.get(adapter.id, auth.fingerprint) : undefined;
+		if (cached) {
+			return {
+				state: {
+					providerId: adapter.id,
+					providerName: adapter.displayName,
+					displayState,
+					status: "ready",
+					report: cached,
+				},
+				fingerprint: auth.fingerprint,
+			};
+		}
+
+		const failureKey = `${adapter.id}:${auth.fingerprint}`;
+		const previousFailure = failureBackoff.get(failureKey);
+		if (!force && previousFailure && previousFailure.until > Date.now()) {
+			return {
+				state: {
+					providerId: adapter.id,
+					providerName: adapter.displayName,
+					displayState,
+					status: "query-failed",
+					message: previousFailure.message,
+				},
+				fingerprint: auth.fingerprint,
+			};
+		}
+		failureBackoff.delete(failureKey);
+		querySequence += 1;
+		const queryId = querySequence;
+		setBoundedMap(latestQueries, failureKey, queryId, MAX_ACCOUNT_STATES);
+
+		try {
+			const remainingMs = Math.max(1, DEFAULT_TIMEOUT_MS - (Date.now() - startedAt));
+			const report = await queryProviderUsage(adapter, auth, signal, remainingMs);
+			if (latestQueries.get(failureKey) === queryId) {
+				cache.set(adapter.id, auth.fingerprint, report);
+				failureBackoff.delete(failureKey);
+			}
+			return {
+				state: {
+					providerId: adapter.id,
+					providerName: adapter.displayName,
+					displayState,
+					status: "ready",
+					report,
+				},
+				fingerprint: auth.fingerprint,
+			};
+		} catch (error) {
+			if (isStaleExtensionContextError(error) || isAbortError(error)) throw error;
+			const message = errorMessage(error);
+			const now = Date.now();
+			for (const [key, failure] of failureBackoff) {
+				if (failure.until <= now) failureBackoff.delete(key);
+			}
+			if (latestQueries.get(failureKey) === queryId) {
+				setBoundedMap(
+					failureBackoff,
+					failureKey,
+					{ until: now + FAILURE_BACKOFF_MS, message },
+					MAX_ACCOUNT_STATES,
+				);
+			}
+			return {
+				state: {
+					providerId: adapter.id,
+					providerName: adapter.displayName,
+					displayState,
+					status: "query-failed",
+					message,
+				},
+				fingerprint: auth.fingerprint,
+			};
+		}
+	};
+
+	const queryCurrentState = async (
+		ctx: ExtensionContext,
+		model: PiModel | undefined,
+		force: boolean,
+		signal: AbortSignal,
+	): Promise<QueryOutcome> => {
+		const adapter = adapterForProvider(model?.provider);
+		if (!adapter) {
+			const providerId = model?.provider ?? "none";
+			transitionCurrentIdentity(`unsupported:${providerId}`, providerId);
+			return {
+				state: {
+					providerId,
+					providerName: providerDisplayName(ctx, providerId),
+					displayState: "current",
+					status: "unsupported",
+					message: model
+						? `Usage reporting is not supported for ${providerDisplayName(ctx, providerId)}.`
+						: "No model is selected.",
+				},
+			};
+		}
+		return queryAdapterState(ctx, adapter, "current", force, signal);
+	};
+
+	const refreshCurrentStatus = async (
+		ctx: ExtensionContext,
+		model: PiModel | undefined,
+		force: boolean,
+	) => {
+		const adapter = adapterForProvider(model?.provider);
+		if (!adapter || !model) {
+			const providerId = model?.provider ?? "none";
+			transitionCurrentIdentity(`unsupported:${providerId}`, providerId);
+			clearStatus(ctx);
+			return;
+		}
+		statusGeneration += 1;
+		const generation = statusGeneration;
+		statusController?.abort();
+		const controller = new AbortController();
+		statusController = controller;
+		activeControllers.add(controller);
+		try {
+			if (!safeSetStatus(ctx, "checking")) return;
+			const outcome = await queryCurrentState(ctx, model, force, controller.signal);
+			if (!sessionActive || generation !== statusGeneration || controller.signal.aborted) return;
+			if (!(await outcomeStillCurrent(ctx, model, generation, outcome, controller.signal))) {
+				if (sessionActive && generation === statusGeneration) {
+					queueMicrotask(() => startStatusRefresh(ctx, ctx.model, false));
+				}
+				return;
+			}
+			publishStatus(ctx, outcome, model, true);
+		} finally {
+			activeControllers.delete(controller);
+			if (statusController === controller) statusController = undefined;
+		}
+	};
+
+	const startStatusRefresh = (
+		ctx: ExtensionContext,
+		model: PiModel | undefined,
+		force: boolean,
+	) => {
+		void refreshCurrentStatus(ctx, model, force).catch((error) => {
+			if (isStaleExtensionContextError(error) || isAbortError(error)) return;
+			safeSetStatus(ctx, "usage error");
+		});
+	};
+
+	const runMenuOperation = async <T>(
+		ctx: ExtensionCommandContext,
+		label: string,
+		parentSignal: AbortSignal,
+		operation: (signal: AbortSignal) => Promise<T>,
+	): Promise<T | undefined> => {
+		if (ctx.mode !== "tui") return operation(parentSignal);
+		const result = await ctx.ui.custom<LoaderResult<T> | null>((tui, theme, _keybindings, done) => {
+			const loader = new BorderedLoader(tui, theme, label);
+			let finished = false;
+			const finish = (value: LoaderResult<T> | null) => {
+				if (finished) return;
+				finished = true;
+				done(value);
+			};
+			loader.onAbort = () => finish(null);
+			const signal = AbortSignal.any([parentSignal, loader.signal]);
+			void operation(signal)
+				.then((value) => finish({ ok: true, value }))
+				.catch((error) => {
+					if (isAbortError(error)) finish(null);
+					else finish({ ok: false, error });
+				});
+			return loader;
+		});
+		if (!result) return undefined;
+		if (!result.ok) throw result.error;
+		return result.value;
+	};
+
+	const outcomeStillCurrent = async (
+		ctx: ExtensionContext,
+		model: PiModel | undefined,
+		generation: number,
+		outcome: QueryOutcome,
+		signal: AbortSignal,
+	): Promise<boolean> => {
+		if (generation !== statusGeneration || modelIdentity(ctx.model) !== modelIdentity(model)) {
+			return false;
+		}
+		const adapter = adapterForProvider(model?.provider);
+		if (outcome.authState === "unavailable") {
+			if (!adapter) return false;
+			try {
+				const auth = await awaitWithDeadline(
+					resolveUsageAuth(ctx, adapter),
+					signal,
+					DEFAULT_TIMEOUT_MS,
+					`revalidating ${adapter.displayName} runtime auth`,
+				);
+				return (
+					generation === statusGeneration &&
+					modelIdentity(ctx.model) === modelIdentity(model) &&
+					auth === undefined
+				);
+			} catch (error) {
+				if (isAbortError(error) || isStaleExtensionContextError(error)) throw error;
+				return false;
+			}
+		}
+		if (!outcome.fingerprint) return true;
+		if (!adapter) return false;
+		try {
+			const auth = await awaitWithDeadline(
+				resolveUsageAuth(ctx, adapter),
+				signal,
+				DEFAULT_TIMEOUT_MS,
+				`revalidating ${adapter.displayName} runtime auth`,
+			);
+			return (
+				generation === statusGeneration &&
+				modelIdentity(ctx.model) === modelIdentity(model) &&
+				auth?.fingerprint === outcome.fingerprint
+			);
+		} catch (error) {
+			if (isAbortError(error) || isStaleExtensionContextError(error)) throw error;
+			return false;
+		}
+	};
+
+	const queryStableCurrent = async (
+		ctx: ExtensionCommandContext,
+		force: boolean,
+		controller: AbortController,
+		label: string,
+	): Promise<StableCurrent | undefined> => {
+		for (let attempt = 0; attempt < 3; attempt += 1) {
+			const model = ctx.model;
+			const generation = statusGeneration;
+			const result = await runMenuOperation(ctx, label, controller.signal, async (signal) => {
+				const outcome = await queryCurrentState(ctx, model, force, signal);
+				return {
+					outcome,
+					stable: await outcomeStillCurrent(ctx, model, generation, outcome, signal),
+				};
+			});
+			if (!result) return undefined;
+			if (result.stable) return { outcome: result.outcome, model };
+			force = false;
+		}
+		ctx.ui.notify("The active model or account kept changing; reopen /usage to retry.", "warning");
+		return undefined;
+	};
+
+	const publishStableCurrent = (ctx: ExtensionCommandContext, current: StableCurrent) => {
+		if (current.model) publishStatus(ctx, current.outcome, current.model, sessionActive);
+		else safeSetStatus(ctx, undefined);
+	};
+
+	const showMenu = async (ctx: ExtensionCommandContext): Promise<void> => {
+		if (!ctx.hasUI) {
+			ctx.ui.notify("/usage requires an interactive Pi mode.", "warning");
+			return;
+		}
+		statusGeneration += 1;
+		statusController?.abort();
+		statusController = undefined;
+		clearStatusTimer();
+		const controller = new AbortController();
+		activeControllers.add(controller);
+		try {
+			let stableCurrent = await queryStableCurrent(
+				ctx,
+				false,
+				controller,
+				"Checking current usage…",
+			);
+			if (!stableCurrent) return;
+			publishStableCurrent(ctx, stableCurrent);
+			let current = stableCurrent.outcome;
+			let visibleStates: ProviderUsageState[] = [current.state];
+
+			while (!controller.signal.aborted) {
+				const action = await ctx.ui.select(formatProviderStates(visibleStates), [...MENU_ACTIONS], {
+					signal: controller.signal,
+				});
+				if (!action || action === CLOSE) return;
+				if (action === REFRESH_CURRENT) {
+					stableCurrent = await queryStableCurrent(
+						ctx,
+						true,
+						controller,
+						"Refreshing current usage…",
+					);
+					if (!stableCurrent) continue;
+					publishStableCurrent(ctx, stableCurrent);
+					current = stableCurrent.outcome;
+					visibleStates = [current.state];
+					continue;
+				}
+				if (action === VIEW_ANOTHER) {
+					const others = configuredAdapters(ctx).filter(
+						(adapter) => adapter.id !== ctx.model?.provider,
+					);
+					if (others.length === 0) {
+						ctx.ui.notify("No other supported provider has configured runtime auth.", "info");
+						continue;
+					}
+					const choice = await ctx.ui.select(
+						"Select a configured provider",
+						others.map((adapter) => adapter.displayName),
+						{ signal: controller.signal },
+					);
+					const adapter = others.find((candidate) => candidate.displayName === choice);
+					if (!adapter) continue;
+					const outcome = await runMenuOperation(
+						ctx,
+						`Checking ${adapter.displayName} usage…`,
+						controller.signal,
+						(signal) => queryAdapterState(ctx, adapter, "configured", false, signal),
+					);
+					if (!outcome) continue;
+					const revalidated = await queryStableCurrent(
+						ctx,
+						false,
+						controller,
+						"Revalidating current usage…",
+					);
+					if (!revalidated) continue;
+					stableCurrent = revalidated;
+					current = revalidated.outcome;
+					visibleStates =
+						outcome.state.providerId === current.state.providerId
+							? [current.state]
+							: [current.state, { ...outcome.state, displayState: "configured" }];
+					continue;
+				}
+				if (action === VIEW_ALL) {
+					const adapters = configuredAdapters(ctx);
+					const currentProviderId = ctx.model?.provider;
+					const settled = await runMenuOperation(
+						ctx,
+						"Checking configured provider usage…",
+						controller.signal,
+						(signal) =>
+							runWithConcurrency(
+								adapters,
+								ALL_PROVIDER_CONCURRENCY,
+								(adapter, _index, workerSignal) =>
+									queryAdapterState(
+										ctx,
+										adapter,
+										adapter.id === currentProviderId ? "current" : "configured",
+										true,
+										workerSignal,
+									),
+								signal,
+							),
+					);
+					if (!settled) continue;
+					const queriedStates: ProviderUsageState[] = settled.map((result, index) => {
+						if (result.status === "fulfilled") {
+							return { ...result.value.state, displayState: "configured" };
+						}
+						const adapter = adapters[index] as UsageProviderAdapter;
+						return {
+							providerId: adapter.id,
+							providerName: adapter.displayName,
+							displayState: "configured",
+							status: "query-failed",
+							message: errorMessage(result.reason),
+						};
+					});
+					const revalidated = await queryStableCurrent(
+						ctx,
+						false,
+						controller,
+						"Revalidating current usage…",
+					);
+					if (!revalidated) continue;
+					stableCurrent = revalidated;
+					current = revalidated.outcome;
+					visibleStates = [
+						current.state,
+						...queriedStates.filter((state) => state.providerId !== current.state.providerId),
+					];
+				}
+			}
+		} finally {
+			controller.abort();
+			activeControllers.delete(controller);
+		}
+	};
+
+	const commandHandler = async (args: string, ctx: ExtensionCommandContext) => {
+		if (args.trim()) {
+			ctx.ui.notify("/usage does not accept arguments; choose an action from its menu.", "warning");
+			return;
+		}
+		try {
+			await showMenu(ctx);
+		} catch (error) {
+			if (isStaleExtensionContextError(error) || isAbortError(error)) return;
+			throw error;
+		}
+	};
+
+	pi.registerCommand("usage", {
+		description: "Show usage for the current runtime account",
+		handler: commandHandler,
+	});
+	pi.registerCommand("codex-status", {
+		description: "Open /usage (temporary compatibility alias)",
+		handler: commandHandler,
+	});
+
+	pi.on("session_start", (_event, ctx) => {
+		sessionActive = true;
+		startStatusRefresh(ctx, ctx.model, false);
+	});
+	pi.on("session_tree", (_event, ctx) => {
+		startStatusRefresh(ctx, ctx.model, false);
+	});
+	pi.on("model_select", (event, ctx) => {
+		startStatusRefresh(ctx, event.model, false);
+	});
+	pi.on("turn_start", (_event, ctx) => {
+		startStatusRefresh(ctx, ctx.model, false);
+	});
+	pi.on("session_shutdown", (_event, ctx) => {
+		sessionActive = false;
+		statusGeneration += 1;
+		clearStatusTimer();
+		for (const controller of activeControllers) controller.abort();
+		activeControllers.clear();
+		statusController = undefined;
+		cache.clear();
+		failureBackoff.clear();
+		latestQueries.clear();
+		activeCurrentIdentity = undefined;
+		safeSetStatus(ctx, undefined);
+	});
+}
+
+function configuredAdapters(ctx: ExtensionContext): UsageProviderAdapter[] {
+	return SUPPORTED_ADAPTERS.filter(
+		(adapter) => adapter.id === ctx.model?.provider || providerIsConfigured(ctx, adapter.id),
+	);
+}
+
+function providerDisplayName(ctx: ExtensionContext, providerId: string): string {
+	try {
+		return sanitizeDisplayText(ctx.modelRegistry.getProviderDisplayName(providerId), 80);
+	} catch {
+		return sanitizeDisplayText(providerId, 80);
+	}
+}
+
+function setBoundedMap<T>(map: Map<string, T>, key: string, value: T, limit: number): void {
+	map.delete(key);
+	while (map.size >= limit) {
+		const oldest = map.keys().next().value;
+		if (oldest === undefined) break;
+		map.delete(oldest);
+	}
+	map.set(key, value);
+}
+
+function modelIdentity(model: PiModel | undefined): string | undefined {
+	return model ? `${model.provider}/${model.id}` : undefined;
+}
+
+function isAbortError(error: unknown): boolean {
+	return error instanceof Error && error.name === "AbortError";
+}
+
+function isTimeoutError(error: unknown): boolean {
+	return error instanceof Error && error.name === "TimeoutError";
+}