fix(cliproxy): poll paste-callback auth urls

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<PasteCallbackStartData> {
+  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<void> {
+  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<string | null> {
+  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);
+    let startData: PasteCallbackStartData;
-    const startResponse = await fetch(buildProxyUrl(target, startPath), {
+    try {
-      ...(provider === 'kiro' ? { method: 'POST' } : {}),
+      startData = await requestPasteCallbackStart(provider, target);
-      headers:
+    } catch (error) {
-        provider === 'kiro'
+      const startError = (error as Error).message;
-          ? buildManagementHeaders(target, { 'Content-Type': 'application/json' })
-          : buildManagementHeaders(target),
-    });
-
-    if (!startResponse.ok) {
-      const startError = `OAuth start failed with status ${startResponse.status}`;
       console.log(fail('Failed to start OAuth flow'));
       warnPossible403Ban(provider, startError);
       return null;
     }
 
-    const startData = (await startResponse.json()) as {
+    const authUrl = await resolvePasteCallbackAuthUrl(target, startData, OAUTH_STATE_TIMEOUT_MS);
-      url?: string;
-      auth_url?: string;
-      status?: string;
-    };
-    const authUrl = startData.url || startData.auth_url;
 
     if (!authUrl) {
       console.log(fail('No authorization URL received'));

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');
+  });
+});