mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-07-26 09:52:11 +03:00
feat(providers): add weekly quota tracking for grok-web (#8127)
* feat(providers): add weekly quota fetcher for grok-web (grok.com SSO) Implements a bespoke QuotaFetcher for the grok-web provider that: - Reads OIDC tokens from ~/.grok/auth.json (local Grok CLI login) - Refreshes tokens via auth.x.ai OIDC if expired - Calls https://cli-chat-proxy.grok.com/v1/billing?format=credits - Returns a single 'weekly' window with creditUsagePercent and resetAt - Caches results with 60s TTL (matching codexQuotaFetcher pattern) - Supports GROK_AUTH_PATH env var override for testing - Registers in chat.ts before registerGenericQuotaFetchers Tests cover: missing auth, successful fetch with header verification, 401-triggered token refresh with retry, 60s cache TTL, and preflight integration. Closes #6444 * chore(quality): rebaseline chat.ts for #8127 own-growth --------- Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>
This commit is contained in:
Apostol Apostolov
committed by
GitHub
GitHub
parent
2a865aaaa7
commit
f3ed4a49d4
@@ -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 791<cap pre-feature). Cohesive request/logging chokepoint wiring; structural shrink of chat.ts tracked in #3501.",
|
||||
"_rebaseline_2026_07_09_6678_routing_strategy_9router": "#6678 (SeaXen) — 9router-parity Routing Strategy settings card + per-provider/combo sticky-round-robin override. Own growth: ProviderDetailPageClient.tsx 784->786 (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,
|
||||
|
||||
468
open-sse/services/grokQuotaFetcher.ts
Normal file
468
open-sse/services/grokQuotaFetcher.ts
Normal file
@@ -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<string, CacheEntry>();
|
||||
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<string, GrokAuthEntry> | null {
|
||||
if (!existsSync(getAuthPath())) return null;
|
||||
try {
|
||||
const raw = JSON.parse(readFileSync(getAuthPath(), "utf8")) as Record<string, GrokAuthEntry>;
|
||||
if (!raw || typeof raw !== "object") return null;
|
||||
return raw;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
function pickAuthEntry(file: Record<string, GrokAuthEntry>): 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<string> {
|
||||
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<ResolvedAuth> {
|
||||
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<ResolvedAuth> {
|
||||
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<UsageSnapshot> {
|
||||
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<string, unknown>,
|
||||
): Promise<QuotaInfo | null> {
|
||||
// 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]);
|
||||
}
|
||||
@@ -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
|
||||
|
||||
290
tests/unit/grok-quota-fetcher.test.ts
Normal file
290
tests/unit/grok-quota-fetcher.test.ts
Normal file
@@ -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<string, string>;
|
||||
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);
|
||||
});
|
||||
Reference in New Issue
Block a user