token-monitor/providers/xai-billing.js

/**
 * xai-billing.js — xAI Management API billing module
 *
 * Wraps the xAI Management API (`https://management-api.x.ai`) for programmatic
 * access to spend data, invoice previews, and prepaid balance.
 *
 * Environment variables required:
 *   XAI_MANAGEMENT_KEY  — management key (separate from API keys)
 *   XAI_TEAM_ID         — team UUID from console.x.ai → Settings → Team
 *
 * Both are exported in ~/.secrets/keys.env.
 *
 * API reference: xai-docs page "developers/rest-api-reference/management/billing"
 *
 * Unit conventions:
 *   - Usage `usd` aggregation returns dollars (floating point)
 *   - Invoice amounts (totalWithCorr, prepaidCredits) are in USD cents (integer strings)
 */

const MANAGEMENT_BASE = 'https://management-api.x.ai';

/**
 * Format a Date as "YYYY-MM-DD HH:MM:SS" (required by analytics timeRange).
 * @param {Date} d
 * @returns {string}
 */
function toAnalyticsDate(d) {
  return d.toISOString().replace('T', ' ').replace(/\.\d+Z$/, '');
}

/**
 * Parse a key label like "xai-vigilio (xai-...3Xxp)" → "xai-vigilio".
 * Falls back to the raw label if pattern doesn't match.
 * @param {string} label
 * @returns {string}
 */
function extractKeyName(label) {
  const m = label.match(/^([^\s(]+)/);
  return m ? m[1] : label;
}

/**
 * Fetch usage data for a date range, grouped by API key and unit type.
 *
 * @param {Object} opts
 * @param {string} opts.managementKey
 * @param {string} opts.teamId
 * @param {Date}   opts.startDate  — inclusive start (will be floored to midnight UTC)
 * @param {Date}   opts.endDate    — inclusive end (will be ceiled to 23:59:59 UTC)
 * @returns {Promise<Array<{keyId, keyName, unitType, usd}>>}  sorted by keyName then unitType
 */
export async function fetchXaiUsage({ managementKey, teamId, startDate, endDate }) {
  if (!managementKey || !teamId) {
    return [];
  }

  // Build start/end in YYYY-MM-DD HH:MM:SS format (Etc/GMT timezone)
  const start = new Date(startDate);
  start.setUTCHours(0, 0, 0, 0);
  const end = new Date(endDate);
  end.setUTCHours(23, 59, 59, 0);

  const body = {
    analyticsRequest: {
      timeRange: {
        startTime: toAnalyticsDate(start),
        endTime: toAnalyticsDate(end),
        timezone: 'Etc/GMT',
      },
      timeUnit: 'TIME_UNIT_NONE',
      values: [{ name: 'usd', aggregation: 'AGGREGATION_SUM' }],
      groupBy: ['api_key_id', 'unit_type'],
      filters: [],
    },
  };

  const response = await fetch(
    `${MANAGEMENT_BASE}/v1/billing/teams/${teamId}/usage`,
    {
      method: 'POST',
      headers: {
        'Authorization': `Bearer ${managementKey}`,
        'Content-Type': 'application/json',
      },
      body: JSON.stringify(body),
    }
  );

  if (!response.ok) {
    const text = await response.text().catch(() => '');
    throw new Error(`xAI usage API ${response.status}: ${text}`);
  }

  const data = await response.json();
  const results = [];

  for (const series of data.timeSeries || []) {
    const [keyId, unitType] = series.group || [];
    const [keyLabel] = series.groupLabels || [];
    const usd = series.dataPoints?.[0]?.values?.[0] ?? 0;
    if (usd === 0) continue;

    results.push({
      keyId,
      keyName: extractKeyName(keyLabel || keyId),
      unitType: unitType || 'unknown',
      usd,
    });
  }

  // Sort by keyName, then unitType
  results.sort((a, b) =>
    a.keyName.localeCompare(b.keyName) || a.unitType.localeCompare(b.unitType)
  );

  return results;
}

/**
 * Fetch the current billing period invoice preview (postpaid).
 *
 * Provides the running total for the current month, prepaid credits,
 * and how much of those credits have been consumed.
 *
 * @param {Object} opts
 * @param {string} opts.managementKey
 * @param {string} opts.teamId
 * @returns {Promise<{
 *   totalCents: number,
 *   prepaidCreditsCents: number,
 *   prepaidUsedCents: number,
 *   billingCycle: { year: number, month: number }
 * }|null>}
 */
export async function fetchXaiInvoicePreview({ managementKey, teamId }) {
  if (!managementKey || !teamId) return null;

  const response = await fetch(
    `${MANAGEMENT_BASE}/v1/billing/teams/${teamId}/postpaid/invoice/preview`,
    { headers: { 'Authorization': `Bearer ${managementKey}` } }
  );

  if (!response.ok) return null;

  const data = await response.json();
  const inv = data.coreInvoice || {};

  return {
    // totalWithCorr is the net amount after credits — in USD cents
    totalCents: parseInt(inv.totalWithCorr?.val ?? inv.amountAfterVat ?? '0', 10),
    // prepaidCredits is negative (money allocated from prepaid pool)
    prepaidCreditsCents: Math.abs(parseInt(inv.prepaidCredits?.val ?? '0', 10)),
    // prepaidCreditsUsed is negative (how much of that pool was consumed)
    prepaidUsedCents: Math.abs(parseInt(inv.prepaidCreditsUsed?.val ?? '0', 10)),
    billingCycle: data.billingCycle || null,
  };
}

/**
 * Aggregate raw usage rows by keyName, summing USD across unit types.
 *
 * @param {Array} rows — output of fetchXaiUsage()
 * @returns {Map<string, {total: number, byUnit: Object<string,number>}>}
 *   keyed by keyName, sorted by total descending
 */
export function aggregateByKey(rows) {
  const map = new Map();

  for (const { keyName, unitType, usd } of rows) {
    if (!map.has(keyName)) map.set(keyName, { total: 0, byUnit: {} });
    const entry = map.get(keyName);
    entry.total += usd;
    entry.byUnit[unitType] = (entry.byUnit[unitType] || 0) + usd;
  }

  // Sort by total descending
  return new Map([...map.entries()].sort((a, b) => b[1].total - a[1].total));
}

/**
 * Render an xAI spend section for the analyze.js text report.
 *
 * @param {Object} opts
 * @param {Array}  opts.usageRows    — from fetchXaiUsage()
 * @param {Object} opts.preview      — from fetchXaiInvoicePreview() (or null)
 * @param {string} opts.startDateStr — human-readable period start
 * @param {string} opts.endDateStr   — human-readable period end
 * @returns {string}
 */
export function renderXaiSection({ usageRows, preview, startDateStr, endDateStr }) {
  const lines = [];
  const period = startDateStr ? ` (${startDateStr} – ${endDateStr})` : '';
  lines.push(`xAI Spend${period}`);

  // Unit type abbreviations for compact display
  const UNIT_ABBREV = {
    'Cached prompt text tokens':  'cached',
    'Prompt text tokens':         'prompt',
    'Reasoning tokens':           'reason',
    'Reasoning text tokens':      'reason',
    'Completion text tokens':     'compl',
    'Image generation':           'image',
    'Generated image':            'image',
    'Batch completion tokens':    'batch',
    'Web searches':               'search',
  };

  const byKey = aggregateByKey(usageRows);

  if (byKey.size === 0) {
    lines.push('  (no spend data — check XAI_MANAGEMENT_KEY and XAI_TEAM_ID)');
  } else {
    for (const [keyName, { total, byUnit }] of byKey) {
      const totalStr = `$${total.toFixed(2)}`;
      const parts = Object.entries(byUnit)
        .filter(([, v]) => v > 0.001)
        .sort((a, b) => b[1] - a[1])
        .map(([unit, v]) => `${UNIT_ABBREV[unit] || unit}: $${v.toFixed(2)}`);
      const detail = parts.length > 0 ? `  [${parts.join('  ')}]` : '';
      lines.push(`  ${keyName.padEnd(20)} ${totalStr.padStart(8)}${detail}`);
    }
  }

  // Grand total from usage rows
  const grandTotal = [...byKey.values()].reduce((s, e) => s + e.total, 0);
  if (byKey.size > 1) {
    lines.push(`  ${'total'.padEnd(20)} ${('$' + grandTotal.toFixed(2)).padStart(8)}`);
  }

  // Invoice preview: prepaid balance
  if (preview) {
    const remaining = (preview.prepaidCreditsCents - preview.prepaidUsedCents) / 100;
    const used = preview.prepaidUsedCents / 100;
    const total = preview.prepaidCreditsCents / 100;
    if (total > 0) {
      lines.push(`  prepaid credit: $${remaining.toFixed(2)} remaining ($${used.toFixed(2)} of $${total.toFixed(2)} used this period)`);
    }
    const cycle = preview.billingCycle;
    if (cycle) {
      lines.push(`  billing cycle:  ${cycle.year}-${String(cycle.month).padStart(2, '0')}`);
    }
  }

  return lines.join('\n');
}