@narumitw/pi-usage 0.53.0 → 0.54.0

package/src/query.ts

@@ -7,18 +7,23 @@ import {
 } from "./oauth-credential-source.js";
 import { normalizeCodexBackendPayload } from "./providers/codex.js";
 import { normalizeGitHubCopilotUsagePayload } from "./providers/github-copilot.js";
+import { normalizeKimiCodingUsagePayload } from "./providers/kimi-coding.js";
 import { normalizeOpenCodeZenPayload } from "./providers/opencode-zen.js";
 import { normalizeOpenRouterKeyPayload } from "./providers/openrouter.js";
+import { normalizeXaiBillingPayload } from "./providers/xai.js";
 import { normalizeZaiQuotaPayload } from "./providers/zai.js";
 import type {
 	CodexBackendPayload,
 	GitHubCopilotUsagePayload,
+	KimiCodingUsagePayload,
 	OpenCodeZenPayload,
 	OpenRouterKeyPayload,
 	PiModel,
 	ResolvedUsageAuth,
 	UsageProviderAdapter,
 	UsageReport,
+	XaiBillingPayload,
+	XaiUserPayload,
 	ZaiQuotaPayload,
 } from "./types.js";
 
@@ -26,11 +31,21 @@ const CODEX_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";
 const GITHUB_COPILOT_USAGE_URL = "https://api.github.com/copilot_internal/user";
 const OPENROUTER_KEY_URL = "https://openrouter.ai/api/v1/key";
 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 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({
+	"X-XAI-Token-Auth": "xai-grok-cli",
+	"x-grok-client-version": "1.0.10",
+	"x-grok-client-mode": "interactive",
+});
 const MAX_SUCCESS_BODY_BYTES = 64 * 1024;
 const MAX_ERROR_BODY_BYTES = 4 * 1024;
 
 export const AUTH_FINGERPRINT_SALT = randomBytes(32);
 
+export type UsageRequestGuard = () => Promise<void>;
+
 export const SUPPORTED_ADAPTERS: readonly UsageProviderAdapter[] = [
 	{
 		id: "openai-codex",
@@ -98,6 +113,22 @@ export const SUPPORTED_ADAPTERS: readonly UsageProviderAdapter[] = [
 			return normalizeOpenCodeZenPayload(payload as OpenCodeZenPayload, Date.now());
 		},
 	},
+	{
+		id: "kimi-coding",
+		displayName: "Kimi For Coding",
+		semantics: { kind: "consumer-subscription", label: "Kimi Coding Plan usage" },
+		async query(auth, signal, timeoutMs) {
+			const payload = await fetchProviderJson(
+				KIMI_CODING_USAGE_URL,
+				auth,
+				signal,
+				timeoutMs,
+				"Kimi Coding usage endpoint",
+				{ redirect: "error" },
+			);
+			return normalizeKimiCodingUsagePayload(payload as KimiCodingUsagePayload, Date.now());
+		},
+	},
 	{
 		id: "zai",
 		displayName: "Z.AI",
@@ -137,10 +168,70 @@ export const SUPPORTED_ADAPTERS: readonly UsageProviderAdapter[] = [
 	},
 ];
 
+// Reviewed contract pins:
+// - Pi xAI provider at https://api.x.ai and OAuth scope
+//   "openid profile email offline_access grok-cli:access api:access" at
+//   e86823096c5bad39e1ca282ec24bc5eb9bec745b, unchanged at
+//   ccfe79ed238674f760c986e3a61493aab794000a.
+// - Grok Build identity, credits routes/structs, required token-auth and version headers, and
+//   client-mode telemetry at 9684fa3cdbf2995e30ea8b9b637f1db008f144fc (client version 1.0.10).
+// - xAI Management API's separate team billing boundary at
+//   723dd2aa22d17be35617463837dc47cda008d90e.
+// x-userid remains attached only to billing to bind the proxy-canonical identity as Grok Build does.
+export const XAI_ADAPTER: UsageProviderAdapter = {
+	id: "xai",
+	displayName: "xAI",
+	semantics: {
+		kind: "consumer-subscription",
+		label: "xAI consumer subscription usage",
+	},
+	publishesStatusline: false,
+	async query(auth, signal, timeoutMs, guard) {
+		if (!guard) throw new Error("xAI usage requires request-boundary revalidation.");
+		const startedAt = Date.now();
+		const clientAuth = {
+			...auth,
+			headers: { ...auth.headers, ...XAI_CLIENT_HEADERS },
+		};
+		await guard();
+		const userPayload = (await fetchProviderJson(
+			XAI_USER_URL,
+			clientAuth,
+			signal,
+			remainingTimeout(timeoutMs, startedAt),
+			"xAI consumer identity endpoint",
+			{ redirect: "error", userAgent: false },
+		)) as XaiUserPayload;
+		await guard();
+		const userId = validatedXaiUserId(userPayload.userId);
+		const billingAuth = {
+			...clientAuth,
+			headers: { ...clientAuth.headers, "x-userid": userId },
+			secrets: [...clientAuth.secrets, userId],
+		};
+		await guard();
+		const billingPayload = (await fetchProviderJson(
+			XAI_BILLING_URL,
+			billingAuth,
+			signal,
+			remainingTimeout(timeoutMs, startedAt),
+			"xAI consumer billing endpoint",
+			{ redirect: "error", userAgent: false },
+		)) as XaiBillingPayload;
+		await guard();
+		return normalizeXaiBillingPayload(billingPayload, userPayload.subscriptionTier, Date.now());
+	},
+};
+
+export function usageAdapters(xaiUsage = true): readonly UsageProviderAdapter[] {
+	return xaiUsage ? [...SUPPORTED_ADAPTERS, XAI_ADAPTER] : SUPPORTED_ADAPTERS;
+}
+
 export function adapterForProvider(
 	providerId: string | undefined,
+	xaiUsage = true,
 ): UsageProviderAdapter | undefined {
-	return SUPPORTED_ADAPTERS.find((adapter) => adapter.id === providerId);
+	return usageAdapters(xaiUsage).find((adapter) => adapter.id === providerId);
 }
 
 export function isStaleExtensionContextError(error: unknown): boolean {
@@ -204,6 +295,13 @@ export async function resolveUsageAuth(
 			offered.offeredCount === 0,
 		);
 	}
+	if (adapter.id === "xai") {
+		const offered = candidateReader
+			? candidateReader(ctx, adapter.id)
+			: fallbackOAuthCredentialCandidates(adapter.id, credentialReader);
+		if (!offered.ok) throw new Error("xAI OAuth credential discovery failed closed.");
+		return resolveXaiUsageAuth(auth, model, salt, offered.candidates);
+	}
 	const authorization = authorizationFrom(auth);
 	if (!authorization) return undefined;
 	const headers = { Authorization: authorization };
@@ -224,9 +322,10 @@ export async function queryProviderUsage(
 	auth: ResolvedUsageAuth,
 	signal: AbortSignal,
 	timeoutMs: number,
+	guard?: UsageRequestGuard,
 ): Promise<UsageReport> {
 	try {
-		return await adapter.query(auth, signal, timeoutMs);
+		return await adapter.query(auth, signal, timeoutMs, guard);
 	} catch (error) {
 		if (isStaleExtensionContextError(error) || isAbortError(error)) throw error;
 		throw new Error(redactUsageError(errorMessage(error), auth.secrets));
@@ -266,6 +365,8 @@ export async function fetchProviderJson(
 	request: {
 		method?: "GET" | "POST";
 		body?: Record<string, unknown>;
+		redirect?: RequestRedirect;
+		userAgent?: boolean;
 	} = {},
 ): Promise<Record<string, unknown>> {
 	const controller = new AbortController();
@@ -279,7 +380,9 @@ export async function fetchProviderJson(
 	}, timeoutMs);
 	try {
 		const headers = { ...auth.headers };
-		if (!hasHeader(headers, "User-Agent")) headers["User-Agent"] = "pi-usage";
+		if (request.userAgent !== false && !hasHeader(headers, "User-Agent")) {
+			headers["User-Agent"] = "pi-usage";
+		}
 		if (request.body && !hasHeader(headers, "Content-Type")) {
 			headers["Content-Type"] = "application/json";
 		}
@@ -287,8 +390,10 @@ export async function fetchProviderJson(
 			method: request.method ?? "GET",
 			headers,
 			...(request.body ? { body: JSON.stringify(request.body) } : {}),
+			...(request.redirect ? { redirect: request.redirect } : {}),
 			signal: controller.signal,
 		});
+		if (response.redirected) throw new Error(`${description} refused a redirected response.`);
 		if (controller.signal.aborted)
 			throw Object.assign(new Error("Usage query aborted."), { name: "AbortError" });
 		const text = await readBoundedResponse(
@@ -296,6 +401,7 @@ export async function fetchProviderJson(
 			response.ok ? MAX_SUCCESS_BODY_BYTES : MAX_ERROR_BODY_BYTES,
 			!response.ok,
 			description,
+			controller.signal,
 		);
 		if (controller.signal.aborted)
 			throw Object.assign(new Error("Usage query aborted."), { name: "AbortError" });
@@ -332,12 +438,16 @@ async function readBoundedResponse(
 	maxBytes: number,
 	truncateOverflow: boolean,
 	description: string,
+	signal: AbortSignal,
 ): Promise<string> {
 	if (!response.body) return "";
 	const reader = response.body.getReader();
 	const chunks: Uint8Array[] = [];
 	let total = 0;
 	let truncated = false;
+	const abort = () => void reader.cancel().catch(() => undefined);
+	if (signal.aborted) abort();
+	else signal.addEventListener("abort", abort, { once: true });
 	try {
 		while (true) {
 			const { done, value } = await reader.read();
@@ -354,6 +464,7 @@ async function readBoundedResponse(
 			total += value.byteLength;
 		}
 	} finally {
+		signal.removeEventListener("abort", abort);
 		reader.releaseLock();
 	}
 	if (truncated && !truncateOverflow) {
@@ -388,6 +499,73 @@ type UsageAuthRegistry = {
 	>;
 };
 
+function resolveXaiUsageAuth(
+	auth: RequestAuth,
+	model: PiModel,
+	salt: Uint8Array,
+	candidates: readonly unknown[],
+): ResolvedUsageAuth {
+	const resolvedAccess = bearerToken(headerValue(auth.headers, "Authorization")) ?? auth.apiKey;
+	if (!resolvedAccess) throw new Error("xAI runtime authentication was incomplete.");
+	let sawOAuth = false;
+	let sawMatchingAccess = false;
+	let sawIncompleteMatch = false;
+	const matches: Array<{ access: string; refresh: string }> = [];
+	for (const candidate of candidates) {
+		try {
+			const credential = asObject(candidate);
+			if (credential?.type !== "oauth") continue;
+			sawOAuth = true;
+			if (credential.access !== resolvedAccess) continue;
+			sawMatchingAccess = true;
+			if (
+				typeof credential.access !== "string" ||
+				!credential.access ||
+				typeof credential.refresh !== "string" ||
+				!credential.refresh ||
+				typeof credential.expires !== "number" ||
+				!Number.isFinite(credential.expires)
+			) {
+				sawIncompleteMatch = true;
+				continue;
+			}
+			matches.push({ access: credential.access, refresh: credential.refresh });
+		} catch {
+			// Malformed candidates never authorize a consumer-proxy request.
+		}
+	}
+	if (sawIncompleteMatch) throw new Error("The matching xAI OAuth credential was incomplete.");
+	if (matches.length > 1) {
+		throw new Error("Multiple OAuth credentials match the active xAI runtime account.");
+	}
+	const match = matches[0];
+	if (!match) {
+		if (!sawOAuth) {
+			throw new Error(
+				"xAI consumer usage requires the OAuth subscription account configured through Pi /login; XAI_API_KEY users can review API spend at console.x.ai.",
+			);
+		}
+		if (sawMatchingAccess) throw new Error("The matching xAI OAuth credential was incomplete.");
+		throw new Error("The active xAI runtime account does not match Pi's stored OAuth account.");
+	}
+	const authorization = `Bearer ${match.access}`;
+	const headers = { Authorization: authorization };
+	return {
+		apiKey: match.access,
+		headers,
+		fingerprint: fingerprintResolvedAuth({ headers }, salt),
+		secrets: [
+			match.access,
+			match.refresh,
+			resolvedAccess,
+			auth.apiKey,
+			headerValue(auth.headers, "Authorization"),
+			authorization,
+		].filter((value): value is string => Boolean(value)),
+		model,
+	};
+}
+
 function resolveGitHubCopilotUsageAuth(
 	auth: RequestAuth,
 	model: PiModel,
@@ -509,6 +687,8 @@ function hasOfficialUrlOrigin(value: string, providerId: string): boolean {
 		if (providerId === "openai-codex") return url.origin === "https://chatgpt.com";
 		if (providerId === "openrouter") return url.origin === "https://openrouter.ai";
 		if (providerId === "opencode-go") return url.origin === "https://opencode.ai";
+		if (providerId === "kimi-coding") return url.origin === "https://api.kimi.com";
+		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";
 		if (providerId === "github-copilot") {
@@ -536,6 +716,19 @@ function hasHeader(headers: Record<string, string>, name: string): boolean {
 	return Object.keys(headers).some((key) => key.toLowerCase() === name.toLowerCase());
 }
 
+function validatedXaiUserId(value: unknown): string {
+	if (typeof value !== "string" || !/^[A-Za-z0-9._~-]{1,128}$/u.test(value)) {
+		throw new Error("xAI consumer identity returned an unsafe canonical user ID.");
+	}
+	return value;
+}
+
+function remainingTimeout(timeoutMs: number, startedAt: number): number {
+	const remaining = timeoutMs - (Date.now() - startedAt);
+	if (remaining <= 0) throw new Error("Timed out while fetching xAI consumer usage.");
+	return remaining;
+}
+
 function zaiMonitorUrl(baseUrl: string | undefined): string {
 	const base = baseUrl?.trim();
 	if (!base) throw new Error("Z.AI model base URL is unavailable.");

package/src/settings.ts

@@ -9,10 +9,12 @@ export const MAX_USAGE_SETTINGS_BYTES = 64 * 1024;
 
 export interface UsageSettings {
 	codexFastMode: boolean;
+	xaiUsage: boolean;
 }
 
 export const DEFAULT_USAGE_SETTINGS: Readonly<UsageSettings> = Object.freeze({
 	codexFastMode: false,
+	xaiUsage: true,
 });
 
 export interface UsageSettingsState {
@@ -52,11 +54,16 @@ export function normalizeUsageSettings(value: unknown): UsageSettings | undefine
 	if (Object.hasOwn(value, "codexFastMode") && typeof value.codexFastMode !== "boolean") {
 		return undefined;
 	}
+	if (Object.hasOwn(value, "xaiUsage") && typeof value.xaiUsage !== "boolean") {
+		return undefined;
+	}
 	return {
 		codexFastMode:
 			typeof value.codexFastMode === "boolean"
 				? value.codexFastMode
 				: DEFAULT_USAGE_SETTINGS.codexFastMode,
+		xaiUsage:
+			typeof value.xaiUsage === "boolean" ? value.xaiUsage : DEFAULT_USAGE_SETTINGS.xaiUsage,
 	};
 }
 

package/src/types.ts

@@ -4,7 +4,7 @@ export type PiModel = NonNullable<ExtensionContext["model"]>;
 export type UsageModel = Pick<PiModel, "id" | "name" | "provider">;
 
 export type UsageSemanticsKind = "consumer-subscription" | "api-key" | "project";
-export type UsageUnit = "percent" | "usd" | "count";
+export type UsageUnit = "percent" | "usd" | "currency" | "count";
 export type UsageDisplayState = "current" | "configured";
 
 export interface UsageSemantics {
@@ -32,6 +32,7 @@ export interface UsageMetric {
 	label: string;
 	value: number | string;
 	unit?: UsageUnit;
+	currency?: string;
 }
 
 export interface UsageReport {
@@ -59,7 +60,12 @@ export interface UsageProviderAdapter {
 	displayName: string;
 	semantics: UsageSemantics;
 	publishesStatusline?: boolean;
-	query(auth: ResolvedUsageAuth, signal: AbortSignal, timeoutMs: number): Promise<UsageReport>;
+	query(
+		auth: ResolvedUsageAuth,
+		signal: AbortSignal,
+		timeoutMs: number,
+		guard?: () => Promise<void>,
+	): Promise<UsageReport>;
 }
 
 export type ProviderUsageState =
@@ -102,6 +108,21 @@ export type ZaiQuotaPayload = {
 	data?: unknown;
 };
 
+export type KimiCodingUsagePayload = {
+	usage?: unknown;
+	limits?: unknown;
+	boosterWallet?: unknown;
+};
+
+export type XaiUserPayload = {
+	userId?: unknown;
+	subscriptionTier?: unknown;
+};
+
+export type XaiBillingPayload = {
+	config?: unknown;
+};
+
 export type CodexBackendPayload = {
 	plan_type?: unknown;
 	rate_limit?: unknown;

package/src/usage-helpers.ts

@@ -1,10 +1,10 @@
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { sanitizeDisplayText } from "./core.js";
-import { providerIsConfigured, SUPPORTED_ADAPTERS } from "./query.js";
+import { providerIsConfigured, usageAdapters } from "./query.js";
 import type { PiModel, UsageProviderAdapter } from "./types.js";
 
-export function configuredAdapters(ctx: ExtensionContext): UsageProviderAdapter[] {
-	return SUPPORTED_ADAPTERS.filter(
+export function configuredAdapters(ctx: ExtensionContext, xaiUsage = true): UsageProviderAdapter[] {
+	return usageAdapters(xaiUsage).filter(
 		(adapter) => adapter.id === ctx.model?.provider || providerIsConfigured(ctx, adapter.id),
 	);
 }

package/src/usage-settings-ui.ts

@@ -0,0 +1,128 @@
+import {
+	type ExtensionCommandContext,
+	getSettingsListTheme,
+} from "@earendil-works/pi-coding-agent";
+import {
+	Container,
+	Key,
+	matchesKey,
+	type SettingItem,
+	SettingsList,
+	Text,
+} from "@earendil-works/pi-tui";
+import { errorMessage } from "./core.js";
+import type { UsageSettings, UsageSettingsRuntime } from "./settings.js";
+
+const OFF = "Off";
+const ON = "On";
+
+type UsageSettingId = keyof UsageSettings;
+
+export async function showUsageSettings(
+	ctx: ExtensionCommandContext,
+	settingsRuntime: UsageSettingsRuntime,
+	parentSignal: AbortSignal,
+	isCurrent: () => boolean,
+	onApplied: (id: UsageSettingId, previous: boolean, next: boolean) => void,
+): Promise<boolean> {
+	if (ctx.mode !== "tui") {
+		if (ctx.hasUI) ctx.ui.notify(`Edit settings manually: ${settingsRuntime.get().path}`, "info");
+		return false;
+	}
+	if (parentSignal.aborted || !isCurrent()) return false;
+
+	return 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: "xaiUsage",
+				label: "xAI usage",
+				description: "Report OAuth subscription allowance and credits.",
+				currentValue: state.kind !== "invalid" && state.settings.xaiUsage ? ON : OFF,
+				values: [OFF, ON],
+			},
+		];
+		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);
+		};
+		settingsList = new SettingsList(
+			items,
+			items.length + 2,
+			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") {
+						const effectivePrevious = settingId === "xaiUsage" ? false : previous;
+						settingsList.updateValue(id, displayValue(settingId, effectivePrevious));
+						if (!signal.aborted && isCurrent()) {
+							ctx.ui.notify("Repair pi-usage.json and reload before changing settings.", "error");
+							tui.requestRender();
+						}
+						return;
+					}
+					try {
+						await settingsRuntime.update({ [settingId]: requested }, signal);
+					} catch (error) {
+						if (signal.aborted || !isCurrent()) return;
+						settingsList.updateValue(id, displayValue(settingId, previous));
+						ctx.ui.notify(`Could not save pi-usage.json: ${errorMessage(error)}`, "error");
+						tui.requestRender();
+						return;
+					}
+					if (previous !== requested) {
+						changed = true;
+						onApplied(settingId, previous, requested);
+					}
+					if (signal.aborted || !isCurrent()) return;
+					settingsList.updateValue(id, displayValue(settingId, requested));
+					tui.requestRender();
+				});
+			},
+			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);
+			},
+		};
+	});
+}
+
+function displayValue(_id: UsageSettingId, enabled: boolean): string {
+	return enabled ? ON : OFF;
+}