From 325d8d861d96ddc3527751e98f7541868a2d46fb Mon Sep 17 00:00:00 2001
From: Tam Nhu Tran <kaitran.ntt@gmail.com>
Date: Sun, 5 Apr 2026 02:27:17 -0400
Subject: [PATCH] fix(kiro): parse Builder ID selector dynamically

---
 src/cliproxy/auth/oauth-process.ts            | 27 +++++++++++++--
 .../oauth-process-error-parser.test.ts        | 33 +++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/src/cliproxy/auth/oauth-process.ts b/src/cliproxy/auth/oauth-process.ts
index ac566af2..74b78d8d 100644
--- a/src/cliproxy/auth/oauth-process.ts
+++ b/src/cliproxy/auth/oauth-process.ts
@@ -191,6 +191,21 @@ export function validateManualCallbackUrl(callbackUrl: string, authUrl: string):
   return null;
 }
 
+export function getKiroBuilderIdSelectionInput(output: string): string | null {
+  const promptMatch = /Select login method/i.exec(output);
+  if (!promptMatch || promptMatch.index === undefined) {
+    return 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);
+  if (!optionMatch) {
+    return null;
+  }
+
+  return `${optionMatch[1]}\n`;
+}
+
 async function promptManualCallbackUrl(
   displayName: string,
   state: ProcessState,
@@ -318,8 +333,16 @@ async function handleStdout(
     state.accumulatedOutput.includes('Select login method')
   ) {
     state.kiroMethodSelectionHandled = true;
-    authProcess.stdin?.write('1\n');
-    log('Auto-selected Kiro Builder ID flow');
+    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;
+    }
+
+    authProcess.stdin?.write(builderIdSelection);
+    log(`Auto-selected Kiro Builder ID flow (${builderIdSelection.trim()})`);
   }
 
   // Parse project list when available
diff --git a/tests/unit/cliproxy/oauth-process-error-parser.test.ts b/tests/unit/cliproxy/oauth-process-error-parser.test.ts
index 6bb84821..322f0a32 100644
--- a/tests/unit/cliproxy/oauth-process-error-parser.test.ts
+++ b/tests/unit/cliproxy/oauth-process-error-parser.test.ts
@@ -2,6 +2,7 @@ import { describe, expect, it } from 'bun:test';
 import {
   extractLikelyAuthFailureFromStderr,
   getExpectedLocalCallback,
+  getKiroBuilderIdSelectionInput,
   validateManualCallbackUrl,
 } from '../../../src/cliproxy/auth/oauth-process';
 
@@ -84,3 +85,35 @@ describe('oauth-process manual callback validation', () => {
     ).toContain('state does not match');
   });
 });
+
+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
+`;
+
+    expect(getKiroBuilderIdSelectionInput(output)).toBe('1\n');
+  });
+
+  it('selects the Builder ID option even when upstream reorders the menu', () => {
+    const output = `
+Select login method
+1. IAM Identity Center
+2. AWS Builder ID
+`;
+
+    expect(getKiroBuilderIdSelectionInput(output)).toBe('2\n');
+  });
+
+  it('returns null when the Builder ID option is not present in the prompt window', () => {
+    const output = `
+Select login method
+1. IAM Identity Center
+2. Google
+`;
+
+    expect(getKiroBuilderIdSelectionInput(output)).toBeNull();
+  });
+});