@oh-my-pi/pi-ai 16.3.14 → 16.4.0

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

@@ -1,5 +1,5 @@
 import { afterEach, describe, expect, it, vi } from "bun:test";
-import { isXAIAccessTokenExpiring, refreshXAIOAuthToken, validateXAIEndpoint, XAIOAuthFlow } from "../xai-oauth";
+import { isXAIAccessTokenExpiring, loginXAIOAuth, refreshXAIOAuthToken, validateXAIEndpoint } from "../xai-oauth";
 
 afterEach(() => {
 	vi.restoreAllMocks();
@@ -11,6 +11,76 @@ function jwtWithExp(exp: number): string {
 	return `${header}.${payload}.sig`;
 }
 
+const DISCOVERY_URL = "https://auth.x.ai/.well-known/openid-configuration";
+const DEVICE_CODE_URL = "https://auth.x.ai/oauth2/device/code";
+const TOKEN_ENDPOINT = "https://auth.x.ai/oauth2/token";
+const CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828";
+const SCOPE = "openid profile email offline_access grok-cli:access api:access";
+
+const DEVICE_AUTHORIZATION = {
+	device_code: "device-code-123",
+	user_code: "ABCD-EFGH",
+	verification_uri: "https://auth.x.ai/activate",
+	verification_uri_complete: "https://auth.x.ai/activate?user_code=ABCD-EFGH",
+	expires_in: 600,
+	interval: 1,
+};
+
+type RecordedRequest = {
+	url: string;
+	init: RequestInit | undefined;
+};
+
+type TokenResponse = {
+	body: unknown;
+	status?: number;
+};
+
+function jsonResponse(body: unknown, status: number = 200): Response {
+	return new Response(JSON.stringify(body), {
+		status,
+		headers: { "Content-Type": "application/json" },
+	});
+}
+
+function createDeviceFlowFetch(tokenResponses: readonly TokenResponse[]) {
+	const requests: RecordedRequest[] = [];
+	let tokenResponseIndex = 0;
+	const fetchMock = vi.fn(async (input: string | URL | Request, init?: RequestInit) => {
+		const url = typeof input === "string" ? input : input instanceof Request ? input.url : input.toString();
+		requests.push({ url, init });
+
+		if (url === DISCOVERY_URL) {
+			return jsonResponse({ token_endpoint: TOKEN_ENDPOINT });
+		}
+		if (url === DEVICE_CODE_URL) {
+			return jsonResponse(DEVICE_AUTHORIZATION);
+		}
+		if (url === TOKEN_ENDPOINT) {
+			const tokenResponse = tokenResponses[tokenResponseIndex];
+			tokenResponseIndex += 1;
+			if (!tokenResponse) {
+				throw new Error(`Unexpected xAI token poll ${tokenResponseIndex}`);
+			}
+			return jsonResponse(tokenResponse.body, tokenResponse.status);
+		}
+		throw new Error(`Unexpected xAI OAuth request: ${url}`);
+	});
+
+	return {
+		fetchMock: fetchMock as unknown as typeof fetch,
+		requests,
+	};
+}
+
+function requestForm(request: RecordedRequest | undefined): URLSearchParams {
+	const body = request?.init?.body;
+	if (!(body instanceof URLSearchParams)) {
+		throw new Error("Expected an application/x-www-form-urlencoded request body");
+	}
+	return body;
+}
+
 describe("isXAIAccessTokenExpiring", () => {
 	it("returns false for an empty string", () => {
 		expect(isXAIAccessTokenExpiring("")).toBe(false);
@@ -63,97 +133,138 @@ describe("refreshXAIOAuthToken", () => {
 	});
 });
 
-describe("XAIOAuthFlow", () => {
-	it("pins the redirect URI to xAI's allowlisted loopback port", () => {
-		const flow = new XAIOAuthFlow({});
-
-		expect(flow.redirectUri).toBe("http://127.0.0.1:56121/callback");
-	});
-
-	it("uses pasted-code login without starting a callback server", async () => {
-		const serveSpy = vi.spyOn(Bun, "serve").mockImplementation(() => {
-			throw new Error("callback server should not start");
-		});
-		let authUrl = "";
-		let tokenRequestBody = "";
-		const progress: string[] = [];
-		const fetchMock = vi.fn(async (input: string | URL | Request, init?: RequestInit) => {
-			const url = typeof input === "string" ? input : input instanceof Request ? input.url : input.toString();
-			if (url.includes("/.well-known/openid-configuration")) {
-				return new Response(
-					JSON.stringify({
-						authorization_endpoint: "https://auth.x.ai/oauth/authorize",
-						token_endpoint: "https://auth.x.ai/oauth/token",
-					}),
-					{ status: 200, headers: { "Content-Type": "application/json" } },
-				);
-			}
-			tokenRequestBody = init?.body instanceof URLSearchParams ? init.body.toString() : String(init?.body ?? "");
-			return new Response(
-				JSON.stringify({
+describe("loginXAIOAuth", () => {
+	it("performs the RFC 8628 device flow and returns the issued credentials", async () => {
+		const now = 1_800_000_000_000;
+		vi.spyOn(Date, "now").mockReturnValue(now);
+		const { fetchMock, requests } = createDeviceFlowFetch([
+			{
+				body: {
 					access_token: "access-token",
 					refresh_token: "refresh-token",
 					expires_in: 3600,
-				}),
-				{ status: 200, headers: { "Content-Type": "application/json" } },
-			);
+				},
+			},
+		]);
+		const authEvents: Array<{ url: string; instructions?: string }> = [];
+		const progress: string[] = [];
+		const onAuth = vi.fn((info: { url: string; instructions?: string }) => {
+			authEvents.push(info);
+		});
+		const onProgress = vi.fn((message: string) => {
+			progress.push(message);
+		});
+		const onManualCodeInput = vi.fn(async () => {
+			throw new Error("device authorization must not request a pasted code");
 		});
 
-		const flow = new XAIOAuthFlow({
-			fetch: fetchMock as unknown as typeof fetch,
-			onAuth: info => {
-				authUrl = info.url;
-			},
-			onManualCodeInput: async () => {
-				const parsed = new URL(authUrl);
-				const redirectUri = parsed.searchParams.get("redirect_uri") ?? "";
-				const state = parsed.searchParams.get("state") ?? "";
-				return `${redirectUri}?code=code-xyz&state=${encodeURIComponent(state)}`;
-			},
-			onProgress: message => progress.push(message),
+		const credentials = await loginXAIOAuth({
+			fetch: fetchMock,
+			onAuth,
+			onProgress,
+			onManualCodeInput,
 		});
 
-		const credentials = await flow.login();
-		const authorizeUrl = new URL(authUrl);
-		const tokenParams = new URLSearchParams(tokenRequestBody);
+		expect(requests.map(request => request.url)).toEqual([DISCOVERY_URL, DEVICE_CODE_URL, TOKEN_ENDPOINT]);
 
-		expect(serveSpy).not.toHaveBeenCalled();
-		expect(authorizeUrl.searchParams.get("redirect_uri")).toBe("http://127.0.0.1:56121/callback");
-		expect(progress).toContain("Waiting for pasted authorization code...");
-		expect(tokenParams.get("code")).toBe("code-xyz");
-		expect(credentials.access).toBe("access-token");
-		expect(credentials.refresh).toBe("refresh-token");
-	});
-});
+		const discoveryRequest = requests[0];
+		expect(discoveryRequest?.init?.method).toBe("GET");
+		expect(new Headers(discoveryRequest?.init?.headers).get("Accept")).toBe("application/json");
 
-describe("XAIOAuthFlow.exchangeToken", () => {
-	it("rejects when the token-exchange response is missing access_token", async () => {
-		const fetchMock = vi.fn(async (input: string | URL) => {
-			const url = typeof input === "string" ? input : input.toString();
-			if (url.includes("/.well-known/openid-configuration")) {
-				return new Response(
-					JSON.stringify({
-						authorization_endpoint: "https://auth.x.ai/oauth/authorize",
-						token_endpoint: "https://auth.x.ai/oauth/token",
-					}),
-					{ status: 200, headers: { "Content-Type": "application/json" } },
-				);
-			}
-			// Token-exchange response deliberately omits `access_token` to exercise
-			// the missing-token rejection path. The value of `refresh_token` here is
-			// a literal test marker, not a real secret — the test verifies
-			// exchangeToken throws before any token would be persisted.
-			return new Response(JSON.stringify({ refresh_token: "stub-refresh-token-for-test-only" }), {
-				status: 200,
-				headers: { "Content-Type": "application/json" },
-			});
+		const deviceRequest = requests[1];
+		expect(deviceRequest?.init?.method).toBe("POST");
+		const deviceHeaders = new Headers(deviceRequest?.init?.headers);
+		expect(deviceHeaders.get("Content-Type")).toBe("application/x-www-form-urlencoded");
+		expect(deviceHeaders.get("Accept")).toBe("application/json");
+		const deviceForm = requestForm(deviceRequest);
+		expect([...deviceForm.keys()].sort()).toEqual(["client_id", "scope"]);
+		expect(Object.fromEntries(deviceForm)).toEqual({
+			client_id: CLIENT_ID,
+			scope: SCOPE,
 		});
 
-		const flow = new XAIOAuthFlow({ fetch: fetchMock as unknown as typeof fetch });
-		await flow.generateAuthUrl("state-abc", "http://127.0.0.1:56121/callback");
+		const tokenRequest = requests[2];
+		expect(tokenRequest?.init?.method).toBe("POST");
+		const tokenHeaders = new Headers(tokenRequest?.init?.headers);
+		expect(tokenHeaders.get("Content-Type")).toBe("application/x-www-form-urlencoded");
+		expect(tokenHeaders.get("Accept")).toBe("application/json");
+		const tokenForm = requestForm(tokenRequest);
+		expect([...tokenForm.keys()].sort()).toEqual(["client_id", "device_code", "grant_type"]);
+		expect(Object.fromEntries(tokenForm)).toEqual({
+			grant_type: "urn:ietf:params:oauth:grant-type:device_code",
+			client_id: CLIENT_ID,
+			device_code: DEVICE_AUTHORIZATION.device_code,
+		});
+
+		expect(authEvents).toEqual([
+			{
+				url: DEVICE_AUTHORIZATION.verification_uri_complete,
+				instructions: `Enter code: ${DEVICE_AUTHORIZATION.user_code}`,
+			},
+		]);
+		expect(authEvents[0]?.instructions).not.toMatch(/hermes/i);
+		expect(onManualCodeInput).not.toHaveBeenCalled();
+		expect(progress).toEqual(["Waiting for xAI device authorization..."]);
+		expect(credentials).toEqual({
+			access: "access-token",
+			refresh: "refresh-token",
+			expires: now + 3_300_000,
+		});
+	});
+
+	it("continues through authorization_pending and slow_down responses", async () => {
+		const sleepSpy = vi.spyOn(Bun, "sleep").mockResolvedValue(undefined);
+		const { fetchMock, requests } = createDeviceFlowFetch([
+			{ status: 400, body: { error: "authorization_pending" } },
+			{ status: 400, body: { error: "slow_down" } },
+			{
+				body: {
+					access_token: "eventual-access-token",
+					refresh_token: "eventual-refresh-token",
+					expires_in: 3600,
+				},
+			},
+		]);
+
+		const credentials = await loginXAIOAuth({ fetch: fetchMock });
+
+		const tokenRequests = requests.filter(request => request.url === TOKEN_ENDPOINT);
+		expect(tokenRequests).toHaveLength(3);
+		expect(tokenRequests.map(request => Object.fromEntries(requestForm(request)))).toEqual([
+			{
+				grant_type: "urn:ietf:params:oauth:grant-type:device_code",
+				client_id: CLIENT_ID,
+				device_code: DEVICE_AUTHORIZATION.device_code,
+			},
+			{
+				grant_type: "urn:ietf:params:oauth:grant-type:device_code",
+				client_id: CLIENT_ID,
+				device_code: DEVICE_AUTHORIZATION.device_code,
+			},
+			{
+				grant_type: "urn:ietf:params:oauth:grant-type:device_code",
+				client_id: CLIENT_ID,
+				device_code: DEVICE_AUTHORIZATION.device_code,
+			},
+		]);
+		expect(sleepSpy.mock.calls).toEqual([[1000], [6000]]);
+		expect(credentials.access).toBe("eventual-access-token");
+		expect(credentials.refresh).toBe("eventual-refresh-token");
+	});
+
+	it("rejects a token response that omits access_token", async () => {
+		const { fetchMock, requests } = createDeviceFlowFetch([
+			{
+				body: {
+					refresh_token: "refresh-token",
+					expires_in: 3600,
+				},
+			},
+		]);
 
-		await expect(flow.exchangeToken("code-xyz", "state-abc", "http://127.0.0.1:56121/callback")).rejects.toThrow(
-			/access_token/,
+		await expect(loginXAIOAuth({ fetch: fetchMock })).rejects.toThrow(
+			/xAI device-code token response missing access_token/,
 		);
+		expect(requests.filter(request => request.url === TOKEN_ENDPOINT)).toHaveLength(1);
 	});
 });

package/src/registry/oauth/device-code.ts

@@ -0,0 +1,92 @@
+import * as AIError from "../../error";
+
+const DEVICE_FLOW_CANCEL_MESSAGE = "Login cancelled";
+const DEVICE_FLOW_TIMEOUT_MESSAGE = "Device flow timed out";
+const DEVICE_FLOW_SLOW_DOWN_TIMEOUT_MESSAGE =
+	"Device flow timed out after one or more slow_down responses. This is often caused by clock drift in WSL or VM environments. Please sync or restart the VM clock and try again.";
+const MINIMUM_DEVICE_FLOW_INTERVAL_MS = 1000;
+const DEFAULT_DEVICE_FLOW_INTERVAL_SECONDS = 5;
+const SLOW_DOWN_INTERVAL_INCREMENT_MS = 5000;
+
+/** Result returned by one OAuth device-code polling attempt. */
+export type OAuthDeviceCodePollResult<T> =
+	| { status: "complete"; value: T }
+	| { status: "pending" }
+	| { status: "slow_down" }
+	| { status: "failed"; message: string };
+
+/** Options for polling an RFC 8628-style OAuth device-code flow. */
+export interface OAuthDeviceCodeFlowOptions<T> {
+	/** Poll the provider once and classify the response. */
+	poll(): OAuthDeviceCodePollResult<T> | Promise<OAuthDeviceCodePollResult<T>>;
+	/** Provider-requested polling cadence; defaults to RFC 8628's five seconds. */
+	intervalSeconds?: number;
+	/** Provider-issued expiry window for the device code. */
+	expiresInSeconds?: number;
+	/** Cancels the flow with the legacy "Login cancelled" error. */
+	signal?: AbortSignal;
+}
+
+async function abortableDeviceFlowSleep(ms: number, signal: AbortSignal | undefined): Promise<void> {
+	if (!signal) {
+		await Bun.sleep(ms);
+		return;
+	}
+	if (signal.aborted) {
+		throw new AIError.LoginCancelledError(DEVICE_FLOW_CANCEL_MESSAGE);
+	}
+
+	const { promise, resolve, reject } = Promise.withResolvers<void>();
+	let timer: Timer | undefined;
+	const onAbort = () => {
+		clearTimeout(timer);
+		reject(new AIError.LoginCancelledError(DEVICE_FLOW_CANCEL_MESSAGE));
+	};
+	timer = setTimeout(() => {
+		signal.removeEventListener("abort", onAbort);
+		resolve();
+	}, ms);
+	signal.addEventListener("abort", onAbort, { once: true });
+	await promise;
+}
+
+/** Poll an OAuth device-code flow until completion, provider failure, timeout, or cancellation. */
+export async function pollOAuthDeviceCodeFlow<T>(options: OAuthDeviceCodeFlowOptions<T>): Promise<T> {
+	const deadline =
+		typeof options.expiresInSeconds === "number"
+			? Date.now() + options.expiresInSeconds * 1000
+			: Number.POSITIVE_INFINITY;
+	let intervalMs = Math.max(
+		MINIMUM_DEVICE_FLOW_INTERVAL_MS,
+		Math.floor((options.intervalSeconds ?? DEFAULT_DEVICE_FLOW_INTERVAL_SECONDS) * 1000),
+	);
+	let slowDownResponses = 0;
+
+	while (Date.now() < deadline) {
+		if (options.signal?.aborted) {
+			throw new AIError.LoginCancelledError(DEVICE_FLOW_CANCEL_MESSAGE);
+		}
+		const result = await options.poll();
+		if (result.status === "complete") {
+			return result.value;
+		}
+		if (result.status === "failed") {
+			throw new AIError.OAuthError(result.message, { kind: "polling" });
+		}
+		if (result.status === "slow_down") {
+			slowDownResponses += 1;
+			intervalMs = Math.max(MINIMUM_DEVICE_FLOW_INTERVAL_MS, intervalMs + SLOW_DOWN_INTERVAL_INCREMENT_MS);
+		}
+
+		const remainingMs = deadline - Date.now();
+		if (remainingMs <= 0) {
+			break;
+		}
+		await abortableDeviceFlowSleep(Math.min(intervalMs, remainingMs), options.signal);
+	}
+
+	throw new AIError.OAuthError(
+		slowDownResponses > 0 ? DEVICE_FLOW_SLOW_DOWN_TIMEOUT_MESSAGE : DEVICE_FLOW_TIMEOUT_MESSAGE,
+		{ kind: "timeout" },
+	);
+}

package/src/registry/oauth/index.ts

@@ -12,99 +12,9 @@ import type {
 	OAuthProviderInterface,
 } from "./types";
 
+export * from "./device-code";
 export type * from "./types";
 
-const DEVICE_FLOW_CANCEL_MESSAGE = "Login cancelled";
-const DEVICE_FLOW_TIMEOUT_MESSAGE = "Device flow timed out";
-const DEVICE_FLOW_SLOW_DOWN_TIMEOUT_MESSAGE =
-	"Device flow timed out after one or more slow_down responses. This is often caused by clock drift in WSL or VM environments. Please sync or restart the VM clock and try again.";
-const MINIMUM_DEVICE_FLOW_INTERVAL_MS = 1000;
-const DEFAULT_DEVICE_FLOW_INTERVAL_SECONDS = 5;
-const SLOW_DOWN_INTERVAL_INCREMENT_MS = 5000;
-
-/** Result returned by one OAuth device-code polling attempt. */
-export type OAuthDeviceCodePollResult<T> =
-	| { status: "complete"; value: T }
-	| { status: "pending" }
-	| { status: "slow_down" }
-	| { status: "failed"; message: string };
-
-/** Options for polling an RFC 8628-style OAuth device-code flow. */
-export interface OAuthDeviceCodeFlowOptions<T> {
-	/** Poll the provider once and classify the response. */
-	poll(): OAuthDeviceCodePollResult<T> | Promise<OAuthDeviceCodePollResult<T>>;
-	/** Provider-requested polling cadence; defaults to RFC 8628's five seconds. */
-	intervalSeconds?: number;
-	/** Provider-issued expiry window for the device code. */
-	expiresInSeconds?: number;
-	/** Cancels the flow with the legacy "Login cancelled" error. */
-	signal?: AbortSignal;
-}
-
-async function abortableDeviceFlowSleep(ms: number, signal: AbortSignal | undefined): Promise<void> {
-	if (!signal) {
-		await Bun.sleep(ms);
-		return;
-	}
-	if (signal.aborted) {
-		throw new AIError.LoginCancelledError(DEVICE_FLOW_CANCEL_MESSAGE);
-	}
-
-	const { promise, resolve, reject } = Promise.withResolvers<void>();
-	let timer: Timer | undefined;
-	const onAbort = () => {
-		if (timer) clearTimeout(timer);
-		reject(new AIError.LoginCancelledError(DEVICE_FLOW_CANCEL_MESSAGE));
-	};
-	timer = setTimeout(() => {
-		signal.removeEventListener("abort", onAbort);
-		resolve();
-	}, ms);
-	signal.addEventListener("abort", onAbort, { once: true });
-	await promise;
-}
-
-/** Poll an OAuth device-code flow until completion, provider failure, timeout, or cancellation. */
-export async function pollOAuthDeviceCodeFlow<T>(options: OAuthDeviceCodeFlowOptions<T>): Promise<T> {
-	const deadline =
-		typeof options.expiresInSeconds === "number"
-			? Date.now() + options.expiresInSeconds * 1000
-			: Number.POSITIVE_INFINITY;
-	let intervalMs = Math.max(
-		MINIMUM_DEVICE_FLOW_INTERVAL_MS,
-		Math.floor((options.intervalSeconds ?? DEFAULT_DEVICE_FLOW_INTERVAL_SECONDS) * 1000),
-	);
-	let slowDownResponses = 0;
-
-	while (Date.now() < deadline) {
-		if (options.signal?.aborted) {
-			throw new AIError.LoginCancelledError(DEVICE_FLOW_CANCEL_MESSAGE);
-		}
-		const result = await options.poll();
-		if (result.status === "complete") {
-			return result.value;
-		}
-		if (result.status === "failed") {
-			throw new AIError.OAuthError(result.message, { kind: "polling" });
-		}
-		if (result.status === "slow_down") {
-			slowDownResponses += 1;
-			intervalMs = Math.max(MINIMUM_DEVICE_FLOW_INTERVAL_MS, intervalMs + SLOW_DOWN_INTERVAL_INCREMENT_MS);
-		}
-
-		const remainingMs = deadline - Date.now();
-		if (remainingMs <= 0) {
-			break;
-		}
-		await abortableDeviceFlowSleep(Math.min(intervalMs, remainingMs), options.signal);
-	}
-
-	throw new AIError.OAuthError(
-		slowDownResponses > 0 ? DEVICE_FLOW_SLOW_DOWN_TIMEOUT_MESSAGE : DEVICE_FLOW_TIMEOUT_MESSAGE,
-		{ kind: "timeout" },
-	);
-}
-
 const builtInOAuthProviders: OAuthProviderInfo[] = PROVIDER_REGISTRY.filter(
 	provider => provider.login && provider.showInLoginList !== false,
 ).map(provider => ({