diff --git a/src/cliproxy/auth/oauth-handler.ts b/src/cliproxy/auth/oauth-handler.ts index 5ef17166..e989ae56 100644 --- a/src/cliproxy/auth/oauth-handler.ts +++ b/src/cliproxy/auth/oauth-handler.ts @@ -45,7 +45,12 @@ import { isHeadlessEnvironment, killProcessOnPort, showStep } from './environmen import { getProviderTokenDir, isAuthenticated, registerAccountFromToken } from './token-manager'; import { executeOAuthProcess } from './oauth-process'; import { importKiroToken } from './kiro-import'; -import { getProxyTarget, buildProxyUrl, buildManagementHeaders } from '../proxy-target-resolver'; +import { + getProxyTarget, + buildProxyUrl, + buildManagementHeaders, + type ProxyTarget, +} from '../proxy-target-resolver'; import { checkNewAccountConflict, warnNewAccountConflict, @@ -54,6 +59,86 @@ import { } from '../account-safety'; import { ensureCliAntigravityResponsibility } from '../antigravity-responsibility'; +interface PasteCallbackStartData { + url?: string; + auth_url?: string; + state?: string; + status?: string; +} + +const PASTE_CALLBACK_AUTH_URL_POLL_INTERVAL_MS = 3000; + +export async function requestPasteCallbackStart( + provider: CLIProxyProvider, + target: ProxyTarget +): Promise { + const startPath = getPasteCallbackStartPath(provider); + const response = await fetch(buildProxyUrl(target, startPath), { + ...(provider === 'kiro' ? { method: 'POST' } : {}), + headers: + provider === 'kiro' + ? buildManagementHeaders(target, { 'Content-Type': 'application/json' }) + : buildManagementHeaders(target), + }); + + if (!response.ok) { + throw new Error(`OAuth start failed with status ${response.status}`); + } + + return (await response.json()) as PasteCallbackStartData; +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +export async function resolvePasteCallbackAuthUrl( + target: ProxyTarget, + startData: PasteCallbackStartData, + timeoutMs: number, + pollIntervalMs: number = PASTE_CALLBACK_AUTH_URL_POLL_INTERVAL_MS +): Promise { + const authUrl = startData.url || startData.auth_url; + if (authUrl) { + return authUrl; + } + + const state = startData.state; + if (!state) { + return null; + } + + const deadline = Date.now() + timeoutMs; + + while (Date.now() < deadline) { + const response = await fetch( + buildProxyUrl(target, `/v0/management/get-auth-status?state=${encodeURIComponent(state)}`), + { headers: buildManagementHeaders(target) } + ); + + if (response.ok) { + const statusData = (await response.json()) as PasteCallbackStartData; + const polledAuthUrl = statusData.url || statusData.auth_url; + + if (polledAuthUrl) { + return polledAuthUrl; + } + + if (statusData.status === 'error' || statusData.status === 'device_code') { + return null; + } + } + + if (Date.now() + pollIntervalMs >= deadline) { + break; + } + + await sleep(pollIntervalMs); + } + + return null; +} + /** * Prompt user to add another account */ @@ -279,28 +364,17 @@ async function handlePasteCallbackMode( // Request auth URL from CLIProxyAPI. // Kiro keeps its legacy start route because CLI auth methods do not share the generic // management auth-url contract used by providers like Claude. - const startPath = getPasteCallbackStartPath(provider); - const startResponse = await fetch(buildProxyUrl(target, startPath), { - ...(provider === 'kiro' ? { method: 'POST' } : {}), - headers: - provider === 'kiro' - ? buildManagementHeaders(target, { 'Content-Type': 'application/json' }) - : buildManagementHeaders(target), - }); - - if (!startResponse.ok) { - const startError = `OAuth start failed with status ${startResponse.status}`; + let startData: PasteCallbackStartData; + try { + startData = await requestPasteCallbackStart(provider, target); + } catch (error) { + const startError = (error as Error).message; console.log(fail('Failed to start OAuth flow')); warnPossible403Ban(provider, startError); return null; } - const startData = (await startResponse.json()) as { - url?: string; - auth_url?: string; - status?: string; - }; - const authUrl = startData.url || startData.auth_url; + const authUrl = await resolvePasteCallbackAuthUrl(target, startData, OAUTH_STATE_TIMEOUT_MS); if (!authUrl) { console.log(fail('No authorization URL received')); diff --git a/tests/unit/cliproxy/oauth-handler-paste-callback.test.ts b/tests/unit/cliproxy/oauth-handler-paste-callback.test.ts new file mode 100644 index 00000000..1b8da8a9 --- /dev/null +++ b/tests/unit/cliproxy/oauth-handler-paste-callback.test.ts @@ -0,0 +1,104 @@ +import { afterEach, describe, expect, it } from 'bun:test'; +import type { ProxyTarget } from '../../../src/cliproxy/proxy-target-resolver'; +import { getCapturedFetchRequests, mockFetch, restoreFetch } from '../../mocks'; + +const remoteTarget: ProxyTarget = { + host: 'proxy.example.com', + port: 8317, + protocol: 'https', + managementKey: 'test-mgmt-key', + isRemote: true, +}; + +afterEach(() => { + restoreFetch(); +}); + +describe('requestPasteCallbackStart', () => { + it('uses management auth-url route for non-kiro providers', async () => { + mockFetch([ + { + url: /\/v0\/management\/anthropic-auth-url\?is_webui=true$/, + response: { auth_url: 'https://auth.example.com/claude' }, + }, + ]); + + const { requestPasteCallbackStart } = await import( + `../../../src/cliproxy/auth/oauth-handler?request-claude-start=${Date.now()}` + ); + const startData = await requestPasteCallbackStart('claude', remoteTarget); + + expect(startData.auth_url).toBe('https://auth.example.com/claude'); + + const [request] = getCapturedFetchRequests(); + expect(request.url).toBe( + 'https://proxy.example.com:8317/v0/management/anthropic-auth-url?is_webui=true' + ); + expect(request.method).toBe('GET'); + expect(request.headers['Authorization']).toBe('Bearer test-mgmt-key'); + expect(request.headers['Content-Type']).toBeUndefined(); + }); + + it('keeps kiro on the legacy start route with POST', async () => { + mockFetch([ + { + url: /\/oauth\/kiro\/start$/, + method: 'POST', + response: { auth_url: 'https://auth.example.com/kiro' }, + }, + ]); + + const { requestPasteCallbackStart } = await import( + `../../../src/cliproxy/auth/oauth-handler?request-kiro-start=${Date.now()}` + ); + const startData = await requestPasteCallbackStart('kiro', remoteTarget); + + expect(startData.auth_url).toBe('https://auth.example.com/kiro'); + + const [request] = getCapturedFetchRequests(); + expect(request.url).toBe('https://proxy.example.com:8317/oauth/kiro/start'); + expect(request.method).toBe('POST'); + expect(request.headers['Authorization']).toBe('Bearer test-mgmt-key'); + expect(request.headers['Content-Type']).toBe('application/json'); + }); +}); + +describe('resolvePasteCallbackAuthUrl', () => { + it('returns the immediate auth URL without polling', async () => { + const { resolvePasteCallbackAuthUrl } = await import( + `../../../src/cliproxy/auth/oauth-handler?resolve-immediate-auth-url=${Date.now()}` + ); + const authUrl = await resolvePasteCallbackAuthUrl( + remoteTarget, + { auth_url: 'https://auth.example.com/direct' }, + 50, + 0 + ); + + expect(authUrl).toBe('https://auth.example.com/direct'); + expect(getCapturedFetchRequests()).toHaveLength(0); + }); + + it('polls management status when the start response only returns state', async () => { + mockFetch([ + { + url: /\/v0\/management\/get-auth-status\?state=state-123$/, + response: { status: 'auth_url', auth_url: 'https://auth.example.com/polled' }, + }, + ]); + + const { resolvePasteCallbackAuthUrl } = await import( + `../../../src/cliproxy/auth/oauth-handler?resolve-polled-auth-url=${Date.now()}` + ); + const authUrl = await resolvePasteCallbackAuthUrl(remoteTarget, { state: 'state-123' }, 50, 0); + + expect(authUrl).toBe('https://auth.example.com/polled'); + + const [request] = getCapturedFetchRequests(); + expect(request.url).toBe( + 'https://proxy.example.com:8317/v0/management/get-auth-status?state=state-123' + ); + expect(request.method).toBe('GET'); + expect(request.headers['Authorization']).toBe('Bearer test-mgmt-key'); + }); +});