@opengeni/xai-subscription 0.1.0

package/src/oauth.ts

@@ -0,0 +1,335 @@
+import {
+  OAUTH_MAX_RESPONSE_BYTES,
+  readResponseJsonBounded,
+  validateHttpUrl,
+} from "@opengeni/network";
+
+import { runBoundedXaiOperation } from "./bounded-operation";
+import {
+  XAI_CLIENT_VERSION,
+  XAI_DEVICE_AUTHORIZATION_URL,
+  XAI_DEVICE_CODE_GRANT_TYPE,
+  XAI_OAUTH_CLIENT_ID,
+  XAI_OAUTH_OPERATION_TIMEOUT_MS,
+  XAI_OAUTH_SCOPES,
+  XAI_TOKEN_URL,
+  XAI_USERINFO_URL,
+} from "./constants";
+import {
+  XaiSubscriptionError,
+  XaiSubscriptionReloginRequired,
+  XaiSubscriptionTransientError,
+} from "./errors";
+
+export type XaiFetch = (input: string | URL | Request, init?: RequestInit) => Promise<Response>;
+
+export type XaiDeviceCode = {
+  deviceCode: string;
+  userCode: string;
+  verificationUri: string;
+  verificationUriComplete: string | null;
+  expiresInSeconds: number;
+  intervalSeconds: number;
+};
+
+export type XaiOAuthTokens = {
+  accessToken: string;
+  refreshToken: string;
+  idToken: string | null;
+  tokenType: string | null;
+  scope: string | null;
+  expiresInSeconds: number;
+};
+
+export type XaiDevicePollResult =
+  | { status: "pending"; intervalSeconds: number }
+  | { status: "slow_down"; intervalSeconds: number }
+  | { status: "expired" }
+  | { status: "denied" }
+  | { status: "authorized"; tokens: XaiOAuthTokens };
+
+export type XaiVerifiedIdentity = {
+  subject: string;
+  email: string | null;
+  emailVerified: boolean | null;
+  name: string | null;
+};
+
+type XaiOAuthOptions = {
+  fetch?: XaiFetch;
+  clientId?: string;
+  deviceAuthorizationUrl?: string;
+  tokenUrl?: string;
+  userinfoUrl?: string;
+  timeoutMs?: number;
+};
+
+function oauthHeaders(): Headers {
+  return new Headers({
+    accept: "application/json",
+    "content-type": "application/x-www-form-urlencoded",
+    "user-agent": `opengeni/${XAI_CLIENT_VERSION}`,
+    "x-grok-client-version": XAI_CLIENT_VERSION,
+    "x-grok-client-surface": "headless",
+  });
+}
+
+function positiveSeconds(value: unknown, fallback: number): number {
+  const number = Number(value);
+  return Number.isFinite(number) && number > 0 ? Math.ceil(number) : fallback;
+}
+
+function requireNonEmptyString(value: unknown, field: string): string {
+  if (typeof value !== "string" || value.length === 0) {
+    throw new XaiSubscriptionError("invalid_response", `xAI OAuth response is missing ${field}`);
+  }
+  return value;
+}
+
+function nullableString(value: unknown): string | null {
+  return typeof value === "string" && value.length > 0 ? value : null;
+}
+
+function validateUserCode(userCode: string): void {
+  if (![...userCode].every((char) => /[A-Za-z0-9-]/.test(char))) {
+    throw new XaiSubscriptionError("invalid_response", "xAI returned an invalid user code");
+  }
+}
+
+function validateVerificationUri(uri: string): string {
+  try {
+    return validateHttpUrl(uri, {
+      allowLoopbackHttp: true,
+      label: "xAI device verification",
+    });
+  } catch {
+    throw new XaiSubscriptionError("invalid_response", "xAI returned an invalid verification URL");
+  }
+}
+
+async function boundedFetchJson(
+  label: string,
+  input: string,
+  init: RequestInit,
+  options: Pick<XaiOAuthOptions, "fetch" | "timeoutMs">,
+): Promise<{ response: Response; body: Record<string, unknown> }> {
+  const fetchImpl = options.fetch ?? fetch;
+  const timeoutMs = options.timeoutMs ?? XAI_OAUTH_OPERATION_TIMEOUT_MS;
+  const fetched = await runBoundedXaiOperation(async (signal) => {
+    const response = await fetchImpl(input, { ...init, signal });
+    const body = await readResponseJsonBounded<Record<string, unknown>>(
+      response,
+      OAUTH_MAX_RESPONSE_BYTES,
+      label,
+      { signal },
+    );
+    return { response, body };
+  }, timeoutMs);
+  if (!fetched.ok) {
+    throw new XaiSubscriptionTransientError(`xAI ${label} ${fetched.reason}`);
+  }
+  return fetched.value;
+}
+
+export async function requestXaiDeviceCode(options: XaiOAuthOptions = {}): Promise<XaiDeviceCode> {
+  const { response, body } = await boundedFetchJson(
+    "device code request",
+    options.deviceAuthorizationUrl ?? XAI_DEVICE_AUTHORIZATION_URL,
+    {
+      method: "POST",
+      headers: oauthHeaders(),
+      body: new URLSearchParams({
+        client_id: options.clientId ?? XAI_OAUTH_CLIENT_ID,
+        scope: XAI_OAUTH_SCOPES.join(" "),
+        referrer: "opengeni",
+      }).toString(),
+    },
+    options,
+  );
+  if (response.status === 404) {
+    throw new XaiSubscriptionError(
+      "not_enabled",
+      "SuperGrok device-code login is not enabled by xAI",
+      404,
+    );
+  }
+  if (!response.ok) {
+    throw new XaiSubscriptionTransientError(
+      `xAI device code request failed (${response.status})`,
+      response.status,
+    );
+  }
+
+  const deviceCode = requireNonEmptyString(body.device_code, "device_code");
+  const userCode = requireNonEmptyString(body.user_code, "user_code");
+  const verificationUri = validateVerificationUri(
+    requireNonEmptyString(body.verification_uri, "verification_uri"),
+  );
+  const verificationUriComplete = nullableString(body.verification_uri_complete);
+  validateUserCode(userCode);
+
+  return {
+    deviceCode,
+    userCode,
+    verificationUri,
+    verificationUriComplete: verificationUriComplete
+      ? validateVerificationUri(verificationUriComplete)
+      : null,
+    expiresInSeconds: positiveSeconds(body.expires_in, 5 * 60),
+    intervalSeconds: Math.max(1, positiveSeconds(body.interval, 5)),
+  };
+}
+
+export async function pollXaiDeviceCode(
+  input: { deviceCode: string; intervalSeconds: number },
+  options: XaiOAuthOptions = {},
+): Promise<XaiDevicePollResult> {
+  const { response, body } = await boundedFetchJson(
+    "device token exchange",
+    options.tokenUrl ?? XAI_TOKEN_URL,
+    {
+      method: "POST",
+      headers: oauthHeaders(),
+      body: new URLSearchParams({
+        grant_type: XAI_DEVICE_CODE_GRANT_TYPE,
+        client_id: options.clientId ?? XAI_OAUTH_CLIENT_ID,
+        device_code: input.deviceCode,
+      }).toString(),
+    },
+    options,
+  );
+
+  if (response.ok) {
+    return { status: "authorized", tokens: parseTokenResponse(body, true) };
+  }
+  const code = nullableString(body.error);
+  if (code === "authorization_pending") {
+    return { status: "pending", intervalSeconds: Math.max(1, input.intervalSeconds) };
+  }
+  if (code === "slow_down") {
+    return { status: "slow_down", intervalSeconds: Math.max(1, input.intervalSeconds) + 5 };
+  }
+  if (code === "access_denied" || code === "authorization_denied") {
+    return { status: "denied" };
+  }
+  if (code === "expired_token") {
+    return { status: "expired" };
+  }
+  throw new XaiSubscriptionTransientError(
+    `xAI device token exchange failed (${response.status})`,
+    response.status,
+  );
+}
+
+export async function refreshXaiToken(
+  refreshToken: string,
+  options: XaiOAuthOptions = {},
+): Promise<XaiOAuthTokens> {
+  const { response, body } = await boundedFetchJson(
+    "token refresh",
+    options.tokenUrl ?? XAI_TOKEN_URL,
+    {
+      method: "POST",
+      headers: oauthHeaders(),
+      body: new URLSearchParams({
+        grant_type: "refresh_token",
+        refresh_token: refreshToken,
+        client_id: options.clientId ?? XAI_OAUTH_CLIENT_ID,
+      }).toString(),
+    },
+    options,
+  );
+  if (!response.ok) {
+    const code = nullableString(body.error);
+    if (
+      response.status === 401 ||
+      code === "invalid_grant" ||
+      code === "invalid_token" ||
+      code === "access_denied"
+    ) {
+      throw new XaiSubscriptionReloginRequired();
+    }
+    throw new XaiSubscriptionTransientError(
+      `xAI token refresh failed (${response.status})`,
+      response.status,
+    );
+  }
+  const parsed = parseTokenResponse(body, false);
+  return {
+    ...parsed,
+    refreshToken: parsed.refreshToken || refreshToken,
+  };
+}
+
+function parseTokenResponse(
+  body: Record<string, unknown>,
+  requireRefresh: boolean,
+): XaiOAuthTokens {
+  const accessToken = requireNonEmptyString(body.access_token, "access_token");
+  const refreshToken = nullableString(body.refresh_token);
+  if (requireRefresh && !refreshToken) {
+    throw new XaiSubscriptionError(
+      "invalid_response",
+      "xAI OAuth response is missing refresh_token",
+    );
+  }
+  return {
+    accessToken,
+    refreshToken: refreshToken ?? "",
+    idToken: nullableString(body.id_token),
+    tokenType: nullableString(body.token_type),
+    scope: nullableString(body.scope),
+    expiresInSeconds: positiveSeconds(body.expires_in, 60 * 60),
+  };
+}
+
+export async function fetchXaiVerifiedIdentity(
+  accessToken: string,
+  options: XaiOAuthOptions = {},
+): Promise<XaiVerifiedIdentity> {
+  const headers = new Headers({
+    accept: "application/json",
+    authorization: `Bearer ${accessToken}`,
+    "user-agent": `opengeni/${XAI_CLIENT_VERSION}`,
+    "x-grok-client-version": XAI_CLIENT_VERSION,
+  });
+  const { response, body } = await boundedFetchJson(
+    "userinfo request",
+    options.userinfoUrl ?? XAI_USERINFO_URL,
+    { method: "GET", headers },
+    options,
+  );
+  if (response.status === 401 || response.status === 403) {
+    throw new XaiSubscriptionReloginRequired();
+  }
+  if (!response.ok) {
+    throw new XaiSubscriptionTransientError(
+      `xAI userinfo request failed (${response.status})`,
+      response.status,
+    );
+  }
+  return {
+    subject: requireNonEmptyString(body.sub, "userinfo sub"),
+    email: nullableString(body.email),
+    emailVerified: typeof body.email_verified === "boolean" ? body.email_verified : null,
+    name: nullableString(body.name),
+  };
+}
+
+export function decodeXaiJwtPayload(jwt: string): Record<string, unknown> | null {
+  const payload = jwt.split(".")[1];
+  if (!payload) return null;
+  try {
+    return JSON.parse(Buffer.from(payload, "base64url").toString("utf8")) as Record<
+      string,
+      unknown
+    >;
+  } catch {
+    return null;
+  }
+}
+
+export function xaiAccessTokenExpiry(accessToken: string): Date | null {
+  const payload = decodeXaiJwtPayload(accessToken);
+  return typeof payload?.exp === "number" ? new Date(payload.exp * 1_000) : null;
+}

package/src/proxy.ts

@@ -0,0 +1,86 @@
+import { OAUTH_MAX_RESPONSE_BYTES, pinnedFetch, readResponseJsonBounded } from "@opengeni/network";
+
+import { runBoundedXaiOperation } from "./bounded-operation";
+import {
+  XAI_CLIENT_MODE,
+  XAI_CLIENT_VERSION,
+  XAI_SUBSCRIPTION_PROXY_BASE_URL,
+  XAI_TOKEN_AUTH_HEADER_VALUE,
+} from "./constants";
+import { XaiSubscriptionReloginRequired, XaiSubscriptionTransientError } from "./errors";
+import type { XaiFetchLike } from "./fetch";
+import type { XaiSubscriptionTokenSnapshot } from "./request-context";
+
+export type XaiProxyAuthContext = {
+  clientVersion?: string;
+  getToken: () => Promise<XaiSubscriptionTokenSnapshot>;
+  refresh: () => Promise<XaiSubscriptionTokenSnapshot>;
+};
+
+const defaultProxyFetch: XaiFetchLike = async (input, init) =>
+  await pinnedFetch(
+    input,
+    init,
+    { environment: "production", integrationsAllowPrivateNetworkTargets: false },
+    { label: "xAI subscription proxy", requireHttpsOutsideLocalTest: true },
+  );
+
+export async function fetchXaiProxyJson<T>(input: {
+  path: string;
+  context: XaiProxyAuthContext;
+  fetch?: XaiFetchLike;
+  timeoutMs?: number;
+  maxBytes?: number;
+  baseUrl?: string;
+  label: string;
+}): Promise<T> {
+  const fetchImpl = input.fetch ?? defaultProxyFetch;
+  const timeoutMs = input.timeoutMs ?? 15_000;
+  const maxBytes = input.maxBytes ?? OAUTH_MAX_RESPONSE_BYTES;
+  const url = `${(input.baseUrl ?? XAI_SUBSCRIPTION_PROXY_BASE_URL).replace(/\/+$/, "")}/${input.path.replace(/^\/+/, "")}`;
+  const request = async (token: XaiSubscriptionTokenSnapshot, signal: AbortSignal) => {
+    const headers = xaiSubscriptionProxyHeaders(token, input.context.clientVersion);
+    return await fetchImpl(url, { method: "GET", redirect: "error", headers, signal });
+  };
+  const fetched = await runBoundedXaiOperation(async (signal) => {
+    let response = await request(await input.context.getToken(), signal);
+    if (response.status === 401) {
+      await response.body?.cancel().catch(() => undefined);
+      response = await request(await input.context.refresh(), signal);
+    }
+    if (response.status === 401 || response.status === 403) {
+      await response.body?.cancel().catch(() => undefined);
+      throw new XaiSubscriptionReloginRequired();
+    }
+    if (!response.ok) {
+      await response.body?.cancel().catch(() => undefined);
+      throw new XaiSubscriptionTransientError(
+        `xAI ${input.label} failed (${response.status})`,
+        response.status,
+      );
+    }
+    return await readResponseJsonBounded<T>(response, maxBytes, `xAI ${input.label}`, { signal });
+  }, timeoutMs);
+  if (!fetched.ok) {
+    throw new XaiSubscriptionTransientError(`xAI ${input.label} ${fetched.reason}`);
+  }
+  return fetched.value;
+}
+
+export function xaiSubscriptionProxyHeaders(
+  token: XaiSubscriptionTokenSnapshot,
+  clientVersion = XAI_CLIENT_VERSION,
+): Headers {
+  return new Headers({
+    accept: "application/json",
+    authorization: `Bearer ${token.accessToken}`,
+    "user-agent": `opengeni/${clientVersion}`,
+    "x-grok-client-version": clientVersion,
+    "x-grok-client-identifier": "opengeni",
+    "x-grok-client-mode": XAI_CLIENT_MODE,
+    "x-authenticateresponse": "authenticate-response",
+    "x-xai-token-auth": XAI_TOKEN_AUTH_HEADER_VALUE,
+    "x-userid": token.userId,
+    "x-grok-user-id": token.userId,
+  });
+}

package/src/quota.ts

@@ -0,0 +1,101 @@
+import type { XaiFetchLike } from "./fetch";
+import { fetchXaiProxyJson, type XaiProxyAuthContext } from "./proxy";
+
+export type XaiUsagePeriod = {
+  type: string | null;
+  start: Date | null;
+  end: Date | null;
+};
+
+export type XaiSubscriptionQuota = {
+  usedPercent: number | null;
+  period: XaiUsagePeriod | null;
+  prepaidBalanceCents: number | null;
+  onDemandCapCents: number | null;
+  onDemandUsedCents: number | null;
+  onDemandEnabled: boolean | null;
+  unifiedBilling: boolean | null;
+  subscriptionTier: string | null;
+  checkedAt: Date;
+};
+
+export async function fetchXaiSubscriptionQuota(input: {
+  context: XaiProxyAuthContext;
+  fetch?: XaiFetchLike;
+  timeoutMs?: number;
+  baseUrl?: string;
+}): Promise<XaiSubscriptionQuota> {
+  const body = await fetchXaiProxyJson<Record<string, unknown>>({
+    path: "billing?format=credits",
+    context: input.context,
+    ...(input.fetch ? { fetch: input.fetch } : {}),
+    ...(input.timeoutMs ? { timeoutMs: input.timeoutMs } : {}),
+    ...(input.baseUrl ? { baseUrl: input.baseUrl } : {}),
+    label: "billing request",
+  });
+  const config = record(body.config);
+  const currentPeriod = record(config?.currentPeriod ?? config?.current_period);
+  const monthlyLimit = cents(config?.monthlyLimit ?? config?.monthly_limit);
+  const used = cents(config?.used);
+  const directPercent = finitePercent(config?.creditUsagePercent ?? config?.credit_usage_percent);
+  const derivedPercent =
+    monthlyLimit !== null && monthlyLimit > 0 && used !== null
+      ? Math.max(0, Math.min(100, (used / monthlyLimit) * 100))
+      : null;
+  const period = currentPeriod
+    ? {
+        type: stringOrNull(currentPeriod.type),
+        start: dateOrNull(currentPeriod.start),
+        end: dateOrNull(currentPeriod.end),
+      }
+    : legacyPeriod(config);
+  return {
+    usedPercent: directPercent ?? derivedPercent,
+    period,
+    prepaidBalanceCents: cents(config?.prepaidBalance ?? config?.prepaid_balance),
+    onDemandCapCents: cents(config?.onDemandCap ?? config?.on_demand_cap),
+    onDemandUsedCents: cents(config?.onDemandUsed ?? config?.on_demand_used),
+    onDemandEnabled: booleanOrNull(body.onDemandEnabled ?? body.on_demand_enabled),
+    unifiedBilling: booleanOrNull(config?.isUnifiedBillingUser ?? config?.is_unified_billing_user),
+    subscriptionTier: stringOrNull(body.subscriptionTier ?? body.subscription_tier),
+    checkedAt: new Date(),
+  };
+}
+
+function record(value: unknown): Record<string, unknown> | null {
+  return value && typeof value === "object" && !Array.isArray(value)
+    ? (value as Record<string, unknown>)
+    : null;
+}
+
+function cents(value: unknown): number | null {
+  if (typeof value === "number" && Number.isSafeInteger(value)) return value;
+  const object = record(value);
+  const candidate = object?.val;
+  return typeof candidate === "number" && Number.isSafeInteger(candidate) ? candidate : null;
+}
+
+function finitePercent(value: unknown): number | null {
+  const candidate = Number(value);
+  return Number.isFinite(candidate) ? Math.max(0, Math.min(100, candidate)) : null;
+}
+
+function stringOrNull(value: unknown): string | null {
+  return typeof value === "string" && value.length > 0 ? value : null;
+}
+
+function booleanOrNull(value: unknown): boolean | null {
+  return typeof value === "boolean" ? value : null;
+}
+
+function dateOrNull(value: unknown): Date | null {
+  if (typeof value !== "string") return null;
+  const date = new Date(value);
+  return Number.isNaN(date.getTime()) ? null : date;
+}
+
+function legacyPeriod(config: Record<string, unknown> | null): XaiUsagePeriod | null {
+  const start = dateOrNull(config?.billingPeriodStart ?? config?.billing_period_start);
+  const end = dateOrNull(config?.billingPeriodEnd ?? config?.billing_period_end);
+  return start || end ? { type: null, start, end } : null;
+}

package/src/request-context.ts

@@ -0,0 +1,31 @@
+import { AsyncLocalStorage } from "node:async_hooks";
+
+export type XaiSubscriptionTokenSnapshot = {
+  accessToken: string;
+  userId: string;
+};
+
+export type XaiHostedSearchOptions = {
+  webSearch?: boolean | Record<string, unknown>;
+  xSearch?: boolean | Record<string, unknown>;
+};
+
+export type XaiFinalContextUsage = {
+  inputTokens: number;
+  outputTokens: number;
+  totalTokens: number;
+};
+
+export type XaiSubscriptionRequestContext = {
+  clientVersion: string;
+  sessionId: string;
+  turnId: string;
+  getToken: () => Promise<XaiSubscriptionTokenSnapshot>;
+  refresh: () => Promise<XaiSubscriptionTokenSnapshot>;
+  resolveModel: (slug: string) => string;
+  hostedSearch?: XaiHostedSearchOptions;
+  onFinalContextUsage?: (usage: XaiFinalContextUsage) => void;
+  nextRequestId?: () => string;
+};
+
+export const xaiSubscriptionRequestStorage = new AsyncLocalStorage<XaiSubscriptionRequestContext>();