From c8fa53b1b7fba1a47acbe22ca06dc135a996db7f Mon Sep 17 00:00:00 2001 From: Hello Date: Fri, 24 Jul 2026 03:37:58 +0800 Subject: [PATCH] fix(quota): restore Grok usage from current CLI auth and billing schema Grok CLI stores the token as nested `key` and reports usage under `config.used.val`, so the Settings usage card never showed monthly credits for normal CLI logins. Also surface the subscription plan label from settings and account details from auth.json. Co-authored-by: Cursor --- .../quota-fetcher/providers/grok.test.ts | 185 +++++++++++++++ .../services/quota-fetcher/providers/grok.ts | 222 ++++++++++++++---- .../services/quota-fetcher/service.test.ts | 5 + 3 files changed, 368 insertions(+), 44 deletions(-) create mode 100644 packages/server/src/services/quota-fetcher/providers/grok.test.ts diff --git a/packages/server/src/services/quota-fetcher/providers/grok.test.ts b/packages/server/src/services/quota-fetcher/providers/grok.test.ts new file mode 100644 index 0000000000..f9372c609f --- /dev/null +++ b/packages/server/src/services/quota-fetcher/providers/grok.test.ts @@ -0,0 +1,185 @@ +import { promises as fs } from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { afterEach, describe, expect, test, vi } from "vitest"; + +import { createTestLogger } from "../../../test-utils/test-logger.js"; +import { GrokQuotaProvider } from "./grok.js"; + +const temporaryDirectories: string[] = []; + +async function createGrokHome(): Promise { + const grokHome = await fs.mkdtemp(path.join(os.tmpdir(), "paseo-grok-usage-")); + temporaryDirectories.push(grokHome); + return grokHome; +} + +afterEach(async () => { + await Promise.all( + temporaryDirectories.splice(0).map(async (directory) => { + await fs.rm(directory, { recursive: true, force: true }); + }), + ); +}); + +describe("GrokQuotaProvider", () => { + test("reads the current Grok CLI credential store and billing counters", async () => { + const grokHome = await createGrokHome(); + await fs.writeFile( + path.join(grokHome, "auth.json"), + JSON.stringify({ + "https://auth.x.ai::account": { + key: "grok_cli_token", + email: "person@example.test", + first_name: "Test", + last_name: "Account", + principal_type: "user", + }, + }), + ); + const fetchApi = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + expect(init?.headers).toMatchObject({ Authorization: "Bearer grok_cli_token" }); + const url = input.toString(); + if (url.endsWith("/v1/settings")) { + return new Response(JSON.stringify({ subscription_tier_display: "SuperGrok Heavy" }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + return new Response( + JSON.stringify({ config: { monthlyLimit: { val: 100 }, used: { val: 25 } } }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + }) as unknown as typeof fetch; + const provider = new GrokQuotaProvider({ + logger: createTestLogger(), + fetch: fetchApi, + grokHome, + }); + + await expect(provider.fetchUsage()).resolves.toMatchObject({ + status: "available", + planLabel: "SuperGrok Heavy", + balances: [expect.objectContaining({ used: 25, remaining: 75, limit: 100 })], + details: [ + { id: "account_email", label: "Account email", value: "person@example.test" }, + { id: "account_name", label: "Account", value: "Test Account" }, + { id: "account_type", label: "Account type", value: "user" }, + ], + }); + }); + + test("refreshes an expired CLI credential before requesting usage", async () => { + const grokHome = await createGrokHome(); + const authPath = path.join(grokHome, "auth.json"); + await fs.writeFile( + authPath, + JSON.stringify({ + "https://auth.x.ai::account": { + key: "expired_token", + expires_at: "2020-01-01T00:00:00.000Z", + refresh_token: "refresh_token", + }, + }), + ); + const refreshAuth = vi.fn(async () => { + await fs.writeFile( + authPath, + JSON.stringify({ + "https://auth.x.ai::account": { + key: "refreshed_token", + expires_at: "2099-01-01T00:00:00.000Z", + refresh_token: "refresh_token", + }, + }), + ); + }); + const fetchApi = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + expect(init?.headers).toMatchObject({ Authorization: "Bearer refreshed_token" }); + if (input.toString().endsWith("/v1/settings")) { + return new Response(JSON.stringify({}), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + return new Response(JSON.stringify({ config: { monthlyLimit: { val: 100 } } }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + }) as unknown as typeof fetch; + const provider = new GrokQuotaProvider({ + logger: createTestLogger(), + fetch: fetchApi, + grokHome, + refreshAuth, + now: () => Date.parse("2026-01-01T00:00:00.000Z"), + }); + + await expect(provider.fetchUsage()).resolves.toMatchObject({ status: "available" }); + expect(refreshAuth).toHaveBeenCalledTimes(1); + expect(fetchApi).toHaveBeenCalledTimes(2); + }); + + test("refreshes and retries once when the billing endpoint rejects a cached credential", async () => { + const grokHome = await createGrokHome(); + const authPath = path.join(grokHome, "auth.json"); + await fs.writeFile( + authPath, + JSON.stringify({ + "https://auth.x.ai::account": { + key: "stale_token", + expires_at: "2099-01-01T00:00:00.000Z", + refresh_token: "refresh_token", + }, + }), + ); + const refreshAuth = vi.fn(async () => { + await fs.writeFile( + authPath, + JSON.stringify({ + "https://auth.x.ai::account": { + key: "fresh_token", + expires_at: "2099-01-01T00:00:00.000Z", + refresh_token: "refresh_token", + }, + }), + ); + }); + let billingCalls = 0; + const fetchApi = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + const url = input.toString(); + if (url.endsWith("/v1/settings")) { + expect(init?.headers).toMatchObject({ Authorization: "Bearer fresh_token" }); + return new Response(JSON.stringify({ subscription_tier_display: "SuperGrok" }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + billingCalls += 1; + if (billingCalls === 1) { + expect(init?.headers).toMatchObject({ Authorization: "Bearer stale_token" }); + return new Response(null, { status: 401 }); + } + expect(init?.headers).toMatchObject({ Authorization: "Bearer fresh_token" }); + return new Response( + JSON.stringify({ config: { monthlyLimit: { val: 100 }, used: { val: 25 } } }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + }) as unknown as typeof fetch; + const provider = new GrokQuotaProvider({ + logger: createTestLogger(), + fetch: fetchApi, + grokHome, + refreshAuth, + }); + + await expect(provider.fetchUsage()).resolves.toMatchObject({ + status: "available", + planLabel: "SuperGrok", + balances: [expect.objectContaining({ used: 25, remaining: 75 })], + }); + expect(refreshAuth).toHaveBeenCalledTimes(1); + expect(fetchApi).toHaveBeenCalledTimes(3); + }); +}); diff --git a/packages/server/src/services/quota-fetcher/providers/grok.ts b/packages/server/src/services/quota-fetcher/providers/grok.ts index 8d660c34c1..c032f21b09 100644 --- a/packages/server/src/services/quota-fetcher/providers/grok.ts +++ b/packages/server/src/services/quota-fetcher/providers/grok.ts @@ -1,18 +1,25 @@ +import { execFile } from "node:child_process"; import { existsSync, promises as fs } from "node:fs"; import { homedir } from "node:os"; -import { join } from "node:path"; +import { join, resolve } from "node:path"; +import { promisify } from "node:util"; import type { Logger } from "pino"; import { z } from "zod"; import type { ProviderUsage, ProviderUsageBalance } from "../../../server/messages.js"; +import { expandTilde } from "../../../utils/path.js"; import type { ProviderApiFetch, ProviderUsageFetcher } from "../provider.js"; import { ApiNumberSchema, - toneFromUsedPct, - usedPctOf, + balanceToneFromRemaining, fetchProviderApi, unavailableUsage, } from "../usage.js"; +const execFileAsync = promisify(execFile); +const GROK_AUTH_REFRESH_TIMEOUT_MS = 10_000; + +const GROK_CLI_PROXY_BASE = "https://cli-chat-proxy.grok.com"; + const GrokUsageResponseSchema = z.object({ config: z .object({ @@ -21,6 +28,11 @@ const GrokUsageResponseSchema = z.object({ val: ApiNumberSchema.optional(), }) .nullish(), + used: z + .object({ + val: ApiNumberSchema.optional(), + }) + .nullish(), }) .nullish(), usage: z @@ -30,13 +42,35 @@ const GrokUsageResponseSchema = z.object({ .nullish(), }); -const GrokAuthSchema = z.object({ - access_token: z.string().optional(), +const GrokSettingsResponseSchema = z.object({ + subscription_tier_display: z.string().trim().min(1).optional(), }); +const GrokAuthEntrySchema = z.object({ + access_token: z.string().trim().min(1).optional(), + key: z.string().trim().min(1).optional(), + email: z.string().trim().email().optional(), + first_name: z.string().trim().min(1).optional(), + last_name: z.string().trim().min(1).optional(), + principal_type: z.string().trim().min(1).optional(), + expires_at: z.string().trim().min(1).optional(), + refresh_token: z.string().trim().min(1).optional(), +}); +const GrokAuthStoreSchema = z.record(z.string(), GrokAuthEntrySchema); + +interface GrokAuth { + token: string; + details: NonNullable; + expiresAt: string | null; + canRefresh: boolean; +} + interface GrokQuotaProviderOptions { logger: Logger; fetch?: ProviderApiFetch; + grokHome?: string; + refreshAuth?: () => Promise; + now?: () => number; } export class GrokQuotaProvider implements ProviderUsageFetcher { @@ -45,29 +79,36 @@ export class GrokQuotaProvider implements ProviderUsageFetcher { private readonly logger: Logger; private readonly fetchApi: ProviderApiFetch; + private readonly grokHome: string; + private readonly refreshAuth: () => Promise; + private readonly now: () => number; constructor(options: GrokQuotaProviderOptions) { this.logger = options.logger; this.fetchApi = options.fetch ?? fetch; + this.grokHome = resolveGrokHome(options.grokHome); + this.refreshAuth = options.refreshAuth ?? refreshGrokCliAuth; + this.now = options.now ?? Date.now; } async fetchUsage(): Promise { - const token = - process.env["GROK_API_KEY"] || process.env["GROK_TOKEN"] || (await this.readGrokToken()); + const environmentToken = process.env["GROK_API_KEY"] || process.env["GROK_TOKEN"]; + let storedAuth = await this.readGrokAuth(); + if (!environmentToken && storedAuth && this.isExpired(storedAuth.expiresAt)) { + storedAuth = await this.refreshStoredAuth(); + } + + let token = environmentToken || storedAuth?.token; if (!token) return unavailableUsage(this); - const res = await fetchProviderApi( - this.fetchApi, - "https://cli-chat-proxy.grok.com/v1/billing", - { - headers: { - Authorization: `Bearer ${token}`, - "X-XAI-Token-Auth": "xai-grok-cli", - Accept: "application/json", - }, - }, - ); + let res = await this.fetchBilling(token); + if (!environmentToken && res.status === 401 && storedAuth?.canRefresh) { + storedAuth = await this.refreshStoredAuth(); + token = storedAuth?.token; + if (!token) return unavailableUsage(this); + res = await this.fetchBilling(token); + } if (!res.ok) { this.logger.debug({ status: res.status }, "Grok usage fetch failed"); @@ -75,45 +116,138 @@ export class GrokQuotaProvider implements ProviderUsageFetcher { } const resp = GrokUsageResponseSchema.parse(await res.json()); - const monthlyLimit = resp.config?.monthlyLimit?.val ?? null; - const creditUsage = resp.usage?.creditUsage ?? null; - const balances: ProviderUsageBalance[] = []; - if (monthlyLimit !== null || creditUsage !== null) { - const remaining = - monthlyLimit !== null && creditUsage !== null - ? Math.max(0, monthlyLimit - creditUsage) - : null; - balances.push({ - id: "monthly_credits", - label: "Monthly credits", - used: creditUsage, - remaining, - limit: monthlyLimit, - unit: "credits", - tone: toneFromUsedPct(usedPctOf(creditUsage, monthlyLimit)), - }); - } - return { providerId: this.providerId, displayName: this.displayName, status: "available", - planLabel: null, + planLabel: await this.fetchPlanLabel(token), windows: [], - balances, - details: [], + balances: resolveGrokBalances(resp), + details: storedAuth?.details ?? [], error: null, }; } - private async readGrokToken(): Promise { - const path = join(homedir(), ".grok", "auth.json"); + private grokAuthHeaders(token: string): Record { + return { + Authorization: `Bearer ${token}`, + "X-XAI-Token-Auth": "xai-grok-cli", + Accept: "application/json", + }; + } + + private fetchBilling(token: string): Promise { + return fetchProviderApi(this.fetchApi, `${GROK_CLI_PROXY_BASE}/v1/billing`, { + headers: this.grokAuthHeaders(token), + }); + } + + private async fetchPlanLabel(token: string): Promise { + try { + const res = await fetchProviderApi(this.fetchApi, `${GROK_CLI_PROXY_BASE}/v1/settings`, { + headers: this.grokAuthHeaders(token), + }); + if (!res.ok) return null; + const settings = GrokSettingsResponseSchema.safeParse(await res.json()); + return settings.success ? (settings.data.subscription_tier_display ?? null) : null; + } catch (error) { + this.logger.debug({ err: error }, "Grok subscription plan fetch failed"); + return null; + } + } + + private async refreshStoredAuth(): Promise { + try { + await this.refreshAuth(); + return this.readGrokAuth(); + } catch (error) { + this.logger.debug({ err: error }, "Grok CLI credential refresh failed"); + return null; + } + } + + private isExpired(expiresAt: string | null): boolean { + if (!expiresAt) return false; + const expiresAtMs = Date.parse(expiresAt); + return Number.isFinite(expiresAtMs) && expiresAtMs <= this.now(); + } + + private async readGrokAuth(): Promise { + const path = join(this.grokHome, "auth.json"); if (!existsSync(path)) return null; try { - const auth = GrokAuthSchema.parse(JSON.parse(await fs.readFile(path, "utf8"))); - return auth.access_token ?? null; + const auth = JSON.parse(await fs.readFile(path, "utf8")); + const directEntry = GrokAuthEntrySchema.safeParse(auth); + const directAuth = resolveGrokAuthEntry(directEntry.success ? directEntry.data : null); + if (directAuth) return directAuth; + + const store = GrokAuthStoreSchema.safeParse(auth); + if (!store.success) return null; + for (const entry of Object.values(store.data)) { + const entryAuth = resolveGrokAuthEntry(entry); + if (entryAuth) return entryAuth; + } + return null; } catch { return null; } } } + +function resolveGrokBalances( + response: z.infer, +): ProviderUsageBalance[] { + const monthlyLimit = response.config?.monthlyLimit?.val ?? null; + const creditUsage = response.config?.used?.val ?? response.usage?.creditUsage ?? null; + if (monthlyLimit === null && creditUsage === null) return []; + + const remaining = + monthlyLimit !== null && creditUsage !== null ? Math.max(0, monthlyLimit - creditUsage) : null; + return [ + { + id: "monthly_credits", + label: "Monthly credits", + used: creditUsage, + remaining, + limit: monthlyLimit, + unit: "credits", + tone: balanceToneFromRemaining(remaining), + }, + ]; +} + +function resolveGrokHome(configuredHome: string | undefined): string { + const grokHome = configuredHome ?? process.env["GROK_HOME"]; + return grokHome ? resolve(expandTilde(grokHome)) : join(homedir(), ".grok"); +} + +function resolveGrokAuthEntry(entry: z.infer | null): GrokAuth | null { + if (!entry) return null; + + const token = entry.access_token ?? entry.key; + if (!token) return null; + + const displayName = [entry.first_name, entry.last_name].filter(Boolean).join(" "); + const details: NonNullable = []; + if (entry.email) { + details.push({ id: "account_email", label: "Account email", value: entry.email }); + } + if (displayName) { + details.push({ id: "account_name", label: "Account", value: displayName }); + } + if (entry.principal_type) { + details.push({ id: "account_type", label: "Account type", value: entry.principal_type }); + } + + return { + token, + details, + expiresAt: entry.expires_at ?? null, + canRefresh: Boolean(entry.refresh_token), + }; +} + +async function refreshGrokCliAuth(): Promise { + // `grok models` is the official non-interactive command that performs the CLI's silent OIDC refresh. + await execFileAsync("grok", ["models"], { timeout: GROK_AUTH_REFRESH_TIMEOUT_MS }); +} diff --git a/packages/server/src/services/quota-fetcher/service.test.ts b/packages/server/src/services/quota-fetcher/service.test.ts index 5dfc184080..dd16d62a90 100644 --- a/packages/server/src/services/quota-fetcher/service.test.ts +++ b/packages/server/src/services/quota-fetcher/service.test.ts @@ -724,6 +724,10 @@ describe("real provider usage fetchers", () => { usage: { creditUsage: 0 }, }), ], + [ + "https://cli-chat-proxy.grok.com/v1/settings", + () => jsonResponse({ subscription_tier_display: "SuperGrok Heavy" }), + ], ]), ); @@ -731,6 +735,7 @@ describe("real provider usage fetchers", () => { expect(grok).toMatchObject({ status: "available", + planLabel: "SuperGrok Heavy", balances: [ expect.objectContaining({ id: "monthly_credits",