Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
41 changes: 41 additions & 0 deletions open-sse/services/usage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
};
Comment on lines +334 to +346

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Since xAI usage is tracked on a monthly basis using getMonthlyProviderTokensForConnection, the usage will reset at the start of each calendar month. Hardcoding resetAt to null is misleading for the user as the dashboard won't display when the usage resets, and the count will suddenly drop to 0 on the 1st of the next month. We should calculate and provide the actual reset timestamp (the start of the next UTC month), mirroring the pattern used in getXiaomiMimoUsage.

    const now = new Date();
    const resetAt = new Date(
      Date.UTC(now.getUTCFullYear(), now.getUTCMonth() + 1, 1)
    ).toISOString();
    return {
      plan: "xAI / Grok (OmniRoute-tracked)",
      quotas: {
        monthly: {
          used,
          total: 0,
          remaining: 100,
          remainingPercentage: 100,
          resetAt,
          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
Expand Down Expand Up @@ -497,6 +534,7 @@ export const USAGE_FETCHER_PROVIDERS = [
"opencode",
"opencode-zen",
"xiaomi-mimo",
"xai",
"vertex",
"vertex-partner",
"codebuddy-cn",
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -1006,6 +1046,7 @@ export const __testing = {
getMiniMaxRemainingPercent,
getMiniMaxUsage,
getXiaomiMimoUsage,
getXaiUsage,
getVertexUsage,
getMiniMaxAuthErrorMessage,
getMiniMaxErrorSummary,
Expand Down
136 changes: 136 additions & 0 deletions tests/unit/xai-usage.test.ts
Original file line number Diff line number Diff line change
@@ -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);
});
});
Loading