9router/tests/unit/xai-usage.test.js

import { describe, it, expect, vi, beforeEach } from "vitest";

vi.mock("../../open-sse/utils/proxyFetch.js", () => ({
  proxyAwareFetch: vi.fn(),
}));

import { proxyAwareFetch } from "../../open-sse/utils/proxyFetch.js";
import { getUsageForProvider } from "../../open-sse/services/usage.js";
import { parseGrokCreditsConfig } from "../../open-sse/services/usage/xai.js";
import { parseQuotaData } from "@/app/(dashboard)/dashboard/usage/components/ProviderLimits/utils.js";

function billingResponse(config) {
  return new Response(JSON.stringify({ config }), {
    status: 200,
    headers: { "Content-Type": "application/json" },
  });
}

function settingsResponse(tier = "SuperGrok") {
  return new Response(JSON.stringify({ subscription_tier_display: tier }), {
    status: 200,
    headers: { "Content-Type": "application/json" },
  });
}

/**
 * Build a minimal grpc-web GetGrokCreditsConfig payload matching the live
 * SuperGrok response shape: usedPercent float + start/end timestamps.
 */
function buildCreditsGrpcFrame({ usedPercent = 18, startSec = 1783673158, endSec = 1784277958 }) {
  // Encode protobuf Timestamp {1: seconds}
  const encodeVarint = (n) => {
    const out = [];
    let v = n >>> 0;
    while (v >= 0x80) {
      out.push((v & 0x7f) | 0x80);
      v >>>= 7;
    }
    out.push(v);
    return Buffer.from(out);
  };
  const encodeKey = (field, wt) => encodeVarint((field << 3) | wt);
  const encodeTimestamp = (seconds) => {
    const body = Buffer.concat([encodeKey(1, 0), encodeVarint(seconds)]);
    return body;
  };
  const encodeLen = (field, bytes) =>
    Buffer.concat([encodeKey(field, 2), encodeVarint(bytes.length), bytes]);
  const encodeFloat = (field, f) => {
    const buf = Buffer.alloc(4);
    buf.writeFloatLE(f, 0);
    return Buffer.concat([encodeKey(field, 5), buf]);
  };

  const startTs = encodeTimestamp(startSec);
  const endTs = encodeTimestamp(endSec);
  const config = Buffer.concat([
    encodeFloat(1, usedPercent),
    encodeLen(4, startTs),
    encodeLen(5, endTs),
  ]);
  const msg = encodeLen(1, config);
  const frame = Buffer.alloc(5 + msg.length);
  frame[0] = 0;
  frame.writeUInt32BE(msg.length, 1);
  msg.copy(frame, 5);
  return frame;
}

function creditsResponse(opts = {}) {
  const frame = buildCreditsGrpcFrame(opts);
  return new Response(frame, {
    status: 200,
    headers: { "Content-Type": "application/grpc-web+proto" },
  });
}

function mockXaiHappyPath({ used = 733, limit = 15000, weeklyPercent = 18 } = {}) {
  proxyAwareFetch.mockImplementation(async (url) => {
    if (String(url).includes("/v1/billing")) {
      return billingResponse({
        monthlyLimit: { val: limit },
        used: { val: used },
        onDemandCap: { val: 0 },
        billingPeriodStart: "2026-07-01T00:00:00+00:00",
        billingPeriodEnd: "2026-08-01T00:00:00+00:00",
      });
    }
    if (String(url).includes("GetGrokCreditsConfig")) {
      return creditsResponse({ usedPercent: weeklyPercent });
    }
    if (String(url).includes("/v1/settings")) {
      return settingsResponse("SuperGrok");
    }
    return new Response("{}", { status: 404 });
  });
}

describe("xAI usage", () => {
  beforeEach(() => {
    vi.clearAllMocks();
  });

  it("parseGrokCreditsConfig extracts weekly used% and reset timestamp", () => {
    const frame = buildCreditsGrpcFrame({
      usedPercent: 18,
      startSec: 1783673158,
      endSec: 1784277958,
    });
    const parsed = parseGrokCreditsConfig(frame);
    expect(parsed).toMatchObject({
      usedPercent: 18,
      remainingPercent: 82,
    });
    expect(parsed.resetAt).toBe(new Date(1784277958 * 1000).toISOString());
    expect(parsed.periodStart).toBe(new Date(1783673158 * 1000).toISOString());
  });

  it("fetches billing + weekly credits + settings in parallel", async () => {
    mockXaiHappyPath();

    const usage = await getUsageForProvider({
      provider: "xai",
      accessToken: "tok-abc",
    });

    const urls = proxyAwareFetch.mock.calls.map((c) => String(c[0]));
    expect(urls).toEqual(
      expect.arrayContaining([
        "https://cli-chat-proxy.grok.com/v1/billing",
        "https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig",
        "https://cli-chat-proxy.grok.com/v1/settings",
      ]),
    );

    expect(usage.plan).toBe("SuperGrok");
    expect(usage.quotas.weekly).toMatchObject({
      used: 18,
      total: 100,
      remaining: 82,
    });
    expect(usage.quotas.api_usage).toMatchObject({
      used: 733,
      total: 15000,
      unit: "credits",
    });
    expect(usage.quotas.api_usage.resetAt).toBe("2026-08-01T00:00:00.000Z");
  });

  it("still returns weekly when billing fails but credits succeed", async () => {
    proxyAwareFetch.mockImplementation(async (url) => {
      if (String(url).includes("/v1/billing")) {
        return new Response("nope", { status: 503 });
      }
      if (String(url).includes("GetGrokCreditsConfig")) {
        return creditsResponse({ usedPercent: 42 });
      }
      if (String(url).includes("/v1/settings")) {
        return settingsResponse("SuperGrok");
      }
      return new Response("{}", { status: 404 });
    });

    const usage = await getUsageForProvider({
      provider: "xai",
      accessToken: "tok",
    });

    expect(usage.plan).toBe("SuperGrok");
    expect(usage.quotas.weekly).toMatchObject({ used: 42, total: 100, remaining: 58 });
    expect(usage.quotas.api_usage).toBeUndefined();
  });

  it("includes on-demand row when onDemandCap > 0", async () => {
    proxyAwareFetch.mockImplementation(async (url) => {
      if (String(url).includes("/v1/billing")) {
        return billingResponse({
          monthlyLimit: { val: 1000 },
          used: { val: 100 },
          onDemandCap: { val: 500 },
          billingPeriodEnd: "2026-08-01T00:00:00+00:00",
        });
      }
      if (String(url).includes("GetGrokCreditsConfig")) {
        return creditsResponse({ usedPercent: 10 });
      }
      if (String(url).includes("/v1/settings")) {
        return settingsResponse("SuperGrok");
      }
      return new Response("{}", { status: 404 });
    });

    const usage = await getUsageForProvider({
      provider: "xai",
      accessToken: "tok",
    });

    expect(usage.quotas.on_demand).toMatchObject({
      used: 0,
      total: 500,
      remainingCredits: 500,
    });
  });

  it("returns auth message on 401 billing", async () => {
    proxyAwareFetch.mockImplementation(async (url) => {
      if (String(url).includes("/v1/billing")) {
        return new Response("unauthorized", { status: 401 });
      }
      return new Response("{}", { status: 404 });
    });

    const usage = await getUsageForProvider({
      provider: "xai",
      accessToken: "tok",
    });

    expect(usage.message).toMatch(/re-authorize/i);
  });

  it("parseQuotaData maps weekly + api_usage labels", () => {
    const rows = parseQuotaData("xai", {
      quotas: {
        weekly: {
          used: 18,
          total: 100,
          remaining: 82,
          remainingPercentage: 82,
          resetAt: "2026-07-17T08:45:58.000Z",
        },
        api_usage: {
          used: 733,
          total: 15000,
          remainingCredits: 14267,
          unit: "credits",
          resetAt: "2026-08-01T00:00:00.000Z",
        },
      },
    });

    expect(rows).toEqual(
      expect.arrayContaining([
        expect.objectContaining({
          name: "Weekly limit",
          used: 18,
          total: 100,
          remaining: 82,
        }),
        expect.objectContaining({
          name: "Api usage",
          used: 733,
          total: 15000,
          unit: "credits",
        }),
      ]),
    );
    const api = rows.find((r) => r.name === "Api usage");
    expect(api.remaining).toBeUndefined();
  });
});