fix(kiro): harden callback replay and auth detection

src/cliproxy/auth/oauth-handler.ts

@@ -11,6 +11,7 @@
  */
 
 import * as fs from 'fs';
+import * as path from 'path';
 import { fail, info, warn, color, ok } from '../../utils/ui';
 import { ensureCLIProxyBinary } from '../binary-manager';
 import { generateConfig } from '../config-generator';
@@ -46,7 +47,12 @@ import {
   normalizeKiroIDCFlow,
 } from './auth-types';
 import { isHeadlessEnvironment, killProcessOnPort, showStep } from './environment-detector';
-import { getProviderTokenDir, isAuthenticated, registerAccountFromToken } from './token-manager';
+import {
+  getProviderTokenDir,
+  isAuthenticated,
+  isTokenFileForProvider,
+  registerAccountFromToken,
+} from './token-manager';
 import { executeOAuthProcess } from './oauth-process';
 import { importKiroToken } from './kiro-import';
 import {
@@ -71,6 +77,12 @@ interface PasteCallbackStartData {
 }
 
 const PASTE_CALLBACK_AUTH_URL_POLL_INTERVAL_MS = 3000;
+const POLLED_AUTH_LOCAL_TOKEN_GRACE_MS = 15 * 1000;
+
+type ProviderTokenSnapshot = {
+  file: string;
+  mtimeMs: number;
+};
 
 export async function requestPasteCallbackStart(
   provider: CLIProxyProvider,
@@ -122,6 +134,134 @@ function sleep(ms: number): Promise<void> {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
+function parseAuthUrlState(url: string | null | undefined): string | null {
+  if (!url) {
+    return null;
+  }
+
+  try {
+    return new URL(url).searchParams.get('state');
+  } catch {
+    return null;
+  }
+}
+
+function listProviderTokenSnapshots(
+  provider: CLIProxyProvider,
+  tokenDir: string
+): ProviderTokenSnapshot[] {
+  if (!fs.existsSync(tokenDir)) {
+    return [];
+  }
+
+  return fs
+    .readdirSync(tokenDir)
+    .filter((file) => file.endsWith('.json'))
+    .map((file): ProviderTokenSnapshot | null => {
+      const filePath = path.join(tokenDir, file);
+      if (!isTokenFileForProvider(filePath, provider)) {
+        return null;
+      }
+
+      return {
+        file,
+        mtimeMs: fs.statSync(filePath).mtimeMs,
+      };
+    })
+    .filter((snapshot): snapshot is ProviderTokenSnapshot => snapshot !== null)
+    .sort((left, right) => right.mtimeMs - left.mtimeMs);
+}
+
+export function findNewTokenSnapshotForManualAuth(
+  provider: CLIProxyProvider,
+  tokenDir: string,
+  knownTokenFiles: ProviderTokenSnapshot[],
+  expectedAccountId?: string
+): ProviderTokenSnapshot | null {
+  const knownTokenMtimes = new Map(
+    knownTokenFiles.map((snapshot) => [snapshot.file, snapshot.mtimeMs])
+  );
+
+  return (
+    listProviderTokenSnapshots(provider, tokenDir).find((snapshot) => {
+      const knownMtime = knownTokenMtimes.get(snapshot.file);
+      if (knownMtime === undefined) {
+        return true;
+      }
+
+      if (!expectedAccountId) {
+        return false;
+      }
+
+      return snapshot.mtimeMs > knownMtime + 1;
+    }) || null
+  );
+}
+
+async function waitForManualCallbackToken(
+  provider: CLIProxyProvider,
+  target: ProxyTarget,
+  tokenDir: string,
+  oauthState: string | null,
+  knownTokenFiles: ProviderTokenSnapshot[],
+  expectedAccountId: string | undefined,
+  timeoutMs: number,
+  pollIntervalMs: number = PASTE_CALLBACK_AUTH_URL_POLL_INTERVAL_MS
+): Promise<{ tokenSnapshot: ProviderTokenSnapshot | null; error?: string }> {
+  const deadline = Date.now() + timeoutMs;
+  let upstreamCompletedAt: number | null = null;
+
+  while (Date.now() < deadline) {
+    const tokenSnapshot = findNewTokenSnapshotForManualAuth(
+      provider,
+      tokenDir,
+      knownTokenFiles,
+      expectedAccountId
+    );
+    if (tokenSnapshot) {
+      return { tokenSnapshot };
+    }
+
+    if (oauthState) {
+      const response = await fetch(
+        buildProxyUrl(
+          target,
+          `/v0/management/get-auth-status?state=${encodeURIComponent(oauthState)}`
+        ),
+        { headers: buildManagementHeaders(target) }
+      );
+
+      if (response.ok) {
+        const data = (await response.json()) as { status?: string; error?: string };
+        if (data.status === 'error') {
+          return {
+            tokenSnapshot: null,
+            error: data.error || 'Authentication failed while waiting for local token persistence',
+          };
+        }
+        if (data.status === 'ok' && upstreamCompletedAt === null) {
+          upstreamCompletedAt = Date.now();
+        }
+      }
+    }
+
+    if (
+      upstreamCompletedAt !== null &&
+      Date.now() - upstreamCompletedAt >= POLLED_AUTH_LOCAL_TOKEN_GRACE_MS
+    ) {
+      break;
+    }
+
+    if (Date.now() + pollIntervalMs >= deadline) {
+      break;
+    }
+
+    await sleep(pollIntervalMs);
+  }
+
+  return { tokenSnapshot: null };
+}
+
 export async function resolvePasteCallbackAuthUrl(
   target: ProxyTarget,
   startData: PasteCallbackStartData,
@@ -397,6 +537,9 @@ async function handlePasteCallbackMode(
       return null;
     }
 
+    const oauthState = startData.state || parseAuthUrlState(authUrl);
+    const knownTokenFiles = listProviderTokenSnapshots(provider, tokenDir);
+
     // Display auth URL in box
     console.log('');
     console.log('  ╔══════════════════════════════════════════════════════════════╗');
@@ -489,15 +632,49 @@ async function handlePasteCallbackMode(
       return null;
     }
 
-    console.log(ok('Authentication successful!'));
+    console.log(info('Callback submitted. Waiting for token exchange...'));
+    const { tokenSnapshot, error: tokenWaitError } = await waitForManualCallbackToken(
+      provider,
+      target,
+      tokenDir,
+      oauthState,
+      knownTokenFiles,
+      expectedAccountId,
+      OAUTH_STATE_TIMEOUT_MS
+    );
+
+    if (tokenWaitError) {
+      console.log(fail(tokenWaitError));
+      warnPossible403Ban(provider, tokenWaitError);
+      return null;
+    }
+
+    if (!tokenSnapshot) {
+      console.log(
+        fail(
+          'Authentication completed upstream, but no new local token was saved for this account. Update CCS/CLIProxy and retry.'
+        )
+      );
+      return null;
+    }
+
     const account = registerAccountFromToken(
       provider,
       tokenDir,
       nickname,
       verbose,
-      expectedAccountId
+      tokenSnapshot.file
     );
 
+    if (!account) {
+      console.log(
+        fail('Authenticated token could not be matched to the requested account. Retry the flow.')
+      );
+      return null;
+    }
+
+    console.log(ok('Authentication successful!'));
+
     // Account safety: check for cross-provider conflicts
     if (account?.email) {
       const conflicts = checkNewAccountConflict(provider, account.email);

src/cliproxy/auth/oauth-process.ts

@@ -198,7 +198,7 @@ export function getKiroBuilderIdSelectionInput(output: string): string | null {
   }
 
   const promptWindow = output.slice(promptMatch.index, promptMatch.index + 600);
-  const optionMatch = /(?:^|\n)\s*(\d+)\s*[\).:-]?\s*(?:AWS\s+)?Builder ID\b/im.exec(promptWindow);
+  const optionMatch = /(?:^|\n)\s*(\d+)\s*[\).:-]?\s*.*\bBuilder ID\b/im.exec(promptWindow);
   if (!optionMatch) {
     return null;
   }
@@ -206,6 +206,33 @@ export function getKiroBuilderIdSelectionInput(output: string): string | null {
   return `${optionMatch[1]}\n`;
 }
 
+export function extractLikelyOAuthAuthorizationUrl(output: string): string | null {
+  const urls = Array.from(output.matchAll(/https?:\/\/[^\s]+/g), (match) => match[0]);
+  let selectedUrl: string | null = null;
+  let selectedScore = 0;
+
+  for (const url of urls) {
+    try {
+      const parsed = new URL(url);
+      let score = 0;
+      if (parsed.searchParams.has('redirect_uri')) score += 4;
+      if (parsed.searchParams.has('state')) score += 2;
+      if (parsed.searchParams.has('code_challenge')) score += 1;
+      if (parsed.pathname.includes('/authorize')) score += 1;
+      if (isLoopbackHost(parsed.hostname)) score -= 3;
+
+      if (score >= selectedScore && score > 0) {
+        selectedUrl = url;
+        selectedScore = score;
+      }
+    } catch {
+      continue;
+    }
+  }
+
+  return selectedUrl;
+}
+
 async function promptManualCallbackUrl(
   displayName: string,
   state: ProcessState,
@@ -332,17 +359,12 @@ async function handleStdout(
     !state.kiroMethodSelectionHandled &&
     state.accumulatedOutput.includes('Select login method')
   ) {
-    state.kiroMethodSelectionHandled = true;
     const builderIdSelection = getKiroBuilderIdSelectionInput(state.accumulatedOutput);
-    if (!builderIdSelection) {
-      console.log(fail('Unable to auto-select Kiro Builder ID from the upstream login menu.'));
-      console.log('    The upstream Kiro prompt format may have changed.');
-      killWithEscalation(authProcess);
-      return;
+    if (builderIdSelection) {
+      state.kiroMethodSelectionHandled = true;
+      authProcess.stdin?.write(builderIdSelection);
+      log(`Auto-selected Kiro Builder ID flow (${builderIdSelection.trim()})`);
     }
-
-    authProcess.stdin?.write(builderIdSelection);
-    log(`Auto-selected Kiro Builder ID flow (${builderIdSelection.trim()})`);
   }
 
   // Parse project list when available
@@ -413,11 +435,11 @@ async function handleStdout(
 
   // Display OAuth URL for all modes (enables VS Code terminal URL detection popup)
   if (!isDeviceCodeFlow && !state.urlDisplayed) {
-    const urlMatch = output.match(/https?:\/\/[^\s]+/);
-    if (urlMatch) {
+    const authUrl = extractLikelyOAuthAuthorizationUrl(state.accumulatedOutput);
+    if (authUrl) {
       console.log('');
       console.log(info(`${options.oauthConfig.displayName} OAuth URL:`));
-      console.log(`    ${urlMatch[0]}`);
+      console.log(`    ${authUrl}`);
       console.log('');
       state.urlDisplayed = true;
 
@@ -426,7 +448,7 @@ async function handleStdout(
         await replayManualCallback(
           options.oauthConfig,
           authProcess,
-          urlMatch[0],
+          authUrl,
           options.verbose,
           state,
           10 * 60 * 1000
@@ -442,11 +464,11 @@ function displayUrlFromStderr(
   state: ProcessState,
   oauthConfig: ProviderOAuthConfig
 ): void {
-  const urlMatch = output.match(/https?:\/\/[^\s]+/);
-  if (urlMatch) {
+  const authUrl = extractLikelyOAuthAuthorizationUrl(output);
+  if (authUrl) {
     console.log('');
     console.log(info(`${oauthConfig.displayName} OAuth URL:`));
-    console.log(`    ${urlMatch[0]}`);
+    console.log(`    ${authUrl}`);
     console.log('');
     state.urlDisplayed = true;
   }

tests/unit/cliproxy/oauth-handler-paste-callback.test.ts

@@ -1,4 +1,7 @@
 import { afterEach, describe, expect, it } from 'bun:test';
+import * as fs from 'fs';
+import * as os from 'os';
+import * as path from 'path';
 import type { ProxyTarget } from '../../../src/cliproxy/proxy-target-resolver';
 import { getCapturedFetchRequests, mockFetch, restoreFetch } from '../../mocks';
 
@@ -130,6 +133,56 @@ describe('resolvePasteCallbackAuthUrl', () => {
   });
 });
 
+describe('findNewTokenSnapshotForManualAuth', () => {
+  it('detects newly created provider token files', async () => {
+    const tokenDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-kiro-manual-auth-'));
+    const existingFile = path.join(tokenDir, 'kiro-existing.json');
+    fs.writeFileSync(existingFile, JSON.stringify({ type: 'kiro', email: 'existing@example.com' }));
+    const existingMtimeMs = fs.statSync(existingFile).mtimeMs;
+
+    const { findNewTokenSnapshotForManualAuth } = await import(
+      `../../../src/cliproxy/auth/oauth-handler?manual-auth-new-token=${Date.now()}`
+    );
+
+    const newFile = path.join(tokenDir, 'kiro-new.json');
+    fs.writeFileSync(newFile, JSON.stringify({ type: 'kiro', email: 'new@example.com' }));
+
+    const snapshot = findNewTokenSnapshotForManualAuth(
+      'kiro',
+      tokenDir,
+      [{ file: 'kiro-existing.json', mtimeMs: existingMtimeMs }]
+    );
+
+    expect(snapshot?.file).toBe('kiro-new.json');
+    fs.rmSync(tokenDir, { recursive: true, force: true });
+  });
+
+  it('treats a modified existing token as the new token during reauth', async () => {
+    const tokenDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-kiro-reauth-'));
+    const tokenFile = path.join(tokenDir, 'kiro-existing.json');
+    fs.writeFileSync(tokenFile, JSON.stringify({ type: 'kiro', email: 'existing@example.com' }));
+    const existingMtimeMs = fs.statSync(tokenFile).mtimeMs;
+
+    const { findNewTokenSnapshotForManualAuth } = await import(
+      `../../../src/cliproxy/auth/oauth-handler?manual-auth-updated-token=${Date.now()}`
+    );
+
+    fs.writeFileSync(tokenFile, JSON.stringify({ type: 'kiro', email: 'existing@example.com', refreshed: true }));
+    const bumpedTime = new Date(existingMtimeMs + 10_000);
+    fs.utimesSync(tokenFile, bumpedTime, bumpedTime);
+
+    const snapshot = findNewTokenSnapshotForManualAuth(
+      'kiro',
+      tokenDir,
+      [{ file: 'kiro-existing.json', mtimeMs: existingMtimeMs }],
+      'kiro-existing.json'
+    );
+
+    expect(snapshot?.file).toBe('kiro-existing.json');
+    fs.rmSync(tokenDir, { recursive: true, force: true });
+  });
+});
+
 describe('getCliAuthNicknameError', () => {
   it('allows omitted nicknames for no-email providers', async () => {
     const { getCliAuthNicknameError } = await import(

tests/unit/cliproxy/oauth-process-error-parser.test.ts

@@ -1,6 +1,7 @@
 import { describe, expect, it } from 'bun:test';
 import {
   extractLikelyAuthFailureFromStderr,
+  extractLikelyOAuthAuthorizationUrl,
   getExpectedLocalCallback,
   getKiroBuilderIdSelectionInput,
   validateManualCallbackUrl,
@@ -89,9 +90,9 @@ describe('oauth-process manual callback validation', () => {
 describe('oauth-process Kiro Builder ID menu parsing', () => {
   it('selects Builder ID when it is the first option', () => {
     const output = `
-Select login method
-1. Builder ID
-2. IAM Identity Center
+? Select login method:
+  1) Use with Builder ID (personal AWS account)
+  2) Use with IDC Account (organization SSO)
 `;
 
     expect(getKiroBuilderIdSelectionInput(output)).toBe('1\n');
@@ -117,3 +118,28 @@ Select login method
     expect(getKiroBuilderIdSelectionInput(output)).toBeNull();
   });
 });
+
+describe('oauth-process OAuth URL extraction', () => {
+  it('prefers the real auth URL over the IDC start URL banner', () => {
+    const authUrl =
+      'https://oidc.us-east-1.amazonaws.com/authorize?response_type=code&client_id=test-client&redirect_uri=http%3A%2F%2F127.0.0.1%3A9876%2Foauth%2Fcallback&state=test-state&code_challenge=test-challenge&code_challenge_method=S256';
+    const output = `
+Using IDC with Start URL: https://d-123.awsapps.com/start
+Region: us-east-1
+URL: ${authUrl}
+`;
+
+    expect(extractLikelyOAuthAuthorizationUrl(output)).toBe(authUrl);
+  });
+
+  it('ignores local callback server URLs when the auth URL is also present', () => {
+    const authUrl =
+      'https://device.sso.us-east-1.amazonaws.com/authorize?response_type=code&client_id=test-client&redirect_uri=http%3A%2F%2F127.0.0.1%3A9876%2Foauth%2Fcallback&state=test-state&code_challenge=test-challenge&code_challenge_method=S256';
+    const output = `
+Callback server started, redirect URI: http://127.0.0.1:9876/oauth/callback
+URL: ${authUrl}
+`;
+
+    expect(extractLikelyOAuthAuthorizationUrl(output)).toBe(authUrl);
+  });
+});