feat(xai): surface Grok usage on quota dashboard via local usageHistory aggregation (#5806)

xAI has no public per-account quota API (the billing console requires a
session cookie, not an API key). Add getXaiUsage(connectionId), mirroring
the existing Xiaomi MiMo self-track pattern: sum tokens routed to the
connection from usage_history via getMonthlyProviderTokensForConnection
and surface them as a cumulative, uncapped quota (unlimited: true,
remaining: 100 — xAI has no fixed monthly cap). Register 'xai' in
USAGE_FETCHER_PROVIDERS and wire a switch case in getUsageForProvider.


Inspired-by: https://github.com/decolua/9router/pull/2150

Co-authored-by: ron <devestacion@gmail.com>
This commit is contained in:
Diego Rodrigues de Sa e Souza
2026-07-03 08:51:01 -03:00
committed by GitHub
parent 5fe225850e
commit b7160e9fc5
3 changed files with 178 additions and 0 deletions

View File

@@ -19,6 +19,7 @@
- **feat(usage):** add on-demand period-scoped usage-data reset (Settings → System Storage) with a purge API and time-window selector.
- **feat(claude-code):** add an opt-in auto-permission classifier compat mode (off/auto/always) for Claude Code, toggleable from the CLI Code settings.
- **feat(providers):** add optional client-identity header profiles for compatible nodes — preset User-Agent/fingerprint headers (e.g. matching a known CLI) merged into the existing customHeaders field.
- **feat(xai):** surface Grok usage on the quota dashboard via local usage-history aggregation. (thanks @DevEstacion)
### 🔧 Bug Fixes

View File

@@ -312,6 +312,43 @@ async function getXiaomiMimoUsage(connectionId: string) {
}
}
/**
* xAI (Grok) — SELF-TRACKED cumulative usage.
*
* xAI has no public per-account quota API (the billing console at console.x.ai
* requires a session cookie, not an API key), so — exactly like the Xiaomi
* MiMo self-track pattern above — OmniRoute sums the tokens it itself routed
* to this connection (from `usage_history`) instead of calling an upstream
* endpoint. Unlike Xiaomi MiMo, xAI has no fixed monthly cap, so the
* aggregate is reported as `unlimited: true` with `remaining: 100` — this
* renders the dashboard's green "100%" badge instead of a meaningless
* progress bar against a `total: 0`.
*/
async function getXaiUsage(connectionId: string) {
if (!connectionId) {
return { message: "xAI: connection id unavailable for self-tracked usage." };
}
try {
const { getMonthlyProviderTokensForConnection } = await import("@/lib/usage/usageStats");
const used = getMonthlyProviderTokensForConnection("xai", connectionId);
return {
plan: "xAI / Grok (OmniRoute-tracked)",
quotas: {
monthly: {
used,
total: 0,
remaining: 100,
remainingPercentage: 100,
resetAt: null,
unlimited: true,
} as UsageQuota,
},
};
} catch (error) {
return { message: `xAI self-tracked usage error: ${(error as Error).message}` };
}
}
/**
* OpenCode Go / OpenCode / OpenCode Zen Usage
* Delegates to the dedicated opencodeQuotaFetcher and shapes the result into
@@ -497,6 +534,7 @@ export const USAGE_FETCHER_PROVIDERS = [
"opencode",
"opencode-zen",
"xiaomi-mimo",
"xai",
"vertex",
"vertex-partner",
"codebuddy-cn",
@@ -578,6 +616,8 @@ export async function getUsageForProvider(
return await getOpencodeUsage(id || "", apiKey || "");
case "xiaomi-mimo":
return await getXiaomiMimoUsage(id || "");
case "xai":
return await getXaiUsage(id || "");
case "codebuddy-cn":
return await getCodeBuddyCnUsage(accessToken, apiKey, providerSpecificData);
default:
@@ -1006,6 +1046,7 @@ export const __testing = {
getMiniMaxRemainingPercent,
getMiniMaxUsage,
getXiaomiMimoUsage,
getXaiUsage,
getVertexUsage,
getMiniMaxAuthErrorMessage,
getMiniMaxErrorSummary,

View File

@@ -0,0 +1,136 @@
/**
* tests/unit/xai-usage.test.ts
*
* xAI (Grok) has no public per-account quota API (the billing console at
* console.x.ai requires a session cookie, not an API key), so — exactly like
* the Xiaomi MiMo self-track pattern — OmniRoute self-tracks it: it sums the
* tokens it routed to the connection from `usage_history` and surfaces them
* as a cumulative, uncapped ("unlimited") usage figure on the quota
* dashboard. These tests cover the aggregation helper + the fetcher shape,
* with a real temp DB, and assert provider + connection scoping (no bleed).
*/
import { describe, it, before, after } from "node:test";
import assert from "node:assert/strict";
import os from "node:os";
import path from "node:path";
import fs from "node:fs";
// DATA_DIR must be set before any module that opens the DB is imported.
const TMP = fs.mkdtempSync(path.join(os.tmpdir(), "omni-xai-usage-"));
process.env.DATA_DIR = TMP;
const core = await import("../../src/lib/db/core.ts");
const { getMonthlyProviderTokensForConnection } = await import(
"../../src/lib/usage/usageStats.ts"
);
const { __testing, USAGE_FETCHER_PROVIDERS, getUsageForProvider } = await import(
"../../open-sse/services/usage.ts"
);
const { getXaiUsage } = __testing;
function insertUsage(
connectionId: string,
provider: string,
tokensIn: number,
tokensOut: number,
timestamp: string
) {
const db = core.getDbInstance();
db.prepare(
`INSERT INTO usage_history (provider, connection_id, tokens_input, tokens_output, timestamp)
VALUES (?, ?, ?, ?, ?)`
).run(provider, connectionId, tokensIn, tokensOut, timestamp);
}
describe("xAI self-tracked usage", () => {
before(() => {
core.getDbInstance(); // trigger migrations
const now = new Date();
const inWindow = now.toISOString();
const outOfWindow = new Date(
Date.UTC(now.getUTCFullYear(), now.getUTCMonth() - 1, 15)
).toISOString();
// in-window usage for conn-x: 2.0M + 0.3M
insertUsage("conn-x", "xai", 2_000_000, 0, inWindow);
insertUsage("conn-x", "xai", 0, 300_000, inWindow);
// out-of-window usage must NOT count toward the current aggregate
insertUsage("conn-x", "xai", 9_000_000, 9_000_000, outOfWindow);
// a different connection must not bleed in
insertUsage("conn-y", "xai", 5_000_000, 0, inWindow);
// a different provider on the same connection must not bleed in
insertUsage("conn-x", "minimax", 8_000_000, 0, inWindow);
});
after(() => {
core.resetDbInstance();
try {
fs.rmSync(TMP, { recursive: true, force: true });
} catch {
// best-effort temp cleanup
}
});
it("registers 'xai' as a usage-fetcher provider", () => {
assert.ok(
(USAGE_FETCHER_PROVIDERS as readonly string[]).includes("xai"),
"xai must be listed in USAGE_FETCHER_PROVIDERS"
);
});
it("aggregates only in-window tokens for the given provider+connection", () => {
// 2.0M + 0.3M = 2.3M; excludes out-of-window, conn-y, and minimax rows.
assert.equal(getMonthlyProviderTokensForConnection("xai", "conn-x"), 2_300_000);
});
it("returns 0 for an unknown connection (fail-open, no bleed)", () => {
assert.equal(getMonthlyProviderTokensForConnection("xai", "conn-none"), 0);
});
it("getXaiUsage returns a cumulative unlimited quota scoped to the connection", async () => {
const r = (await getXaiUsage("conn-x")) as {
plan?: string;
quotas?: Record<
string,
{
used: number;
total: number;
remaining?: number;
remainingPercentage?: number;
unlimited: boolean;
resetAt: string | null;
}
>;
message?: string;
};
assert.ok(r.quotas, `expected quotas, got message: ${r.message}`);
const m = r.quotas!.monthly;
assert.ok(m, "cumulative window present");
assert.equal(m.used, 2_300_000);
assert.equal(m.unlimited, true, "xAI has no fixed monthly cap");
assert.equal(m.remaining, 100, "unlimited rows report remaining: 100 (matches upstream UX)");
});
it("getXaiUsage does not bleed a different connection's usage", async () => {
const r = (await getXaiUsage("conn-y")) as {
quotas?: { monthly?: { used: number } };
};
assert.equal(r.quotas?.monthly?.used, 5_000_000);
});
it("getXaiUsage returns a message when connection id is missing", async () => {
const r = (await getXaiUsage("")) as { message?: string; quotas?: unknown };
assert.ok(r.message && !r.quotas, "no quota without a connection id");
});
it("getUsageForProvider('xai', ...) delegates to getXaiUsage", async () => {
const r = (await getUsageForProvider({
id: "conn-x",
provider: "xai",
} as Parameters<typeof getUsageForProvider>[0])) as {
quotas?: { monthly?: { used: number; unlimited: boolean } };
};
assert.equal(r.quotas?.monthly?.used, 2_300_000);
assert.equal(r.quotas?.monthly?.unlimited, true);
});
});