feat(ai): SuperGrok usage tracking via billing API

packages/ai/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- Added SuperGrok (`xai-oauth`) usage tracking for weekly credits, product limits, and positive on-demand caps.
+
 ## [17.0.6] - 2026-07-20
 
 ### Fixed

packages/ai/src/auth-storage.ts

@@ -11,7 +11,7 @@ import { Database, type Statement } from "bun:sqlite";
 import { createHash } from "node:crypto";
 import * as fs from "node:fs/promises";
 import * as path from "node:path";
-import { getAgentDbPath, logger } from "@oh-my-pi/pi-utils";
+import { $env, getAgentDbPath, logger } from "@oh-my-pi/pi-utils";
 import type { ApiKeyResolver } from "./auth-retry";
 import * as AIError from "./error";
 import { isUsageLimitOutcome } from "./error/rate-limit";
@@ -57,6 +57,7 @@ import {
 	listCodexResetCredits,
 } from "./usage/openai-codex-reset";
 import { opencodeGoUsageProvider } from "./usage/opencode-go";
+import { xaiOauthUsageProvider } from "./usage/xai-oauth";
 import { zaiRankingStrategy, zaiUsageProvider } from "./usage/zai";
 
 const USAGE_RANKING_METRIC_EPSILON = 1e-9;
@@ -596,6 +597,7 @@ const DEFAULT_USAGE_PROVIDERS: UsageProvider[] = [
 	opencodeGoUsageProvider,
 	githubCopilotUsageProvider,
 	cursorUsageProvider,
+	xaiOauthUsageProvider,
 ];
 
 const DEFAULT_USAGE_PROVIDER_MAP = new Map<Provider, UsageProvider>(
@@ -3227,6 +3229,29 @@ export class AuthStorage {
 				entries = dedupedEntries;
 			}
 
+			// SuperGrok billing only accepts OAuth bearers. Catalog envVars for
+			// xai-oauth are [XAI_OAUTH_TOKEN, XAI_API_KEY], so the generic path
+			// would (a) build api_key usage requests from stored keys / the paid
+			// API env var and (b) never fall through to XAI_OAUTH_TOKEN when a
+			// non-OAuth row is the only stored credential. Skip api_key material
+			// and only env-fallback to the dedicated OAuth bearer.
+			if (providerId === "xai-oauth") {
+				let hasUsableStoredOAuthCredential = false;
+				for (const entry of entries) {
+					if (entry.credential.type !== "oauth") continue;
+					const request = this.#buildUsageRequestForOauth(provider, entry.credential, baseUrl);
+					if (providerImpl.supports && !providerImpl.supports(request)) continue;
+					requests.push(request);
+					hasUsableStoredOAuthCredential = true;
+				}
+				const oauthToken = $env.XAI_OAUTH_TOKEN?.trim();
+				if (!hasUsableStoredOAuthCredential && oauthToken) {
+					const request = this.#buildUsageRequest(provider, { type: "oauth", accessToken: oauthToken }, baseUrl);
+					if (!providerImpl.supports || providerImpl.supports(request)) requests.push(request);
+				}
+				continue;
+			}
+
 			if (entries.length === 0) {
 				const runtimeKey = this.#runtimeOverrides.get(providerId);
 				const envKey = getEnvApiKey(providerId);

packages/ai/src/index.ts

@@ -39,6 +39,7 @@ export * from "./usage/ollama";
 export * from "./usage/openai-codex";
 export * from "./usage/openai-codex-reset";
 export * from "./usage/opencode-go";
+export * from "./usage/xai-oauth";
 export * from "./usage/zai";
 export * from "./utils/anthropic-auth";
 export * from "./utils/event-stream";

packages/ai/src/registry/oauth/__tests__/xai-oauth.test.ts

@@ -1,19 +1,35 @@
 import { afterEach, describe, expect, it, vi } from "bun:test";
-import { isXAIAccessTokenExpiring, loginXAIOAuth, refreshXAIOAuthToken, validateXAIEndpoint } from "../xai-oauth";
+import {
+	buildXAICliBillingUrl,
+	extractXAIAccessTokenSubject,
+	fetchXAIOAuthIdentity,
+	getXAICliBillingHeaders,
+	isXAIAccessTokenExpiring,
+	loginXAIOAuth,
+	parseXAIAccessTokenPayload,
+	refreshXAIOAuthToken,
+	validateXAIBillingEndpoint,
+	validateXAIEndpoint,
+} from "../xai-oauth";
 
 afterEach(() => {
 	vi.restoreAllMocks();
 });
 
 function jwtWithExp(exp: number): string {
+	return jwtWithPayload({ exp });
+}
+
+function jwtWithPayload(payload: Record<string, unknown>): string {
 	const header = Buffer.from(JSON.stringify({ alg: "HS256", typ: "JWT" })).toString("base64url");
-	const payload = Buffer.from(JSON.stringify({ exp })).toString("base64url");
-	return `${header}.${payload}.sig`;
+	const encodedPayload = Buffer.from(JSON.stringify(payload)).toString("base64url");
+	return `${header}.${encodedPayload}.sig`;
 }
 
 const DISCOVERY_URL = "https://auth.x.ai/.well-known/openid-configuration";
 const DEVICE_CODE_URL = "https://auth.x.ai/oauth2/device/code";
 const TOKEN_ENDPOINT = "https://auth.x.ai/oauth2/token";
+const USERINFO_URL = "https://auth.x.ai/oauth2/userinfo";
 const CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828";
 const SCOPE = "openid profile email offline_access grok-cli:access api:access";
 
@@ -43,7 +59,10 @@ function jsonResponse(body: unknown, status: number = 200): Response {
 	});
 }
 
-function createDeviceFlowFetch(tokenResponses: readonly TokenResponse[]) {
+function createDeviceFlowFetch(
+	tokenResponses: readonly TokenResponse[],
+	userinfoResponse: TokenResponse = { body: {} },
+) {
 	const requests: RecordedRequest[] = [];
 	let tokenResponseIndex = 0;
 	const fetchMock = vi.fn(async (input: string | URL | Request, init?: RequestInit) => {
@@ -64,6 +83,9 @@ function createDeviceFlowFetch(tokenResponses: readonly TokenResponse[]) {
 			}
 			return jsonResponse(tokenResponse.body, tokenResponse.status);
 		}
+		if (url === USERINFO_URL) {
+			return jsonResponse(userinfoResponse.body, userinfoResponse.status);
+		}
 		throw new Error(`Unexpected xAI OAuth request: ${url}`);
 	});
 
@@ -101,6 +123,109 @@ describe("isXAIAccessTokenExpiring", () => {
 	});
 });
 
+describe("xAI OAuth helpers", () => {
+	it("parses JWT payloads and extracts subjects", () => {
+		const token = jwtWithPayload({ sub: " subject-123 ", exp: 1_900_000_000 });
+
+		expect(parseXAIAccessTokenPayload(token)).toEqual({ sub: " subject-123 ", exp: 1_900_000_000 });
+		expect(extractXAIAccessTokenSubject(token)).toBe("subject-123");
+		expect(parseXAIAccessTokenPayload("not-a-jwt")).toBeNull();
+		expect(extractXAIAccessTokenSubject("not-a-jwt")).toBeUndefined();
+	});
+
+	it("builds the billing URL and CLI-aligned headers", () => {
+		expect(buildXAICliBillingUrl()).toBe("https://cli-chat-proxy.grok.com/v1/billing?format=credits");
+		expect(buildXAICliBillingUrl("tokens")).toBe("https://cli-chat-proxy.grok.com/v1/billing?format=tokens");
+		expect(getXAICliBillingHeaders({ accessToken: "access-token" })).toEqual({
+			Authorization: "Bearer access-token",
+			Accept: "application/json",
+			"X-XAI-Token-Auth": "xai-grok-cli",
+		});
+	});
+
+	it("pins SuperGrok billing URLs to https grok.com hosts", () => {
+		expect(validateXAIBillingEndpoint("https://cli-chat-proxy.grok.com/v1/billing")).toBe(
+			"https://cli-chat-proxy.grok.com/v1/billing",
+		);
+		expect(() => validateXAIBillingEndpoint("https://auth.x.ai/v1/billing")).toThrow(/Invalid xAI billing_url/);
+		expect(() => validateXAIBillingEndpoint("http://cli-chat-proxy.grok.com/v1/billing")).toThrow(
+			/Invalid xAI billing_url/,
+		);
+		expect(() => validateXAIBillingEndpoint("https://evil.com/v1/billing")).toThrow(/Invalid xAI billing_url/);
+	});
+
+	it("normalizes OIDC userinfo identity", async () => {
+		const requests: RecordedRequest[] = [];
+		const fetchMock = vi.fn(async (input: string | URL | Request, init?: RequestInit) => {
+			requests.push({
+				url: typeof input === "string" ? input : input instanceof Request ? input.url : input.toString(),
+				init,
+			});
+			return jsonResponse({ sub: "profile-sub", email: "User@Example.com", name: "User" });
+		});
+
+		await expect(fetchXAIOAuthIdentity("access-token", fetchMock as unknown as typeof fetch)).resolves.toEqual({
+			accountId: "profile-sub",
+			email: "user@example.com",
+			name: "User",
+		});
+		expect(requests[0]?.url).toBe(USERINFO_URL);
+		expect(new Headers(requests[0]?.init?.headers)).toEqual(
+			new Headers({ Authorization: "Bearer access-token", Accept: "application/json" }),
+		);
+		expect(requests[0]?.init?.redirect).toBe("error");
+	});
+
+	it("combines caller cancellation with the 15-second userinfo timeout", async () => {
+		const timeoutControllers: AbortController[] = [];
+		const timeoutSpy = vi.spyOn(AbortSignal, "timeout").mockImplementation(timeoutMs => {
+			expect(timeoutMs).toBe(15_000);
+			const controller = new AbortController();
+			timeoutControllers.push(controller);
+			return controller.signal;
+		});
+		const requests: RecordedRequest[] = [];
+		const fetchMock = vi.fn(async (input: string | URL | Request, init?: RequestInit) => {
+			requests.push({
+				url: typeof input === "string" ? input : input instanceof Request ? input.url : input.toString(),
+				init,
+			});
+			const { promise, reject } = Promise.withResolvers<Response>();
+			const requestSignal = init?.signal;
+			if (!requestSignal) {
+				reject(new Error("expected userinfo request signal"));
+			} else if (requestSignal.aborted) {
+				reject(requestSignal.reason);
+			} else {
+				requestSignal.addEventListener("abort", () => reject(requestSignal.reason), { once: true });
+			}
+			return promise;
+		});
+
+		const callerController = new AbortController();
+		const callerCancelled = fetchXAIOAuthIdentity(
+			"access-token",
+			fetchMock as unknown as typeof fetch,
+			callerController.signal,
+		);
+		callerController.abort();
+		await expect(callerCancelled).resolves.toBeNull();
+		expect(requests[0]?.init?.signal).not.toBe(callerController.signal);
+
+		expect(timeoutSpy).toHaveBeenCalledWith(15_000);
+		const timeoutCancelled = fetchXAIOAuthIdentity(
+			"access-token",
+			fetchMock as unknown as typeof fetch,
+			new AbortController().signal,
+		);
+		const timeoutController = timeoutControllers[1];
+		expect(timeoutController).toBeDefined();
+		timeoutController?.abort();
+		await expect(timeoutCancelled).resolves.toBeNull();
+		expect(requests[1]?.init?.signal).not.toBe(timeoutController?.signal);
+	});
+});
+
 describe("validateXAIEndpoint", () => {
 	it("rejects non-HTTPS URLs", () => {
 		expect(() => validateXAIEndpoint("http://x.ai/token", "token_endpoint")).toThrow(/Invalid xAI token_endpoint/);
@@ -131,6 +256,35 @@ describe("refreshXAIOAuthToken", () => {
 		);
 		expect(fetchMock).not.toHaveBeenCalled();
 	});
+
+	it("persists refreshed OAuth identity from OIDC userinfo", async () => {
+		const accessToken = jwtWithPayload({ sub: "jwt-sub" });
+		const requests: RecordedRequest[] = [];
+		const fetchMock = vi.fn(async (input: string | URL | Request, init?: RequestInit) => {
+			const url = typeof input === "string" ? input : input instanceof Request ? input.url : input.toString();
+			requests.push({ url, init });
+			if (url === DISCOVERY_URL) return jsonResponse({ token_endpoint: TOKEN_ENDPOINT });
+			if (url === TOKEN_ENDPOINT) {
+				return jsonResponse({
+					access_token: accessToken,
+					expires_in: 3600,
+				});
+			}
+			if (url === USERINFO_URL) return jsonResponse({ sub: "profile-sub", email: "User@Example.com" });
+			throw new Error(`Unexpected xAI OAuth request: ${url}`);
+		});
+
+		await expect(
+			refreshXAIOAuthToken("old-refresh-token", fetchMock as unknown as typeof fetch),
+		).resolves.toMatchObject({
+			access: accessToken,
+			refresh: "old-refresh-token",
+			accountId: "profile-sub",
+			email: "user@example.com",
+		});
+		expect(requests.map(request => request.url)).toEqual([DISCOVERY_URL, TOKEN_ENDPOINT, USERINFO_URL]);
+		expect(requests.find(request => request.url === TOKEN_ENDPOINT)?.init?.redirect).toBe("error");
+	});
 });
 
 describe("loginXAIOAuth", () => {
@@ -165,7 +319,12 @@ describe("loginXAIOAuth", () => {
 			onManualCodeInput,
 		});
 
-		expect(requests.map(request => request.url)).toEqual([DISCOVERY_URL, DEVICE_CODE_URL, TOKEN_ENDPOINT]);
+		expect(requests.map(request => request.url)).toEqual([
+			DISCOVERY_URL,
+			DEVICE_CODE_URL,
+			TOKEN_ENDPOINT,
+			USERINFO_URL,
+		]);
 
 		const discoveryRequest = requests[0];
 		expect(discoveryRequest?.init?.method).toBe("GET");
@@ -230,6 +389,7 @@ describe("loginXAIOAuth", () => {
 
 		const tokenRequests = requests.filter(request => request.url === TOKEN_ENDPOINT);
 		expect(tokenRequests).toHaveLength(3);
+		expect(tokenRequests.every(request => request.init?.redirect === "error")).toBe(true);
 		expect(tokenRequests.map(request => Object.fromEntries(requestForm(request)))).toEqual([
 			{
 				grant_type: "urn:ietf:params:oauth:grant-type:device_code",
@@ -252,6 +412,54 @@ describe("loginXAIOAuth", () => {
 		expect(credentials.refresh).toBe("eventual-refresh-token");
 	});
 
+	it("retains the JWT subject and succeeds when OIDC userinfo fails", async () => {
+		const accessToken = jwtWithPayload({ sub: "jwt-sub" });
+		const controller = new AbortController();
+		const { fetchMock, requests } = createDeviceFlowFetch(
+			[
+				{
+					body: {
+						access_token: accessToken,
+						refresh_token: "refresh-token",
+						expires_in: 3600,
+					},
+				},
+			],
+			{ body: { error: "userinfo unavailable" }, status: 503 },
+		);
+
+		const credentials = await loginXAIOAuth({ fetch: fetchMock, signal: controller.signal });
+
+		expect(credentials).toMatchObject({
+			access: accessToken,
+			refresh: "refresh-token",
+			accountId: "jwt-sub",
+		});
+		expect(requests.at(-1)?.url).toBe(USERINFO_URL);
+		expect(requests.at(-1)?.init?.signal).not.toBe(controller.signal);
+	});
+
+	it("keeps the JWT subject when userinfo returns only an email", async () => {
+		const accessToken = jwtWithPayload({ sub: "jwt-sub" });
+		const { fetchMock } = createDeviceFlowFetch(
+			[
+				{
+					body: {
+						access_token: accessToken,
+						refresh_token: "refresh-token",
+						expires_in: 3600,
+					},
+				},
+			],
+			{ body: { email: "User@Example.com" } },
+		);
+
+		await expect(loginXAIOAuth({ fetch: fetchMock })).resolves.toMatchObject({
+			accountId: "jwt-sub",
+			email: "user@example.com",
+		});
+	});
+
 	it("rejects a token response that omits access_token", async () => {
 		const { fetchMock, requests } = createDeviceFlowFetch([
 			{