diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index f35650cadc1..53184017f66 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added SuperGrok OAuth subscription-credit usage reporting to `omp usage` through Grok's billing RPC. + ## [17.0.2] - 2026-07-17 ### Fixed diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 631561c2ae5..674f0947f78 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -57,6 +57,7 @@ import { listCodexResetCredits, } from "./usage/openai-codex-reset"; import { opencodeGoUsageProvider } from "./usage/opencode-go"; +import { xaiOAuthUsageProvider } from "./usage/xai"; import { zaiRankingStrategy, zaiUsageProvider } from "./usage/zai"; const USAGE_RANKING_METRIC_EPSILON = 1e-9; @@ -587,6 +588,7 @@ const DEFAULT_USAGE_PROVIDERS: UsageProvider[] = [ opencodeGoUsageProvider, githubCopilotUsageProvider, cursorUsageProvider, + xaiOAuthUsageProvider, ]; const DEFAULT_USAGE_PROVIDER_MAP = new Map( diff --git a/packages/ai/src/usage/xai.ts b/packages/ai/src/usage/xai.ts new file mode 100644 index 00000000000..3aaee7017aa --- /dev/null +++ b/packages/ai/src/usage/xai.ts @@ -0,0 +1,265 @@ +import type { + UsageFetchContext, + UsageFetchParams, + UsageLimit, + UsageProvider, + UsageReport, + UsageStatus, +} from "../usage"; + +const XAI_OAUTH_PROVIDER = "xai-oauth"; +const GROK_BILLING_URL = "https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig"; +const EMPTY_GRPC_WEB_MESSAGE = new Uint8Array(5); +const MIN_UNIX_TIMESTAMP_SECONDS = 1_700_000_000; +const MAX_UNIX_TIMESTAMP_SECONDS = 2_100_000_000; +const MAX_PROTOBUF_DEPTH = 4; + +interface ProtobufFixed32Field { + path: number[]; + value: number; + order: number; +} + +interface ProtobufVarintField { + path: number[]; + value: number; +} + +interface ProtobufScan { + fixed32Fields: ProtobufFixed32Field[]; + varintFields: ProtobufVarintField[]; + nextOrder: number; +} + +export interface GrokUsageSnapshot { + usedPercent: number; + resetsAt?: number; +} + +function readVarint(bytes: Uint8Array, cursor: { index: number }): number | undefined { + let value = 0; + let multiplier = 1; + for (let count = 0; cursor.index < bytes.length && count < 10; count++) { + const byte = bytes[cursor.index++]; + value += (byte & 0x7f) * multiplier; + if ((byte & 0x80) === 0) return Number.isSafeInteger(value) ? value : undefined; + multiplier *= 128; + } + return undefined; +} + +function scanProtobuf(bytes: Uint8Array, depth: number, path: number[], scan: ProtobufScan): void { + const cursor = { index: 0 }; + while (cursor.index < bytes.length) { + const fieldStart = cursor.index; + const key = readVarint(bytes, cursor); + if (key === undefined || key === 0) { + cursor.index = fieldStart + 1; + continue; + } + const fieldNumber = Math.floor(key / 8); + const wireType = key & 0x07; + const fieldPath = [...path, fieldNumber]; + switch (wireType) { + case 0: { + const value = readVarint(bytes, cursor); + if (value === undefined) cursor.index = fieldStart + 1; + else scan.varintFields.push({ path: fieldPath, value }); + break; + } + case 1: + if (cursor.index + 8 > bytes.length) return; + cursor.index += 8; + break; + case 2: { + const length = readVarint(bytes, cursor); + if (length === undefined || length < 0 || cursor.index + length > bytes.length) { + cursor.index = fieldStart + 1; + break; + } + const end = cursor.index + length; + if (depth < MAX_PROTOBUF_DEPTH) { + scanProtobuf(bytes.subarray(cursor.index, end), depth + 1, fieldPath, scan); + } + cursor.index = end; + break; + } + case 5: { + if (cursor.index + 4 > bytes.length) return; + const view = new DataView(bytes.buffer, bytes.byteOffset + cursor.index, 4); + scan.fixed32Fields.push({ path: fieldPath, value: view.getFloat32(0, true), order: scan.nextOrder++ }); + cursor.index += 4; + break; + } + default: + cursor.index = fieldStart + 1; + } + } +} + +function looksLikeProtobuf(bytes: Uint8Array): boolean { + if (bytes.length === 0) return false; + const fieldNumber = bytes[0] >> 3; + const wireType = bytes[0] & 0x07; + return fieldNumber > 0 && (wireType === 0 || wireType === 1 || wireType === 2 || wireType === 5); +} + +function decodeGrpcWebFrames(bytes: Uint8Array): { payloads: Uint8Array[]; trailers: Map } | null { + const payloads: Uint8Array[] = []; + const trailers = new Map(); + let index = 0; + while (index < bytes.length) { + if (index + 5 > bytes.length) return null; + const flags = bytes[index]; + const length = + bytes[index + 1] * 0x1000000 + bytes[index + 2] * 0x10000 + bytes[index + 3] * 0x100 + bytes[index + 4]; + const start = index + 5; + const end = start + length; + if (end > bytes.length) return null; + if ((flags & 0x80) === 0) { + payloads.push(bytes.subarray(start, end)); + } else { + const text = new TextDecoder().decode(bytes.subarray(start, end)); + for (const line of text.split(/\r?\n/)) { + const separator = line.indexOf(":"); + if (separator <= 0) continue; + const key = line.slice(0, separator).trim().toLowerCase(); + const rawValue = line.slice(separator + 1).trim(); + let value = rawValue; + try { + value = decodeURIComponent(rawValue); + } catch { + // Keep malformed percent-encoding verbatim so the upstream error remains visible. + } + trailers.set(key, value); + } + } + index = end; + } + return { payloads, trailers }; +} + +function assertGrpcStatus(status: string | null | undefined, message?: string | null): void { + if (status === null || status === undefined || status === "" || status === "0") return; + throw new Error(`Grok billing RPC failed with status ${status}${message ? `: ${message}` : ""}`); +} + +function samePath(path: number[], expected: readonly number[]): boolean { + return path.length === expected.length && path.every((value, index) => value === expected[index]); +} + +export function parseGrokUsageResponse(bytes: Uint8Array, nowMs: number = Date.now()): GrokUsageSnapshot { + const decoded = decodeGrpcWebFrames(bytes); + let payloads: Uint8Array[]; + if (decoded) { + assertGrpcStatus(decoded.trailers.get("grpc-status"), decoded.trailers.get("grpc-message")); + if (decoded.payloads.length === 0) throw new Error("Grok billing returned no valid protobuf payload"); + payloads = decoded.payloads; + } else if (looksLikeProtobuf(bytes)) { + payloads = [bytes]; + } else { + throw new Error("Grok billing returned no valid protobuf payload"); + } + + const scan: ProtobufScan = { fixed32Fields: [], varintFields: [], nextOrder: 0 }; + for (const payload of payloads) scanProtobuf(payload, 0, [], scan); + + const percentField = scan.fixed32Fields + .filter( + field => field.path.at(-1) === 1 && Number.isFinite(field.value) && field.value >= 0 && field.value <= 100, + ) + .sort((left, right) => left.path.length - right.path.length || left.order - right.order)[0]; + const resetCandidates = scan.varintFields + .filter(field => field.value >= MIN_UNIX_TIMESTAMP_SECONDS && field.value <= MAX_UNIX_TIMESTAMP_SECONDS) + .map(field => ({ ...field, resetsAt: field.value * 1000 })) + .filter(field => field.resetsAt > nowMs); + const preferredReset = resetCandidates + .filter(field => samePath(field.path, [1, 5, 1])) + .sort((left, right) => left.resetsAt - right.resetsAt)[0]; + const fallbackReset = resetCandidates.sort((left, right) => left.resetsAt - right.resetsAt)[0]; + const resetsAt = preferredReset?.resetsAt ?? fallbackReset?.resetsAt; + const hasUsagePeriod = scan.varintFields.some( + field => + (field.path.length >= 2 && field.path[0] === 1 && field.path[1] === 6) || + (samePath(field.path, [1, 8, 1]) && (field.value === 1 || field.value === 2)), + ); + const usedPercent = + percentField?.value ?? (scan.fixed32Fields.length === 0 && resetsAt && hasUsagePeriod ? 0 : undefined); + if (usedPercent === undefined) throw new Error("Could not parse Grok billing usage"); + return { usedPercent, ...(resetsAt !== undefined ? { resetsAt } : {}) }; +} + +function resolveStatus(usedFraction: number): UsageStatus { + if (usedFraction >= 1) return "exhausted"; + if (usedFraction >= 0.9) return "warning"; + return "ok"; +} + +async function fetchXaiOAuthUsage(params: UsageFetchParams, ctx: UsageFetchContext): Promise { + const accessToken = params.credential.accessToken; + if (params.provider !== XAI_OAUTH_PROVIDER || params.credential.type !== "oauth" || !accessToken) return null; + const response = await ctx.fetch(GROK_BILLING_URL, { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + Origin: "https://grok.com", + Referer: "https://grok.com/?_s=usage", + Accept: "*/*", + "Content-Type": "application/grpc-web+proto", + "x-grpc-web": "1", + "x-user-agent": "connect-es/2.1.1", + "User-Agent": "oh-my-pi/xai", + }, + body: EMPTY_GRPC_WEB_MESSAGE, + signal: params.signal, + }); + if (!response.ok) { + const detail = (await response.text()).slice(0, 400); + throw new Error(`Grok billing request failed with HTTP ${response.status}${detail ? `: ${detail}` : ""}`); + } + assertGrpcStatus(response.headers.get("grpc-status"), response.headers.get("grpc-message")); + const snapshot = parseGrokUsageResponse(new Uint8Array(await response.arrayBuffer())); + const usedFraction = snapshot.usedPercent / 100; + const remainingPercent = Math.max(0, 100 - snapshot.usedPercent); + const limit: UsageLimit = { + id: "xai-oauth:grok-credits", + label: "Grok Credits", + scope: { + provider: XAI_OAUTH_PROVIDER, + ...(params.credential.accountId ? { accountId: params.credential.accountId } : {}), + windowId: "grok-credits", + shared: true, + }, + window: { + id: "grok-credits", + label: "Current Period", + ...(snapshot.resetsAt !== undefined ? { resetsAt: snapshot.resetsAt } : {}), + }, + amount: { + used: snapshot.usedPercent, + limit: 100, + remaining: remainingPercent, + usedFraction, + remainingFraction: remainingPercent / 100, + unit: "percent", + }, + status: resolveStatus(usedFraction), + }; + return { + provider: XAI_OAUTH_PROVIDER, + fetchedAt: Date.now(), + limits: [limit], + metadata: { + ...(params.credential.email ? { email: params.credential.email } : {}), + ...(params.credential.accountId ? { accountId: params.credential.accountId } : {}), + }, + }; +} + +export const xaiOAuthUsageProvider: UsageProvider = { + id: XAI_OAUTH_PROVIDER, + supports: params => + params.provider === XAI_OAUTH_PROVIDER && params.credential.type === "oauth" && !!params.credential.accessToken, + fetchUsage: fetchXaiOAuthUsage, + validatesCredentials: true, +}; diff --git a/packages/ai/test/xai-usage.test.ts b/packages/ai/test/xai-usage.test.ts new file mode 100644 index 00000000000..def0e4b40eb --- /dev/null +++ b/packages/ai/test/xai-usage.test.ts @@ -0,0 +1,159 @@ +import { describe, expect, it } from "bun:test"; +import type { FetchImpl } from "@oh-my-pi/pi-ai/types"; +import type { UsageFetchContext, UsageFetchParams } from "@oh-my-pi/pi-ai/usage"; +import { parseGrokUsageResponse, xaiOAuthUsageProvider } from "@oh-my-pi/pi-ai/usage/xai"; + +function concatBytes(...chunks: Uint8Array[]): Uint8Array { + const length = chunks.reduce((total, chunk) => total + chunk.length, 0); + const result = new Uint8Array(length); + let offset = 0; + for (const chunk of chunks) { + result.set(chunk, offset); + offset += chunk.length; + } + return result; +} + +function encodeVarint(value: number): Uint8Array { + const bytes: number[] = []; + let remaining = value; + do { + let byte = remaining % 128; + remaining = Math.floor(remaining / 128); + if (remaining > 0) byte |= 0x80; + bytes.push(byte); + } while (remaining > 0); + return Uint8Array.from(bytes); +} + +function encodeMessageField(fieldNumber: number, payload: Uint8Array): Uint8Array { + return concatBytes(encodeVarint(fieldNumber * 8 + 2), encodeVarint(payload.length), payload); +} + +function encodeVarintField(fieldNumber: number, value: number): Uint8Array { + return concatBytes(encodeVarint(fieldNumber * 8), encodeVarint(value)); +} + +function encodeFixed32Field(fieldNumber: number, value: number): Uint8Array { + const bytes = new Uint8Array(5); + bytes[0] = fieldNumber * 8 + 5; + new DataView(bytes.buffer).setFloat32(1, value, true); + return bytes; +} + +function grpcWebFrame(payload: Uint8Array, flags: number = 0): Uint8Array { + const frame = new Uint8Array(payload.length + 5); + frame[0] = flags; + new DataView(frame.buffer).setUint32(1, payload.length, false); + frame.set(payload, 5); + return frame; +} + +function grokUsageResponse(usedPercent: number, preferredResetSeconds: number): Uint8Array { + const fallbackReset = encodeMessageField(7, encodeVarintField(1, preferredResetSeconds - 3600)); + const preferredReset = encodeMessageField(5, encodeVarintField(1, preferredResetSeconds)); + const payload = encodeMessageField( + 1, + concatBytes(encodeFixed32Field(1, usedPercent), fallbackReset, preferredReset), + ); + const trailers = grpcWebFrame(new TextEncoder().encode("grpc-status: 0\r\n"), 0x80); + return concatBytes(grpcWebFrame(payload), trailers); +} + +function oauthParams(): UsageFetchParams { + return { + provider: "xai-oauth", + credential: { + type: "oauth", + accessToken: "grok-test-token", + email: "grok@example.com", + accountId: "grok-account", + }, + }; +} + +describe("xAI Grok OAuth usage provider", () => { + it("fetches and normalizes Grok subscription credits", async () => { + const resetSeconds = 2_000_000_000; + let requestedUrl = ""; + let requestedInit: RequestInit | undefined; + const fetch: FetchImpl = async (input, init) => { + requestedUrl = String(input); + requestedInit = init; + return new Response(grokUsageResponse(92.5, resetSeconds), { + status: 200, + headers: { "content-type": "application/grpc-web+proto" }, + }); + }; + const report = await xaiOAuthUsageProvider.fetchUsage(oauthParams(), { fetch }); + + expect(requestedUrl).toBe("https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig"); + expect(requestedInit?.method).toBe("POST"); + expect(Array.from(requestedInit?.body as Uint8Array)).toEqual([0, 0, 0, 0, 0]); + const headers = new Headers(requestedInit?.headers); + expect(headers.get("authorization")).toBe("Bearer grok-test-token"); + expect(headers.get("origin")).toBe("https://grok.com"); + expect(headers.get("referer")).toBe("https://grok.com/?_s=usage"); + expect(headers.get("content-type")).toBe("application/grpc-web+proto"); + expect(headers.get("x-grpc-web")).toBe("1"); + expect(headers.get("x-user-agent")).toBe("connect-es/2.1.1"); + expect(report?.provider).toBe("xai-oauth"); + expect(report?.metadata).toEqual({ email: "grok@example.com", accountId: "grok-account" }); + expect(report?.limits).toHaveLength(1); + expect(report?.limits[0]).toMatchObject({ + id: "xai-oauth:grok-credits", + label: "Grok Credits", + status: "warning", + window: { id: "grok-credits", label: "Current Period", resetsAt: resetSeconds * 1000 }, + amount: { + used: 92.5, + limit: 100, + remaining: 7.5, + usedFraction: 0.925, + remainingFraction: 0.075, + unit: "percent", + }, + }); + }); + + it("prefers the documented reset path over an earlier fallback timestamp", () => { + const resetSeconds = 2_000_000_000; + expect(parseGrokUsageResponse(grokUsageResponse(25, resetSeconds), 1_900_000_000_000)).toEqual({ + usedPercent: 25, + resetsAt: resetSeconds * 1000, + }); + }); + + it("rejects non-zero gRPC trailers", () => { + const payload = encodeMessageField(1, encodeFixed32Field(1, 25)); + const trailers = grpcWebFrame( + new TextEncoder().encode("grpc-status: 16\r\ngrpc-message: unauthenticated\r\n"), + 0x80, + ); + expect(() => parseGrokUsageResponse(concatBytes(grpcWebFrame(payload), trailers))).toThrow( + /Grok billing RPC failed with status 16: unauthenticated/, + ); + expect(() => parseGrokUsageResponse(trailers)).toThrow(/Grok billing RPC failed with status 16: unauthenticated/); + }); + + it("rejects malformed responses", () => { + expect(() => parseGrokUsageResponse(Uint8Array.of(0))).toThrow(/no valid protobuf payload/); + }); + + it("supports only xAI OAuth credentials with an access token", () => { + expect(xaiOAuthUsageProvider.supports?.(oauthParams())).toBe(true); + expect( + xaiOAuthUsageProvider.supports?.({ provider: "xai", credential: { type: "api_key", apiKey: "xai-key" } }), + ).toBe(false); + expect(xaiOAuthUsageProvider.supports?.({ provider: "xai-oauth", credential: { type: "oauth" } })).toBe(false); + }); + + it("surfaces non-OK billing responses", async () => { + const ctx: UsageFetchContext = { + fetch: async () => new Response("denied", { status: 403 }), + }; + await expect(xaiOAuthUsageProvider.fetchUsage(oauthParams(), ctx)).rejects.toThrow( + /Grok billing request failed with HTTP 403: denied/, + ); + }); +}); diff --git a/packages/coding-agent/src/cli/usage-cli.ts b/packages/coding-agent/src/cli/usage-cli.ts index 52df8bd1e57..a1aa5f8e766 100644 --- a/packages/coding-agent/src/cli/usage-cli.ts +++ b/packages/coding-agent/src/cli/usage-cli.ts @@ -759,9 +759,8 @@ function collectStoredAccounts(authStorage: AuthStorage): UsageAccountIdentity[] * * `hasUsageProvider` is injected (in practice {@link AuthStorage.usageProviderFor}) * so custom/broker resolvers stay authoritative — no provider list is duplicated - * here. An explicit `--provider` request bypasses the cull, so - * `omp usage --provider xai` can still confirm the stored credential has no - * usage endpoint. + * here. An explicit `--provider` request bypasses the cull so credentials for + * providers without usage reporting can still be inspected directly. */ export function selectReportableAccounts( accounts: UsageAccountIdentity[],