9router/open-sse/services/usage/xai.js

/**
 * xAI (Grok) OAuth usage handler
 *
 * SuperGrok quota is split across two upstream surfaces (OAuth only):
 *
 * 1) Monthly / API usage allotment (JSON)
 *    GET https://cli-chat-proxy.grok.com/v1/billing
 *    {
 *      "config": {
 *        "monthlyLimit": { "val": 15000 },
 *        "used": { "val": 733 },
 *        "onDemandCap": { "val": 0 },
 *        "billingPeriodStart": "...",
 *        "billingPeriodEnd": "..."
 *      }
 *    }
 *
 * 2) Weekly SuperGrok limit (grpc-web protobuf)
 *    POST https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig
 *    Empty request frame; response message1 contains:
 *      usedPercent (float32), window start/end timestamps, nested windows.
 *    This is what the grok.com usage page labels "Weekly limit" / "Resets …".
 *
 * Plan label comes from cli-chat-proxy settings:
 *    GET https://cli-chat-proxy.grok.com/v1/settings → subscription_tier_display
 *
 * Note: grok.com/rest/rate-limits is a short chat-window (e.g. 2h query count)
 * behind Cloudflare browser cookies — not usable with pure OAuth bearer.
 */

import { proxyAwareFetch } from "../../utils/proxyFetch.js";
import { U, parseResetTime, toFiniteNumber } from "./shared.js";

// Empty grpc-web request frame: flag(0) + length(0) + no payload.
const GRPC_WEB_EMPTY_FRAME = Buffer.from([0x00, 0x00, 0x00, 0x00, 0x00]);

function moneyVal(wrapper) {
  if (wrapper == null) return null;
  if (typeof wrapper === "number") return toFiniteNumber(wrapper, null);
  if (typeof wrapper === "object" && wrapper.val != null) {
    return toFiniteNumber(wrapper.val, null);
  }
  return null;
}

function authHeaders(accessToken, extra = {}) {
  return {
    Authorization: `Bearer ${accessToken}`,
    ...extra,
  };
}

function readVarint(buf, offset) {
  let val = 0;
  let shift = 0;
  let i = offset;
  while (i < buf.length) {
    const b = buf[i++];
    val |= (b & 0x7f) << shift;
    if ((b & 0x80) === 0) return { value: val >>> 0, offset: i };
    shift += 7;
    if (shift > 35) break;
  }
  return null;
}

/**
 * Minimal protobuf decoder for GetGrokCreditsConfig.
 * Only understands varint / fixed32 / fixed64 / length-delimited.
 */
function parseProtobufFields(buf) {
  const fields = [];
  let i = 0;
  while (i < buf.length) {
    const key = readVarint(buf, i);
    if (!key) break;
    i = key.offset;
    const field = key.value >>> 3;
    const wt = key.value & 7;

    if (wt === 0) {
      const v = readVarint(buf, i);
      if (!v) break;
      i = v.offset;
      fields.push({ field, type: "varint", value: v.value });
    } else if (wt === 1) {
      if (i + 8 > buf.length) break;
      fields.push({ field, type: "fixed64", value: buf.subarray(i, i + 8) });
      i += 8;
    } else if (wt === 5) {
      if (i + 4 > buf.length) break;
      fields.push({
        field,
        type: "fixed32",
        value: buf.readFloatLE(i),
      });
      i += 4;
    } else if (wt === 2) {
      const ln = readVarint(buf, i);
      if (!ln) break;
      i = ln.offset;
      if (i + ln.value > buf.length) break;
      fields.push({
        field,
        type: "bytes",
        value: buf.subarray(i, i + ln.value),
      });
      i += ln.value;
    } else {
      break;
    }
  }
  return fields;
}

function parseTimestamp(bytes) {
  if (!bytes || !bytes.length) return null;
  const fields = parseProtobufFields(bytes);
  const seconds = fields.find((f) => f.field === 1 && f.type === "varint")?.value;
  if (!Number.isFinite(seconds) || seconds <= 0) return null;
  return new Date(seconds * 1000).toISOString();
}

/**
 * Parse grpc-web response bytes from GetGrokCreditsConfig.
 * Returns { usedPercent, resetAt, periodStart } or null.
 */
export function parseGrokCreditsConfig(raw) {
  if (!raw || !raw.length) return null;
  const buf = Buffer.isBuffer(raw) ? raw : Buffer.from(raw);

  // grpc-web data frame: 1-byte flag + 4-byte big-endian length + message
  if (buf.length < 5) return null;
  const flag = buf[0];
  // Data frames have flag 0; ignore trailer frames (flag 0x80).
  if (flag !== 0) return null;
  const msgLen = buf.readUInt32BE(1);
  if (msgLen <= 0 || 5 + msgLen > buf.length) return null;
  const msg = buf.subarray(5, 5 + msgLen);

  // Response is typically { 1: CreditsConfig }
  const top = parseProtobufFields(msg);
  const configBytes = top.find((f) => f.field === 1 && f.type === "bytes")?.value || msg;
  const fields = parseProtobufFields(configBytes);

  const usedPercentRaw = fields.find((f) => f.field === 1 && f.type === "fixed32")?.value;
  const periodStart = parseTimestamp(fields.find((f) => f.field === 4 && f.type === "bytes")?.value);
  const periodEnd = parseTimestamp(fields.find((f) => f.field === 5 && f.type === "bytes")?.value);

  if (usedPercentRaw == null || !Number.isFinite(usedPercentRaw)) return null;

  const usedPercent = Math.max(0, Math.min(100, usedPercentRaw));
  return {
    usedPercent,
    remainingPercent: Math.max(0, 100 - usedPercent),
    periodStart,
    resetAt: periodEnd,
  };
}

async function fetchBilling(accessToken, billingUrl, proxyOptions) {
  const response = await proxyAwareFetch(billingUrl, {
    method: "GET",
    headers: authHeaders(accessToken, { Accept: "application/json" }),
  }, proxyOptions);

  if (response.status === 401 || response.status === 403) {
    return { error: "auth", status: response.status };
  }
  if (!response.ok) {
    return { error: "http", status: response.status };
  }

  const data = await response.json().catch(() => null);
  if (!data || typeof data !== "object") {
    return { error: "json" };
  }
  return { data };
}

async function fetchWeeklyCredits(accessToken, creditsUrl, proxyOptions) {
  if (!creditsUrl) return null;
  try {
    const response = await proxyAwareFetch(creditsUrl, {
      method: "POST",
      headers: authHeaders(accessToken, {
        "Content-Type": "application/grpc-web+proto",
        "x-grpc-web": "1",
        "x-user-agent": "connect-es/2.1.1",
        Accept: "*/*",
        Origin: "https://grok.com",
        Referer: "https://grok.com/?_s=usage",
      }),
      body: GRPC_WEB_EMPTY_FRAME,
    }, proxyOptions);

    if (!response.ok) return null;
    const ab = await response.arrayBuffer();
    return parseGrokCreditsConfig(Buffer.from(ab));
  } catch {
    return null;
  }
}

async function fetchPlanLabel(accessToken, settingsUrl, proxyOptions) {
  if (!settingsUrl) return null;
  try {
    const response = await proxyAwareFetch(settingsUrl, {
      method: "GET",
      headers: authHeaders(accessToken, { Accept: "application/json" }),
    }, proxyOptions);
    if (!response.ok) return null;
    const data = await response.json().catch(() => null);
    const label = data?.subscription_tier_display;
    return typeof label === "string" && label.trim() ? label.trim() : null;
  } catch {
    return null;
  }
}

/**
 * @param {string} accessToken - xAI OAuth access token
 * @param {object|null} proxyOptions
 */
export async function getXaiUsage(accessToken, proxyOptions = null) {
  if (!accessToken) {
    return { message: "xAI usage unavailable: no access token. Re-authorize the connection." };
  }

  const cfg = U("xai") || {};
  const billingUrl = cfg.url;
  if (!billingUrl) {
    return { message: "xAI usage endpoint is not configured." };
  }

  try {
    const [billingResult, weekly, planLabel] = await Promise.all([
      fetchBilling(accessToken, billingUrl, proxyOptions),
      fetchWeeklyCredits(accessToken, cfg.creditsUrl, proxyOptions),
      fetchPlanLabel(accessToken, cfg.settingsUrl, proxyOptions),
    ]);

    if (billingResult.error === "auth") {
      return { message: "xAI OAuth token expired or unauthorized. Please re-authorize." };
    }

    const quotas = {};
    let periodStart = null;
    let periodEnd = null;
    let onDemandCap = 0;

    if (billingResult.data) {
      const config =
        billingResult.data.config && typeof billingResult.data.config === "object"
          ? billingResult.data.config
          : billingResult.data;
      const monthlyLimit = moneyVal(config.monthlyLimit);
      const used = moneyVal(config.used);
      onDemandCap = moneyVal(config.onDemandCap) ?? 0;
      periodEnd = parseResetTime(config.billingPeriodEnd);
      periodStart = parseResetTime(config.billingPeriodStart);

      // Absolute credit counts — do NOT put remaining credits on `remaining`
      // (QuotaTable treats remaining as a 0-100 percentage; same pitfall as Qoder).
      if (monthlyLimit != null && monthlyLimit > 0) {
        const usedSafe = Math.max(0, used ?? 0);
        quotas.api_usage = {
          used: usedSafe,
          total: monthlyLimit,
          remainingCredits: Math.max(0, monthlyLimit - usedSafe),
          unit: "credits",
          resetAt: periodEnd,
          unlimited: false,
        };
      }

      if (onDemandCap > 0) {
        quotas.on_demand = {
          used: 0,
          total: onDemandCap,
          remainingCredits: onDemandCap,
          unit: "credits",
          resetAt: periodEnd,
          unlimited: false,
        };
      }
    }

    // Weekly SuperGrok window — percentage-based like Claude/Codex windows.
    if (weekly) {
      quotas.weekly = {
        used: weekly.usedPercent,
        total: 100,
        remaining: weekly.remainingPercent,
        remainingPercentage: weekly.remainingPercent,
        resetAt: weekly.resetAt || null,
        unlimited: false,
      };
      if (!periodStart && weekly.periodStart) periodStart = weekly.periodStart;
    }

    if (Object.keys(quotas).length === 0) {
      const statusHint =
        billingResult.error === "http"
          ? ` Billing API temporarily unavailable (${billingResult.status}).`
          : "";
      return {
        plan: planLabel || "xAI",
        message: `xAI connected. No quota allotment reported for this account.${statusHint}`,
        periodStart,
        periodEnd,
        quotas: {},
      };
    }

    return {
      plan: planLabel || "xAI",
      periodStart,
      periodEnd,
      onDemandCap,
      quotas,
    };
  } catch (error) {
    return { message: `xAI connected. Unable to fetch billing: ${error.message}` };
  }
}