refactor(ai): share device code polling

packages/ai/CHANGELOG.md

@@ -9,7 +9,7 @@
 
 ### Added
 
-- Added first-class OAuth device-code callback metadata for login flows and GitHub Copilot OAuth.
+- Added first-class OAuth device-code callback metadata, shared polling support, and GitHub Copilot OAuth integration.
 
 ### Fixed
 

packages/ai/src/utils/oauth/device-code.ts

@@ -0,0 +1,96 @@
+export type OAuthDeviceCodeAuthorization = {
+	userCode: string;
+	verificationUri: string;
+	intervalSeconds: number;
+	expiresInSeconds?: number;
+};
+
+export type OAuthDeviceCodePollResult<T> =
+	| { status: "pending" }
+	| { status: "slow_down"; intervalSeconds?: number }
+	| { status: "complete"; value: T }
+	| { status: "failed"; message: string };
+
+export type OAuthDeviceCodePollOptions<T> = {
+	authorization: OAuthDeviceCodeAuthorization;
+	poll: () => Promise<OAuthDeviceCodePollResult<T>>;
+	signal?: AbortSignal;
+	cancelMessage?: string;
+	timeoutMessage?: string;
+	slowDownTimeoutMessage?: string;
+	initialIntervalMultiplier?: number;
+	slowDownIntervalMultiplier?: number;
+	slowDownIntervalIncrementSeconds?: number;
+	minimumIntervalMs?: number;
+};
+
+function abortableSleep(ms: number, signal: AbortSignal | undefined, cancelMessage: string): Promise<void> {
+	return new Promise((resolve, reject) => {
+		if (signal?.aborted) {
+			reject(new Error(cancelMessage));
+			return;
+		}
+
+		const onAbort = () => {
+			clearTimeout(timeout);
+			reject(new Error(cancelMessage));
+		};
+		const timeout = setTimeout(() => {
+			signal?.removeEventListener("abort", onAbort);
+			resolve();
+		}, ms);
+
+		signal?.addEventListener("abort", onAbort, { once: true });
+	});
+}
+
+export async function pollOAuthDeviceCodeFlow<T>(options: OAuthDeviceCodePollOptions<T>): Promise<T> {
+	const cancelMessage = options.cancelMessage ?? "Login cancelled";
+	const minimumIntervalMs = options.minimumIntervalMs ?? 1000;
+	const slowDownIntervalIncrementSeconds = options.slowDownIntervalIncrementSeconds ?? 5;
+	const initialIntervalMultiplier = options.initialIntervalMultiplier ?? 1;
+	const slowDownIntervalMultiplier = options.slowDownIntervalMultiplier ?? initialIntervalMultiplier;
+	const timeoutMessage = options.timeoutMessage ?? "Device flow timed out";
+
+	const deadline =
+		typeof options.authorization.expiresInSeconds === "number"
+			? Date.now() + options.authorization.expiresInSeconds * 1000
+			: Number.POSITIVE_INFINITY;
+	let intervalMs = Math.max(minimumIntervalMs, Math.floor(options.authorization.intervalSeconds * 1000));
+	let intervalMultiplier = initialIntervalMultiplier;
+	let slowDownResponses = 0;
+
+	while (Date.now() < deadline) {
+		if (options.signal?.aborted) {
+			throw new Error(cancelMessage);
+		}
+
+		const remainingMs = deadline - Date.now();
+		const waitMs = Math.min(Math.ceil(intervalMs * intervalMultiplier), remainingMs);
+		await abortableSleep(waitMs, options.signal, cancelMessage);
+
+		const result = await options.poll();
+		if (result.status === "complete") {
+			return result.value;
+		}
+		if (result.status === "pending") {
+			continue;
+		}
+		if (result.status === "slow_down") {
+			slowDownResponses += 1;
+			intervalMs =
+				typeof result.intervalSeconds === "number" && result.intervalSeconds > 0
+					? result.intervalSeconds * 1000
+					: Math.max(minimumIntervalMs, intervalMs + slowDownIntervalIncrementSeconds * 1000);
+			intervalMultiplier = slowDownIntervalMultiplier;
+			continue;
+		}
+		throw new Error(result.message);
+	}
+
+	if (slowDownResponses > 0 && options.slowDownTimeoutMessage) {
+		throw new Error(options.slowDownTimeoutMessage);
+	}
+
+	throw new Error(timeoutMessage);
+}

packages/ai/src/utils/oauth/github-copilot.ts

@@ -144,95 +144,56 @@ async function startDeviceFlow(domain: string): Promise<DeviceCodeResponse> {
 	};
 }
 
-/**
- * Sleep that can be interrupted by an AbortSignal
- */
-function abortableSleep(ms: number, signal?: AbortSignal): Promise<void> {
-	return new Promise((resolve, reject) => {
-		if (signal?.aborted) {
-			reject(new Error("Login cancelled"));
-			return;
-		}
-
-		const timeout = setTimeout(resolve, ms);
-
-		signal?.addEventListener(
-			"abort",
-			() => {
-				clearTimeout(timeout);
-				reject(new Error("Login cancelled"));
-			},
-			{ once: true },
-		);
-	});
-}
-
-async function pollForGitHubAccessToken(
-	domain: string,
-	deviceCode: string,
-	intervalSeconds: number,
-	expiresIn: number,
-	signal?: AbortSignal,
-) {
+async function pollForGitHubAccessToken(domain: string, device: DeviceCodeResponse, signal?: AbortSignal) {
 	const urls = getUrls(domain);
-	const deadline = Date.now() + expiresIn * 1000;
-	let intervalMs = Math.max(1000, Math.floor(intervalSeconds * 1000));
-	let intervalMultiplier = INITIAL_POLL_INTERVAL_MULTIPLIER;
-	let slowDownResponses = 0;
-
-	while (Date.now() < deadline) {
-		if (signal?.aborted) {
-			throw new Error("Login cancelled");
-		}
-
-		const remainingMs = deadline - Date.now();
-		const waitMs = Math.min(Math.ceil(intervalMs * intervalMultiplier), remainingMs);
-		await abortableSleep(waitMs, signal);
-
-		const raw = await fetchJson(urls.accessTokenUrl, {
-			method: "POST",
-			headers: {
-				Accept: "application/json",
-				"Content-Type": "application/x-www-form-urlencoded",
-				"User-Agent": "GitHubCopilotChat/0.35.0",
-			},
-			body: new URLSearchParams({
-				client_id: CLIENT_ID,
-				device_code: deviceCode,
-				grant_type: "urn:ietf:params:oauth:grant-type:device_code",
-			}),
-		});
-
-		if (raw && typeof raw === "object" && typeof (raw as DeviceTokenSuccessResponse).access_token === "string") {
-			return (raw as DeviceTokenSuccessResponse).access_token;
-		}
-
-		if (raw && typeof raw === "object" && typeof (raw as DeviceTokenErrorResponse).error === "string") {
-			const { error, error_description: description, interval } = raw as DeviceTokenErrorResponse;
-			if (error === "authorization_pending") {
-				continue;
-			}
-
-			if (error === "slow_down") {
-				slowDownResponses += 1;
-				intervalMs =
-					typeof interval === "number" && interval > 0 ? interval * 1000 : Math.max(1000, intervalMs + 5000);
-				intervalMultiplier = SLOW_DOWN_POLL_INTERVAL_MULTIPLIER;
-				continue;
-			}
-
-			const descriptionSuffix = description ? `: ${description}` : "";
-			throw new Error(`Device flow failed: ${error}${descriptionSuffix}`);
-		}
-	}
-
-	if (slowDownResponses > 0) {
-		throw new Error(
+	return pollOAuthDeviceCodeFlow<string>({
+		authorization: {
+			userCode: device.user_code,
+			verificationUri: device.verification_uri,
+			intervalSeconds: device.interval,
+			expiresInSeconds: device.expires_in,
+		},
+		signal,
+		initialIntervalMultiplier: INITIAL_POLL_INTERVAL_MULTIPLIER,
+		slowDownIntervalMultiplier: SLOW_DOWN_POLL_INTERVAL_MULTIPLIER,
+		slowDownTimeoutMessage:
 			"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.",
-		);
-	}
+		poll: async () => {
+			const raw = await fetchJson(urls.accessTokenUrl, {
+				method: "POST",
+				headers: {
+					Accept: "application/json",
+					"Content-Type": "application/x-www-form-urlencoded",
+					"User-Agent": "GitHubCopilotChat/0.35.0",
+				},
+				body: new URLSearchParams({
+					client_id: CLIENT_ID,
+					device_code: device.device_code,
+					grant_type: "urn:ietf:params:oauth:grant-type:device_code",
+				}),
+			});
 
-	throw new Error("Device flow timed out");
+			if (raw && typeof raw === "object" && typeof (raw as DeviceTokenSuccessResponse).access_token === "string") {
+				return { status: "complete", value: (raw as DeviceTokenSuccessResponse).access_token };
+			}
+
+			if (raw && typeof raw === "object" && typeof (raw as DeviceTokenErrorResponse).error === "string") {
+				const { error, error_description: description, interval } = raw as DeviceTokenErrorResponse;
+				if (error === "authorization_pending") {
+					return { status: "pending" };
+				}
+
+				if (error === "slow_down") {
+					return { status: "slow_down", intervalSeconds: interval };
+				}
+
+				const descriptionSuffix = description ? `: ${description}` : "";
+				return { status: "failed", message: `Device flow failed: ${error}${descriptionSuffix}` };
+			}
+
+			return { status: "failed", message: "Invalid device token response" };
+		},
+	});
 }
 
 /**
@@ -361,13 +322,7 @@ export async function loginGitHubCopilot(options: {
 		options.onAuth(device.verification_uri, `Enter code: ${device.user_code}`);
 	}
 
-	const githubAccessToken = await pollForGitHubAccessToken(
-		domain,
-		device.device_code,
-		device.interval,
-		device.expires_in,
-		options.signal,
-	);
+	const githubAccessToken = await pollForGitHubAccessToken(domain, device, options.signal);
 	const credentials = await refreshGitHubCopilotToken(githubAccessToken, enterpriseDomain ?? undefined);
 
 	// Enable all models after successful login

packages/ai/src/utils/oauth/index.ts

@@ -9,6 +9,7 @@
 
 // Anthropic
 export { anthropicOAuthProvider, loginAnthropic, refreshAnthropicToken } from "./anthropic.ts";
+export * from "./device-code.ts";
 // GitHub Copilot
 export {
 	getGitHubCopilotBaseUrl,

packages/ai/test/oauth-device-code.test.ts

@@ -0,0 +1,64 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { pollOAuthDeviceCodeFlow } from "../src/utils/oauth/device-code.js";
+
+describe("OAuth device-code polling", () => {
+	afterEach(() => {
+		vi.useRealTimers();
+	});
+
+	it("waits before the first poll and returns the completed value", async () => {
+		vi.useFakeTimers();
+		vi.setSystemTime(new Date("2026-03-09T00:00:00Z"));
+
+		const pollTimes: number[] = [];
+		const poll = vi.fn(async () => {
+			pollTimes.push(Date.now());
+			return pollTimes.length === 1
+				? { status: "pending" as const }
+				: { status: "complete" as const, value: "token" };
+		});
+
+		const resultPromise = pollOAuthDeviceCodeFlow({
+			authorization: {
+				userCode: "ABCD-EFGH",
+				verificationUri: "https://example.com/device",
+				intervalSeconds: 2,
+				expiresInSeconds: 30,
+			},
+			poll,
+			initialIntervalMultiplier: 1.5,
+		});
+
+		await vi.advanceTimersByTimeAsync(2999);
+		expect(pollTimes).toEqual([]);
+
+		await vi.advanceTimersByTimeAsync(1);
+		expect(pollTimes).toEqual([new Date("2026-03-09T00:00:03Z").getTime()]);
+
+		await vi.advanceTimersByTimeAsync(3000);
+		await expect(resultPromise).resolves.toBe("token");
+		expect(pollTimes).toEqual([
+			new Date("2026-03-09T00:00:03Z").getTime(),
+			new Date("2026-03-09T00:00:06Z").getTime(),
+		]);
+	});
+
+	it("cancels an in-flight wait", async () => {
+		vi.useFakeTimers();
+		const controller = new AbortController();
+
+		const resultPromise = pollOAuthDeviceCodeFlow({
+			authorization: {
+				userCode: "ABCD-EFGH",
+				verificationUri: "https://example.com/device",
+				intervalSeconds: 5,
+				expiresInSeconds: 30,
+			},
+			poll: async () => ({ status: "pending" }),
+			signal: controller.signal,
+		});
+
+		controller.abort();
+		await expect(resultPromise).rejects.toThrow("Login cancelled");
+	});
+});