@rise-so/pi-grok-build 0.1.0

package/src/xai-grok-build.ts

@@ -0,0 +1,549 @@
+// Core of the xAI Grok Build integration: model metadata, PKCE, OIDC
+// discovery, loopback callback server, token exchange/refresh, and credential
+// narrowing. OAuth only — API keys are deliberately unsupported in this
+// project (the subscription token is the whole point). This module has ZERO
+// pi-ai runtime
+// imports so it can be shared by both the pi-ai adapter (provider.ts) and the
+// pi extension adapter (extension.ts) — the extension loader's jiti alias
+// mangles pi-ai subpath imports, so anything the extension reaches must stay
+// free of them.
+//
+// The OAuth flow borrows the Grok CLI's public client id, scope, and fixed
+// loopback port (56121) because xAI has no public client registration. xAI can
+// break this flow at any time, and a concurrent Grok CLI login on the same
+// machine collides on the port — accepted for a personal tool.
+
+import { createHash, randomBytes } from "node:crypto";
+import { createServer, type Server } from "node:http";
+// Type-only: erased at runtime, so the extension path never resolves pi-ai.
+import type { Model } from "@earendil-works/pi-ai";
+
+export const XAI_PROVIDER_ID = "xai-grok-build";
+export const XAI_MODEL_ID = "grok-build-0.1";
+export const XAI_BASE_URL = "https://api.x.ai/v1";
+
+const XAI_DISCOVERY_URL = "https://auth.x.ai/.well-known/openid-configuration";
+const XAI_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828";
+const XAI_SCOPE = "openid profile email offline_access grok-cli:access api:access";
+const XAI_REDIRECT_HOST = "127.0.0.1";
+const XAI_REDIRECT_PORT = 56_121;
+const XAI_REDIRECT_PATH = "/callback";
+const XAI_REFRESH_LEAD_MS = 5 * 60 * 1000;
+
+const MODEL_COMPAT = {
+  sendSessionIdHeader: false,
+  supportsLongCacheRetention: false,
+} as const;
+
+// Catalog notes:
+// - Context windows and per-Mtok pricing for grok-build-0.1 and grok-4.3 are
+//   confirmed by xAI's live GET /v1/models (see `pnpm smoke:models`). Live
+//   pricing doubles past a 200k-token threshold; pi's cost model has no
+//   long-context tier, so the base tier is recorded.
+// - maxTokens values are third-party registry metadata, unverified against
+//   live output caps: pi-ai sends max_output_tokens only when stream options
+//   set maxTokens, which they normally don't.
+// - grok-composer-2.5-fast does NOT appear in the live /models catalog but
+//   accepts /responses calls with a subscription OAuth token. Context/output
+//   metadata comes from third-party registries; pricing was supplied by the
+//   maintainer (not machine-verifiable against /models since it's unlisted).
+//
+// The credential also sees media models (grok-imagine-image,
+// grok-imagine-image-quality, grok-imagine-video, grok-imagine-video-1.5).
+// They use xAI's /images and /videos endpoints, not /responses, so they are
+// deliberately NOT registered here — see README "Media models".
+
+export const GROK_BUILD_MODEL = {
+  id: XAI_MODEL_ID,
+  name: "Grok Build 0.1",
+  api: "openai-responses",
+  provider: XAI_PROVIDER_ID,
+  baseUrl: XAI_BASE_URL,
+  reasoning: false,
+  input: ["text", "image"],
+  cost: {
+    input: 1,
+    output: 2,
+    cacheRead: 0.2,
+    cacheWrite: 0,
+  },
+  contextWindow: 256_000,
+  maxTokens: 256_000,
+  compat: MODEL_COMPAT,
+} satisfies Model<"openai-responses">;
+
+export const GROK_43_MODEL = {
+  id: "grok-4.3",
+  name: "Grok 4.3",
+  api: "openai-responses",
+  provider: XAI_PROVIDER_ID,
+  baseUrl: XAI_BASE_URL,
+  reasoning: true,
+  input: ["text", "image"],
+  cost: {
+    input: 1.25,
+    output: 2.5,
+    cacheRead: 0.2,
+    cacheWrite: 0,
+  },
+  contextWindow: 1_000_000,
+  maxTokens: 65_536,
+  compat: MODEL_COMPAT,
+} satisfies Model<"openai-responses">;
+
+export const GROK_COMPOSER_MODEL = {
+  id: "grok-composer-2.5-fast",
+  name: "Grok Composer 2.5 Fast",
+  api: "openai-responses",
+  provider: XAI_PROVIDER_ID,
+  baseUrl: XAI_BASE_URL,
+  reasoning: false,
+  input: ["text"],
+  cost: {
+    input: 3,
+    output: 15,
+    cacheRead: 0.5,
+    cacheWrite: 0,
+  },
+  contextWindow: 200_000,
+  maxTokens: 32_768,
+  compat: MODEL_COMPAT,
+} satisfies Model<"openai-responses">;
+
+export const GROK_MODELS = [GROK_BUILD_MODEL, GROK_43_MODEL, GROK_COMPOSER_MODEL] as const;
+
+export type XaiOptions = {
+  /**
+   * Authorization-URL `referrer` param identifying the host application.
+   * Defaults to "pi-grok-build"; the pi extension adapter passes "pi".
+   */
+  referrer?: string;
+};
+
+/**
+ * Everything the login flow needs from its host. Deliberately smaller than
+ * either pi surface: xAI's flow never prompts, only notifies.
+ */
+export type XaiLoginHost = {
+  authUrl(url: string, instructions: string): void;
+  progress(message: string): void;
+  signal?: AbortSignal;
+};
+
+/**
+ * `type`, not `interface` — needs an implicit index signature to be assignable
+ * to pi's OAuthCredential / OAuthCredentials (`[key: string]: unknown`).
+ * Exactly these four fields, nothing more: pi coding-agent's refresh path does
+ * a FULL REPLACE of the stored credential, so anything refresh omits is gone
+ * forever. Identity fields (idToken/email/subject) are deliberately not
+ * persisted — nothing reads them and the refresh grant may omit id_token.
+ */
+export type XaiOAuthCredential = {
+  access: string;
+  refresh: string;
+  expires: number;
+  tokenEndpoint: string;
+};
+
+/** pi hands back `Record<string, unknown>`. Narrow at the boundary; fail fast. */
+export function requireXaiOAuthCredential(value: Record<string, unknown>): XaiOAuthCredential {
+  const access = stringValue(value["access"]);
+  const refresh = stringValue(value["refresh"]);
+  const tokenEndpoint = stringValue(value["tokenEndpoint"]);
+  const expires = value["expires"];
+  if (
+    access === undefined ||
+    refresh === undefined ||
+    tokenEndpoint === undefined ||
+    typeof expires !== "number"
+  ) {
+    throw new Error("Stored xAI credential is incomplete; log in again");
+  }
+  return { access, refresh, expires, tokenEndpoint };
+}
+
+type XaiDiscovery = {
+  authorizationEndpoint: string;
+  tokenEndpoint: string;
+};
+
+type XaiCallbackResult =
+  | { type: "success"; code: string; state: string }
+  | { type: "error"; error: string; state: string };
+
+type XaiCallbackServer = {
+  redirectUri: string;
+  close: () => void;
+  waitForCallback: () => Promise<XaiCallbackResult>;
+};
+
+type XaiTokenData = {
+  accessToken: string;
+  refreshToken?: string;
+  expiresIn: number;
+};
+
+export async function loginXai(
+  host: XaiLoginHost,
+  opts: XaiOptions = {},
+): Promise<XaiOAuthCredential> {
+  const pkce = generatePkce();
+  const state = randomToken();
+  const nonce = randomToken();
+
+  const discovery = await discoverXaiOAuth(host.signal);
+  const callbackServer = await startXaiCallbackServer(state);
+
+  try {
+    const authUrl = buildAuthorizeUrl({
+      authorizationEndpoint: discovery.authorizationEndpoint,
+      redirectUri: callbackServer.redirectUri,
+      codeChallenge: pkce.challenge,
+      state,
+      nonce,
+      referrer: opts.referrer ?? "pi-grok-build",
+    });
+
+    host.authUrl(
+      authUrl,
+      "Complete xAI login in your browser. This local process is waiting for the OAuth callback.",
+    );
+    host.progress("Waiting for xAI authentication callback...");
+
+    // The callback server only settles on requests carrying our state, so no
+    // re-check is needed here.
+    const result = await waitForCallbackOrAbort(callbackServer.waitForCallback(), host.signal);
+    if (result.type === "error") {
+      throw new Error(`xAI authentication failed: ${result.error}`);
+    }
+
+    host.progress("Exchanging xAI authorization code...");
+    const token = await exchangeXaiCode(
+      result.code,
+      callbackServer.redirectUri,
+      pkce.verifier,
+      discovery.tokenEndpoint,
+      host.signal,
+    );
+    return tokenToCredential(token, discovery.tokenEndpoint);
+  } finally {
+    callbackServer.close();
+  }
+}
+
+export async function refreshXai(credential: XaiOAuthCredential): Promise<XaiOAuthCredential> {
+  const token = await postTokenForm(credential.tokenEndpoint, {
+    grant_type: "refresh_token",
+    client_id: XAI_CLIENT_ID,
+    refresh_token: credential.refresh,
+  });
+  return tokenToCredential(token, credential.tokenEndpoint, credential.refresh);
+}
+
+async function discoverXaiOAuth(signal?: AbortSignal): Promise<XaiDiscovery> {
+  const response = await fetch(XAI_DISCOVERY_URL, {
+    headers: { Accept: "application/json" },
+    ...(signal !== undefined ? { signal } : {}),
+  });
+  const text = await response.text();
+  if (!response.ok) {
+    throw new Error(`xAI discovery failed (${response.status}): ${text || response.statusText}`);
+  }
+  const body = parseJsonObject(text, "xAI discovery response");
+  const authorizationEndpoint = validateXaiOAuthEndpoint(
+    body["authorization_endpoint"],
+    "authorization_endpoint",
+  );
+  const tokenEndpoint = validateXaiOAuthEndpoint(body["token_endpoint"], "token_endpoint");
+  return { authorizationEndpoint, tokenEndpoint };
+}
+
+function buildAuthorizeUrl(input: {
+  authorizationEndpoint: string;
+  redirectUri: string;
+  codeChallenge: string;
+  state: string;
+  nonce: string;
+  referrer: string;
+}): string {
+  const url = new URL(input.authorizationEndpoint);
+  url.searchParams.set("response_type", "code");
+  url.searchParams.set("client_id", XAI_CLIENT_ID);
+  url.searchParams.set("redirect_uri", input.redirectUri);
+  url.searchParams.set("scope", XAI_SCOPE);
+  url.searchParams.set("code_challenge", input.codeChallenge);
+  url.searchParams.set("code_challenge_method", "S256");
+  url.searchParams.set("state", input.state);
+  url.searchParams.set("nonce", input.nonce);
+  url.searchParams.set("plan", "generic");
+  url.searchParams.set("referrer", input.referrer);
+  return url.toString();
+}
+
+async function startXaiCallbackServer(expectedState: string): Promise<XaiCallbackServer> {
+  const redirectUri = `http://${XAI_REDIRECT_HOST}:${XAI_REDIRECT_PORT}${XAI_REDIRECT_PATH}`;
+
+  return new Promise((resolve, reject) => {
+    let settled = false;
+    let settleWait: ((value: XaiCallbackResult) => void) | undefined;
+
+    const waitPromise = new Promise<XaiCallbackResult>((resolveWait) => {
+      settleWait = (value) => {
+        if (settled) return;
+        settled = true;
+        resolveWait(value);
+      };
+    });
+
+    const server = createServer((req, res) => {
+      try {
+        const url = new URL(req.url ?? "", redirectUri);
+        if (url.pathname !== XAI_REDIRECT_PATH) {
+          res.writeHead(404, { "Content-Type": "text/html; charset=utf-8" });
+          res.end(oauthPage("Authentication failed", "Callback route not found."));
+          return;
+        }
+
+        const error = stringValue(url.searchParams.get("error"));
+        const state = stringValue(url.searchParams.get("state"));
+        const code = stringValue(url.searchParams.get("code"));
+
+        // Only a request carrying the expected state may settle the wait: the
+        // fixed loopback port is shared with the Grok CLI, so a stray request
+        // or stale browser tab must not kill a legitimate login in flight.
+        // OAuth requires xAI to echo our state on both success and error
+        // redirects, so the real callback always passes this gate.
+        if (state !== expectedState) {
+          res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
+          res.end(oauthPage("Authentication failed", "OAuth state mismatch."));
+          return;
+        }
+
+        if (error !== undefined) {
+          res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
+          res.end(oauthPage("Authentication failed", `xAI returned: ${error}`));
+          settleWait?.({ type: "error", error, state });
+          return;
+        }
+
+        if (code === undefined) {
+          res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
+          res.end(oauthPage("Authentication failed", "Missing authorization code."));
+          settleWait?.({ type: "error", error: "missing authorization code", state });
+          return;
+        }
+
+        res.writeHead(200, { "Content-Type": "text/html; charset=utf-8" });
+        res.end(
+          oauthPage(
+            "Authentication successful",
+            "xAI authentication completed. You can close this window.",
+          ),
+        );
+        settleWait?.({ type: "success", code, state });
+      } catch {
+        // An unparseable request cannot be our callback; answer it and keep
+        // waiting. The login timeout still bounds the wait.
+        res.writeHead(500, { "Content-Type": "text/html; charset=utf-8" });
+        res.end(
+          oauthPage("Authentication failed", "Internal error while processing OAuth callback."),
+        );
+      }
+    });
+
+    const onListenError = (error: NodeJS.ErrnoException) => {
+      if (error.code === "EADDRINUSE") {
+        reject(
+          new Error(
+            `Port ${XAI_REDIRECT_PORT} is already in use — likely a concurrent Grok CLI ` +
+              "login (the OAuth client shares its fixed loopback port). Close it and retry.",
+          ),
+        );
+        return;
+      }
+      reject(error);
+    };
+
+    server.once("error", onListenError);
+    server.listen(XAI_REDIRECT_PORT, XAI_REDIRECT_HOST, () => {
+      server.off("error", onListenError);
+      resolve({
+        redirectUri,
+        close: () => closeServer(server),
+        waitForCallback: () => waitPromise,
+      });
+    });
+  });
+}
+
+async function waitForCallbackOrAbort<T>(promise: Promise<T>, signal?: AbortSignal): Promise<T> {
+  const timeoutSignal = AbortSignal.timeout(5 * 60 * 1000);
+  const combinedSignal =
+    signal !== undefined ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal;
+
+  if (combinedSignal.aborted) throw new Error("xAI OAuth login cancelled");
+
+  return Promise.race([
+    promise,
+    new Promise<T>((_, reject) => {
+      combinedSignal.addEventListener(
+        "abort",
+        () => reject(new Error("xAI OAuth login cancelled or timed out")),
+        { once: true },
+      );
+    }),
+  ]);
+}
+
+async function exchangeXaiCode(
+  code: string,
+  redirectUri: string,
+  verifier: string,
+  tokenEndpoint: string,
+  signal?: AbortSignal,
+): Promise<XaiTokenData> {
+  return postTokenForm(
+    tokenEndpoint,
+    {
+      grant_type: "authorization_code",
+      code,
+      redirect_uri: redirectUri,
+      client_id: XAI_CLIENT_ID,
+      code_verifier: verifier,
+    },
+    signal,
+  );
+}
+
+async function postTokenForm(
+  tokenEndpoint: string,
+  fields: Record<string, string>,
+  signal?: AbortSignal,
+): Promise<XaiTokenData> {
+  const endpoint = validateXaiOAuthEndpoint(tokenEndpoint, "token_endpoint");
+  const response = await fetch(endpoint, {
+    method: "POST",
+    headers: { "Content-Type": "application/x-www-form-urlencoded", Accept: "application/json" },
+    body: new URLSearchParams(fields),
+    ...(signal !== undefined ? { signal } : {}),
+  });
+  const text = await response.text();
+  if (!response.ok) {
+    throw new Error(
+      `xAI token request failed (${response.status}): ${text || response.statusText}`,
+    );
+  }
+  return parseTokenResponse(text);
+}
+
+function parseTokenResponse(text: string): XaiTokenData {
+  const body = parseJsonObject(text, "xAI token response");
+  const accessToken = requiredString(body["access_token"], "access_token");
+  const refreshToken = stringValue(body["refresh_token"]);
+  const expiresIn = numberValue(body["expires_in"]);
+  if (expiresIn === undefined) {
+    throw new Error("xAI token response missing expires_in");
+  }
+  assertBearer(stringValue(body["token_type"]));
+  return {
+    accessToken,
+    ...(refreshToken !== undefined ? { refreshToken } : {}),
+    expiresIn,
+  };
+}
+
+// RFC 6749 §7.1: token_type values are case-insensitive; xAI sends "Bearer".
+function assertBearer(tokenType: string | undefined): void {
+  if (tokenType !== undefined && tokenType.toLowerCase() !== "bearer") {
+    throw new Error(`xAI returned unsupported token_type: ${tokenType}`);
+  }
+}
+
+function tokenToCredential(
+  token: XaiTokenData,
+  tokenEndpoint: string,
+  fallbackRefreshToken?: string,
+): XaiOAuthCredential {
+  const refresh = token.refreshToken ?? fallbackRefreshToken;
+  if (refresh === undefined) {
+    throw new Error("xAI token response missing refresh_token");
+  }
+  return {
+    access: token.accessToken,
+    refresh,
+    expires: Date.now() + token.expiresIn * 1000 - XAI_REFRESH_LEAD_MS,
+    tokenEndpoint,
+  };
+}
+
+function generatePkce(): { verifier: string; challenge: string } {
+  const verifier = randomBytes(96).toString("base64url");
+  const challenge = createHash("sha256").update(verifier).digest("base64url");
+  return { verifier, challenge };
+}
+
+function randomToken(): string {
+  return randomBytes(16).toString("hex");
+}
+
+function validateXaiOAuthEndpoint(value: unknown, field: string): string {
+  const endpoint = requiredString(value, field);
+  const url = new URL(endpoint);
+  if (url.protocol !== "https:") {
+    throw new Error(`xAI discovery ${field} must use https`);
+  }
+  const host = url.hostname.toLowerCase();
+  if (host !== "x.ai" && !host.endsWith(".x.ai")) {
+    throw new Error(`xAI discovery ${field} host is not on x.ai`);
+  }
+  return endpoint;
+}
+
+function parseJsonObject(text: string, label: string): Record<string, unknown> {
+  const parsed: unknown = JSON.parse(text);
+  if (!isRecord(parsed)) throw new Error(`${label} must be a JSON object`);
+  return parsed;
+}
+
+function requiredString(value: unknown, field: string): string {
+  const result = stringValue(value);
+  if (result === undefined) throw new Error(`xAI response missing ${field}`);
+  return result;
+}
+
+function numberValue(value: unknown): number | undefined {
+  return typeof value === "number" && Number.isFinite(value) ? value : undefined;
+}
+
+function stringValue(value: unknown): string | undefined {
+  if (typeof value !== "string") return undefined;
+  const trimmed = value.trim();
+  return trimmed.length > 0 ? trimmed : undefined;
+}
+
+function oauthPage(title: string, message: string): string {
+  const escapedTitle = escapeHtml(title);
+  const escapedMessage = escapeHtml(message);
+  return `<!doctype html><html lang="en"><head><meta charset="utf-8"><title>${escapedTitle}</title></head><body><main><h1>${escapedTitle}</h1><p>${escapedMessage}</p></main></body></html>`;
+}
+
+function escapeHtml(value: string): string {
+  return value
+    .replaceAll("&", "&amp;")
+    .replaceAll("<", "&lt;")
+    .replaceAll(">", "&gt;")
+    .replaceAll('"', "&quot;")
+    .replaceAll("'", "&#39;");
+}
+
+function closeServer(server: Server): void {
+  server.close((error) => {
+    if (error) {
+      // Nothing actionable for the caller: the login flow has already settled.
+    }
+  });
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}