@oh-my-pi/pi-ai 16.3.13 → 16.3.15

package/src/registry/oauth/xai-oauth.ts

@@ -1,31 +1,22 @@
-// Ported from NousResearch/hermes-agent (MIT) — hermes_cli/auth.py xAI sections (L93-111, L2979-3160, L5286-5469).
+// Device authorization and token refresh adapted from NousResearch/hermes-agent (MIT).
 
 /**
- * xAI Grok (SuperGrok or X Premium+) OAuth flow.
+ * xAI Grok OAuth device authorization flow.
  *
- * Manual-code PKCE flow using `127.0.0.1:56121/callback` as the allowlisted
- * redirect URI. One token unlocks Grok-4.x
- * chat, Grok Imagine image generation, and Grok Voice TTS via subsequent
- * commits. Endpoint discovery is hardened against MITM via
- * {@link validateXAIEndpoint}: any non-HTTPS or non-`x.ai`/`*.x.ai` host is
- * rejected on every call site, not just the first.
+ * Requests an RFC 8628 device code, opens xAI's verification page, and polls
+ * the discovered token endpoint until the user approves the login.
  */
 
 import * as AIError from "../../error";
 import type { FetchImpl } from "../../types";
-import { OAuthCallbackFlow, type OAuthCallbackFlowOptions } from "./callback-server";
-import { generatePKCE } from "./pkce";
+import { type OAuthDeviceCodePollResult, pollOAuthDeviceCodeFlow } from "./device-code";
 import type { OAuthController, OAuthCredentials } from "./types";
 
-// Hermes hermes_cli/auth.py L93-111
 const XAI_OAUTH_ISSUER = "https://auth.x.ai";
 const XAI_OAUTH_DISCOVERY_URL = `${XAI_OAUTH_ISSUER}/.well-known/openid-configuration`;
+const XAI_OAUTH_DEVICE_CODE_URL = `${XAI_OAUTH_ISSUER}/oauth2/device/code`;
 const XAI_OAUTH_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828";
 const XAI_OAUTH_SCOPE = "openid profile email offline_access grok-cli:access api:access";
-const XAI_OAUTH_REDIRECT_HOST = "127.0.0.1";
-const XAI_OAUTH_REDIRECT_PORT = 56121;
-const XAI_OAUTH_REDIRECT_PATH = "/callback";
-const XAI_OAUTH_DOCS_URL = "https://hermes-agent.nousresearch.com/docs/guides/xai-grok-oauth";
 
 // Mirrors the 5-min skew used by anthropic.ts:160 — keeps every provider on the
 // same conservative client-side expiry window.
@@ -35,18 +26,27 @@ const DISCOVERY_TIMEOUT_MS = 15_000;
 const TOKEN_REQUEST_TIMEOUT_MS = 20_000;
 
 interface XAIOAuthDiscovery {
-	authorization_endpoint: string;
 	token_endpoint: string;
 }
 
+interface XAIDeviceAuthorization {
+	deviceCode: string;
+	userCode: string;
+	verificationUriComplete: string;
+	expiresInSeconds: number;
+	intervalSeconds: number;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null;
+}
+
 /**
- * Validate an xAI OIDC discovery endpoint against scheme + host.
+ * Validate an xAI OIDC endpoint against its scheme and host.
  *
- * Hermes `_xai_validate_oauth_endpoint` L2997-3035. The discovery response is
- * long-lived and cached in {@link OAuthCredentials}; a single MITM during
- * initial login could substitute a malicious `token_endpoint` that would then
- * receive every future refresh_token. Rejecting non-HTTPS or non-`x.ai` /
- * `*.x.ai` hosts pins the cached endpoint to the xAI auth origin.
+ * The discovery response is long-lived and its token endpoint receives every
+ * future refresh token. Rejecting non-HTTPS or non-`x.ai` / `*.x.ai` hosts
+ * pins that endpoint to the xAI auth origin.
  *
  * @throws Error with message `Invalid xAI <field>: <url>` when the URL fails
  *         either scheme or host validation.
@@ -68,11 +68,7 @@ export function validateXAIEndpoint(url: string, field: string): string {
 	return url;
 }
 
-/**
- * Fetch xAI's OIDC discovery document and validate both endpoints.
- *
- * Hermes `_xai_oauth_discovery` L3038-3084.
- */
+/** Fetch xAI's OIDC discovery document and validate the token endpoint. */
 async function xaiOAuthDiscovery(
 	timeoutMs: number = DISCOVERY_TIMEOUT_MS,
 	fetchOverride?: FetchImpl,
@@ -111,37 +107,29 @@ async function xaiOAuthDiscovery(
 			{ kind: "validation", provider: "xai", cause: error },
 		);
 	}
-	if (!payload || typeof payload !== "object") {
+	if (!isRecord(payload)) {
 		throw new AIError.OAuthError("xAI OIDC discovery response was not a JSON object.", {
 			kind: "validation",
 			provider: "xai",
 		});
 	}
-	const obj = payload as Record<string, unknown>;
-	const authorizationEndpoint =
-		typeof obj.authorization_endpoint === "string" ? obj.authorization_endpoint.trim() : "";
-	const tokenEndpoint = typeof obj.token_endpoint === "string" ? obj.token_endpoint.trim() : "";
-	if (!authorizationEndpoint || !tokenEndpoint) {
-		throw new AIError.OAuthError("xAI OIDC discovery response was missing required endpoints.", {
+	const tokenEndpoint = typeof payload.token_endpoint === "string" ? payload.token_endpoint.trim() : "";
+	if (!tokenEndpoint) {
+		throw new AIError.OAuthError("xAI OIDC discovery response was missing token_endpoint.", {
 			kind: "validation",
 			provider: "xai",
 		});
 	}
-	validateXAIEndpoint(authorizationEndpoint, "authorization_endpoint");
 	validateXAIEndpoint(tokenEndpoint, "token_endpoint");
-	return {
-		authorization_endpoint: authorizationEndpoint,
-		token_endpoint: tokenEndpoint,
-	};
+	return { token_endpoint: tokenEndpoint };
 }
 
 /**
  * Check whether a JWT access token is at or past its `exp` claim (with an
  * optional refresh-skew margin).
  *
- * Hermes `_xai_access_token_is_expiring` L2979-2994. Returns `false` for any
- * malformed input — this is a refresh-trigger check, not a validation, so
- * non-JWTs ("no token in cache") must NOT trigger a spurious refresh.
+ * Returns `false` for malformed input because this is a refresh-trigger check,
+ * not token validation.
  */
 export function isXAIAccessTokenExpiring(jwt: string, skewSeconds: number = 0): boolean {
 	try {
@@ -151,7 +139,8 @@ export function isXAIAccessTokenExpiring(jwt: string, skewSeconds: number = 0):
 		const payloadPart = parts[1];
 		if (!payloadPart) return false;
 		const decoded = Buffer.from(payloadPart, "base64url").toString("utf8");
-		const payload = JSON.parse(decoded) as { exp?: unknown };
+		const payload: unknown = JSON.parse(decoded);
+		if (!isRecord(payload)) return false;
 		const exp = payload.exp;
 		if (typeof exp !== "number" || !Number.isFinite(exp)) return false;
 		const now = Math.floor(Date.now() / 1000);
@@ -162,161 +151,232 @@ export function isXAIAccessTokenExpiring(jwt: string, skewSeconds: number = 0):
 	}
 }
 
-interface BuildXAIAuthorizeUrlOptions {
-	authorizationEndpoint: string;
-	redirectUri: string;
-	codeChallenge: string;
-	state: string;
-	nonce: string;
-}
-
-/**
- * Build the xAI authorization URL.
- *
- * Hermes `_xai_oauth_build_authorize_url` L5286-5312. `plan=generic` opts the
- * consent screen into xAI's generic OAuth plan tier; without it,
- * `accounts.x.ai` rejects loopback OAuth from non-allowlisted clients.
- * `referrer=oh-my-pi` lets xAI attribute oh-my-pi-originated logins in their
- * OAuth server logs (Hermes uses `referrer=hermes-agent`; oh-my-pi mirrors the
- * pattern with its own attribution string).
- */
-function buildXAIAuthorizeUrl(opts: BuildXAIAuthorizeUrlOptions): string {
-	const params = new URLSearchParams({
-		response_type: "code",
-		client_id: XAI_OAUTH_CLIENT_ID,
-		redirect_uri: opts.redirectUri,
-		scope: XAI_OAUTH_SCOPE,
-		code_challenge: opts.codeChallenge,
-		code_challenge_method: "S256",
-		state: opts.state,
-		nonce: opts.nonce,
-		plan: "generic",
-		referrer: "oh-my-pi",
-	});
-	return `${opts.authorizationEndpoint}?${params.toString()}`;
-}
-
-/**
- * xAI Grok OAuth code flow (Hermes `_xai_oauth_loopback_login` L5315-5469).
- */
-export class XAIOAuthFlow extends OAuthCallbackFlow {
-	#verifier: string = "";
-	#fetch: FetchImpl;
-
-	constructor(ctrl: OAuthController) {
-		super(ctrl, {
-			preferredPort: XAI_OAUTH_REDIRECT_PORT,
-			callbackPath: XAI_OAUTH_REDIRECT_PATH,
-			callbackHostname: XAI_OAUTH_REDIRECT_HOST,
-			redirectUri: `http://${XAI_OAUTH_REDIRECT_HOST}:${XAI_OAUTH_REDIRECT_PORT}${XAI_OAUTH_REDIRECT_PATH}`,
-			manualInputOnly: true,
-		} satisfies OAuthCallbackFlowOptions);
-		this.#fetch = ctrl.fetch ?? fetch;
+function parseXAIDeviceAuthorization(payload: unknown): XAIDeviceAuthorization {
+	if (!isRecord(payload)) {
+		throw new AIError.OAuthError("xAI device-code response was not a JSON object.", {
+			kind: "validation",
+			provider: "xai",
+		});
 	}
 
-	async generateAuthUrl(state: string, redirectUri: string): Promise<{ url: string; instructions?: string }> {
-		const pkce = await generatePKCE();
-		this.#verifier = pkce.verifier;
-		const nonce = crypto.randomUUID().replace(/-/g, "");
-
-		const discovery = await xaiOAuthDiscovery(DISCOVERY_TIMEOUT_MS, this.#fetch);
-		const url = buildXAIAuthorizeUrl({
-			authorizationEndpoint: discovery.authorization_endpoint,
-			redirectUri,
-			codeChallenge: pkce.challenge,
-			state,
-			nonce,
+	const deviceCode = typeof payload.device_code === "string" ? payload.device_code.trim() : "";
+	const userCode = typeof payload.user_code === "string" ? payload.user_code.trim() : "";
+	const verificationUri = typeof payload.verification_uri === "string" ? payload.verification_uri.trim() : "";
+	const verificationUriComplete =
+		typeof payload.verification_uri_complete === "string" ? payload.verification_uri_complete.trim() : "";
+	const expiresInSeconds = payload.expires_in;
+	const intervalSeconds = payload.interval;
+	if (
+		!deviceCode ||
+		!userCode ||
+		!verificationUri ||
+		!verificationUriComplete ||
+		typeof expiresInSeconds !== "number" ||
+		!Number.isFinite(expiresInSeconds) ||
+		expiresInSeconds <= 0 ||
+		typeof intervalSeconds !== "number" ||
+		!Number.isFinite(intervalSeconds) ||
+		intervalSeconds <= 0
+	) {
+		throw new AIError.OAuthError("xAI device-code response missing or invalid required fields.", {
+			kind: "validation",
+			provider: "xai",
 		});
-
-		return {
-			url,
-			instructions: `Complete login in your browser for xAI Grok (SuperGrok or X Premium+). Docs: ${XAI_OAUTH_DOCS_URL}`,
-		};
 	}
 
-	async exchangeToken(code: string, _state: string, redirectUri: string): Promise<OAuthCredentials> {
-		const discovery = await xaiOAuthDiscovery(DISCOVERY_TIMEOUT_MS, this.#fetch);
-		const tokenEndpoint = validateXAIEndpoint(discovery.token_endpoint, "token_endpoint");
+	validateXAIEndpoint(verificationUri, "verification_uri");
+	validateXAIEndpoint(verificationUriComplete, "verification_uri_complete");
+	return {
+		deviceCode,
+		userCode,
+		verificationUriComplete,
+		expiresInSeconds,
+		intervalSeconds,
+	};
+}
 
-		const body = new URLSearchParams({
-			grant_type: "authorization_code",
-			client_id: XAI_OAUTH_CLIENT_ID,
-			code,
-			redirect_uri: redirectUri,
-			code_verifier: this.#verifier,
+function parseXAITokenResponse(payload: unknown, label: string, refreshTokenFallback?: string): OAuthCredentials {
+	if (!isRecord(payload)) {
+		throw new AIError.OAuthError(`${label} was not a JSON object`, {
+			kind: "validation",
+			provider: "xai",
 		});
+	}
+	const accessToken = typeof payload.access_token === "string" ? payload.access_token : "";
+	const responseRefreshToken = typeof payload.refresh_token === "string" ? payload.refresh_token : "";
+	const refreshToken = responseRefreshToken || refreshTokenFallback || "";
+	const expiresInSeconds = payload.expires_in;
+	if (!accessToken) {
+		throw new AIError.OAuthError(`${label} missing access_token`, {
+			kind: "validation",
+			provider: "xai",
+		});
+	}
+	if (!refreshToken) {
+		throw new AIError.OAuthError(`${label} missing refresh_token`, {
+			kind: "validation",
+			provider: "xai",
+		});
+	}
+	if (typeof expiresInSeconds !== "number" || !Number.isFinite(expiresInSeconds)) {
+		throw new AIError.OAuthError(`${label} missing expires_in`, {
+			kind: "validation",
+			provider: "xai",
+		});
+	}
+	return {
+		access: accessToken,
+		refresh: refreshToken,
+		expires: Date.now() + expiresInSeconds * 1000 - ACCESS_TOKEN_CLIENT_SKEW_MS,
+	};
+}
 
-		const response = await this.#fetch(tokenEndpoint, {
+async function requestXAIDeviceAuthorization(
+	fetchImpl: FetchImpl,
+	signal?: AbortSignal,
+): Promise<XAIDeviceAuthorization> {
+	let response: Response;
+	try {
+		const timeoutSignal = AbortSignal.timeout(TOKEN_REQUEST_TIMEOUT_MS);
+		response = await fetchImpl(XAI_OAUTH_DEVICE_CODE_URL, {
 			method: "POST",
 			headers: {
 				"Content-Type": "application/x-www-form-urlencoded",
 				Accept: "application/json",
 			},
-			body,
-			signal: AbortSignal.timeout(TOKEN_REQUEST_TIMEOUT_MS),
+			body: new URLSearchParams({
+				client_id: XAI_OAUTH_CLIENT_ID,
+				scope: XAI_OAUTH_SCOPE,
+			}),
+			signal: signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal,
 		});
+	} catch (error) {
+		if (signal?.aborted) throw new AIError.LoginCancelledError();
+		throw new AIError.OAuthError(
+			`xAI device-code request failed: ${error instanceof Error ? error.message : String(error)}`,
+			{ kind: "device-auth", provider: "xai", cause: error },
+		);
+	}
 
-		if (!response.ok) {
-			let detail = "";
-			try {
-				detail = (await response.text()).trim();
-			} catch {
-				// Ignore body-read failures; the status code is the diagnostic.
-			}
-			throw new AIError.OAuthError(`xAI token exchange failed: ${response.status}${detail ? ` ${detail}` : ""}`, {
-				kind: "token-exchange",
-				provider: "xai",
-				status: response.status,
-			});
-		}
-
-		let tokenData: { access_token?: unknown; refresh_token?: unknown; expires_in?: unknown };
+	if (!response.ok) {
+		let detail = "";
 		try {
-			tokenData = (await response.json()) as typeof tokenData;
-		} catch (error) {
-			throw new AIError.OAuthError(
-				`xAI token exchange returned invalid JSON: ${error instanceof Error ? error.message : String(error)}`,
-				{ kind: "validation", provider: "xai", cause: error },
-			);
+			detail = (await response.text()).trim();
+		} catch {
+			// Ignore body-read failures; the status code is the diagnostic.
 		}
+		throw new AIError.OAuthError(`xAI device-code request failed: ${response.status}${detail ? ` ${detail}` : ""}`, {
+			kind: "device-auth",
+			provider: "xai",
+			status: response.status,
+		});
+	}
 
-		if (typeof tokenData.access_token !== "string" || !tokenData.access_token) {
-			throw new AIError.OAuthError("xAI token exchange response missing access_token", {
-				kind: "validation",
-				provider: "xai",
-			});
-		}
-		if (typeof tokenData.refresh_token !== "string" || !tokenData.refresh_token) {
-			throw new AIError.OAuthError("xAI token exchange response missing refresh_token", {
-				kind: "validation",
-				provider: "xai",
-			});
-		}
-		if (typeof tokenData.expires_in !== "number" || !Number.isFinite(tokenData.expires_in)) {
-			throw new AIError.OAuthError("xAI token exchange response missing expires_in", {
-				kind: "validation",
-				provider: "xai",
-			});
-		}
+	let payload: unknown;
+	try {
+		payload = await response.json();
+	} catch (error) {
+		throw new AIError.OAuthError(
+			`xAI device-code response returned invalid JSON: ${error instanceof Error ? error.message : String(error)}`,
+			{ kind: "validation", provider: "xai", cause: error },
+		);
+	}
+	return parseXAIDeviceAuthorization(payload);
+}
+
+async function pollXAIDeviceToken(
+	tokenEndpoint: string,
+	deviceCode: string,
+	fetchImpl: FetchImpl,
+	signal?: AbortSignal,
+): Promise<OAuthDeviceCodePollResult<OAuthCredentials>> {
+	let response: Response;
+	try {
+		const timeoutSignal = AbortSignal.timeout(TOKEN_REQUEST_TIMEOUT_MS);
+		response = await fetchImpl(tokenEndpoint, {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/x-www-form-urlencoded",
+				Accept: "application/json",
+			},
+			body: new URLSearchParams({
+				grant_type: "urn:ietf:params:oauth:grant-type:device_code",
+				client_id: XAI_OAUTH_CLIENT_ID,
+				device_code: deviceCode,
+			}),
+			signal: signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal,
+		});
+	} catch (error) {
+		if (signal?.aborted) throw new AIError.LoginCancelledError();
+		throw new AIError.OAuthError(
+			`xAI device-code token polling failed: ${error instanceof Error ? error.message : String(error)}`,
+			{ kind: "polling", provider: "xai", cause: error },
+		);
+	}
 
+	let payload: unknown;
+	try {
+		payload = await response.json();
+	} catch (error) {
+		throw new AIError.OAuthError(
+			`xAI device-code token polling returned invalid JSON: ${
+				error instanceof Error ? error.message : String(error)
+			}`,
+			{ kind: "polling", provider: "xai", status: response.status, cause: error },
+		);
+	}
+
+	if (response.ok) {
 		return {
-			access: tokenData.access_token,
-			refresh: tokenData.refresh_token,
-			expires: Date.now() + tokenData.expires_in * 1000 - ACCESS_TOKEN_CLIENT_SKEW_MS,
+			status: "complete",
+			value: parseXAITokenResponse(payload, "xAI device-code token response"),
 		};
 	}
+	if (!isRecord(payload)) {
+		throw new AIError.OAuthError(`xAI device-code token polling failed: ${response.status}`, {
+			kind: "polling",
+			provider: "xai",
+			status: response.status,
+		});
+	}
+
+	const errorCode = typeof payload.error === "string" ? payload.error : "";
+	if (errorCode === "authorization_pending") return { status: "pending" };
+	if (errorCode === "slow_down") return { status: "slow_down" };
+
+	const errorDescription = typeof payload.error_description === "string" ? payload.error_description : "";
+	const detail = errorDescription || errorCode || String(response.status);
+	throw new AIError.OAuthError(`xAI device-code token polling failed: ${detail}`, {
+		kind: "polling",
+		provider: "xai",
+		status: response.status,
+	});
 }
 
+/** Log in to xAI Grok with the RFC 8628 device authorization grant. */
 export async function loginXAIOAuth(ctrl: OAuthController): Promise<OAuthCredentials> {
-	return new XAIOAuthFlow(ctrl).login();
+	const fetchImpl = ctrl.fetch ?? fetch;
+	const discovery = await xaiOAuthDiscovery(DISCOVERY_TIMEOUT_MS, fetchImpl);
+	const device = await requestXAIDeviceAuthorization(fetchImpl, ctrl.signal);
+	ctrl.onAuth?.({
+		url: device.verificationUriComplete,
+		instructions: `Enter code: ${device.userCode}`,
+	});
+	ctrl.onProgress?.("Waiting for xAI device authorization...");
+
+	return pollOAuthDeviceCodeFlow({
+		poll: () => pollXAIDeviceToken(discovery.token_endpoint, device.deviceCode, fetchImpl, ctrl.signal),
+		intervalSeconds: device.intervalSeconds,
+		expiresInSeconds: device.expiresInSeconds,
+		signal: ctrl.signal,
+	});
 }
 
 /**
  * Refresh an xAI OAuth access token using a stored refresh_token.
  *
- * Hermes `refresh_xai_oauth_pure` L3087-3160. Re-runs OIDC discovery and
- * re-validates the cached `token_endpoint` on the refresh hot path so a
- * cached-but-poisoned endpoint cannot silently leak a refresh_token.
+ * Re-runs OIDC discovery and re-validates the token endpoint before sending
+ * the stored refresh token.
  */
 export async function refreshXAIOAuthToken(refreshToken: string, fetchOverride?: FetchImpl): Promise<OAuthCredentials> {
 	const fetchImpl = fetchOverride ?? fetch;
@@ -357,34 +417,14 @@ export async function refreshXAIOAuthToken(refreshToken: string, fetchOverride?:
 		});
 	}
 
-	let data: { access_token?: unknown; refresh_token?: unknown; expires_in?: unknown };
+	let payload: unknown;
 	try {
-		data = (await response.json()) as typeof data;
+		payload = await response.json();
 	} catch (error) {
 		throw new AIError.OAuthError(
 			`xAI token refresh returned invalid JSON: ${error instanceof Error ? error.message : String(error)}`,
 			{ kind: "validation", provider: "xai", cause: error },
 		);
 	}
-
-	if (typeof data.access_token !== "string" || !data.access_token) {
-		throw new AIError.OAuthError("xAI token refresh response missing access_token", {
-			kind: "validation",
-			provider: "xai",
-		});
-	}
-	if (typeof data.expires_in !== "number" || !Number.isFinite(data.expires_in)) {
-		throw new AIError.OAuthError("xAI token refresh response missing expires_in", {
-			kind: "validation",
-			provider: "xai",
-		});
-	}
-
-	const newRefresh = typeof data.refresh_token === "string" && data.refresh_token ? data.refresh_token : refreshToken;
-
-	return {
-		access: data.access_token,
-		refresh: newRefresh,
-		expires: Date.now() + data.expires_in * 1000 - ACCESS_TOKEN_CLIENT_SKEW_MS,
-	};
+	return parseXAITokenResponse(payload, "xAI token refresh response", refreshToken);
 }

package/src/registry/xai-oauth.ts

@@ -14,5 +14,4 @@ export const xaiOauthProvider = {
 		const { refreshXAIOAuthToken } = await import("./oauth/xai-oauth");
 		return refreshXAIOAuthToken(credentials.refresh);
 	},
-	pasteCodeFlow: true,
 } as const satisfies ProviderDefinition;

package/src/types.ts

@@ -388,9 +388,9 @@ export interface StreamOptions {
 	 */
 	sessionId?: string;
 	/**
-	 * Optional prompt-cache identity. When set, OpenAI Responses-compatible
-	 * providers use this for `prompt_cache_key` while keeping `sessionId` for
-	 * provider routing / conversation headers.
+	 * Optional prompt-cache identity. OpenAI-family providers use this for
+	 * `prompt_cache_key` payloads and cache-affinity headers such as
+	 * `x-grok-conv-id`; when omitted, they fall back to `sessionId`.
 	 */
 	promptCacheKey?: string;
 	/**