@narumitw/pi-usage 0.58.0 → 0.59.0

package/src/providers/vercel-ai-gateway.ts

@@ -0,0 +1,43 @@
+import type { UsageMetric, UsageReport, VercelAIGatewayCreditsPayload } from "../types.js";
+
+const DECIMAL_AMOUNT = /^(?:0|[1-9]\d*)(?:\.\d+)?$/u;
+
+export function normalizeVercelAIGatewayCreditsPayload(
+	payload: VercelAIGatewayCreditsPayload,
+	capturedAt: number,
+): UsageReport {
+	const balance = decimalAmount(payload.balance, "balance");
+	const totalUsed = decimalAmount(payload.total_used, "total used");
+	const metrics: UsageMetric[] = [
+		{
+			id: "credit-balance",
+			label: "Credit balance",
+			value: balance,
+			unit: "currency",
+			currency: "USD",
+		},
+		{
+			id: "lifetime-spend",
+			label: "Lifetime spend",
+			value: totalUsed,
+			unit: "currency",
+			currency: "USD",
+		},
+	];
+	return {
+		providerId: "vercel-ai-gateway",
+		providerName: "Vercel AI Gateway",
+		capturedAt,
+		source: "vercel-ai-gateway-credits",
+		semantics: { kind: "api-key", label: "AI Gateway credits and lifetime spend" },
+		buckets: [],
+		metrics,
+	};
+}
+
+function decimalAmount(value: unknown, label: string): string {
+	if (typeof value !== "string" || value.length > 64 || !DECIMAL_AMOUNT.test(value)) {
+		throw new Error(`Vercel AI Gateway ${label} was not a valid nonnegative amount.`);
+	}
+	return value;
+}

package/src/query.ts

@@ -5,6 +5,7 @@ import {
 	fallbackOAuthCredentialCandidates,
 	type OAuthCredentialCandidateReader,
 } from "./oauth-credential-source.js";
+import { normalizeBasetenBillingUsagePayload } from "./providers/baseten.js";
 import { normalizeCodexBackendPayload } from "./providers/codex.js";
 import { normalizeDeepSeekBalancePayload } from "./providers/deepseek.js";
 import {
@@ -14,17 +15,27 @@ import {
 } from "./providers/fireworks.js";
 import { normalizeGitHubCopilotUsagePayload } from "./providers/github-copilot.js";
 import { normalizeKimiCodingUsagePayload } from "./providers/kimi-coding.js";
+import {
+	type MiniMaxProviderId,
+	miniMaxUsageKind,
+	normalizeMiniMaxUsagePayload,
+} from "./providers/minimax.js";
+import { normalizeMoonshotBalancePayload } from "./providers/moonshot.js";
 import { normalizeOpenCodeZenPayload } from "./providers/opencode-zen.js";
 import { normalizeOpenRouterKeyPayload } from "./providers/openrouter.js";
+import { normalizeVercelAIGatewayCreditsPayload } from "./providers/vercel-ai-gateway.js";
 import { normalizeXaiBillingPayload } from "./providers/xai.js";
 import { normalizeZaiQuotaPayload } from "./providers/zai.js";
 import type {
+	BasetenBillingUsagePayload,
 	CodexBackendPayload,
 	DeepSeekBalancePayload,
 	FireworksAccountsPayload,
 	FireworksBillingSummaryPayload,
 	GitHubCopilotUsagePayload,
 	KimiCodingUsagePayload,
+	MiniMaxUsagePayload,
+	MoonshotBalancePayload,
 	OpenCodeZenPayload,
 	OpenRouterKeyPayload,
 	PiModel,
@@ -32,11 +43,14 @@ import type {
 	UsageProviderAdapter,
 	UsageQuerySettings,
 	UsageReport,
+	VercelAIGatewayCreditsPayload,
 	XaiBillingPayload,
 	XaiUserPayload,
 	ZaiQuotaPayload,
 } from "./types.js";
 
+const BASETEN_BILLING_USAGE_URL = "https://api.baseten.co/v1/billing/usage_summary";
+const BASETEN_USAGE_WINDOW_DAYS = 30;
 const CODEX_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";
 const DEEPSEEK_BALANCE_URL = "https://api.deepseek.com/user/balance";
 const FIREWORKS_BILLING_SUMMARY_ORIGIN = "https://api.fireworks.ai";
@@ -44,8 +58,18 @@ const FIREWORKS_SPEND_WINDOW_DAYS = 30;
 const FIREWORKS_MAX_ACCOUNT_PAGES = 5;
 const GITHUB_COPILOT_USAGE_URL = "https://api.github.com/copilot_internal/user";
 const OPENROUTER_KEY_URL = "https://openrouter.ai/api/v1/key";
+const VERCEL_AI_GATEWAY_CREDITS_URL = "https://ai-gateway.vercel.sh/v1/credits";
 const OPENCODE_GO_USAGE_URL = "https://opencode.ai/zen/go/v1/usage";
 const KIMI_CODING_USAGE_URL = "https://api.kimi.com/coding/v1/usages";
+const MINIMAX_API_ROOTS = Object.freeze({
+	minimax: "https://api.minimax.io",
+	"minimax-cn": "https://api.minimaxi.com",
+});
+const MOONSHOT_BALANCE_URLS = Object.freeze({
+	moonshotai: "https://api.moonshot.ai/v1/users/me/balance",
+	"moonshotai-cn": "https://api.moonshot.cn/v1/users/me/balance",
+});
+const SHARED_MOONSHOT_ENV_VAR = "MOONSHOT_API_KEY";
 const XAI_USER_URL = "https://cli-chat-proxy.grok.com/v1/user?include=subscription";
 const XAI_BILLING_URL = "https://cli-chat-proxy.grok.com/v1/billing?format=credits";
 const XAI_CLIENT_HEADERS = Object.freeze({
@@ -61,6 +85,27 @@ export const AUTH_FINGERPRINT_SALT = randomBytes(32);
 export type UsageRequestGuard = () => Promise<void>;
 
 export const SUPPORTED_ADAPTERS: readonly UsageProviderAdapter[] = [
+	{
+		id: "baseten",
+		displayName: "Baseten",
+		semantics: { kind: "api-key", label: "Organization Model APIs spend" },
+		async query(auth, signal, timeoutMs, guard) {
+			if (!guard) throw new Error("Baseten billing usage requires request-boundary revalidation.");
+			const startedAt = Date.now();
+			await guard();
+			const windowAt = Date.now();
+			const payload = (await fetchProviderJson(
+				basetenBillingUsageUrl(windowAt),
+				auth,
+				signal,
+				remainingTimeout(timeoutMs, startedAt, "fetching Baseten billing usage"),
+				"Baseten billing usage endpoint",
+				{ redirect: "error" },
+			)) as BasetenBillingUsagePayload;
+			await guard();
+			return normalizeBasetenBillingUsagePayload(payload, Date.now());
+		},
+	},
 	{
 		id: "openai-codex",
 		displayName: "OpenAI Codex",
@@ -133,6 +178,27 @@ export const SUPPORTED_ADAPTERS: readonly UsageProviderAdapter[] = [
 			return normalizeOpenRouterKeyPayload(payload as OpenRouterKeyPayload, Date.now());
 		},
 	},
+	{
+		id: "vercel-ai-gateway",
+		displayName: "Vercel AI Gateway",
+		semantics: { kind: "api-key", label: "AI Gateway credits and lifetime spend" },
+		async query(auth, signal, timeoutMs, guard) {
+			if (!guard)
+				throw new Error("Vercel AI Gateway usage requires request-boundary revalidation.");
+			const startedAt = Date.now();
+			await guard();
+			const payload = (await fetchProviderJson(
+				VERCEL_AI_GATEWAY_CREDITS_URL,
+				auth,
+				signal,
+				remainingTimeout(timeoutMs, startedAt, "fetching Vercel AI Gateway credits"),
+				"Vercel AI Gateway credits endpoint",
+				{ redirect: "error" },
+			)) as VercelAIGatewayCreditsPayload;
+			await guard();
+			return normalizeVercelAIGatewayCreditsPayload(payload, Date.now());
+		},
+	},
 	{
 		id: "fireworks",
 		displayName: "Fireworks",
@@ -193,6 +259,38 @@ export const SUPPORTED_ADAPTERS: readonly UsageProviderAdapter[] = [
 			return normalizeKimiCodingUsagePayload(payload as KimiCodingUsagePayload, Date.now());
 		},
 	},
+	{
+		id: "minimax",
+		displayName: "MiniMax",
+		semantics: { kind: "consumer-subscription", label: "MiniMax usage" },
+		async query(auth, signal, timeoutMs, guard) {
+			return queryMiniMaxUsage("minimax", auth, signal, timeoutMs, guard);
+		},
+	},
+	{
+		id: "minimax-cn",
+		displayName: "MiniMax CN",
+		semantics: { kind: "consumer-subscription", label: "MiniMax usage" },
+		async query(auth, signal, timeoutMs, guard) {
+			return queryMiniMaxUsage("minimax-cn", auth, signal, timeoutMs, guard);
+		},
+	},
+	{
+		id: "moonshotai",
+		displayName: "Moonshot AI",
+		semantics: { kind: "api-key", label: "Moonshot API account balance" },
+		async query(auth, signal, timeoutMs, guard) {
+			return queryMoonshotBalance("moonshotai", auth, signal, timeoutMs, guard);
+		},
+	},
+	{
+		id: "moonshotai-cn",
+		displayName: "Moonshot AI CN",
+		semantics: { kind: "api-key", label: "Moonshot API account balance" },
+		async query(auth, signal, timeoutMs, guard) {
+			return queryMoonshotBalance("moonshotai-cn", auth, signal, timeoutMs, guard);
+		},
+	},
 	{
 		id: "zai",
 		displayName: "Z.AI",
@@ -329,11 +427,14 @@ export async function resolveUsageAuth(
 		if (!result.ok) throw new Error(redactUsageError(result.error));
 		return authorizationFrom(result) ? result : undefined;
 	};
-	if (adapter.id !== "deepseek") modelAuth = await resolveCurrentModelAuth();
+	const resolveSelectedAuthLast = ["deepseek", "minimax", "minimax-cn"].includes(adapter.id);
+	if (!resolveSelectedAuthLast) modelAuth = await resolveCurrentModelAuth();
 	if (typeof registry.getProviderAuth !== "function") {
 		throw new Error("pi-usage requires Pi 0.81.0 or newer to validate resolved provider auth.");
 	}
+	if (!moonshotProviderAuthIsAllowed(ctx, adapter.id)) return undefined;
 	const providerResult = await registry.getProviderAuth(adapter.id);
+	if (!moonshotProviderAuthIsAllowed(ctx, adapter.id)) return undefined;
 	if (
 		providerResult?.auth.baseUrl &&
 		!hasOfficialUrlOrigin(providerResult.auth.baseUrl, adapter.id)
@@ -342,9 +443,9 @@ export async function resolveUsageAuth(
 			`${adapter.displayName} usage cannot send a proxy-resolved credential to the official usage endpoint.`,
 		);
 	}
-	// DeepSeek reads selected-model auth last so a rotation during provider-origin validation
-	// cannot leave the earlier credential queued for the balance request.
-	if (adapter.id === "deepseek") modelAuth = await resolveCurrentModelAuth();
+	// Providers with credential-change retries read selected-model auth last so a rotation during
+	// provider-origin validation cannot leave the earlier credential queued for the usage request.
+	if (resolveSelectedAuthLast) modelAuth = await resolveCurrentModelAuth();
 	const auth = modelAuth ?? providerResult?.auth;
 	if (!auth) return undefined;
 	if (adapter.id === "github-copilot") {
@@ -422,12 +523,53 @@ export async function queryProviderUsage(
 
 export function providerIsConfigured(ctx: ExtensionContext, providerId: string): boolean {
 	try {
-		return ctx.modelRegistry.getProviderAuthStatus(providerId).configured;
+		const status = ctx.modelRegistry.getProviderAuthStatus(providerId);
+		return (
+			status.configured &&
+			moonshotProviderAuthSourceIsAllowed(ctx, providerId, status.source, status.label)
+		);
+	} catch {
+		return (
+			!isMoonshotSiblingProvider(ctx, providerId) && candidateModels(ctx, providerId).length > 0
+		);
+	}
+}
+
+function moonshotProviderAuthIsAllowed(ctx: ExtensionContext, providerId: string): boolean {
+	if (!isMoonshotSiblingProvider(ctx, providerId)) return true;
+	try {
+		const status = ctx.modelRegistry.getProviderAuthStatus(providerId);
+		return moonshotProviderAuthSourceIsAllowed(ctx, providerId, status.source, status.label);
 	} catch {
-		return candidateModels(ctx, providerId).length > 0;
+		return false;
 	}
 }
 
+function moonshotProviderAuthSourceIsAllowed(
+	ctx: ExtensionContext,
+	providerId: string,
+	source: string | undefined,
+	label: string | undefined,
+): boolean {
+	if (!isMoonshotSiblingProvider(ctx, providerId)) return true;
+	if (source === undefined) return false;
+	if (source !== "environment") return true;
+	return (
+		label !== undefined &&
+		!label
+			.split(",")
+			.map((name) => name.trim())
+			.includes(SHARED_MOONSHOT_ENV_VAR)
+	);
+}
+
+function isMoonshotSiblingProvider(ctx: ExtensionContext, providerId: string): boolean {
+	return (
+		(providerId === "moonshotai" || providerId === "moonshotai-cn") &&
+		ctx.model?.provider !== providerId
+	);
+}
+
 function candidateModels(ctx: ExtensionContext, providerId: string): PiModel[] {
 	const candidates: PiModel[] = [];
 	const seen = new Set<string>();
@@ -772,12 +914,20 @@ function hasOfficialOrigin(model: PiModel, providerId: string): boolean {
 function hasOfficialUrlOrigin(value: string, providerId: string): boolean {
 	try {
 		const url = new URL(value);
+		if (providerId === "baseten") {
+			return ["https://inference.baseten.co", "https://api.baseten.co"].includes(url.origin);
+		}
 		if (providerId === "openai-codex") return url.origin === "https://chatgpt.com";
 		if (providerId === "deepseek") return url.origin === "https://api.deepseek.com";
 		if (providerId === "fireworks") return url.origin === FIREWORKS_BILLING_SUMMARY_ORIGIN;
 		if (providerId === "openrouter") return url.origin === "https://openrouter.ai";
+		if (providerId === "vercel-ai-gateway") return url.origin === "https://ai-gateway.vercel.sh";
 		if (providerId === "opencode-go") return url.origin === "https://opencode.ai";
 		if (providerId === "kimi-coding") return url.origin === "https://api.kimi.com";
+		if (providerId === "minimax") return url.origin === "https://api.minimax.io";
+		if (providerId === "minimax-cn") return url.origin === "https://api.minimaxi.com";
+		if (providerId === "moonshotai") return url.origin === "https://api.moonshot.ai";
+		if (providerId === "moonshotai-cn") return url.origin === "https://api.moonshot.cn";
 		if (providerId === "xai") return url.origin === "https://api.x.ai";
 		if (providerId === "zai") return url.origin === "https://api.z.ai";
 		if (providerId === "zai-coding-cn") return url.origin === "https://open.bigmodel.cn";
@@ -813,6 +963,64 @@ function validatedXaiUserId(value: unknown): string {
 	return value;
 }
 
+function basetenBillingUsageUrl(windowAt: number): string {
+	const url = new URL(BASETEN_BILLING_USAGE_URL);
+	url.searchParams.set(
+		"start_date",
+		new Date(windowAt - BASETEN_USAGE_WINDOW_DAYS * 24 * 60 * 60 * 1_000).toISOString(),
+	);
+	url.searchParams.set("end_date", new Date(windowAt).toISOString());
+	return url.toString();
+}
+
+async function queryMiniMaxUsage(
+	providerId: MiniMaxProviderId,
+	auth: ResolvedUsageAuth,
+	signal: AbortSignal,
+	timeoutMs: number,
+	guard: UsageRequestGuard | undefined,
+): Promise<UsageReport> {
+	if (!guard) throw new Error("MiniMax usage requires request-boundary revalidation.");
+	const apiKey = bearerToken(headerValue(auth.headers, "Authorization")) ?? auth.apiKey;
+	if (!apiKey) throw new Error("MiniMax runtime API key was unavailable.");
+	const kind = miniMaxUsageKind(apiKey);
+	const path = kind === "account-balance" ? "/account/query_balance" : "/v1/token_plan/remains";
+	const startedAt = Date.now();
+	await guard();
+	const payload = (await fetchProviderJson(
+		`${MINIMAX_API_ROOTS[providerId]}${path}`,
+		auth,
+		signal,
+		remainingTimeout(timeoutMs, startedAt, "fetching MiniMax usage"),
+		"MiniMax usage endpoint",
+		{ redirect: "error" },
+	)) as MiniMaxUsagePayload;
+	await guard();
+	return normalizeMiniMaxUsagePayload(providerId, kind, payload, Date.now());
+}
+
+async function queryMoonshotBalance(
+	providerId: "moonshotai" | "moonshotai-cn",
+	auth: ResolvedUsageAuth,
+	signal: AbortSignal,
+	timeoutMs: number,
+	guard: UsageRequestGuard | undefined,
+): Promise<UsageReport> {
+	if (!guard) throw new Error("Moonshot AI balance requires request-boundary revalidation.");
+	const startedAt = Date.now();
+	await guard();
+	const payload = (await fetchProviderJson(
+		MOONSHOT_BALANCE_URLS[providerId],
+		auth,
+		signal,
+		remainingTimeout(timeoutMs, startedAt, "fetching Moonshot AI balance"),
+		"Moonshot AI balance endpoint",
+		{ redirect: "error" },
+	)) as MoonshotBalancePayload;
+	await guard();
+	return normalizeMoonshotBalancePayload(providerId, payload, Date.now());
+}
+
 function remainingTimeout(
 	timeoutMs: number,
 	startedAt: number,

package/src/types.ts

@@ -89,6 +89,12 @@ export type ProviderUsageState =
 			message: string;
 	  };
 
+export type BasetenBillingUsagePayload = {
+	dedicated_usage?: unknown;
+	model_apis_usage?: unknown;
+	training_usage?: unknown;
+};
+
 export type DeepSeekBalancePayload = {
 	is_available?: unknown;
 	balance_infos?: unknown;
@@ -119,6 +125,30 @@ export type OpenRouterKeyPayload = {
 	data?: unknown;
 };
 
+export type MiniMaxUsagePayload = {
+	available_amount?: unknown;
+	balance_alert_switch?: unknown;
+	balance_alert_threshold?: unknown;
+	base_resp?: unknown;
+	cash_balance?: unknown;
+	credit_balance?: unknown;
+	model_remains?: unknown;
+	owed_amount?: unknown;
+	voucher_balance?: unknown;
+};
+
+export type MoonshotBalancePayload = {
+	code?: unknown;
+	data?: unknown;
+	scode?: unknown;
+	status?: unknown;
+};
+
+export type VercelAIGatewayCreditsPayload = {
+	balance?: unknown;
+	total_used?: unknown;
+};
+
 export type OpenCodeZenPayload = {
 	usage?: unknown;
 };

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?.();
@@ -281,7 +275,17 @@ export default function usageExtension(
 				},
 			};
 		}
-		const requiresRequestBoundaryGuard = ["deepseek", "fireworks", "xai"].includes(adapter.id);
+		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 ||
@@ -345,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
@@ -359,9 +363,11 @@ 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();
 						}
@@ -393,7 +399,7 @@ export default function usageExtension(
 		} catch (error) {
 			if (isStaleExtensionContextError(error) || isAbortError(error)) throw error;
 			if (
-				deepSeekAuthChanged &&
+				retryableAuthChanged &&
 				authRetry === 0 &&
 				!signal.aborted &&
 				!requestContextChanged() &&