From 1783a56a4e717e7d1004c2f1687da03dba4c3a4d Mon Sep 17 00:00:00 2001 From: Vegard Stikbakke Date: Tue, 19 May 2026 19:35:09 +0200 Subject: [PATCH] refactor(ai): share device code polling --- packages/ai/CHANGELOG.md | 2 +- packages/ai/src/utils/oauth/device-code.ts | 96 ++++++++++++ packages/ai/src/utils/oauth/github-copilot.ts | 141 ++++++------------ packages/ai/src/utils/oauth/index.ts | 1 + packages/ai/test/oauth-device-code.test.ts | 64 ++++++++ 5 files changed, 210 insertions(+), 94 deletions(-) create mode 100644 packages/ai/src/utils/oauth/device-code.ts create mode 100644 packages/ai/test/oauth-device-code.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 1db0e8a9f..782adee11 100644 --- a/packages/ai/CHANGELOG.md +++ b/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 diff --git a/packages/ai/src/utils/oauth/device-code.ts b/packages/ai/src/utils/oauth/device-code.ts new file mode 100644 index 000000000..68ef1717c --- /dev/null +++ b/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 = + | { status: "pending" } + | { status: "slow_down"; intervalSeconds?: number } + | { status: "complete"; value: T } + | { status: "failed"; message: string }; + +export type OAuthDeviceCodePollOptions = { + authorization: OAuthDeviceCodeAuthorization; + poll: () => Promise>; + 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 { + 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(options: OAuthDeviceCodePollOptions): Promise { + 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); +} diff --git a/packages/ai/src/utils/oauth/github-copilot.ts b/packages/ai/src/utils/oauth/github-copilot.ts index 095b09511..01b4861b9 100644 --- a/packages/ai/src/utils/oauth/github-copilot.ts +++ b/packages/ai/src/utils/oauth/github-copilot.ts @@ -144,95 +144,56 @@ async function startDeviceFlow(domain: string): Promise { }; } -/** - * Sleep that can be interrupted by an AbortSignal - */ -function abortableSleep(ms: number, signal?: AbortSignal): Promise { - 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({ + 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 diff --git a/packages/ai/src/utils/oauth/index.ts b/packages/ai/src/utils/oauth/index.ts index 3a3a01b16..55910322b 100644 --- a/packages/ai/src/utils/oauth/index.ts +++ b/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, diff --git a/packages/ai/test/oauth-device-code.test.ts b/packages/ai/test/oauth-device-code.test.ts new file mode 100644 index 000000000..199fb06df --- /dev/null +++ b/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"); + }); +});