@narumitw/pi-usage 0.54.0 → 0.58.0

package/src/usage.ts

@@ -56,6 +56,7 @@ import {
 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;
@@ -103,23 +104,26 @@ export default function usageExtension(
 	let sessionActive = false;
 	let statusGeneration = 0;
 	let sessionGeneration = 0;
-	let xaiSettingsGeneration = 0;
 	let statusRefreshTimer: ReturnType<typeof setTimeout> | undefined;
+	let statusCountdownTimer: ReturnType<typeof setTimeout> | undefined;
 	let statusController: AbortController | undefined;
 	let fastRuntime: ReturnType<typeof registerCodexFastMode>;
 
-	const xaiUsageEnabled = () => {
-		const state = settingsRuntime.get();
-		return state.kind !== "invalid" && state.settings.xaiUsage;
-	};
-	const activeAdapterForProvider = (providerId: string | undefined) =>
-		adapterForProvider(providerId, xaiUsageEnabled());
-
-	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);
@@ -134,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;
@@ -155,13 +159,14 @@ export default function usageExtension(
 		model: PiModel,
 		shouldSchedule: boolean,
 	) => {
-		if (activeAdapterForProvider(model.provider)?.publishesStatusline === false) {
-			clearStatusTimer();
+		clearStatusCountdownTimer();
+		if (adapterForProvider(model.provider)?.publishesStatusline === false) {
+			clearStatusRefreshTimer();
 			safeSetStatus(ctx, undefined);
 			return;
 		}
 		if (outcome.state.status === "unsupported") {
-			clearStatusTimer();
+			clearStatusRefreshTimer();
 			safeSetStatus(ctx, undefined);
 			return;
 		}
@@ -176,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) => {
@@ -210,18 +248,22 @@ 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 expectedXaiSettingsGeneration = xaiSettingsGeneration;
 		const expectedSessionId = ctx.sessionManager.getSessionId();
 		const expectedModelIdentity = modelIdentity(ctx.model);
+		const expectedFireworksAccountId =
+			adapter.id === "fireworks" ? settingsRuntime.get().settings.fireworksAccountId : undefined;
+		const querySettings =
+			adapter.id === "fireworks" ? { fireworksAccountId: expectedFireworksAccountId } : undefined;
 		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) {
@@ -239,16 +281,14 @@ export default function usageExtension(
 				},
 			};
 		}
-		if (
-			adapter.id === "xai" &&
-			(expectedSessionGeneration !== sessionGeneration ||
-				expectedXaiSettingsGeneration !== xaiSettingsGeneration ||
-				!xaiUsageEnabled() ||
-				ctx.sessionManager.getSessionId() !== expectedSessionId ||
-				modelIdentity(ctx.model) !== expectedModelIdentity)
-		) {
-			throw abortError();
-		}
+		const requiresRequestBoundaryGuard = ["deepseek", "fireworks", "xai"].includes(adapter.id);
+		const requestContextChanged = () =>
+			expectedSessionGeneration !== sessionGeneration ||
+			ctx.sessionManager.getSessionId() !== expectedSessionId ||
+			modelIdentity(ctx.model) !== expectedModelIdentity ||
+			(adapter.id === "fireworks" &&
+				settingsRuntime.get().settings.fireworksAccountId !== expectedFireworksAccountId);
+		if (requiresRequestBoundaryGuard && requestContextChanged()) throw abortError();
 		if (!auth) {
 			if (displayState === "current") {
 				transitionCurrentIdentity(`${adapter.id}:unavailable`, adapter.id);
@@ -264,11 +304,15 @@ export default function usageExtension(
 				authState: "unavailable",
 			};
 		}
+		const queryFingerprint =
+			adapter.id === "fireworks"
+				? `${auth.fingerprint}:account:${expectedFireworksAccountId ?? "auto"}`
+				: auth.fingerprint;
 		if (displayState === "current") {
-			transitionCurrentIdentity(`${adapter.id}:${auth.fingerprint}`, adapter.id);
+			transitionCurrentIdentity(`${adapter.id}:${queryFingerprint}`, adapter.id);
 		}
 
-		const cached = !force ? cache.get(adapter.id, auth.fingerprint) : undefined;
+		const cached = !force ? cache.get(adapter.id, queryFingerprint) : undefined;
 		if (cached) {
 			return {
 				state: {
@@ -282,7 +326,7 @@ export default function usageExtension(
 			};
 		}
 
-		const failureKey = `${adapter.id}:${auth.fingerprint}`;
+		const failureKey = `${adapter.id}:${queryFingerprint}`;
 		const previousFailure = failureBackoff.get(failureKey);
 		if (!force && previousFailure && previousFailure.until > Date.now()) {
 			return {
@@ -301,44 +345,39 @@ 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 guard =
-				adapter.id === "xai"
-					? async () => {
-							if (
-								signal.aborted ||
-								expectedSessionGeneration !== sessionGeneration ||
-								expectedXaiSettingsGeneration !== xaiSettingsGeneration ||
-								!xaiUsageEnabled() ||
-								ctx.sessionManager.getSessionId() !== expectedSessionId ||
-								modelIdentity(ctx.model) !== expectedModelIdentity
-							) {
-								throw abortError();
-							}
-							const revalidated = await awaitWithDeadline(
-								resolveUsageAuth(ctx, adapter, undefined, credentialReader, credentialCandidates),
-								signal,
-								Math.max(1, DEFAULT_TIMEOUT_MS - (Date.now() - startedAt)),
-								"revalidating xAI runtime auth",
-							);
-							if (
-								signal.aborted ||
-								expectedSessionGeneration !== sessionGeneration ||
-								expectedXaiSettingsGeneration !== xaiSettingsGeneration ||
-								!xaiUsageEnabled() ||
-								ctx.sessionManager.getSessionId() !== expectedSessionId ||
-								modelIdentity(ctx.model) !== expectedModelIdentity ||
-								revalidated?.fingerprint !== auth.fingerprint
-							) {
-								throw abortError();
+			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);
+					}
+				: undefined;
+			const report = await queryProviderUsage(
+				adapter,
+				auth,
+				signal,
+				remainingMs,
+				guard,
+				querySettings,
+			);
 			if (guard) await guard();
 			if (latestQueries.get(failureKey) === queryId) {
-				cache.set(adapter.id, auth.fingerprint, report);
+				cache.set(adapter.id, queryFingerprint, report);
 				failureBackoff.delete(failureKey);
 			}
 			return {
@@ -353,6 +392,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) {
@@ -385,7 +442,7 @@ export default function usageExtension(
 		force: boolean,
 		signal: AbortSignal,
 	): Promise<QueryOutcome> => {
-		const adapter = activeAdapterForProvider(model?.provider);
+		const adapter = adapterForProvider(model?.provider);
 		if (!adapter) {
 			const providerId = model?.provider ?? "none";
 			transitionCurrentIdentity(`unsupported:${providerId}`, providerId);
@@ -395,12 +452,9 @@ export default function usageExtension(
 					providerName: providerDisplayName(ctx, providerId),
 					displayState: "current",
 					status: "unsupported",
-					message:
-						providerId === "xai" && !xaiUsageEnabled()
-							? "xAI usage is disabled. Open Settings to enable it."
-							: model
-								? `Usage reporting is not supported for ${providerDisplayName(ctx, providerId)}.`
-								: "No model is selected.",
+					message: model
+						? `Usage reporting is not supported for ${providerDisplayName(ctx, providerId)}.`
+						: "No model is selected.",
 				},
 			};
 		}
@@ -412,7 +466,7 @@ export default function usageExtension(
 		model: PiModel | undefined,
 		force: boolean,
 	) => {
-		const adapter = activeAdapterForProvider(model?.provider);
+		const adapter = adapterForProvider(model?.provider);
 		if (!adapter || !model) {
 			const providerId = model?.provider ?? "none";
 			transitionCurrentIdentity(`unsupported:${providerId}`, providerId);
@@ -425,6 +479,7 @@ export default function usageExtension(
 		}
 		statusGeneration += 1;
 		const generation = statusGeneration;
+		clearStatusCountdownTimer();
 		statusController?.abort();
 		const controller = new AbortController();
 		statusController = controller;
@@ -494,7 +549,7 @@ export default function usageExtension(
 		if (generation !== statusGeneration || modelIdentity(ctx.model) !== modelIdentity(model)) {
 			return false;
 		}
-		const adapter = activeAdapterForProvider(model?.provider);
+		const adapter = adapterForProvider(model?.provider);
 		if (outcome.authState === "unavailable") {
 			if (!adapter) return false;
 			try {
@@ -569,7 +624,7 @@ export default function usageExtension(
 		const menuGeneration = statusGeneration;
 		statusController?.abort();
 		statusController = undefined;
-		clearStatusTimer();
+		clearStatusTimers();
 		const controller = new AbortController();
 		activeControllers.add(controller);
 		try {
@@ -668,7 +723,7 @@ export default function usageExtension(
 					providers: () => ({
 						kind: "actions",
 						title: "Select a configured provider",
-						items: configuredAdapters(ctx, xaiUsageEnabled())
+						items: configuredAdapters(ctx)
 							.filter((adapter) => adapter.id !== ctx.model?.provider)
 							.map((adapter) => ({
 								id: adapter.id,
@@ -742,14 +797,14 @@ export default function usageExtension(
 							settingsRuntime,
 							controller.signal,
 							() => statusGeneration === menuGeneration && !controller.signal.aborted,
-							(id, _previous, next) => {
-								if (id !== "xaiUsage") return;
-								xaiSettingsGeneration += 1;
-								invalidateProviderState("xai");
-								if (!next) {
-									for (const active of activeControllers) {
-										if (active !== controller) active.abort();
-									}
+							(id) => {
+								if (
+									id === "codexStatusResetCountdown" &&
+									stableCurrent &&
+									statusGeneration === menuGeneration &&
+									!controller.signal.aborted
+								) {
+									publishStableCurrent(ctx, stableCurrent);
 								}
 							},
 						);
@@ -946,7 +1001,7 @@ export default function usageExtension(
 						return { kind: "stay" };
 					},
 					another: async () => {
-						const others = configuredAdapters(ctx, xaiUsageEnabled()).filter(
+						const others = configuredAdapters(ctx).filter(
 							(adapter) => adapter.id !== ctx.model?.provider,
 						);
 						if (others.length === 0) {
@@ -956,7 +1011,7 @@ export default function usageExtension(
 						return { kind: "to", screen: "providers" };
 					},
 					provider: async ({ itemId }) => {
-						const adapter = configuredAdapters(ctx, xaiUsageEnabled()).find(
+						const adapter = configuredAdapters(ctx).find(
 							(candidate) => candidate.id === itemId && candidate.id !== ctx.model?.provider,
 						);
 						if (!adapter) return { kind: "back" };
@@ -984,7 +1039,7 @@ export default function usageExtension(
 						return { kind: "back" };
 					},
 					all: async () => {
-						const adapters = configuredAdapters(ctx, xaiUsageEnabled());
+						const adapters = configuredAdapters(ctx);
 						const currentProviderId = ctx.model?.provider;
 						const settled = await runMenuOperation(
 							ctx,
@@ -1065,16 +1120,21 @@ export default function usageExtension(
 			}
 		},
 	});
-	pi.on("session_start", (_event, ctx) => {
+	pi.on("session_start", async (_event, ctx) => {
 		sessionGeneration += 1;
-		xaiSettingsGeneration += 1;
 		statusGeneration += 1;
-		clearStatusTimer();
+		clearStatusTimers();
 		for (const controller of activeControllers) controller.abort();
 		activeControllers.clear();
 		statusController = undefined;
 		sessionActive = true;
-		startStatusRefresh(ctx, ctx.model, false);
+		const ownerGeneration = sessionGeneration;
+		try {
+			await fastRuntime.prepareSession(ctx);
+		} catch (error) {
+			if (isStaleExtensionContextError(error) || ownerGeneration !== sessionGeneration) return;
+			throw error;
+		}
 	});
 	pi.on("session_tree", (_event, ctx) => {
 		startStatusRefresh(ctx, ctx.model, false);
@@ -1088,9 +1148,8 @@ export default function usageExtension(
 	pi.on("session_shutdown", (_event, ctx) => {
 		sessionActive = false;
 		sessionGeneration += 1;
-		xaiSettingsGeneration += 1;
 		statusGeneration += 1;
-		clearStatusTimer();
+		clearStatusTimers();
 		for (const controller of activeControllers) controller.abort();
 		activeControllers.clear();
 		statusController = undefined;
@@ -1101,7 +1160,10 @@ export default function usageExtension(
 		safeSetStatus(ctx, undefined);
 	});
 
-	fastRuntime = registerCodexFastMode(pi, settingsRuntime, (ctx) =>
-		startStatusRefresh(ctx, ctx.model, false),
+	fastRuntime = registerCodexFastMode(
+		pi,
+		settingsRuntime,
+		(ctx) => startStatusRefresh(ctx, ctx.model, false),
+		{ registerSessionStart: false },
 	);
 }