diff --git a/config/quality/file-size-baseline.json b/config/quality/file-size-baseline.json index 2ecfb8bad7..4c9502248e 100644 --- a/config/quality/file-size-baseline.json +++ b/config/quality/file-size-baseline.json @@ -305,7 +305,8 @@ "_rebaseline_2026_06_28_5275_correlation_id_extract": "Extraction of the safe CorrelationId subset of #5275 (hartmark) — request correlation id stored in call_logs (migration 109) and returned via the X-Correlation-Id response header, WITHOUT the combo/resilience or build/lazy-loading changes (those stay in #5275). Own growth: callLogs.ts 975->985 (correlation_id column on CallLogSummaryRow + read/map), usageHistory.ts 983->988 (correlationId metadata normalize), chat.ts 1575->1632 (withCorrelationId response wiring + combo-failure log carrying correlationId), chatHelpers.ts new 811 (withCorrelationId helper + reqId threading; was 791786 (single ProviderAccountRoutingCard mount + import), auth.ts 2448->2458 (providerStrategies override resolution: fallbackStrategy/stickyRoundRobinLimit per-provider cascade in getProviderCredentials). Both additive, zero unrelated refactor; new UI/logic lives in new files (ProviderAccountRoutingCard.tsx, RoutingStrategyCard.tsx, rrState.ts::resolveComboStickyRoundRobinLimit). chat.ts value below reflects the current release tip (grown by other concurrent PRs, e.g. #6640), not this PR own change.", "_rebaseline_2026_07_22_8213_chat_abandoned_target_abort": "PR #8213 (hartmark, fix/gemini-tpm-quota-cooldown-wait) own growth: src/sse/handlers/chat.ts 1794->1860 (+66, measured against the PR's own merge-base — the release tip separately carries an unrelated -5 net shrink from #8013's antigravity callable-catalog alignment, which this PR's branch does not include and this entry does not cover). Adds resolveDispatchClientRawRequest(): merges a per-target modelAbortSignal into clientRawRequest.signal (via mergeAbortSignals) so a combo target abandoned by comboTargetTimeoutMs actually observes its own abort and reaches its cleanup path, instead of hanging forever inside withRateLimit/acquireAccountSemaphore and leaking a permanent 'pending' dashboard entry (live incident, log id 1784418258231-14961a). Also wires combo-exhausted rejection logging to capture request body + attempted models via the new rejectedRequestUsage helper. Irreducible additions at the existing chat dispatch chokepoint. Covered by the PR's own combo-config + integration test additions.", - "src/sse/handlers/chat.ts": 1861, + "_rebaseline_2026_07_23_8127_grok_weekly_quota": "#8127 (@apoapostolov) own growth: src/sse/handlers/chat.ts 1861->1865 (+4) — weekly quota tracking for grok-web wires a quota-fetch hook at the existing dispatch chokepoint. Thin wiring mirroring adjacent provider-quota branches; not extractable. Covered by tests/unit/grok-quota-fetcher.test.ts.", + "src/sse/handlers/chat.ts": 1865, "_rebaseline_2026_07_20_7779_routingcombo_thread": "PR #7779 own growth: chatHelpers.ts 876->877 (+1, thread routingComboId into executeChatWithBreaker for compression-combo assignment). Frozen so it can only shrink.", "src/sse/handlers/chatHelpers.ts": 878, "src/sse/services/auth.ts": 2462, diff --git a/open-sse/services/grokQuotaFetcher.ts b/open-sse/services/grokQuotaFetcher.ts new file mode 100644 index 0000000000..810d90d425 --- /dev/null +++ b/open-sse/services/grokQuotaFetcher.ts @@ -0,0 +1,468 @@ +/** + * grokQuotaFetcher.ts — Grok Web Weekly Quota Fetcher + * + * Implements QuotaFetcher for the grok-web provider (grok.com SSO/cookie + * connections). Since grok-web uses cookie auth rather than an accessToken + * stored in the provider connection, this fetcher reads the account-level + * OIDC tokens from the local Grok CLI (~/.grok/auth.json) and calls the + * same billing API that powers the grok CLI /usage command and third-party + * quota widgets (pi-grok-usage, hermes-grok-usage). + * + * The billing endpoint returns a single weekly credit-usage percentage + * that covers all Grok products (Chat, Build, Imagine, API) under the + * SuperGrok / X Premium+ unified weekly pool. + * + * Registration: call registerGrokWebQuotaFetcher() once at server startup, + * before registerGenericQuotaFetchers() — because the generic path doesn't + * know how to resolve grok OIDC auth from a cookie-based connection. + */ + +import { registerQuotaFetcher, registerQuotaWindows, type QuotaInfo } from "./quotaPreflight.ts"; +import { registerMonitorFetcher } from "./quotaMonitor.ts"; +import { existsSync, readFileSync, writeFileSync } from "node:fs"; +import { homedir } from "node:os"; +import { join } from "node:path"; + +// ─── Constants ─────────────────────────────────────────────────────────────── + +const BILLING_URL = "https://cli-chat-proxy.grok.com/v1/billing?format=credits"; +const DEFAULT_ISSUER = "https://auth.x.ai"; +const FETCH_COOLDOWN_MS = 5 * 60 * 1000; // 5 minutes — matches grok-web rate limit cooldown +const REQUEST_TIMEOUT_MS = 10_000; +const EXPIRY_SKEW_MS = 60_000; // refresh a bit before actual expiry +/** + * Resolve the Grok auth.json path. + * Controlled by the GROK_AUTH_PATH env var for testing; defaults to ~/.grok/auth.json. + */ +function getAuthPath(): string { + const override = (process.env.GROK_AUTH_PATH || "").trim(); + return override || join(homedir(), ".grok", "auth.json"); +} + +/** + * Stable window identifier for the grok-web weekly quota. + * Surfaced in the dashboard as "Weekly" under provider limits. + */ +export const GROK_WINDOW_WEEKLY = "weekly"; + +// ─── Types ─────────────────────────────────────────────────────────────────── + +interface GrokAuthEntry { + key?: string; + refresh_token?: string; + expires_at?: string; + email?: string; + auth_mode?: string; + oidc_issuer?: string; + oidc_client_id?: string; +} + +interface ResolvedAuth { + entryId: string; + token: string; + refreshToken?: string; + email?: string; + expiresAtMs?: number; + issuer: string; + clientId?: string; +} + +interface BillingConfig { + currentPeriod?: { + type?: string; + start?: string; + end?: string; + }; + creditUsagePercent?: number; + onDemandCap?: { val?: number }; + onDemandUsed?: { val?: number }; + prepaidBalance?: { val?: number }; + productUsage?: Array<{ product?: string; usagePercent?: number }>; + isUnifiedBillingUser?: boolean; + billingPeriodStart?: string; + billingPeriodEnd?: string; +} + +interface BillingResponse { + config?: BillingConfig; +} + +interface UsageSnapshot { + percent: number; + periodLabel: string; + resetLabel: string; + endIso?: string; + email?: string; + products: Array<{ product: string; usagePercent?: number }>; + onDemandUsed: number; + onDemandCap: number; + prepaidBalance: number; + fetchedAt: number; +} + +// ─── In-memory cache (per-connection, keyed by connectionId) ───────────────── + +interface CacheEntry { + quota: QuotaInfo | null; + error: string | null; + fetchedAt: number; +} + +const quotaCache = new Map(); +const CACHE_TTL_MS = 60_000; // 60 seconds — matches codexQuotaFetcher + +const _cacheCleanup = setInterval(() => { + const now = Date.now(); + for (const [key, entry] of quotaCache) { + if (now - entry.fetchedAt > CACHE_TTL_MS * 5) { + quotaCache.delete(key); + } + } +}, 5 * 60_000); + +if (typeof _cacheCleanup === "object" && "unref" in _cacheCleanup) { + (_cacheCleanup as { unref?: () => void }).unref?.(); +} + +// ─── Auth helpers (mirrors pi-grok-usage / hermes-grok-usage patterns) ──────── + +function isAllowedXaiUrl(raw: string): boolean { + try { + const url = new URL(raw); + return url.protocol === "https:" && (url.hostname === "x.ai" || url.hostname.endsWith(".x.ai")); + } catch { + return false; + } +} + +function readAuthFile(): Record | null { + if (!existsSync(getAuthPath())) return null; + try { + const raw = JSON.parse(readFileSync(getAuthPath(), "utf8")) as Record; + if (!raw || typeof raw !== "object") return null; + return raw; + } catch { + return null; + } +} + +function pickAuthEntry(file: Record): ResolvedAuth | null { + const now = Date.now(); + const scored = Object.entries(file) + .map(([entryId, e]) => { + const token = typeof e?.key === "string" ? e.key.trim() : ""; + const exp = e?.expires_at ? Date.parse(e.expires_at) : Number.POSITIVE_INFINITY; + const expired = Number.isFinite(exp) ? exp <= now + EXPIRY_SKEW_MS : false; + const hasRefresh = typeof e?.refresh_token === "string" && e.refresh_token.length > 0; + return { entryId, e, token, exp, expired, hasRefresh }; + }) + .filter((x) => x.token.length > 0 || x.hasRefresh); + + if (scored.length === 0) return null; + + scored.sort((a, b) => { + if (a.expired !== b.expired) return a.expired ? 1 : -1; + if (a.hasRefresh !== b.hasRefresh) return a.hasRefresh ? -1 : 1; + return b.exp - a.exp; + }); + + const best = scored[0]; + const issuer = (best.e.oidc_issuer || DEFAULT_ISSUER).replace(/\/$/, ""); + return { + entryId: best.entryId, + token: best.token, + refreshToken: best.e.refresh_token?.trim() || undefined, + email: best.e.email, + expiresAtMs: Number.isFinite(best.exp) ? best.exp : undefined, + issuer, + clientId: best.e.oidc_client_id?.trim() || undefined, + }; +} + +function writeRefreshedTokens( + entryId: string, + update: { access: string; refresh?: string; expiresAtIso: string }, +): void { + const file = readAuthFile(); + if (!file || !file[entryId]) return; + file[entryId] = { + ...file[entryId], + key: update.access, + ...(update.refresh ? { refresh_token: update.refresh } : {}), + expires_at: update.expiresAtIso, + }; + try { + writeFileSync(getAuthPath(), `${JSON.stringify(file, null, 2)}\n`, { encoding: "utf8" }); + } catch { + // Non-fatal: token is still valid in-memory this session + } +} + +async function discoverTokenEndpoint(issuer: string, signal: AbortSignal): Promise { + const discoveryUrl = `${issuer.replace(/\/$/, "")}/.well-known/openid-configuration`; + if (!isAllowedXaiUrl(discoveryUrl)) { + throw new Error("invalid oidc issuer"); + } + const res = await fetch(discoveryUrl, { + headers: { Accept: "application/json" }, + signal, + }); + if (!res.ok) throw new Error(`token refresh failed (discovery HTTP ${res.status})`); + const json = (await res.json()) as { token_endpoint?: string }; + const endpoint = String(json.token_endpoint || ""); + if (!isAllowedXaiUrl(endpoint)) throw new Error("invalid token endpoint"); + return endpoint; +} + +async function refreshAccessToken(auth: ResolvedAuth, signal: AbortSignal): Promise { + if (!auth.refreshToken) { + throw new Error("auth expired — run `grok login`"); + } + if (!auth.clientId) { + throw new Error("auth missing client id — run `grok login`"); + } + + const tokenEndpoint = await discoverTokenEndpoint(auth.issuer, signal); + const res = await fetch(tokenEndpoint, { + method: "POST", + headers: { + "Content-Type": "application/x-www-form-urlencoded", + Accept: "application/json", + "User-Agent": "omniroute-grok-usage/1.0", + }, + body: new URLSearchParams({ + grant_type: "refresh_token", + client_id: auth.clientId, + refresh_token: auth.refreshToken, + }).toString(), + signal, + }); + + if (!res.ok) { + throw new Error(`token refresh failed (HTTP ${res.status})`); + } + + const payload = (await res.json()) as { + access_token?: string; + refresh_token?: string; + expires_in?: number; + }; + + const access = String(payload.access_token || "").trim(); + if (!access) throw new Error("token refresh failed (empty access token)"); + + const expiresInSec = Number(payload.expires_in || 3600); + const expiresAtMs = Date.now() + Math.max(60, expiresInSec) * 1000; + const expiresAtIso = new Date(expiresAtMs).toISOString(); + const refresh = String(payload.refresh_token || auth.refreshToken).trim(); + + try { + writeRefreshedTokens(auth.entryId, { + access, + refresh, + expiresAtIso, + }); + } catch { + // Non-fatal + } + + return { + ...auth, + token: access, + refreshToken: refresh, + expiresAtMs, + }; +} + +function needsRefresh(auth: ResolvedAuth): boolean { + if (!auth.token) return true; + if (auth.expiresAtMs == null) return false; + return auth.expiresAtMs <= Date.now() + EXPIRY_SKEW_MS; +} + +async function resolveAuth(signal: AbortSignal): Promise { + const file = readAuthFile(); + if (!file) throw new Error("no grok auth — run `grok login`"); + const auth = pickAuthEntry(file); + if (!auth) throw new Error("no usable Grok credentials — run `grok login`"); + if (needsRefresh(auth)) { + return refreshAccessToken(auth, signal); + } + return auth; +} + +// ─── Billing fetch ─────────────────────────────────────────────────────────── + +function periodShort(type?: string): string { + if (!type) return ""; + if (type.includes("WEEKLY")) return "weekly"; + if (type.includes("MONTHLY")) return "monthly"; + if (type.includes("DAILY")) return "daily"; + return "period"; +} + +function resetLocalLabel(endIso?: string): string { + if (!endIso) return ""; + const end = new Date(endIso); + if (Number.isNaN(end.getTime())) return ""; + const weekday = end.toLocaleDateString(undefined, { weekday: "short" }); + const hour = end.toLocaleTimeString(undefined, { + hour: "2-digit", + minute: "2-digit", + hour12: false, + }); + return `${weekday} ${hour.replace(/^24:/, "00:")}`; +} + +async function fetchBilling(token: string, signal: AbortSignal): Promise { + const res = await fetch(BILLING_URL, { + method: "GET", + headers: { + Authorization: `Bearer ${token}`, + Accept: "application/json", + "User-Agent": "omniroute-grok-usage/1.0", + "x-grok-client-mode": "cli", + }, + signal, + }); + + if (res.status === 401 || res.status === 403) { + throw new Error(`auth ${res.status}`); + } + if (!res.ok) { + throw new Error(`HTTP ${res.status}`); + } + + const data = (await res.json()) as BillingResponse; + const cfg = data.config ?? {}; + const percent = Number(cfg.creditUsagePercent ?? 0); + const endIso = cfg.currentPeriod?.end ?? cfg.billingPeriodEnd; + const products = (cfg.productUsage ?? []) + .filter((p) => p.product) + .map((p) => ({ product: String(p.product), usagePercent: p.usagePercent })); + + return { + percent: Number.isFinite(percent) ? percent : 0, + periodLabel: periodShort(cfg.currentPeriod?.type), + resetLabel: resetLocalLabel(endIso), + endIso, + products, + onDemandUsed: Number(cfg.onDemandUsed?.val ?? 0), + onDemandCap: Number(cfg.onDemandCap?.val ?? 0), + prepaidBalance: Number(cfg.prepaidBalance?.val ?? 0), + fetchedAt: Date.now(), + }; +} + +// ─── Core Fetcher ──────────────────────────────────────────────────────────── + +/** + * Convert a UsageSnapshot into the preflight QuotaInfo contract. + * Grok has a single weekly window; percentUsed = creditUsagePercent / 100. + * The resetAt comes from currentPeriod.end. + */ +function snapshotToQuotaInfo(snap: UsageSnapshot): QuotaInfo { + const percentUsed = Math.max(0, Math.min(1, snap.percent / 100)); + return { + used: Math.round(snap.percent), + total: 100, + percentUsed, + resetAt: snap.endIso ?? null, + windows: { + [GROK_WINDOW_WEEKLY]: { + percentUsed, + resetAt: snap.endIso ?? null, + }, + }, + }; +} + +/** + * Fetch current weekly quota for a grok-web connection. + * + * Unlike most provider fetchers, this does NOT use the connection's stored + * accessToken/cookie — instead it reads the local Grok CLI OIDC tokens + * (~/.grok/auth.json) which represent the same xAI account. This is because + * grok-web uses cookie auth (SSO) but the billing endpoint requires a Bearer + * token, and both auth paths share the same underlying account. + * + * If no local Grok CLI auth is found, returns null (fail-open — don't block + * grok-web requests on missing quota data). + * + * @param connectionId - Connection ID from the DB (used as cache key) + * @param _connection - Ignored for grok-web; auth is resolved from local CLI + * @returns QuotaInfo with a single "weekly" window, or null + */ +export async function fetchGrokWebQuota( + connectionId: string, + _connection?: Record, +): Promise { + // Check cache first + const cached = quotaCache.get(connectionId); + if (cached && Date.now() - cached.fetchedAt < CACHE_TTL_MS) { + return cached.quota; + } + + try { + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS); + + try { + let auth: ResolvedAuth; + try { + auth = await resolveAuth(controller.signal); + } catch { + // No local grok auth — fail open + quotaCache.set(connectionId, { quota: null, error: null, fetchedAt: Date.now() }); + return null; + } + + try { + const snap = await fetchBilling(auth.token, controller.signal); + const quota = snapshotToQuotaInfo(snap); + quotaCache.set(connectionId, { quota, error: null, fetchedAt: Date.now() }); + return quota; + } catch (err) { + const msg = err instanceof Error ? err.message : ""; + // One retry with token refresh on auth errors + if (msg.startsWith("auth ") && auth.refreshToken) { + auth = await refreshAccessToken(auth, controller.signal); + const snap = await fetchBilling(auth.token, controller.signal); + const quota = snapshotToQuotaInfo(snap); + quotaCache.set(connectionId, { quota, error: null, fetchedAt: Date.now() }); + return quota; + } + // Non-auth error — fail open + quotaCache.set(connectionId, { quota: null, error: msg, fetchedAt: Date.now() }); + return null; + } + } finally { + clearTimeout(timeout); + } + } catch { + quotaCache.set(connectionId, { quota: null, error: null, fetchedAt: Date.now() }); + return null; + } +} + +// ─── Invalidation ──────────────────────────────────────────────────────────── + +/** + * Force-invalidate the cache for a connection. + */ +export function invalidateGrokWebQuotaCache(connectionId: string): void { + quotaCache.delete(connectionId); +} + +// ─── Registration ──────────────────────────────────────────────────────────── + +/** + * Register the grok-web quota fetcher with the preflight and monitor systems. + * Call this once at server startup (in chat.ts), before registerGenericQuotaFetchers() + * so the generic path doesn't try to override this bespoke fetcher. + */ +export function registerGrokWebQuotaFetcher(): void { + registerQuotaFetcher("grok-web", fetchGrokWebQuota); + registerMonitorFetcher("grok-web", fetchGrokWebQuota); + registerQuotaWindows("grok-web", [GROK_WINDOW_WEEKLY]); +} diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 6fe83ca4f9..161321117e 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -133,6 +133,7 @@ import { registerCrofUsageFetcher } from "@omniroute/open-sse/services/crofUsage import { registerDeepseekQuotaFetcher } from "@omniroute/open-sse/services/deepseekQuotaFetcher.ts"; import { registerOpenrouterQuotaFetcher } from "@omniroute/open-sse/services/openrouterQuotaFetcher.ts"; import { registerOpencodeQuotaFetcher } from "@omniroute/open-sse/services/opencodeQuotaFetcher.ts"; +import { registerGrokWebQuotaFetcher } from "@omniroute/open-sse/services/grokQuotaFetcher.ts"; import { registerGenericQuotaFetchers } from "@omniroute/open-sse/services/genericQuotaFetcher.ts"; import "@omniroute/open-sse/services/quotaTrackersBatch.ts"; import { @@ -164,6 +165,14 @@ registerOpenrouterQuotaFetcher(); // quota-aware preflight switching between connections. (#2852) registerOpencodeQuotaFetcher(); +// Register Grok Web quota fetcher. +// Reads account-level OIDC tokens from ~/.grok/auth.json (the local Grok CLI +// login) to surface the weekly credit-usage percentage in the dashboard. +// This runs before registerGenericQuotaFetchers so the bespoke fetcher takes +// precedence over the generic path (which can't resolve grok OIDC auth from +// cookie-based connections). +registerGrokWebQuotaFetcher(); + // Register the generic quota fetcher for every other provider that has a // usage implementation in usage.ts but no bespoke preflight fetcher. This is // what lets the per-window cutoff modal in Dashboard › Limits actually diff --git a/tests/unit/grok-quota-fetcher.test.ts b/tests/unit/grok-quota-fetcher.test.ts new file mode 100644 index 0000000000..0844a870c4 --- /dev/null +++ b/tests/unit/grok-quota-fetcher.test.ts @@ -0,0 +1,290 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { mkdtempSync, writeFileSync, rmSync } from "node:fs"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; + +import { + fetchGrokWebQuota, + invalidateGrokWebQuotaCache, + registerGrokWebQuotaFetcher, +} from "../../open-sse/services/grokQuotaFetcher.ts"; +import { preflightQuota } from "../../open-sse/services/quotaPreflight.ts"; +import { + clearQuotaMonitors, + getActiveMonitorCount, + startQuotaMonitor, + stopQuotaMonitor, +} from "../../open-sse/services/quotaMonitor.ts"; +import { clearSessions, touchSession } from "../../open-sse/services/sessionManager.ts"; + +const originalFetch = globalThis.fetch; +const originalAuthPath = process.env.GROK_AUTH_PATH; + +/** Create a temporary auth.json for testing. */ +function createTestAuth( + overrides: Partial<{ + key: string; + refresh_token: string; + expires_at: string; + email: string; + oidc_issuer: string; + oidc_client_id: string; + }> = {}, +): string { + const dir = mkdtempSync(join(tmpdir(), "grok-auth-test-")); + const entry = { + key: "test-access-token", + refresh_token: "test-refresh-token", + expires_at: new Date(Date.now() + 3600_000).toISOString(), + email: "test@example.com", + auth_mode: "device_code", + oidc_issuer: "https://auth.x.ai", + oidc_client_id: "test-client-id", + ...overrides, + }; + const authFile = join(dir, "auth.json"); + writeFileSync(authFile, JSON.stringify({ "test-entry": entry }, null, 2) + "\n"); + return authFile; +} + +function cleanupTestAuth(authPath: string): void { + try { + rmSync(authPath, { force: true }); + rmSync(join(authPath, ".."), { force: true }); + } catch { + // best effort cleanup + } +} + +test.afterEach(() => { + globalThis.fetch = originalFetch; + if (originalAuthPath !== undefined) { + process.env.GROK_AUTH_PATH = originalAuthPath; + } else { + delete process.env.GROK_AUTH_PATH; + } + clearQuotaMonitors(); + clearSessions(); +}); + +test("fetchGrokWebQuota returns null when no auth.json exists", async () => { + // Point to a nonexistent path + process.env.GROK_AUTH_PATH = "/nonexistent/grok/auth.json"; + const quota = await fetchGrokWebQuota(`missing-auth-${Date.now()}`); + assert.equal(quota, null); +}); + +test("fetchGrokWebQuota reads auth and calls billing endpoint", async () => { + const authFile = createTestAuth(); + process.env.GROK_AUTH_PATH = authFile; + + const connectionId = `grok-test-${Date.now()}`; + const calls: Array<{ url: string; init: RequestInit }> = []; + + globalThis.fetch = async (url, init) => { + calls.push({ url: typeof url === "string" ? url : "stringified", init }); + const urlStr = typeof url === "string" ? url : ""; + // OIDC discovery + if (urlStr.includes(".well-known/openid-configuration")) { + return new Response( + JSON.stringify({ token_endpoint: "https://auth.x.ai/oauth2/token" }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + // Billing endpoint + if (urlStr.includes("billing")) { + return new Response( + JSON.stringify({ + config: { + creditUsagePercent: 15.3, + currentPeriod: { + type: "WEEKLY", + start: "2026-07-17T06:34:33.775Z", + end: "2026-07-24T06:34:33.775Z", + }, + productUsage: [ + { product: "Api", usagePercent: 15 }, + { product: "GrokBuild", usagePercent: 0 }, + ], + onDemandCap: { val: 0 }, + onDemandUsed: { val: 0 }, + prepaidBalance: { val: 0 }, + isUnifiedBillingUser: true, + }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + return new Response("not found", { status: 404 }); + }; + + const quota = await fetchGrokWebQuota(connectionId); + + assert.notEqual(quota, null, "expected non-null quota"); + assert.equal(quota.percentUsed, 0.153, "15.3% → 0.153"); + assert.equal(quota.used, 15, "used = 15"); + assert.equal(quota.total, 100, "total = 100"); + assert.ok(quota.resetAt?.includes("2026-07-24"), "resetAt matches period end"); + assert.ok(quota.windows?.weekly, "weekly window present"); + assert.equal(quota.windows.weekly.percentUsed, 0.153); + + // Verify the billing call had proper headers + const billingCall = calls.find((c) => (c.url as string).includes("billing")); + assert.ok(billingCall, "billing call made"); + const headers = billingCall.init.headers as Record; + assert.equal(headers.Authorization, "Bearer test-access-token"); + assert.equal(headers["x-grok-client-mode"], "cli"); + + invalidateGrokWebQuotaCache(connectionId); + cleanupTestAuth(authFile); +}); + +test("fetchGrokWebQuota refreshes token on 401 and retries", async () => { + const authFile = createTestAuth(); + process.env.GROK_AUTH_PATH = authFile; + + const connectionId = `grok-refresh-${Date.now()}`; + let callCount = 0; + let tokenRefreshCount = 0; + + globalThis.fetch = async (url, init) => { + const urlStr = typeof url === "string" ? url : ""; + + // OIDC discovery + if (urlStr.includes(".well-known/openid-configuration")) { + return new Response( + JSON.stringify({ token_endpoint: "https://auth.x.ai/oauth2/token" }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + + // Token refresh + if (urlStr.includes("oauth2/token")) { + tokenRefreshCount++; + return new Response( + JSON.stringify({ + access_token: "refreshed-access-token", + refresh_token: "new-refresh-token", + expires_in: 3600, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + + // Billing: first call returns 401 (triggers refresh), second call succeeds + if (urlStr.includes("billing")) { + callCount++; + if (callCount === 1) { + return new Response("Unauthorized", { status: 401 }); + } + return new Response( + JSON.stringify({ + config: { + creditUsagePercent: 42, + currentPeriod: { + type: "WEEKLY", + end: new Date(Date.now() + 7 * 24 * 3600_000).toISOString(), + }, + }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + + return new Response("not found", { status: 404 }); + }; + + const quota = await fetchGrokWebQuota(connectionId); + + assert.notEqual(quota, null, "expected quota after refresh"); + assert.equal(quota.percentUsed, 0.42, "42% after refresh"); + assert.equal(tokenRefreshCount, 1, "token was refreshed once"); + assert.equal(callCount, 2, "billing called twice (first fails, second succeeds)"); + + invalidateGrokWebQuotaCache(connectionId); + cleanupTestAuth(authFile); +}); + +test("fetchGrokWebQuota caches results for 60s", async () => { + const authFile = createTestAuth(); + process.env.GROK_AUTH_PATH = authFile; + + const connectionId = `grok-cache-${Date.now()}`; + let fetchCount = 0; + + globalThis.fetch = async (url, init) => { + const urlStr = typeof url === "string" ? url : ""; + if (urlStr.includes("billing")) { + fetchCount++; + return new Response( + JSON.stringify({ + config: { + creditUsagePercent: 30, + currentPeriod: { type: "WEEKLY", end: new Date(Date.now() + 7 * 24 * 3600_000).toISOString() }, + }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + if (urlStr.includes(".well-known")) { + return new Response( + JSON.stringify({ token_endpoint: "https://auth.x.ai/oauth2/token" }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + return new Response("not found", { status: 404 }); + }; + + // First call — should fetch + const q1 = await fetchGrokWebQuota(connectionId); + assert.equal(fetchCount, 1, "first call hits API"); + + // Second call — should use cache + const q2 = await fetchGrokWebQuota(connectionId); + assert.equal(fetchCount, 1, "second call uses cache"); + assert.equal(q1.percentUsed, q2.percentUsed); + + invalidateGrokWebQuotaCache(connectionId); + cleanupTestAuth(authFile); +}); + +test("registerGrokWebQuotaFetcher registers the fetcher for grok-web", async () => { + registerGrokWebQuotaFetcher(); + + // preflightQuota should find the registered fetcher and call it + // Point auth to a test file so it returns something + const authFile = createTestAuth(); + process.env.GROK_AUTH_PATH = authFile; + + globalThis.fetch = async (url, init) => { + const urlStr = typeof url === "string" ? url : ""; + if (urlStr.includes(".well-known")) { + return new Response( + JSON.stringify({ token_endpoint: "https://auth.x.ai/oauth2/token" }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + if (urlStr.includes("billing")) { + return new Response( + JSON.stringify({ + config: { + creditUsagePercent: 50, + currentPeriod: { type: "WEEKLY", end: new Date(Date.now() + 7 * 24 * 3600_000).toISOString() }, + }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + return new Response("not found", { status: 404 }); + }; + + const connectionId = `grok-reg-test-${Date.now()}`; + const result = await preflightQuota("grok-web", connectionId, { + provider: "grok-web", + id: connectionId, + }); + assert.equal(result.proceed, true, "preflight should proceed with headroom"); + + cleanupTestAuth(authFile); +});