@narumitw/pi-usage 0.59.0 → 0.60.1

package/src/settings.ts

@@ -4,6 +4,7 @@ import { chmod, mkdir, open, rename, rm, writeFile } from "node:fs/promises";
 import { basename, dirname, join } from "node:path";
 import { getAgentDir } from "@earendil-works/pi-coding-agent";
 import { isFireworksAccountId } from "./providers/fireworks.js";
+import { isBoundedTargetId } from "./usage-targets.js";
 
 export const USAGE_SETTINGS_FILE = "pi-usage.json";
 export const MAX_USAGE_SETTINGS_BYTES = 64 * 1024;
@@ -11,12 +12,13 @@ export const MAX_USAGE_SETTINGS_BYTES = 64 * 1024;
 export interface UsageSettings {
 	codexFastMode: boolean;
 	codexStatusResetCountdown: boolean;
-	fireworksAccountId?: string;
+	selectedTargets: Record<string, string>;
 }
 
 export const DEFAULT_USAGE_SETTINGS: Readonly<UsageSettings> = Object.freeze({
 	codexFastMode: false,
 	codexStatusResetCountdown: true,
+	selectedTargets: Object.freeze({}),
 });
 
 export interface UsageSettingsState {
@@ -27,6 +29,8 @@ export interface UsageSettingsState {
 	issue?: string;
 }
 
+export type UsageTargetPublicationCheck = () => Promise<void>;
+
 export interface UsageSettingsRuntime {
 	get(): Readonly<UsageSettingsState>;
 	reload(signal?: AbortSignal): Promise<Readonly<UsageSettingsState>>;
@@ -34,6 +38,12 @@ export interface UsageSettingsRuntime {
 		patch: Partial<UsageSettings>,
 		signal?: AbortSignal,
 	): Promise<Readonly<UsageSettingsState>>;
+	updateSelectedTarget(
+		providerId: string,
+		targetId: string,
+		signal?: AbortSignal,
+		checkPublishedSelection?: UsageTargetPublicationCheck,
+	): Promise<Readonly<UsageSettingsState>>;
 	flush(): Promise<void>;
 }
 
@@ -68,6 +78,12 @@ export function normalizeUsageSettings(value: unknown): UsageSettings | undefine
 	) {
 		return undefined;
 	}
+	const selectedTargets = normalizeSelectedTargets(value.selectedTargets);
+	if (Object.hasOwn(value, "selectedTargets") && !selectedTargets) return undefined;
+	const effectiveTargets = { ...(selectedTargets ?? {}) };
+	if (!effectiveTargets.fireworks && isFireworksAccountId(value.fireworksAccountId)) {
+		effectiveTargets.fireworks = value.fireworksAccountId;
+	}
 	return {
 		codexFastMode:
 			typeof value.codexFastMode === "boolean"
@@ -77,9 +93,7 @@ export function normalizeUsageSettings(value: unknown): UsageSettings | undefine
 			typeof value.codexStatusResetCountdown === "boolean"
 				? value.codexStatusResetCountdown
 				: DEFAULT_USAGE_SETTINGS.codexStatusResetCountdown,
-		...(isFireworksAccountId(value.fireworksAccountId)
-			? { fireworksAccountId: value.fireworksAccountId }
-			: {}),
+		selectedTargets: effectiveTargets,
 	};
 }
 
@@ -169,6 +183,39 @@ export function createUsageSettingsRuntime(
 				state = saved;
 				return structuredClone(state);
 			}),
+		updateSelectedTarget: (providerId, targetId, signal, checkPublishedSelection) =>
+			enqueue(async () => {
+				const transaction = await saveUsageTargetSelection(
+					path,
+					providerId,
+					targetId,
+					operations,
+					signal,
+				);
+				try {
+					await checkPublishedSelection?.();
+					throwIfAborted(signal);
+				} catch (error) {
+					try {
+						await restoreUsageSettingsState(
+							path,
+							transaction.saved,
+							transaction.previous,
+							operations,
+						);
+						state = transaction.previous;
+					} catch (rollbackError) {
+						state = await loadUsageSettings(path);
+						throw new AggregateError(
+							[error, rollbackError],
+							"Target selection changed after publication and pi-usage.json rollback failed",
+						);
+					}
+					throw error;
+				}
+				state = transaction.saved;
+				return structuredClone(state);
+			}),
 		flush: () => queue,
 	};
 }
@@ -178,16 +225,63 @@ async function saveUsageSettingsPatch(
 	patch: Partial<UsageSettings>,
 	operations: UsageSettingsFileOperations,
 	signal?: AbortSignal,
+): Promise<UsageSettingsState> {
+	return saveUsageSettingsDocument(
+		path,
+		(document) => {
+			for (const [key, value] of Object.entries(patch)) {
+				if (value === undefined) delete document[key];
+				else document[key] = value;
+			}
+		},
+		operations,
+		signal,
+	);
+}
+
+async function saveUsageTargetSelection(
+	path: string,
+	providerId: string,
+	targetId: string,
+	operations: UsageSettingsFileOperations,
+	signal?: AbortSignal,
+): Promise<{ saved: UsageSettingsState; previous: UsageSettingsState }> {
+	if (!isProviderId(providerId) || !isBoundedTargetId(targetId)) {
+		throw new Error("Refusing to save an invalid usage target selection");
+	}
+	const previous = await loadUsageSettings(path, signal);
+	const saved = await saveUsageSettingsDocument(
+		path,
+		(document) => {
+			document.selectedTargets = {
+				...(normalizeSelectedTargets(document.selectedTargets) ?? {}),
+				[providerId]: targetId,
+			};
+			if (providerId === "fireworks") delete document.fireworksAccountId;
+		},
+		operations,
+		signal,
+		previous,
+	);
+	return { saved, previous };
+}
+
+async function saveUsageSettingsDocument(
+	path: string,
+	mutate: (document: Record<string, unknown>) => void,
+	operations: UsageSettingsFileOperations,
+	signal?: AbortSignal,
+	expected?: UsageSettingsState,
 ): Promise<UsageSettingsState> {
 	const latest = await loadUsageSettings(path, signal);
 	if (latest.kind === "invalid") {
 		throw new Error("Cannot overwrite an invalid pi-usage.json; repair it and reload first");
 	}
-	const document = { ...latest.document };
-	for (const [key, value] of Object.entries(patch)) {
-		if (value === undefined) delete document[key];
-		else document[key] = value;
+	if (expected && !sameUsageSettingsDocument(latest, expected)) {
+		throw new Error("pi-usage.json changed while saving; retry the action");
 	}
+	const document = { ...latest.document };
+	mutate(document);
 	const settings = normalizeUsageSettings(document);
 	if (!settings) throw new Error("Refusing to save invalid pi-usage settings");
 	const directory = dirname(path);
@@ -218,6 +312,44 @@ async function saveUsageSettingsPatch(
 	return { kind: "loaded", path, settings, document };
 }
 
+async function restoreUsageSettingsState(
+	path: string,
+	published: UsageSettingsState,
+	previous: UsageSettingsState,
+	operations: UsageSettingsFileOperations,
+): Promise<void> {
+	if (previous.kind === "missing") {
+		const current = await loadUsageSettings(path);
+		if (!sameUsageSettingsDocument(current, published)) {
+			throw new Error("pi-usage.json changed before target selection rollback");
+		}
+		await rm(path);
+		return;
+	}
+	if (previous.kind !== "loaded" || !previous.document) {
+		throw new Error("Cannot restore invalid prior pi-usage.json settings");
+	}
+	await saveUsageSettingsDocument(
+		path,
+		(document) => {
+			for (const key of Object.keys(document)) delete document[key];
+			Object.assign(document, previous.document);
+		},
+		operations,
+		undefined,
+		published,
+	);
+}
+
+function sameUsageSettingsDocument(
+	left: Pick<UsageSettingsState, "kind" | "document">,
+	right: Pick<UsageSettingsState, "kind" | "document">,
+): boolean {
+	return (
+		left.kind === right.kind && JSON.stringify(left.document) === JSON.stringify(right.document)
+	);
+}
+
 async function chmodPrivate(path: string): Promise<void> {
 	await chmod(path, 0o600);
 }
@@ -233,3 +365,18 @@ function isRecord(value: unknown): value is Record<string, unknown> {
 function isNodeError(error: unknown): error is NodeJS.ErrnoException {
 	return error instanceof Error && "code" in error;
 }
+
+function normalizeSelectedTargets(value: unknown): Record<string, string> | undefined {
+	if (value === undefined) return {};
+	if (!isRecord(value)) return undefined;
+	const targets: Record<string, string> = {};
+	for (const [providerId, targetId] of Object.entries(value)) {
+		if (!isProviderId(providerId) || !isBoundedTargetId(targetId)) return undefined;
+		targets[providerId] = targetId;
+	}
+	return targets;
+}
+
+function isProviderId(value: string): boolean {
+	return /^[A-Za-z0-9][A-Za-z0-9._~-]{0,127}$/u.test(value);
+}

package/src/types.ts

@@ -53,23 +53,51 @@ export interface ResolvedUsageAuth {
 	fingerprint: string;
 	secrets: string[];
 	model: PiModel;
+	auth?: {
+		apiKey?: string;
+		headers?: Record<string, string | null>;
+		baseUrl?: string;
+	};
+	env?: Record<string, string>;
+	source?: string;
+	effectiveBaseUrl?: string;
 }
 
 export interface UsageQuerySettings {
 	fireworksAccountId?: string;
 }
 
+export type UsageRequestGuard = () => Promise<void>;
+
+export interface UsageProviderTarget {
+	id: string;
+	label: string;
+	description?: string;
+}
+
+export interface UsageTargetResolver {
+	singularLabel: string;
+	pluralLabel: string;
+	list(
+		auth: ResolvedUsageAuth,
+		signal: AbortSignal,
+		timeoutMs: number,
+		guard: UsageRequestGuard,
+	): Promise<readonly UsageProviderTarget[]>;
+}
+
 export interface UsageProviderAdapter {
 	id: string;
 	displayName: string;
 	semantics: UsageSemantics;
 	publishesStatusline?: boolean;
+	targets?: UsageTargetResolver;
 	query(
 		auth: ResolvedUsageAuth,
 		signal: AbortSignal,
 		timeoutMs: number,
-		guard?: () => Promise<void>,
-		settings?: Readonly<UsageQuerySettings>,
+		guard?: UsageRequestGuard,
+		targetId?: string,
 	): Promise<UsageReport>;
 }
 
@@ -81,6 +109,15 @@ export type ProviderUsageState =
 			status: "ready";
 			report: UsageReport;
 	  }
+	| {
+			providerId: string;
+			providerName: string;
+			displayState: UsageDisplayState;
+			status: "selection-required";
+			singularLabel: string;
+			pluralLabel: string;
+			choices: readonly UsageProviderTarget[];
+	  }
 	| {
 			providerId: string;
 			providerName: string;
@@ -157,6 +194,17 @@ export type ZaiQuotaPayload = {
 	data?: unknown;
 };
 
+export type ZaiSubscriptionPayload = {
+	code?: unknown;
+	success?: unknown;
+	data?: unknown;
+};
+
+export interface ZaiPlanInfo {
+	name: string;
+	renewsAt?: string;
+}
+
 export type KimiCodingUsagePayload = {
 	usage?: unknown;
 	limits?: unknown;

package/src/usage-settings-ui.ts

@@ -11,16 +11,12 @@ 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";
+import type { 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 };
+type UsageSettingId = "codexFastMode" | "codexStatusResetCountdown";
 
 export async function showUsageSettings(
 	ctx: ExtensionCommandContext,
@@ -33,187 +29,96 @@ export async function showUsageSettings(
 		if (ctx.hasUI) ctx.ui.notify(`Edit settings manually: ${settingsRuntime.get().path}`, "info");
 		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 (
+		(await ctx.ui.custom<boolean>((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 items: SettingItem[] = [
+				{
+					id: "codexFastMode",
+					label: "Codex Fast mode",
+					description: "Use faster Codex routing at increased plan allowance consumption.",
+					currentValue: state.settings.codexFastMode ? ON : OFF,
+					values: [OFF, ON],
+				},
+				{
+					id: "codexStatusResetCountdown",
+					label: "Codex reset countdown",
+					description: "Show time remaining until each Codex usage limit resets.",
+					currentValue: state.settings.codexStatusResetCountdown ? ON : OFF,
+					values: [OFF, ON],
+				},
+			];
+			const container = new Container();
+			container.addChild(new Text(theme.fg("accent", theme.bold("pi-usage Settings")), 1, 1));
 
-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",
-				label: "Codex Fast mode",
-				description: "Use faster Codex routing at increased plan allowance consumption.",
-				currentValue: state.settings.codexFastMode ? ON : OFF,
-				values: [OFF, ON],
-			},
-			{
-				id: "codexStatusResetCountdown",
-				label: "Codex reset countdown",
-				description: "Show time remaining until each Codex usage limit resets.",
-				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));
-
-		let settingsList: SettingsList;
-		const cancel = () => {
-			if (closing) return;
-			closing = true;
-			localController.abort();
-			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");
+			let settingsList: SettingsList;
+			const cancel = () => {
+				if (closing) return;
+				closing = true;
+				localController.abort();
+				done(changed);
+			};
+			const queueUpdate = (id: UsageSettingId, requested: boolean, display: string) => {
+				saveQueue = saveQueue.then(async () => {
+					const previous = settingsRuntime.get().settings[id];
+					if (settingsRuntime.get().kind === "invalid") {
+						settingsList.updateValue(id, previous ? ON : OFF);
+						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, previous ? ON : OFF);
+						ctx.ui.notify(`Could not save pi-usage.json: ${errorMessage(error)}`, "error");
 						tui.requestRender();
+						return;
+					}
+					if (previous !== requested) {
+						changed = true;
+						onApplied(id);
 					}
-					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");
+					settingsList.updateValue(id, display);
 					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,
-			items.length + 2,
-			getSettingsListTheme(),
-			(id, value) => {
-				if (closing || signal.aborted || !isCurrent()) return;
-				if (id === "fireworksAccountId") {
-					if (value === EDIT) {
-						saveQueue = saveQueue.then(() => {
-							if (closing || signal.aborted || !isCurrent()) return;
-							closing = true;
-							done({ changed, editFireworksAccount: true });
-						});
-					}
-					return;
-				}
-				const settingId = id as "codexFastMode" | "codexStatusResetCountdown";
-				queueUpdate(settingId, value !== OFF, value);
-			},
-			cancel,
-		);
-		container.addChild(settingsList);
+				});
+			};
+			settingsList = new SettingsList(
+				items,
+				items.length + 2,
+				getSettingsListTheme(),
+				(id, value) => {
+					if (closing || signal.aborted || !isCurrent()) return;
+					queueUpdate(id as UsageSettingId, value !== OFF, value);
+				},
+				cancel,
+			);
+			container.addChild(settingsList);
 
-		parentSignal.addEventListener("abort", cancel, { once: true });
-		return {
-			render: (width: number) => container.render(width),
-			invalidate: () => container.invalidate(),
-			handleInput(data: string) {
-				if (closing) return;
-				if (matchesKey(data, Key.ctrl("c"))) cancel();
-				else settingsList.handleInput(data);
-				tui.requestRender();
-			},
-			dispose() {
-				localController.abort();
-				parentSignal.removeEventListener("abort", cancel);
-			},
-		};
-	});
-}
-
-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;
+			parentSignal.addEventListener("abort", cancel, { once: true });
+			return {
+				render: (width: number) => container.render(width),
+				invalidate: () => container.invalidate(),
+				handleInput(data: string) {
+					if (closing) return;
+					if (matchesKey(data, Key.ctrl("c"))) cancel();
+					else settingsList.handleInput(data);
+					tui.requestRender();
+				},
+				dispose() {
+					localController.abort();
+					parentSignal.removeEventListener("abort", cancel);
+				},
+			};
+		})) ?? false
+	);
 }