@narumitw/pi-usage 0.53.0 → 0.57.0

package/src/usage.ts

@@ -1,3 +1,5 @@
+// This orchestrator intentionally remains over 1,000 lines because its menu, query generations,
+// account cache, cancellation, and session lifecycle share one consistency boundary.
 import { randomUUID } from "node:crypto";
 import type {
 	ExtensionAPI,
@@ -20,7 +22,13 @@ import {
 	resetOptionExpiration,
 	resolveCodexResetAuth,
 } from "./codex-resets.js";
-import { awaitWithDeadline, errorMessage, runWithConcurrency, UsageCache } from "./core.js";
+import {
+	abortError,
+	awaitWithDeadline,
+	errorMessage,
+	runWithConcurrency,
+	UsageCache,
+} from "./core.js";
 import { formatProviderStates, formatUsageStatusline } from "./format.js";
 import { createOAuthCredentialCandidateReader } from "./oauth-credential-source.js";
 import {
@@ -45,8 +53,10 @@ import {
 	providerDisplayName,
 	setBoundedMap,
 } from "./usage-helpers.js";
+import { showUsageSettings } from "./usage-settings-ui.js";
 
 const CACHE_TTL_MS = 5 * 60 * 1000;
+const STATUS_COUNTDOWN_REFRESH_MS = 60 * 1000;
 const DEFAULT_TIMEOUT_MS = 15_000;
 const ALL_PROVIDER_CONCURRENCY = 2;
 const FAILURE_BACKOFF_MS = 30_000;
@@ -57,6 +67,7 @@ const REFRESH_CURRENT = "Refresh current usage";
 const VIEW_ANOTHER = "View another configured provider…";
 const VIEW_ALL = "View all configured providers…";
 const CLOSE = "Close";
+const SETTINGS = "Settings";
 const REDEEM_CODEX_RESET = "Redeem usage limit reset…";
 
 type UsageExtensionDependencies = {
@@ -92,15 +103,27 @@ export default function usageExtension(
 	let activeCurrentIdentity: string | undefined;
 	let sessionActive = false;
 	let statusGeneration = 0;
+	let sessionGeneration = 0;
 	let statusRefreshTimer: ReturnType<typeof setTimeout> | undefined;
+	let statusCountdownTimer: ReturnType<typeof setTimeout> | undefined;
 	let statusController: AbortController | undefined;
 	let fastRuntime: ReturnType<typeof registerCodexFastMode>;
 
-	const clearStatusTimer = () => {
+	const clearStatusRefreshTimer = () => {
 		if (statusRefreshTimer) clearTimeout(statusRefreshTimer);
 		statusRefreshTimer = undefined;
 	};
 
+	const clearStatusCountdownTimer = () => {
+		if (statusCountdownTimer) clearTimeout(statusCountdownTimer);
+		statusCountdownTimer = undefined;
+	};
+
+	const clearStatusTimers = () => {
+		clearStatusRefreshTimer();
+		clearStatusCountdownTimer();
+	};
+
 	const safeSetStatus = (ctx: ExtensionContext, value: string | undefined): boolean => {
 		try {
 			ctx.ui.setStatus(STATUS_KEY, value);
@@ -115,12 +138,12 @@ export default function usageExtension(
 		statusGeneration += 1;
 		statusController?.abort();
 		statusController = undefined;
-		clearStatusTimer();
+		clearStatusTimers();
 		safeSetStatus(ctx, undefined);
 	};
 
 	const scheduleStatusRefresh = (ctx: ExtensionContext, model: PiModel) => {
-		clearStatusTimer();
+		clearStatusRefreshTimer();
 		const generation = statusGeneration;
 		statusRefreshTimer = setTimeout(() => {
 			statusRefreshTimer = undefined;
@@ -136,13 +159,14 @@ export default function usageExtension(
 		model: PiModel,
 		shouldSchedule: boolean,
 	) => {
+		clearStatusCountdownTimer();
 		if (adapterForProvider(model.provider)?.publishesStatusline === false) {
-			clearStatusTimer();
+			clearStatusRefreshTimer();
 			safeSetStatus(ctx, undefined);
 			return;
 		}
 		if (outcome.state.status === "unsupported") {
-			clearStatusTimer();
+			clearStatusRefreshTimer();
 			safeSetStatus(ctx, undefined);
 			return;
 		}
@@ -157,10 +181,43 @@ export default function usageExtension(
 			}
 			return;
 		}
-		const rawValue = formatUsageStatusline(outcome.state.report, model);
+		const showCodexResetCountdown =
+			outcome.state.report.providerId === "openai-codex" &&
+			settingsRuntime.get().settings.codexStatusResetCountdown;
+		const now = Date.now();
+		const rawValue = formatUsageStatusline(
+			outcome.state.report,
+			model,
+			now,
+			showCodexResetCountdown,
+		);
 		const value = rawValue ? fastRuntime.decorateStatus(model, rawValue) : undefined;
 		if (!safeSetStatus(ctx, value)) return;
 		if (shouldSchedule && sessionActive) scheduleStatusRefresh(ctx, model);
+		if (
+			sessionActive &&
+			showCodexResetCountdown &&
+			outcome.state.report.buckets.some(
+				(bucket) =>
+					bucket.resetsAt !== undefined &&
+					Number.isFinite(bucket.resetsAt) &&
+					bucket.resetsAt * 1_000 > now,
+			)
+		) {
+			const generation = statusGeneration;
+			statusCountdownTimer = setTimeout(() => {
+				statusCountdownTimer = undefined;
+				if (
+					!sessionActive ||
+					generation !== statusGeneration ||
+					modelIdentity(ctx.model) !== modelIdentity(model)
+				) {
+					return;
+				}
+				publishStatus(ctx, outcome, model, false);
+			}, STATUS_COUNTDOWN_REFRESH_MS);
+			statusCountdownTimer.unref?.();
+		}
 	};
 
 	const invalidateProviderState = (providerId: string) => {
@@ -191,14 +248,18 @@ export default function usageExtension(
 		displayState: UsageDisplayState,
 		force: boolean,
 		signal: AbortSignal,
+		authRetry = 0,
+		deadlineAt = Date.now() + DEFAULT_TIMEOUT_MS,
 	): Promise<QueryOutcome> => {
-		const startedAt = Date.now();
+		const expectedSessionGeneration = sessionGeneration;
+		const expectedSessionId = ctx.sessionManager.getSessionId();
+		const expectedModelIdentity = modelIdentity(ctx.model);
 		let auth: ResolvedUsageAuth | undefined;
 		try {
 			auth = await awaitWithDeadline(
 				resolveUsageAuth(ctx, adapter, undefined, credentialReader, credentialCandidates),
 				signal,
-				DEFAULT_TIMEOUT_MS,
+				Math.max(1, deadlineAt - Date.now()),
 				`resolving ${adapter.displayName} runtime auth`,
 			);
 		} catch (error) {
@@ -216,6 +277,12 @@ export default function usageExtension(
 				},
 			};
 		}
+		const requiresRequestBoundaryGuard = adapter.id === "deepseek" || adapter.id === "xai";
+		const requestContextChanged = () =>
+			expectedSessionGeneration !== sessionGeneration ||
+			ctx.sessionManager.getSessionId() !== expectedSessionId ||
+			modelIdentity(ctx.model) !== expectedModelIdentity;
+		if (requiresRequestBoundaryGuard && requestContextChanged()) throw abortError();
 		if (!auth) {
 			if (displayState === "current") {
 				transitionCurrentIdentity(`${adapter.id}:unavailable`, adapter.id);
@@ -268,9 +335,30 @@ export default function usageExtension(
 		const queryId = querySequence;
 		setBoundedMap(latestQueries, failureKey, queryId, MAX_ACCOUNT_STATES);
 
+		let deepSeekAuthChanged = false;
 		try {
-			const remainingMs = Math.max(1, DEFAULT_TIMEOUT_MS - (Date.now() - startedAt));
-			const report = await queryProviderUsage(adapter, auth, signal, remainingMs);
+			const remainingMs = Math.max(1, deadlineAt - Date.now());
+			const guard = requiresRequestBoundaryGuard
+				? async () => {
+						if (signal.aborted || requestContextChanged()) throw abortError();
+						const revalidated = await awaitWithDeadline(
+							resolveUsageAuth(ctx, adapter, undefined, credentialReader, credentialCandidates),
+							signal,
+							Math.max(1, deadlineAt - Date.now()),
+							`revalidating ${adapter.displayName} runtime auth`,
+						);
+						if (signal.aborted || requestContextChanged()) throw abortError();
+						if (revalidated?.fingerprint !== auth.fingerprint) {
+							if (adapter.id === "deepseek") {
+								deepSeekAuthChanged = true;
+								throw new Error("DeepSeek runtime credential changed during the balance query.");
+							}
+							throw abortError();
+						}
+					}
+				: undefined;
+			const report = await queryProviderUsage(adapter, auth, signal, remainingMs, guard);
+			if (guard) await guard();
 			if (latestQueries.get(failureKey) === queryId) {
 				cache.set(adapter.id, auth.fingerprint, report);
 				failureBackoff.delete(failureKey);
@@ -287,6 +375,24 @@ export default function usageExtension(
 			};
 		} catch (error) {
 			if (isStaleExtensionContextError(error) || isAbortError(error)) throw error;
+			if (
+				deepSeekAuthChanged &&
+				authRetry === 0 &&
+				!signal.aborted &&
+				!requestContextChanged() &&
+				Date.now() < deadlineAt
+			) {
+				if (latestQueries.get(failureKey) === queryId) latestQueries.delete(failureKey);
+				return queryAdapterState(
+					ctx,
+					adapter,
+					displayState,
+					true,
+					signal,
+					authRetry + 1,
+					deadlineAt,
+				);
+			}
 			const message = errorMessage(error);
 			const now = Date.now();
 			for (const [key, failure] of failureBackoff) {
@@ -356,6 +462,7 @@ export default function usageExtension(
 		}
 		statusGeneration += 1;
 		const generation = statusGeneration;
+		clearStatusCountdownTimer();
 		statusController?.abort();
 		const controller = new AbortController();
 		statusController = controller;
@@ -500,7 +607,7 @@ export default function usageExtension(
 		const menuGeneration = statusGeneration;
 		statusController?.abort();
 		statusController = undefined;
-		clearStatusTimer();
+		clearStatusTimers();
 		const controller = new AbortController();
 		activeControllers.add(controller);
 		try {
@@ -533,6 +640,7 @@ export default function usageExtension(
 				| "reset-error";
 			type Action =
 				| "refresh"
+				| "settings"
 				| "toggle-fast"
 				| "another"
 				| "all"
@@ -560,6 +668,7 @@ export default function usageExtension(
 							lines: [...formatProviderStates(visibleStates).split("\n"), ...fastLines],
 							items: [
 								{ id: "refresh", label: REFRESH_CURRENT, action: "refresh" },
+								{ id: "settings", label: SETTINGS, action: "settings" },
 								...(fastAvailability.kind === "available"
 									? [
 											{
@@ -665,6 +774,37 @@ export default function usageExtension(
 					}),
 				},
 				actions: {
+					settings: async () => {
+						await showUsageSettings(
+							ctx,
+							settingsRuntime,
+							controller.signal,
+							() => statusGeneration === menuGeneration && !controller.signal.aborted,
+							(id) => {
+								if (
+									id === "codexStatusResetCountdown" &&
+									stableCurrent &&
+									statusGeneration === menuGeneration &&
+									!controller.signal.aborted
+								) {
+									publishStableCurrent(ctx, stableCurrent);
+								}
+							},
+						);
+						fastState = settingsRuntime.get();
+						const revalidated = await queryStableCurrent(
+							ctx,
+							false,
+							controller,
+							"Applying usage settings…",
+						);
+						if (!revalidated) return { kind: "stay" };
+						stableCurrent = revalidated;
+						current = revalidated.outcome;
+						visibleStates = [current.state];
+						publishStableCurrent(ctx, revalidated);
+						return { kind: "stay" };
+					},
 					"toggle-fast": async () => {
 						const availability = fastRuntime.availability(ctx.model);
 						if (availability.kind !== "available" || fastState.kind === "invalid") {
@@ -964,8 +1104,9 @@ export default function usageExtension(
 		},
 	});
 	pi.on("session_start", (_event, ctx) => {
+		sessionGeneration += 1;
 		statusGeneration += 1;
-		clearStatusTimer();
+		clearStatusTimers();
 		for (const controller of activeControllers) controller.abort();
 		activeControllers.clear();
 		statusController = undefined;
@@ -983,8 +1124,9 @@ export default function usageExtension(
 	});
 	pi.on("session_shutdown", (_event, ctx) => {
 		sessionActive = false;
+		sessionGeneration += 1;
 		statusGeneration += 1;
-		clearStatusTimer();
+		clearStatusTimers();
 		for (const controller of activeControllers) controller.abort();
 		activeControllers.clear();
 		statusController = undefined;