diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index f1f3f2af36e..028bb5d264f 100644 --- a/packages/ai/CHANGELOG.md +++ b/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 diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index ae36ea537d2..2c2ecfba63f 100644 --- a/packages/ai/src/auth-storage.ts +++ b/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( @@ -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); diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index e5599d876e4..875d25c7a02 100644 --- a/packages/ai/src/index.ts +++ b/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"; diff --git a/packages/ai/src/registry/oauth/__tests__/xai-oauth.test.ts b/packages/ai/src/registry/oauth/__tests__/xai-oauth.test.ts index ac5012ba017..66eb3109df8 100644 --- a/packages/ai/src/registry/oauth/__tests__/xai-oauth.test.ts +++ b/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 { 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(); + 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([ { diff --git a/packages/ai/src/registry/oauth/xai-oauth.ts b/packages/ai/src/registry/oauth/xai-oauth.ts index d31da2750a4..59330961841 100644 --- a/packages/ai/src/registry/oauth/xai-oauth.ts +++ b/packages/ai/src/registry/oauth/xai-oauth.ts @@ -15,8 +15,12 @@ import type { OAuthController, OAuthCredentials } from "./types"; const XAI_OAUTH_ISSUER = "https://auth.x.ai"; const XAI_OAUTH_DISCOVERY_URL = `${XAI_OAUTH_ISSUER}/.well-known/openid-configuration`; const XAI_OAUTH_DEVICE_CODE_URL = `${XAI_OAUTH_ISSUER}/oauth2/device/code`; +const XAI_OAUTH_USERINFO_URL = `${XAI_OAUTH_ISSUER}/oauth2/userinfo`; const XAI_OAUTH_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828"; const XAI_OAUTH_SCOPE = "openid profile email offline_access grok-cli:access api:access"; +const XAI_CLI_BILLING_BASE_URL = "https://cli-chat-proxy.grok.com"; +const XAI_CLI_BILLING_PATH = "/v1/billing"; +const XAI_CLI_BILLING_FORMAT = "credits"; // Mirrors the 5-min skew used by anthropic.ts:160 — keeps every provider on the // same conservative client-side expiry window. @@ -51,6 +55,15 @@ function isRecord(value: unknown): value is Record { * @throws Error with message `Invalid xAI : ` when the URL fails * either scheme or host validation. */ +function isXaiAuthHostname(host: string): boolean { + return host === "x.ai" || host.endsWith(".x.ai"); +} + +/** SuperGrok CLI billing proxy host (`cli-chat-proxy.grok.com`), not the OIDC issuer. */ +function isXaiBillingHostname(host: string): boolean { + return host === "grok.com" || host.endsWith(".grok.com"); +} + export function validateXAIEndpoint(url: string, field: string): string { let parsed: URL; try { @@ -62,7 +75,28 @@ export function validateXAIEndpoint(url: string, field: string): string { throw new AIError.OAuthError(`Invalid xAI ${field}: ${url}`, { kind: "validation", provider: "xai" }); } const host = parsed.hostname.toLowerCase(); - if (!host || (host !== "x.ai" && !host.endsWith(".x.ai"))) { + if (!host || !isXaiAuthHostname(host)) { + throw new AIError.OAuthError(`Invalid xAI ${field}: ${url}`, { kind: "validation", provider: "xai" }); + } + return url; +} + +/** + * Pin SuperGrok billing URLs to HTTPS `grok.com` / `*.grok.com`. + * The CLI billing proxy is intentionally not on `*.x.ai`. + */ +export function validateXAIBillingEndpoint(url: string, field: string = "billing_url"): string { + let parsed: URL; + try { + parsed = new URL(url); + } catch { + throw new AIError.OAuthError(`Invalid xAI ${field}: ${url}`, { kind: "validation", provider: "xai" }); + } + if (parsed.protocol !== "https:") { + throw new AIError.OAuthError(`Invalid xAI ${field}: ${url}`, { kind: "validation", provider: "xai" }); + } + const host = parsed.hostname.toLowerCase(); + if (!host || !isXaiBillingHostname(host)) { throw new AIError.OAuthError(`Invalid xAI ${field}: ${url}`, { kind: "validation", provider: "xai" }); } return url; @@ -124,6 +158,22 @@ async function xaiOAuthDiscovery( return { token_endpoint: tokenEndpoint }; } +/** Decode an xAI access-token JWT payload without verifying its signature. */ +export function parseXAIAccessTokenPayload(jwt: string): Record | null { + try { + if (typeof jwt !== "string" || !jwt.includes(".")) return null; + const parts = jwt.split("."); + if (parts.length < 2) return null; + const payloadPart = parts[1]; + if (!payloadPart) return null; + const decoded = Buffer.from(payloadPart, "base64url").toString("utf8"); + const payload = JSON.parse(decoded) as unknown; + return isRecord(payload) && !Array.isArray(payload) ? payload : null; + } catch { + return null; + } +} + /** * Check whether a JWT access token is at or past its `exp` claim (with an * optional refresh-skew margin). @@ -132,25 +182,100 @@ async function xaiOAuthDiscovery( * not token validation. */ export function isXAIAccessTokenExpiring(jwt: string, skewSeconds: number = 0): boolean { + const payload = parseXAIAccessTokenPayload(jwt); + if (!payload) return false; + const exp = payload.exp; + if (typeof exp !== "number" || !Number.isFinite(exp)) return false; + const now = Math.floor(Date.now() / 1000); + const skew = Math.max(0, Math.floor(skewSeconds)); + return exp <= now + skew; +} + +/** Extract the stable xAI subject UUID from an access token. */ +export function extractXAIAccessTokenSubject(jwt: string): string | undefined { + const sub = parseXAIAccessTokenPayload(jwt)?.sub; + return typeof sub === "string" && sub.trim() ? sub.trim() : undefined; +} + +export interface XAIOAuthIdentity { + accountId?: string; + email?: string; + name?: string; +} + +/** Fetch optional OIDC userinfo for a valid xAI access token. */ +export async function fetchXAIOAuthIdentity( + accessToken: string, + fetchOverride?: FetchImpl, + signal?: AbortSignal, +): Promise { + const token = accessToken.trim(); + if (!token) return null; + const fetchImpl = fetchOverride ?? fetch; try { - if (typeof jwt !== "string" || !jwt.includes(".")) return false; - const parts = jwt.split("."); - if (parts.length < 2) return false; - const payloadPart = parts[1]; - if (!payloadPart) return false; - const decoded = Buffer.from(payloadPart, "base64url").toString("utf8"); - const payload: unknown = JSON.parse(decoded); - if (!isRecord(payload)) return false; - const exp = payload.exp; - if (typeof exp !== "number" || !Number.isFinite(exp)) return false; - const now = Math.floor(Date.now() / 1000); - const skew = Math.max(0, Math.floor(skewSeconds)); - return exp <= now + skew; + const response = await fetchImpl(XAI_OAUTH_USERINFO_URL, { + method: "GET", + headers: { + Authorization: `Bearer ${token}`, + Accept: "application/json", + }, + redirect: "error", + signal: signal + ? AbortSignal.any([signal, AbortSignal.timeout(DISCOVERY_TIMEOUT_MS)]) + : AbortSignal.timeout(DISCOVERY_TIMEOUT_MS), + }); + if (!response.ok) return null; + const payload = (await response.json()) as unknown; + if (!isRecord(payload) || Array.isArray(payload)) return null; + const sub = typeof payload.sub === "string" && payload.sub.trim() ? payload.sub.trim() : undefined; + const email = typeof payload.email === "string" && payload.email.trim() ? payload.email.trim() : undefined; + const name = typeof payload.name === "string" && payload.name.trim() ? payload.name.trim() : undefined; + if (!sub && !email && !name) return null; + return { + ...(sub ? { accountId: sub } : {}), + ...(email ? { email: email.toLowerCase() } : {}), + ...(name ? { name } : {}), + }; } catch { - return false; + return null; } } +async function withXAIOAuthIdentity( + credentials: OAuthCredentials, + fetchOverride?: FetchImpl, + signal?: AbortSignal, +): Promise { + const identity = await fetchXAIOAuthIdentity(credentials.access, fetchOverride, signal); + const accountId = identity?.accountId ?? credentials.accountId ?? extractXAIAccessTokenSubject(credentials.access); + const email = identity?.email ?? credentials.email; + return { + ...credentials, + ...(accountId ? { accountId } : {}), + ...(email ? { email } : {}), + }; +} + +/** Build the SuperGrok CLI billing URL. */ +export function buildXAICliBillingUrl(format: string = XAI_CLI_BILLING_FORMAT): string { + const url = new URL(XAI_CLI_BILLING_PATH, XAI_CLI_BILLING_BASE_URL); + url.searchParams.set("format", format); + return validateXAIBillingEndpoint(url.toString()); +} + +/** + * Headers for SuperGrok CLI billing (`cli-chat-proxy.grok.com`). + * Official Grok CLI also sends `X-XAI-Token-Auth: xai-grok-cli` on this host; + * include it so billing stays on the same product gate as chat inference. + */ +export function getXAICliBillingHeaders(options: { accessToken: string }): Record { + return { + Authorization: `Bearer ${options.accessToken}`, + Accept: "application/json", + "X-XAI-Token-Auth": "xai-grok-cli", + }; +} + function parseXAIDeviceAuthorization(payload: unknown): XAIDeviceAuthorization { if (!isRecord(payload)) { throw new AIError.OAuthError("xAI device-code response was not a JSON object.", { @@ -304,6 +429,7 @@ async function pollXAIDeviceToken( client_id: XAI_OAUTH_CLIENT_ID, device_code: deviceCode, }), + redirect: "error", signal: signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal, }); } catch (error) { @@ -364,12 +490,13 @@ export async function loginXAIOAuth(ctrl: OAuthController): Promise pollXAIDeviceToken(discovery.token_endpoint, device.deviceCode, fetchImpl, ctrl.signal), intervalSeconds: device.intervalSeconds, expiresInSeconds: device.expiresInSeconds, signal: ctrl.signal, }); + return withXAIOAuthIdentity(credentials, fetchImpl, ctrl.signal); } /** @@ -400,6 +527,7 @@ export async function refreshXAIOAuthToken(refreshToken: string, fetchOverride?: Accept: "application/json", }, body, + redirect: "error", signal: AbortSignal.timeout(TOKEN_REQUEST_TIMEOUT_MS), }); @@ -426,5 +554,6 @@ export async function refreshXAIOAuthToken(refreshToken: string, fetchOverride?: { kind: "validation", provider: "xai", cause: error }, ); } - return parseXAITokenResponse(payload, "xAI token refresh response", refreshToken); + const credentials = parseXAITokenResponse(payload, "xAI token refresh response", refreshToken); + return withXAIOAuthIdentity(credentials, fetchImpl); } diff --git a/packages/ai/src/usage/xai-oauth.ts b/packages/ai/src/usage/xai-oauth.ts new file mode 100644 index 00000000000..98a6bf52e20 --- /dev/null +++ b/packages/ai/src/usage/xai-oauth.ts @@ -0,0 +1,256 @@ +/** + * SuperGrok (`xai-oauth`) subscription usage provider. + * + * Reads weekly credit and product utilization from the Grok CLI billing + * endpoint. Only OAuth access credentials are accepted; paid API keys are a + * separate product and must never be sent here. + */ + +import { + buildXAICliBillingUrl, + extractXAIAccessTokenSubject, + fetchXAIOAuthIdentity, + getXAICliBillingHeaders, +} from "../registry/oauth/xai-oauth"; +import type { + UsageAmount, + UsageFetchContext, + UsageFetchParams, + UsageLimit, + UsageProvider, + UsageReport, + UsageStatus, + UsageWindow, +} from "../usage"; +import { isRecord } from "../utils"; +import { toNumber } from "./shared"; + +const PROVIDER_ID = "xai-oauth"; +const WEEK_MS = 7 * 24 * 60 * 60 * 1000; + +interface XaiBillingPeriod { + start: string; + end: string; + type: string; +} + +interface XaiProductUsage { + product: string; + usagePercent: number; +} + +interface XaiBillingConfig { + currentPeriod: XaiBillingPeriod; + creditUsagePercent: number; + productUsage: XaiProductUsage[]; + onDemandCap?: number; + onDemandUsed?: number; +} + +function parseIsoMs(value: string): number | undefined { + const parsed = Date.parse(value); + return Number.isFinite(parsed) ? parsed : undefined; +} + +function parsePercent(value: unknown): number | undefined { + const percent = toNumber(value); + return percent !== undefined && percent >= 0 && percent <= 100 ? percent : undefined; +} + +function parseOnDemandAmount(value: unknown): number | undefined { + if (!isRecord(value)) return undefined; + const amount = toNumber(value.val); + return amount !== undefined && amount >= 0 ? amount : undefined; +} + +function buildPercentAmount(usagePercent: number): UsageAmount { + const usedFraction = usagePercent / 100; + return { + used: usagePercent, + limit: 100, + remaining: 100 - usagePercent, + usedFraction, + remainingFraction: 1 - usedFraction, + unit: "percent", + }; +} + +function buildUsageStatus(usedFraction: number): UsageStatus { + if (usedFraction >= 1) return "exhausted"; + if (usedFraction >= 0.9) return "warning"; + return "ok"; +} + +function slugifyProduct(product: string): string { + return product + .trim() + .toLowerCase() + .replace(/[^a-z0-9]+/g, "-") + .replace(/^-+|-+$/g, ""); +} + +function buildPeriodWindow(period: XaiBillingPeriod): UsageWindow { + return { + id: "1w", + label: "Weekly", + durationMs: WEEK_MS, + resetsAt: parseIsoMs(period.end), + }; +} + +function parseBillingConfig(payload: unknown): XaiBillingConfig | null { + if (!isRecord(payload) || !isRecord(payload.config)) return null; + const raw = payload.config; + if (!isRecord(raw.currentPeriod)) return null; + + const start = typeof raw.currentPeriod.start === "string" ? parseIsoMs(raw.currentPeriod.start) : undefined; + const end = typeof raw.currentPeriod.end === "string" ? parseIsoMs(raw.currentPeriod.end) : undefined; + const type = typeof raw.currentPeriod.type === "string" ? raw.currentPeriod.type : ""; + // Keep recently-ended weekly windows so /usage still renders across period + // rollover while the billing API is mid-refresh. Reject only inverted ranges + // and non-weekly period types. + if (start === undefined || end === undefined || end <= start || !type.toUpperCase().includes("WEEK")) { + return null; + } + + const creditUsagePercent = parsePercent(raw.creditUsagePercent); + if (creditUsagePercent === undefined) return null; + + const productUsage: XaiProductUsage[] = []; + if (raw.productUsage !== undefined) { + if (!Array.isArray(raw.productUsage)) return null; + for (const item of raw.productUsage) { + if (!isRecord(item)) continue; + const product = typeof item.product === "string" ? item.product.trim() : ""; + const usagePercent = parsePercent(item.usagePercent); + if (!product || usagePercent === undefined) continue; + productUsage.push({ product, usagePercent }); + } + } + + return { + currentPeriod: { + start: raw.currentPeriod.start as string, + end: raw.currentPeriod.end as string, + type, + }, + creditUsagePercent, + productUsage, + onDemandCap: parseOnDemandAmount(raw.onDemandCap), + onDemandUsed: parseOnDemandAmount(raw.onDemandUsed), + }; +} + +function buildLimits(config: XaiBillingConfig, accountId: string | undefined): UsageLimit[] { + const window = buildPeriodWindow(config.currentPeriod); + const scope = { + provider: PROVIDER_ID, + ...(accountId ? { accountId } : {}), + windowId: window.id, + shared: true as const, + }; + const overall = buildPercentAmount(config.creditUsagePercent); + const limits: UsageLimit[] = [ + { + id: `${PROVIDER_ID}:credits:1w`, + label: "SuperGrok Weekly Credits", + scope, + window, + amount: overall, + status: buildUsageStatus(overall.usedFraction ?? 0), + }, + ]; + + for (const item of config.productUsage) { + const amount = buildPercentAmount(item.usagePercent); + const slug = slugifyProduct(item.product); + if (!slug) continue; + limits.push({ + id: `${PROVIDER_ID}:product:${slug}:1w`, + label: `${item.product === "GrokBuild" ? "Grok Build" : item.product === "Api" ? "API" : item.product} (Weekly)`, + scope, + window, + amount, + status: buildUsageStatus(amount.usedFraction ?? 0), + }); + } + if (config.onDemandCap !== undefined && config.onDemandCap > 0 && config.onDemandUsed !== undefined) { + const usedFraction = Math.min(config.onDemandUsed / config.onDemandCap, 1); + limits.push({ + id: `${PROVIDER_ID}:on-demand`, + label: "On-demand", + scope: { + provider: PROVIDER_ID, + ...(accountId ? { accountId } : {}), + shared: true, + }, + amount: { + used: config.onDemandUsed, + limit: config.onDemandCap, + remaining: Math.max(0, config.onDemandCap - config.onDemandUsed), + usedFraction, + remainingFraction: 1 - usedFraction, + unit: "unknown", + }, + status: buildUsageStatus(usedFraction), + }); + } + + return limits; +} + +export const xaiOauthUsageProvider: UsageProvider = { + id: PROVIDER_ID, + + supports(params: UsageFetchParams): boolean { + return params.provider === PROVIDER_ID && params.credential.type === "oauth" && !!params.credential.accessToken; + }, + + async fetchUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise { + if (params.provider !== PROVIDER_ID || params.credential.type !== "oauth") return null; + const accessToken = params.credential.accessToken?.trim(); + if (!accessToken) return null; + if (params.credential.expiresAt !== undefined && params.credential.expiresAt <= Date.now()) return null; + + let accountId = params.credential.accountId?.trim() || extractXAIAccessTokenSubject(accessToken); + let email = params.credential.email?.trim().toLowerCase(); + if (!email) { + try { + const identity = await fetchXAIOAuthIdentity(accessToken, ctx.fetch, params.signal); + email = identity?.email?.trim().toLowerCase() || undefined; + accountId ??= identity?.accountId?.trim() || undefined; + } catch { + // Identity enrichment is best effort; billing remains authoritative. + } + } + + const url = buildXAICliBillingUrl(); + let payload: unknown; + try { + const response = await ctx.fetch(url, { + headers: getXAICliBillingHeaders({ accessToken }), + redirect: "error", + signal: params.signal, + }); + if (!response.ok) return null; + payload = await response.json(); + } catch { + return null; + } + + const config = parseBillingConfig(payload); + if (!config) return null; + return { + provider: PROVIDER_ID, + fetchedAt: Date.now(), + limits: buildLimits(config, accountId), + metadata: { + endpoint: url, + source: "cli-chat-proxy.grok.com/v1/billing", + ...(accountId ? { accountId } : {}), + ...(email ? { email } : {}), + }, + raw: payload, + }; + }, +}; diff --git a/packages/ai/test/auth-storage-xai-oauth-usage.test.ts b/packages/ai/test/auth-storage-xai-oauth-usage.test.ts new file mode 100644 index 00000000000..d20317c7768 --- /dev/null +++ b/packages/ai/test/auth-storage-xai-oauth-usage.test.ts @@ -0,0 +1,144 @@ +import { describe, expect, it } from "bun:test"; +import { type AuthCredentialStore, AuthStorage, type StoredAuthCredential } from "@oh-my-pi/pi-ai/auth-storage"; +import type { UsageFetchParams, UsageProvider } from "@oh-my-pi/pi-ai/usage"; +import { withEnv } from "./helpers"; + +function makeStore(credentials: StoredAuthCredential[] = []): AuthCredentialStore { + const cache = new Map(); + return { + close() {}, + listAuthCredentials() { + return credentials; + }, + updateAuthCredential() {}, + deleteAuthCredential() {}, + tryDisableAuthCredentialIfMatches() { + return false; + }, + replaceAuthCredentialsForProvider() { + return []; + }, + upsertAuthCredentialForProvider() { + return []; + }, + deleteAuthCredentialsForProvider() {}, + getCache(key) { + const entry = cache.get(key); + if (!entry || entry.expiresAtSec * 1000 <= Date.now()) return null; + return entry.value; + }, + setCache(key, value, expiresAtSec) { + cache.set(key, { value, expiresAtSec }); + }, + cleanExpiredCache() {}, + }; +} + +function captureUsageProvider(calls: UsageFetchParams[]): UsageProvider { + return { + id: "xai-oauth", + supports: params => params.credential.type === "oauth" && !!params.credential.accessToken, + async fetchUsage(params) { + calls.push(params); + return { + provider: "xai-oauth", + fetchedAt: Date.now(), + limits: [], + }; + }, + }; +} + +describe("xAI OAuth environment usage", () => { + it("treats dedicated XAI_OAUTH_TOKEN as an OAuth bearer for SuperGrok usage", async () => { + const calls: UsageFetchParams[] = []; + await withEnv({ XAI_OAUTH_TOKEN: "oauth-bearer", XAI_API_KEY: undefined }, async () => { + const storage = new AuthStorage(makeStore(), { + usageProviderResolver: provider => (provider === "xai-oauth" ? captureUsageProvider(calls) : undefined), + }); + await storage.reload(); + + await storage.fetchUsageReports(); + }); + + expect(calls).toHaveLength(1); + expect(calls[0]?.credential).toEqual({ type: "oauth", accessToken: "oauth-bearer" }); + }); + + it("prefers stored OAuth credentials to XAI_OAUTH_TOKEN", async () => { + const calls: UsageFetchParams[] = []; + await withEnv({ XAI_OAUTH_TOKEN: "env-oauth-bearer" }, async () => { + const storage = new AuthStorage( + makeStore([ + { + id: 1, + provider: "xai-oauth", + credential: { + type: "oauth", + access: "stored-oauth-bearer", + refresh: "stored-refresh-token", + expires: Date.now() + 3_600_000, + }, + disabledCause: null, + }, + ]), + { + usageProviderResolver: provider => (provider === "xai-oauth" ? captureUsageProvider(calls) : undefined), + }, + ); + await storage.reload(); + + await storage.fetchUsageReports(); + }); + + expect(calls).toHaveLength(1); + expect(calls[0]?.credential).toEqual({ + type: "oauth", + accessToken: "stored-oauth-bearer", + refreshToken: "stored-refresh-token", + expiresAt: expect.any(Number), + }); + }); + + it("uses XAI_OAUTH_TOKEN when stored xAI credentials contain only an API key", async () => { + const calls: UsageFetchParams[] = []; + await withEnv({ XAI_OAUTH_TOKEN: "env-oauth-bearer" }, async () => { + const storage = new AuthStorage( + makeStore([ + { + id: 1, + provider: "xai-oauth", + credential: { + type: "api_key", + key: "stored-api-key", + }, + disabledCause: null, + }, + ]), + { + usageProviderResolver: provider => (provider === "xai-oauth" ? captureUsageProvider(calls) : undefined), + }, + ); + await storage.reload(); + + await storage.fetchUsageReports(); + }); + + expect(calls).toHaveLength(1); + expect(calls[0]?.credential).toEqual({ type: "oauth", accessToken: "env-oauth-bearer" }); + }); + + it("does not send shared XAI_API_KEY to the SuperGrok usage endpoint", async () => { + const calls: UsageFetchParams[] = []; + await withEnv({ XAI_OAUTH_TOKEN: undefined, XAI_API_KEY: "paid-api-key" }, async () => { + const storage = new AuthStorage(makeStore(), { + usageProviderResolver: provider => (provider === "xai-oauth" ? captureUsageProvider(calls) : undefined), + }); + await storage.reload(); + + await storage.fetchUsageReports(); + }); + + expect(calls).toEqual([]); + }); +}); diff --git a/packages/ai/test/xai-oauth-usage.test.ts b/packages/ai/test/xai-oauth-usage.test.ts new file mode 100644 index 00000000000..7647974e8be --- /dev/null +++ b/packages/ai/test/xai-oauth-usage.test.ts @@ -0,0 +1,177 @@ +import { describe, expect, it } from "bun:test"; +import { buildXAICliBillingUrl } from "@oh-my-pi/pi-ai/oauth/xai-oauth"; +import type { FetchImpl } from "@oh-my-pi/pi-ai/types"; +import type { UsageFetchParams } from "@oh-my-pi/pi-ai/usage"; +import { xaiOauthUsageProvider } from "@oh-my-pi/pi-ai/usage/xai-oauth"; + +const USER_ID = "cf12ecb5-cca4-4ba0-9f02-298071a2d052"; + +const accessTokenFixture = (() => { + const header = Buffer.from(JSON.stringify({ alg: "none", typ: "JWT" })).toString("base64url"); + const body = Buffer.from(JSON.stringify({ sub: USER_ID })).toString("base64url"); + return `${header}.${body}.sig`; +})(); + +function makeBillingPayload(overrides?: Record) { + const periodEnd = new Date(Date.now() + 2 * 24 * 60 * 60 * 1000).toISOString(); + const periodStart = new Date(Date.now() - 5 * 24 * 60 * 60 * 1000).toISOString(); + return { + config: { + creditUsagePercent: 18, + currentPeriod: { + end: periodEnd, + start: periodStart, + type: "USAGE_PERIOD_TYPE_WEEKLY", + }, + productUsage: [ + { product: "GrokBuild", usagePercent: 16 }, + { product: "Api", usagePercent: 2 }, + ], + ...overrides, + }, + }; +} + +function makeCredential(overrides?: Partial): UsageFetchParams["credential"] { + return { + type: "oauth", + accessToken: accessTokenFixture, + refreshToken: "refresh-fixture", + expiresAt: Date.now() + 3_600_000, + ...overrides, + }; +} + +function capturingFetch(payload: unknown): { + fetch: FetchImpl; + calls: Array<{ url: string; headers: Record; redirect?: RequestInit["redirect"] }>; +} { + const calls: Array<{ url: string; headers: Record; redirect?: RequestInit["redirect"] }> = []; + const fetch: FetchImpl = async (input, init) => { + const headers: Record = {}; + const raw = init?.headers; + if (raw && typeof raw === "object" && !Array.isArray(raw)) { + for (const [key, value] of Object.entries(raw as Record)) { + headers[key.toLowerCase()] = value; + } + } + const url = String(input); + calls.push({ url, headers, redirect: init?.redirect }); + if (url.includes("/oauth2/userinfo")) { + return new Response(JSON.stringify({ sub: USER_ID, email: "user@example.com" }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + return new Response(JSON.stringify(payload), { + status: 200, + headers: { "content-type": "application/json" }, + }); + }; + return { fetch, calls }; +} + +describe("xai-oauth usage provider", () => { + it("accepts stored OAuth credentials but never shared API-key fallbacks", () => { + expect(xaiOauthUsageProvider.supports?.({ provider: "xai-oauth", credential: makeCredential() })).toBe(true); + expect( + xaiOauthUsageProvider.supports?.({ + provider: "xai-oauth", + credential: { type: "api_key", apiKey: accessTokenFixture }, + }), + ).toBe(false); + }); + + it("maps weekly credit and product usage with CLI-aligned billing headers", async () => { + const { fetch, calls } = capturingFetch(makeBillingPayload()); + const report = await xaiOauthUsageProvider.fetchUsage( + { provider: "xai-oauth", credential: makeCredential() }, + { fetch: fetch }, + ); + + expect(report?.limits.map(limit => limit.id)).toEqual([ + "xai-oauth:credits:1w", + "xai-oauth:product:grokbuild:1w", + "xai-oauth:product:api:1w", + ]); + expect(report?.limits[0]?.amount.usedFraction).toBeCloseTo(0.18, 5); + expect(report?.metadata?.accountId).toBe(USER_ID); + expect(report?.metadata?.email).toBe("user@example.com"); + + const billingCall = calls.find(call => call.url.includes("/v1/billing")); + expect(billingCall?.url).toBe(buildXAICliBillingUrl()); + expect(billingCall?.headers).toEqual({ + authorization: `Bearer ${accessTokenFixture}`, + accept: "application/json", + "x-xai-token-auth": "xai-grok-cli", + }); + expect(billingCall?.redirect).toBe("error"); + }); + + it("uses a stored email without an extra userinfo request", async () => { + const { fetch, calls } = capturingFetch(makeBillingPayload()); + const report = await xaiOauthUsageProvider.fetchUsage( + { + provider: "xai-oauth", + credential: makeCredential({ accountId: "stored-account", email: "stored@example.com" }), + }, + { fetch: fetch }, + ); + + expect(report?.metadata?.accountId).toBe("stored-account"); + expect(report?.metadata?.email).toBe("stored@example.com"); + expect(calls.some(call => call.url.includes("/oauth2/userinfo"))).toBe(false); + }); + + it("maps a positive on-demand cap", async () => { + const report = await xaiOauthUsageProvider.fetchUsage( + { + provider: "xai-oauth", + credential: makeCredential(), + }, + { fetch: capturingFetch(makeBillingPayload({ onDemandCap: { val: 50 }, onDemandUsed: { val: 10 } })).fetch }, + ); + + const onDemand = report?.limits.find(limit => limit.id === "xai-oauth:on-demand"); + expect(onDemand?.amount.used).toBe(10); + expect(onDemand?.amount.limit).toBe(50); + expect(onDemand?.amount.usedFraction).toBeCloseTo(0.2, 5); + }); + + it("still reports usage when the weekly period has just ended", async () => { + const periodEnd = new Date(Date.now() - 60_000).toISOString(); + const periodStart = new Date(Date.now() - 8 * 24 * 60 * 60 * 1000).toISOString(); + const report = await xaiOauthUsageProvider.fetchUsage( + { provider: "xai-oauth", credential: makeCredential() }, + { + fetch: capturingFetch( + makeBillingPayload({ + currentPeriod: { + end: periodEnd, + start: periodStart, + type: "USAGE_PERIOD_TYPE_WEEKLY", + }, + }), + ).fetch, + }, + ); + + expect(report?.limits[0]?.id).toBe("xai-oauth:credits:1w"); + expect(report?.limits[0]?.window?.resetsAt).toBe(Date.parse(periodEnd)); + }); + + it("skips expired OAuth tokens and returns null for rejected billing", async () => { + const expired = await xaiOauthUsageProvider.fetchUsage( + { provider: "xai-oauth", credential: makeCredential({ expiresAt: Date.now() - 1 }) }, + { fetch: capturingFetch(makeBillingPayload()).fetch }, + ); + expect(expired).toBeNull(); + + const denied: FetchImpl = async () => new Response("denied", { status: 403 }); + const report = await xaiOauthUsageProvider.fetchUsage( + { provider: "xai-oauth", credential: makeCredential() }, + { fetch: denied }, + ); + expect(report).toBeNull(); + }); +});