From b7160e9fc5b8100b01df2a653734317df6763a62 Mon Sep 17 00:00:00 2001
From: Diego Rodrigues de Sa e Souza
 <8016841+diegosouzapw@users.noreply.github.com>
Date: Fri, 3 Jul 2026 08:51:01 -0300
Subject: [PATCH] feat(xai): surface Grok usage on quota dashboard via local
 usageHistory aggregation (#5806)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 CHANGELOG.md                 |   1 +
 open-sse/services/usage.ts   |  41 +++++++++++
 tests/unit/xai-usage.test.ts | 136 +++++++++++++++++++++++++++++++++++
 3 files changed, 178 insertions(+)
 create mode 100644 tests/unit/xai-usage.test.ts

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f436ac87e9..e105c5ef79 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -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
 
diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts
index 76d5fdda87..39f8296082 100644
--- a/open-sse/services/usage.ts
+++ b/open-sse/services/usage.ts
@@ -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,
diff --git a/tests/unit/xai-usage.test.ts b/tests/unit/xai-usage.test.ts
new file mode 100644
index 0000000000..5e027f3c2d
--- /dev/null
+++ b/tests/unit/xai-usage.test.ts
@@ -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);
+  });
+});