mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-04 06:12:10 +03:00
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:
committed by
GitHub
parent
5fe225850e
commit
b7160e9fc5
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
136
tests/unit/xai-usage.test.ts
Normal file
136
tests/unit/xai-usage.test.ts
Normal 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);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user