@narumitw/pi-usage 0.49.3 → 0.51.0

package/src/usage.ts

@@ -1,24 +1,34 @@
+import { randomUUID } from "node:crypto";
 import type {
 	ExtensionAPI,
 	ExtensionCommandContext,
 	ExtensionContext,
 } from "@earendil-works/pi-coding-agent";
+import { FAST_USAGE_WARNING, registerCodexFastMode } from "./codex-fast-runtime.js";
 import {
-	awaitWithDeadline,
-	errorMessage,
-	runWithConcurrency,
-	sanitizeDisplayText,
-	UsageCache,
-} from "./core.js";
+	type CodexResetAvailability,
+	type CodexResetOption,
+	type CodexResetOutcome,
+	codexResetActionDescription,
+	codexResetCount,
+	consumeCodexResetCredit,
+	formatCodexResetOutcome,
+	genericCodexResetOption,
+	listCodexResetCredits,
+	resetConfirmationLines,
+	resetLabel,
+	resetOptionExpiration,
+	resolveCodexResetAuth,
+} from "./codex-resets.js";
+import { awaitWithDeadline, errorMessage, runWithConcurrency, UsageCache } from "./core.js";
 import { formatProviderStates, formatUsageStatusline } from "./format.js";
 import {
 	adapterForProvider,
 	isStaleExtensionContextError,
-	providerIsConfigured,
 	queryProviderUsage,
 	resolveUsageAuth,
-	SUPPORTED_ADAPTERS,
 } from "./query.js";
+import { createUsageSettingsRuntime, type UsageSettingsRuntime } from "./settings.js";
 import type {
 	PiModel,
 	ProviderUsageState,
@@ -26,6 +36,14 @@ import type {
 	UsageDisplayState,
 	UsageProviderAdapter,
 } from "./types.js";
+import {
+	configuredAdapters,
+	isAbortError,
+	isTimeoutError,
+	modelIdentity,
+	providerDisplayName,
+	setBoundedMap,
+} from "./usage-helpers.js";
 
 const CACHE_TTL_MS = 5 * 60 * 1000;
 const DEFAULT_TIMEOUT_MS = 15_000;
@@ -38,6 +56,13 @@ const REFRESH_CURRENT = "Refresh current usage";
 const VIEW_ANOTHER = "View another configured provider…";
 const VIEW_ALL = "View all configured providers…";
 const CLOSE = "Close";
+const REDEEM_CODEX_RESET = "Redeem usage limit reset…";
+
+type UsageExtensionDependencies = {
+	credentialReader?: (providerId: string) => unknown;
+	createRedemptionId?: () => string;
+	settingsRuntime?: UsageSettingsRuntime;
+};
 
 type QueryOutcome = {
 	state: ProviderUsageState;
@@ -50,7 +75,13 @@ type StableCurrent = {
 	model: PiModel | undefined;
 };
 
-export default function usageExtension(pi: ExtensionAPI) {
+export default function usageExtension(
+	pi: ExtensionAPI,
+	dependencies: UsageExtensionDependencies = {},
+) {
+	const credentialReader = dependencies.credentialReader;
+	const createRedemptionId = dependencies.createRedemptionId ?? randomUUID;
+	const settingsRuntime = dependencies.settingsRuntime ?? createUsageSettingsRuntime();
 	const cache = new UsageCache(CACHE_TTL_MS);
 	const failureBackoff = new Map<string, { until: number; message: string }>();
 	const latestQueries = new Map<string, number>();
@@ -61,6 +92,7 @@ export default function usageExtension(pi: ExtensionAPI) {
 	let statusGeneration = 0;
 	let statusRefreshTimer: ReturnType<typeof setTimeout> | undefined;
 	let statusController: AbortController | undefined;
+	let fastRuntime: ReturnType<typeof registerCodexFastMode>;
 
 	const clearStatusTimer = () => {
 		if (statusRefreshTimer) clearTimeout(statusRefreshTimer);
@@ -118,11 +150,22 @@ export default function usageExtension(pi: ExtensionAPI) {
 			}
 			return;
 		}
-		const value = formatUsageStatusline(outcome.state.report, model);
+		const rawValue = formatUsageStatusline(outcome.state.report, model);
+		const value = rawValue ? fastRuntime.decorateStatus(model, rawValue) : undefined;
 		if (!safeSetStatus(ctx, value)) return;
 		if (shouldSchedule && sessionActive) scheduleStatusRefresh(ctx, model);
 	};
 
+	const invalidateProviderState = (providerId: string) => {
+		cache.clearProvider(providerId);
+		for (const key of failureBackoff.keys()) {
+			if (key.startsWith(`${providerId}:`)) failureBackoff.delete(key);
+		}
+		for (const key of latestQueries.keys()) {
+			if (key.startsWith(`${providerId}:`)) latestQueries.delete(key);
+		}
+	};
+
 	const transitionCurrentIdentity = (nextIdentity: string, providerId: string) => {
 		if (!activeCurrentIdentity || activeCurrentIdentity === nextIdentity) {
 			activeCurrentIdentity = nextIdentity;
@@ -130,14 +173,7 @@ export default function usageExtension(pi: ExtensionAPI) {
 		}
 		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);
-			}
+			if (id) invalidateProviderState(id);
 		}
 		activeCurrentIdentity = nextIdentity;
 	};
@@ -346,12 +382,14 @@ export default function usageExtension(pi: ExtensionAPI) {
 		label: string,
 		parentSignal: AbortSignal,
 		operation: (signal: AbortSignal) => Promise<T>,
+		cancellable = true,
 	): Promise<T | undefined> => {
 		const { runTask } = await import("@narumitw/pi-tui-kit");
 		if (parentSignal.aborted) return undefined;
 		const result = await runTask(ctx, {
 			label,
 			signal: parentSignal,
+			cancellable,
 			onError: () => undefined,
 			task: ({ signal }) => operation(signal),
 		});
@@ -465,25 +503,86 @@ export default function usageExtension(pi: ExtensionAPI) {
 			publishStableCurrent(ctx, stableCurrent);
 			let current = stableCurrent.outcome;
 			let visibleStates: ProviderUsageState[] = [current.state];
+			let fastState = settingsRuntime.get();
+			let resetAvailability: CodexResetAvailability | undefined;
+			let selectedReset: CodexResetOption | undefined;
+			let resetAuthFingerprint: string | undefined;
+			let resetModelIdentity: string | undefined;
+			let redemptionId: string | undefined;
+			let resetOutcome: CodexResetOutcome | undefined;
+			let resetFailure: string | undefined;
 			const { defineMenu, runMenu } = await import("@narumitw/pi-tui-kit");
 			if (controller.signal.aborted || statusGeneration !== menuGeneration) return;
-			type Screen = "main" | "providers";
-			type Action = "refresh" | "another" | "all" | "provider";
+			type Screen =
+				| "main"
+				| "providers"
+				| "reset-picker"
+				| "reset-confirm"
+				| "reset-result"
+				| "reset-error";
+			type Action =
+				| "refresh"
+				| "toggle-fast"
+				| "another"
+				| "all"
+				| "provider"
+				| "open-resets"
+				| "select-reset"
+				| "cancel-reset"
+				| "consume-reset"
+				| "back-to-usage"
+				| "back-to-resets";
 			const menu = defineMenu<undefined, Screen, Action, ExtensionCommandContext>({
 				start: "main",
 				screens: {
-					main: () => ({
-						kind: "actions",
-						title: "Provider usage",
-						lines: formatProviderStates(visibleStates).split("\n"),
-						items: [
-							{ id: "refresh", label: REFRESH_CURRENT, action: "refresh" },
-							{ id: "another", label: VIEW_ANOTHER, action: "another" },
-							{ id: "all", label: VIEW_ALL, action: "all" },
-							{ id: "close", label: CLOSE, close: true },
-						],
-						hint: "close",
-					}),
+					main: () => {
+						const fastAvailability = fastRuntime.availability(ctx.model);
+						const fastLines =
+							fastAvailability.kind === "available"
+								? [`Fast mode: ${fastAvailability.enabled ? "On" : "Off"}`, FAST_USAGE_WARNING]
+								: fastAvailability.kind === "unavailable"
+									? [`Fast mode: Unavailable · ${fastAvailability.reason}`]
+									: [];
+						return {
+							kind: "actions",
+							title: "Provider usage",
+							lines: [...formatProviderStates(visibleStates).split("\n"), ...fastLines],
+							items: [
+								{ id: "refresh", label: REFRESH_CURRENT, action: "refresh" },
+								...(fastAvailability.kind === "available"
+									? [
+											{
+												id: "toggle-fast",
+												label: fastAvailability.enabled
+													? "Turn Fast mode off"
+													: "Turn Fast mode on",
+												description:
+													fastState.kind === "invalid"
+														? "Repair pi-usage.json and reload before changing Fast mode."
+														: FAST_USAGE_WARNING,
+												disabled: fastState.kind === "invalid",
+												action: "toggle-fast" as const,
+											},
+										]
+									: []),
+								...(current.state.status === "ready" && current.state.providerId === "openai-codex"
+									? [
+											{
+												id: "open-resets",
+												label: REDEEM_CODEX_RESET,
+												description: codexResetActionDescription(current.state.report),
+												disabled: codexResetCount(current.state.report) === 0,
+												action: "open-resets" as const,
+											},
+										]
+									: []),
+								{ id: "another", label: VIEW_ANOTHER, action: "another" },
+								{ id: "all", label: VIEW_ALL, action: "all" },
+								{ id: "close", label: CLOSE, close: true },
+							],
+							hint: "close",
+						};
+					},
 					providers: () => ({
 						kind: "actions",
 						title: "Select a configured provider",
@@ -496,8 +595,229 @@ export default function usageExtension(pi: ExtensionAPI) {
 							})),
 						hint: "back",
 					}),
+					"reset-picker": () => ({
+						kind: "choice",
+						title: "Usage limit resets",
+						lines: [
+							`${resetAvailability?.availableCount ?? 0} ${resetLabel(resetAvailability?.availableCount ?? 0)} available.`,
+						],
+						items: (resetAvailability?.options ?? []).map((option, index) => ({
+							id: `reset-${index}`,
+							label: option.title,
+							description: resetOptionExpiration(option),
+							details: [option.description],
+						})),
+						action: "select-reset",
+						initialItemId: "reset-0",
+						hint: "back",
+					}),
+					"reset-confirm": () => ({
+						kind: "actions",
+						title: "Use this reset?",
+						lines: resetConfirmationLines(selectedReset),
+						items: [
+							{ id: "cancel-reset", label: "No, go back", action: "cancel-reset" },
+							{ id: "consume-reset", label: "Yes, use reset", action: "consume-reset" },
+						],
+						hint: "back",
+					}),
+					"reset-result": () => ({
+						kind: "actions",
+						title: "Usage limit resets",
+						lines: [
+							formatCodexResetOutcome(
+								resetOutcome,
+								current.state.status === "ready"
+									? codexResetCount(current.state.report)
+									: undefined,
+							),
+						],
+						items: [
+							{
+								id: "back-to-usage",
+								label: "Back to usage",
+								action: "back-to-usage",
+							},
+							{ id: "close", label: CLOSE, close: true },
+						],
+						hint: "back",
+					}),
+					"reset-error": () => ({
+						kind: "actions",
+						title: "Usage limit resets",
+						lines: [resetFailure ?? "Couldn't reset usage. Please try again."],
+						items: [
+							{ id: "consume-reset", label: "Try again", action: "consume-reset" },
+							{ id: "back-to-resets", label: "Back", action: "back-to-resets" },
+						],
+						hint: "back",
+					}),
 				},
 				actions: {
+					"toggle-fast": async () => {
+						const availability = fastRuntime.availability(ctx.model);
+						if (availability.kind !== "available" || fastState.kind === "invalid") {
+							return { kind: "rejected" };
+						}
+						const changed = await fastRuntime.toggle(ctx, !availability.enabled, controller.signal);
+						if (!changed) return { kind: "rejected" };
+						fastState = settingsRuntime.get();
+						return { kind: "stay" };
+					},
+					"open-resets": async () => {
+						const summaryCount =
+							current.state.status === "ready" && current.state.providerId === "openai-codex"
+								? codexResetCount(current.state.report)
+								: undefined;
+						try {
+							const loaded = await runMenuOperation(
+								ctx,
+								"Checking usage limit resets…",
+								controller.signal,
+								async (signal) => {
+									const expectedModel = modelIdentity(ctx.model);
+									const auth = await awaitWithDeadline(
+										resolveCodexResetAuth(ctx, undefined, credentialReader),
+										signal,
+										DEFAULT_TIMEOUT_MS,
+										"resolving current Codex reset authentication",
+									);
+									let availability: CodexResetAvailability;
+									try {
+										availability = await listCodexResetCredits(auth, signal, DEFAULT_TIMEOUT_MS);
+									} catch (error) {
+										if (isAbortError(error) || summaryCount === undefined || summaryCount <= 0) {
+											throw error;
+										}
+										availability = {
+											availableCount: summaryCount,
+											options: [genericCodexResetOption()],
+										};
+									}
+									const revalidated = await awaitWithDeadline(
+										resolveCodexResetAuth(ctx, undefined, credentialReader),
+										signal,
+										DEFAULT_TIMEOUT_MS,
+										"revalidating current Codex reset authentication",
+									);
+									if (
+										modelIdentity(ctx.model) !== expectedModel ||
+										revalidated.fingerprint !== auth.fingerprint
+									) {
+										throw new Error(
+											"The active Codex model or account changed while loading usage limit resets.",
+										);
+									}
+									return { availability, auth, expectedModel };
+								},
+							);
+							if (!loaded) return { kind: "stay" };
+							resetAvailability = loaded.availability;
+							resetAuthFingerprint = loaded.auth.fingerprint;
+							resetModelIdentity = loaded.expectedModel;
+							selectedReset = undefined;
+							redemptionId = undefined;
+							resetOutcome = undefined;
+							resetFailure = undefined;
+							return {
+								kind: "to",
+								screen: loaded.availability.availableCount > 0 ? "reset-picker" : "reset-result",
+							};
+						} catch (error) {
+							if (isAbortError(error) || isStaleExtensionContextError(error)) {
+								return { kind: "stay" };
+							}
+							ctx.ui.notify(`Couldn't load usage limit resets: ${errorMessage(error)}`, "error");
+							return { kind: "stay" };
+						}
+					},
+					"select-reset": ({ itemId }) => {
+						const index = Number(itemId.replace(/^reset-/u, ""));
+						const option = Number.isSafeInteger(index)
+							? resetAvailability?.options[index]
+							: undefined;
+						if (!option) return { kind: "rejected" };
+						selectedReset = option;
+						redemptionId = undefined;
+						resetFailure = undefined;
+						return { kind: "to", screen: "reset-confirm" };
+					},
+					"cancel-reset": () => ({ kind: "back" }),
+					"consume-reset": async () => {
+						if (!selectedReset || !resetAuthFingerprint || !resetModelIdentity) {
+							return { kind: "rejected" };
+						}
+						try {
+							redemptionId ??= createRedemptionId();
+							const attemptId = redemptionId;
+							const option = selectedReset;
+							const expectedFingerprint = resetAuthFingerprint;
+							const expectedModel = resetModelIdentity;
+							const result = await runMenuOperation(
+								ctx,
+								"Resetting your usage…",
+								controller.signal,
+								async (signal) => {
+									const auth = await awaitWithDeadline(
+										resolveCodexResetAuth(ctx, undefined, credentialReader),
+										signal,
+										DEFAULT_TIMEOUT_MS,
+										"revalidating current Codex reset authentication",
+									);
+									if (
+										modelIdentity(ctx.model) !== expectedModel ||
+										auth.fingerprint !== expectedFingerprint
+									) {
+										throw new Error(
+											"The active Codex model or account changed; the reset was not used.",
+										);
+									}
+									const outcome = await consumeCodexResetCredit(
+										auth,
+										option,
+										attemptId,
+										signal,
+										DEFAULT_TIMEOUT_MS,
+									);
+									invalidateProviderState("openai-codex");
+									const model = ctx.model;
+									if (modelIdentity(model) !== expectedModel) return { outcome };
+									const refreshed = await queryCurrentState(ctx, model, true, signal);
+									const stable = await outcomeStillCurrent(
+										ctx,
+										model,
+										menuGeneration,
+										refreshed,
+										signal,
+									);
+									return stable ? { outcome, refreshed, model } : { outcome };
+								},
+								false,
+							);
+							if (!result) return { kind: "close" };
+							resetOutcome = result.outcome;
+							resetFailure = undefined;
+							if (result.refreshed && result.model) {
+								stableCurrent = { outcome: result.refreshed, model: result.model };
+								current = result.refreshed;
+								visibleStates = [current.state];
+								publishStableCurrent(ctx, stableCurrent);
+							}
+							return { kind: "to", screen: "reset-result" };
+						} catch (error) {
+							if (isAbortError(error) || isStaleExtensionContextError(error)) {
+								return { kind: "close" };
+							}
+							resetFailure = `Couldn't reset usage: ${errorMessage(error)}. Try again with the same request.`;
+							return { kind: "to", screen: "reset-error" };
+						}
+					},
+					"back-to-usage": () => ({ kind: "to", screen: "main" }),
+					"back-to-resets": () => {
+						redemptionId = undefined;
+						resetFailure = undefined;
+						return { kind: "to", screen: "reset-picker" };
+					},
 					refresh: async () => {
 						const refreshed = await queryStableCurrent(
 							ctx,
@@ -614,25 +934,30 @@ export default function usageExtension(pi: ExtensionAPI) {
 		}
 	};
 
-	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,
+		handler: async (args, ctx) => {
+			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.on("session_start", (_event, ctx) => {
+		statusGeneration += 1;
+		clearStatusTimer();
+		for (const controller of activeControllers) controller.abort();
+		activeControllers.clear();
+		statusController = undefined;
 		sessionActive = true;
 		startStatusRefresh(ctx, ctx.model, false);
 	});
@@ -658,40 +983,8 @@ export default function usageExtension(pi: ExtensionAPI) {
 		activeCurrentIdentity = undefined;
 		safeSetStatus(ctx, undefined);
 	});
-}
 
-function configuredAdapters(ctx: ExtensionContext): UsageProviderAdapter[] {
-	return SUPPORTED_ADAPTERS.filter(
-		(adapter) => adapter.id === ctx.model?.provider || providerIsConfigured(ctx, adapter.id),
+	fastRuntime = registerCodexFastMode(pi, settingsRuntime, (ctx) =>
+		startStatusRefresh(ctx, ctx.model, false),
 	);
 }
-
-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";
-}