@narumitw/pi-usage 0.57.0 → 0.59.0

package/src/usage-settings-ui.ts

@@ -11,12 +11,16 @@ import {
 	Text,
 } from "@earendil-works/pi-tui";
 import { errorMessage } from "./core.js";
+import { isFireworksAccountId } from "./providers/fireworks.js";
 import type { UsageSettings, UsageSettingsRuntime } from "./settings.js";
 
+const AUTO = "Auto";
+const EDIT = "Edit…";
 const OFF = "Off";
 const ON = "On";
 
 type UsageSettingId = keyof UsageSettings;
+type SettingsScreenResult = { changed: boolean; editFireworksAccount: boolean };
 
 export async function showUsageSettings(
 	ctx: ExtensionCommandContext,
@@ -29,15 +33,38 @@ export async function showUsageSettings(
 		if (ctx.hasUI) ctx.ui.notify(`Edit settings manually: ${settingsRuntime.get().path}`, "info");
 		return false;
 	}
-	if (parentSignal.aborted || !isCurrent()) return false;
+	let changed = false;
+	while (!parentSignal.aborted && isCurrent()) {
+		const result = await showSettingsList(ctx, settingsRuntime, parentSignal, isCurrent, onApplied);
+		if (!result) return changed;
+		changed ||= result.changed;
+		if (!result.editFireworksAccount) return changed;
+		changed ||= await editFireworksAccount(
+			ctx,
+			settingsRuntime,
+			parentSignal,
+			isCurrent,
+			onApplied,
+		);
+	}
+	return changed;
+}
 
-	return ctx.ui.custom<boolean>((tui, theme, _keybindings, done) => {
+async function showSettingsList(
+	ctx: ExtensionCommandContext,
+	settingsRuntime: UsageSettingsRuntime,
+	parentSignal: AbortSignal,
+	isCurrent: () => boolean,
+	onApplied: (id: UsageSettingId) => void,
+): Promise<SettingsScreenResult | undefined> {
+	return ctx.ui.custom<SettingsScreenResult>((tui, theme, _keybindings, done) => {
 		const localController = new AbortController();
 		const signal = AbortSignal.any([parentSignal, localController.signal]);
 		let changed = false;
 		let closing = false;
 		let saveQueue = Promise.resolve();
 		const state = settingsRuntime.get();
+		const fireworksValue = state.settings.fireworksAccountId ?? AUTO;
 		const items: SettingItem[] = [
 			{
 				id: "codexFastMode",
@@ -53,6 +80,15 @@ export async function showUsageSettings(
 				currentValue: state.settings.codexStatusResetCountdown ? ON : OFF,
 				values: [OFF, ON],
 			},
+			{
+				id: "fireworksAccountId",
+				label: "Fireworks account",
+				description: "Select Edit to enter a visible account slug, or submit blank to clear it.",
+				currentValue: fireworksValue,
+				values: state.settings.fireworksAccountId
+					? [state.settings.fireworksAccountId, EDIT]
+					: [AUTO, EDIT],
+			},
 		];
 		const container = new Container();
 		container.addChild(new Text(theme.fg("accent", theme.bold("pi-usage Settings")), 1, 1));
@@ -62,7 +98,40 @@ export async function showUsageSettings(
 			if (closing) return;
 			closing = true;
 			localController.abort();
-			done(changed);
+			done({ changed, editFireworksAccount: false });
+		};
+		const queueUpdate = (
+			id: UsageSettingId,
+			requested: UsageSettings[UsageSettingId],
+			display: string,
+		) => {
+			saveQueue = saveQueue.then(async () => {
+				const previous = settingsRuntime.get().settings[id];
+				if (settingsRuntime.get().kind === "invalid") {
+					settingsList.updateValue(id, displaySetting(id, previous));
+					if (!signal.aborted && isCurrent()) {
+						ctx.ui.notify("Repair pi-usage.json and reload before changing settings.", "error");
+						tui.requestRender();
+					}
+					return;
+				}
+				try {
+					await settingsRuntime.update({ [id]: requested }, signal);
+				} catch (error) {
+					if (signal.aborted || !isCurrent()) return;
+					settingsList.updateValue(id, displaySetting(id, previous));
+					ctx.ui.notify(`Could not save pi-usage.json: ${errorMessage(error)}`, "error");
+					tui.requestRender();
+					return;
+				}
+				if (previous !== requested) {
+					changed = true;
+					onApplied(id);
+				}
+				if (signal.aborted || !isCurrent()) return;
+				settingsList.updateValue(id, display);
+				tui.requestRender();
+			});
 		};
 		settingsList = new SettingsList(
 			items,
@@ -70,35 +139,18 @@ export async function showUsageSettings(
 			getSettingsListTheme(),
 			(id, value) => {
 				if (closing || signal.aborted || !isCurrent()) return;
-				const settingId = id as UsageSettingId;
-				const requested = value !== OFF;
-				saveQueue = saveQueue.then(async () => {
-					const previous = settingsRuntime.get().settings[settingId];
-					if (settingsRuntime.get().kind === "invalid") {
-						settingsList.updateValue(id, displayValue(previous));
-						if (!signal.aborted && isCurrent()) {
-							ctx.ui.notify("Repair pi-usage.json and reload before changing settings.", "error");
-							tui.requestRender();
-						}
-						return;
+				if (id === "fireworksAccountId") {
+					if (value === EDIT) {
+						saveQueue = saveQueue.then(() => {
+							if (closing || signal.aborted || !isCurrent()) return;
+							closing = true;
+							done({ changed, editFireworksAccount: true });
+						});
 					}
-					try {
-						await settingsRuntime.update({ [settingId]: requested }, signal);
-					} catch (error) {
-						if (signal.aborted || !isCurrent()) return;
-						settingsList.updateValue(id, displayValue(previous));
-						ctx.ui.notify(`Could not save pi-usage.json: ${errorMessage(error)}`, "error");
-						tui.requestRender();
-						return;
-					}
-					if (previous !== requested) {
-						changed = true;
-						onApplied(settingId);
-					}
-					if (signal.aborted || !isCurrent()) return;
-					settingsList.updateValue(id, displayValue(requested));
-					tui.requestRender();
-				});
+					return;
+				}
+				const settingId = id as "codexFastMode" | "codexStatusResetCountdown";
+				queueUpdate(settingId, value !== OFF, value);
 			},
 			cancel,
 		);
@@ -122,6 +174,46 @@ export async function showUsageSettings(
 	});
 }
 
-function displayValue(enabled: boolean): string {
-	return enabled ? ON : OFF;
+async function editFireworksAccount(
+	ctx: ExtensionCommandContext,
+	settingsRuntime: UsageSettingsRuntime,
+	signal: AbortSignal,
+	isCurrent: () => boolean,
+	onApplied: (id: UsageSettingId) => void,
+): Promise<boolean> {
+	while (!signal.aborted && isCurrent()) {
+		const state = settingsRuntime.get();
+		if (state.kind === "invalid") {
+			ctx.ui.notify("Repair pi-usage.json and reload before changing settings.", "error");
+			return false;
+		}
+		const entered = await ctx.ui.input(
+			"Fireworks account slug · submit blank for Auto",
+			state.settings.fireworksAccountId ?? "Example: acme",
+			{ signal },
+		);
+		if (signal.aborted || !isCurrent() || entered === undefined) return false;
+		const normalized = entered.trim();
+		const requested = normalized || undefined;
+		if (requested !== undefined && !isFireworksAccountId(requested)) {
+			ctx.ui.notify("Enter a URL-safe Fireworks account slug.", "warning");
+			continue;
+		}
+		if (requested === state.settings.fireworksAccountId) return false;
+		try {
+			await settingsRuntime.update({ fireworksAccountId: requested }, signal);
+		} catch (error) {
+			if (signal.aborted || !isCurrent()) return false;
+			ctx.ui.notify(`Could not save pi-usage.json: ${errorMessage(error)}`, "error");
+			return false;
+		}
+		onApplied("fireworksAccountId");
+		return true;
+	}
+	return false;
+}
+
+function displaySetting(id: UsageSettingId, value: UsageSettings[UsageSettingId]): string {
+	if (id === "fireworksAccountId") return typeof value === "string" ? value : AUTO;
+	return value ? ON : OFF;
 }

package/src/usage.ts

@@ -207,13 +207,7 @@ export default function usageExtension(
 			const generation = statusGeneration;
 			statusCountdownTimer = setTimeout(() => {
 				statusCountdownTimer = undefined;
-				if (
-					!sessionActive ||
-					generation !== statusGeneration ||
-					modelIdentity(ctx.model) !== modelIdentity(model)
-				) {
-					return;
-				}
+				if (!sessionActive || generation !== statusGeneration) return;
 				publishStatus(ctx, outcome, model, false);
 			}, STATUS_COUNTDOWN_REFRESH_MS);
 			statusCountdownTimer.unref?.();
@@ -254,6 +248,10 @@ export default function usageExtension(
 		const expectedSessionGeneration = sessionGeneration;
 		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(
@@ -277,11 +275,23 @@ export default function usageExtension(
 				},
 			};
 		}
-		const requiresRequestBoundaryGuard = adapter.id === "deepseek" || adapter.id === "xai";
+		const requiresRequestBoundaryGuard = [
+			"baseten",
+			"deepseek",
+			"fireworks",
+			"minimax",
+			"minimax-cn",
+			"moonshotai",
+			"moonshotai-cn",
+			"vercel-ai-gateway",
+			"xai",
+		].includes(adapter.id);
 		const requestContextChanged = () =>
 			expectedSessionGeneration !== sessionGeneration ||
 			ctx.sessionManager.getSessionId() !== expectedSessionId ||
-			modelIdentity(ctx.model) !== expectedModelIdentity;
+			modelIdentity(ctx.model) !== expectedModelIdentity ||
+			(adapter.id === "fireworks" &&
+				settingsRuntime.get().settings.fireworksAccountId !== expectedFireworksAccountId);
 		if (requiresRequestBoundaryGuard && requestContextChanged()) throw abortError();
 		if (!auth) {
 			if (displayState === "current") {
@@ -298,11 +308,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: {
@@ -316,7 +330,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 {
@@ -335,7 +349,7 @@ export default function usageExtension(
 		const queryId = querySequence;
 		setBoundedMap(latestQueries, failureKey, queryId, MAX_ACCOUNT_STATES);
 
-		let deepSeekAuthChanged = false;
+		let retryableAuthChanged = false;
 		try {
 			const remainingMs = Math.max(1, deadlineAt - Date.now());
 			const guard = requiresRequestBoundaryGuard
@@ -349,18 +363,27 @@ export default function usageExtension(
 						);
 						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.");
+							if (["deepseek", "minimax", "minimax-cn"].includes(adapter.id)) {
+								retryableAuthChanged = true;
+								throw new Error(
+									`${adapter.displayName} runtime credential changed during the usage query.`,
+								);
 							}
 							throw abortError();
 						}
 					}
 				: undefined;
-			const report = await queryProviderUsage(adapter, auth, signal, remainingMs, guard);
+			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 {
@@ -376,7 +399,7 @@ export default function usageExtension(
 		} catch (error) {
 			if (isStaleExtensionContextError(error) || isAbortError(error)) throw error;
 			if (
-				deepSeekAuthChanged &&
+				retryableAuthChanged &&
 				authRetry === 0 &&
 				!signal.aborted &&
 				!requestContextChanged() &&
@@ -1103,7 +1126,7 @@ export default function usageExtension(
 			}
 		},
 	});
-	pi.on("session_start", (_event, ctx) => {
+	pi.on("session_start", async (_event, ctx) => {
 		sessionGeneration += 1;
 		statusGeneration += 1;
 		clearStatusTimers();
@@ -1111,7 +1134,13 @@ export default function usageExtension(
 		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);
@@ -1137,7 +1166,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 },
 	);
 }