fix(ui): improve Kiro dashboard auth flow

ui/src/components/account/add-account-dialog.tsx

@@ -39,11 +39,15 @@ import {
 } from '@/components/account/antigravity-responsibility-constants';
 import {
   DEFAULT_KIRO_AUTH_METHOD,
+  DEFAULT_KIRO_IDC_FLOW,
+  getKiroEffectiveFlowType,
+  getKiroEffectiveStartEndpoint,
   getKiroAuthMethodOption,
+  isKiroSocialAuthMethod,
   isDeviceCodeProvider,
   KIRO_AUTH_METHOD_OPTIONS,
 } from '@/lib/provider-config';
-import type { KiroAuthMethod } from '@/lib/provider-config';
+import type { KiroAuthMethod, KiroIDCFlow } from '@/lib/provider-config';
 import { useTranslation } from 'react-i18next';
 import { toast } from 'sonner';
 
@@ -81,6 +85,9 @@ export function AddAccountDialog({
   const [powerUserModeEnabled, setPowerUserModeEnabled] = useState(false);
   const [powerUserModeLoading, setPowerUserModeLoading] = useState(false);
   const [kiroAuthMethod, setKiroAuthMethod] = useState<KiroAuthMethod>(DEFAULT_KIRO_AUTH_METHOD);
+  const [kiroIDCStartUrl, setKiroIDCStartUrl] = useState('');
+  const [kiroIDCRegion, setKiroIDCRegion] = useState('');
+  const [kiroIDCFlow, setKiroIDCFlow] = useState<KiroIDCFlow>(DEFAULT_KIRO_IDC_FLOW);
   const { t } = useTranslation();
   const wasAuthenticatingRef = useRef(false);
   const powerUserModeRequestIdRef = useRef(0);
@@ -97,9 +104,19 @@ export function AddAccountDialog({
   const isGeminiRiskAcknowledged = normalizeRiskPhrase(riskAcknowledgementText) === RISK_ACK_PHRASE;
   const defaultDeviceCode = isDeviceCodeProvider(provider);
   const kiroMethodOption = getKiroAuthMethodOption(kiroAuthMethod);
-  const isDeviceCode = isKiro ? kiroMethodOption.flowType === 'device_code' : defaultDeviceCode;
+  const isKiroIdc = isKiro && kiroAuthMethod === 'idc';
+  const isKiroSocial = isKiro && isKiroSocialAuthMethod(kiroAuthMethod);
+  const selectedKiroFlowType = isKiro
+    ? getKiroEffectiveFlowType(kiroAuthMethod, kiroIDCFlow)
+    : undefined;
+  const selectedKiroStartEndpoint = isKiro
+    ? getKiroEffectiveStartEndpoint(kiroAuthMethod)
+    : undefined;
+  const isDeviceCode = isKiro ? selectedKiroFlowType === 'device_code' : defaultDeviceCode;
   const isPending = authFlow.isAuthenticating || kiroImportMutation.isPending;
   const nicknameTrimmed = nickname.trim();
+  const kiroIDCStartUrlTrimmed = kiroIDCStartUrl.trim();
+  const kiroIDCRegionTrimmed = kiroIDCRegion.trim();
   const errorMessage = localError || authFlow.error;
 
   const fetchPowerUserModeState = useCallback(async (): Promise<boolean> => {
@@ -168,6 +185,9 @@ export function AddAccountDialog({
     setPowerUserModeEnabled(false);
     setPowerUserModeLoading(false);
     setKiroAuthMethod(DEFAULT_KIRO_AUTH_METHOD);
+    setKiroIDCStartUrl('');
+    setKiroIDCRegion('');
+    setKiroIDCFlow(DEFAULT_KIRO_IDC_FLOW);
     powerUserModeRequestIdRef.current += 1;
     powerUserModeLoadErrorShownRef.current = false;
     wasAuthenticatingRef.current = false;
@@ -282,12 +302,19 @@ export function AddAccountDialog({
       return;
     }
     setLocalError(null);
+    if (isKiroIdc && !kiroIDCStartUrlTrimmed) {
+      setLocalError('IDC Start URL is required for Kiro IAM Identity Center login.');
+      return;
+    }
     wasAuthenticatingRef.current = true;
     authFlow.startAuth(provider, {
       nickname: nicknameTrimmed || undefined,
       kiroMethod: isKiro ? kiroAuthMethod : undefined,
-      flowType: isKiro ? kiroMethodOption.flowType : undefined,
+      kiroIDCStartUrl: isKiroIdc ? kiroIDCStartUrlTrimmed : undefined,
-      startEndpoint: isKiro ? kiroMethodOption.startEndpoint : undefined,
+      kiroIDCRegion: isKiroIdc && kiroIDCRegionTrimmed ? kiroIDCRegionTrimmed : undefined,
+      kiroIDCFlow: isKiroIdc ? kiroIDCFlow : undefined,
+      flowType: isKiro ? selectedKiroFlowType : undefined,
+      startEndpoint: isKiro ? selectedKiroStartEndpoint : undefined,
       riskAcknowledgement: requiresAgyResponsibilityFlow
         ? {
             version: ANTIGRAVITY_ACK_VERSION,
@@ -399,6 +426,77 @@ export function AddAccountDialog({
                 </SelectContent>
               </Select>
               <p className="text-xs text-muted-foreground">{kiroMethodOption.description}</p>
+              {isKiroSocial && (
+                <p className="text-xs text-muted-foreground">
+                  If your browser does not return automatically after login, CCS can accept the
+                  final
+                  <span className="mx-1 rounded bg-muted px-1.5 py-0.5 font-mono text-[11px]">
+                    kiro://...
+                  </span>
+                  callback URL in the next step.
+                </p>
+              )}
+            </div>
+          )}
+
+          {isKiroIdc && !showAuthUI && (
+            <div className="space-y-4 rounded-lg border border-border/60 bg-muted/20 p-4">
+              <div className="space-y-2">
+                <Label htmlFor="kiro-idc-start-url">IDC Start URL</Label>
+                <Input
+                  id="kiro-idc-start-url"
+                  value={kiroIDCStartUrl}
+                  onChange={(e) => {
+                    setKiroIDCStartUrl(e.target.value);
+                    setLocalError(null);
+                  }}
+                  placeholder="https://d-xxx.awsapps.com/start"
+                  disabled={isPending}
+                />
+                <p className="text-xs text-muted-foreground">
+                  Required for organization IAM Identity Center login.
+                </p>
+              </div>
+
+              <div className="space-y-2">
+                <Label htmlFor="kiro-idc-region">IDC Region</Label>
+                <Input
+                  id="kiro-idc-region"
+                  value={kiroIDCRegion}
+                  onChange={(e) => {
+                    setKiroIDCRegion(e.target.value);
+                    setLocalError(null);
+                  }}
+                  placeholder="us-east-1"
+                  disabled={isPending}
+                />
+                <p className="text-xs text-muted-foreground">
+                  Optional. Leave blank to use the upstream default region.
+                </p>
+              </div>
+
+              <div className="space-y-2">
+                <Label htmlFor="kiro-idc-flow">IDC Flow</Label>
+                <Select
+                  value={kiroIDCFlow}
+                  onValueChange={(value) => {
+                    setKiroIDCFlow(value as KiroIDCFlow);
+                    setLocalError(null);
+                  }}
+                >
+                  <SelectTrigger id="kiro-idc-flow">
+                    <SelectValue placeholder="Select IDC flow" />
+                  </SelectTrigger>
+                  <SelectContent>
+                    <SelectItem value="authcode">Authorization Code</SelectItem>
+                    <SelectItem value="device">Device Code</SelectItem>
+                  </SelectContent>
+                </Select>
+                <p className="text-xs text-muted-foreground">
+                  Auth Code opens a browser and may need the final callback URL pasted back. Device
+                  Code shows a verification code instead.
+                </p>
+              </div>
             </div>
           )}
 
@@ -438,7 +536,9 @@ export function AddAccountDialog({
                 <p className="text-xs text-muted-foreground mt-1">
                   {authFlow.isDeviceCodeFlow
                     ? t('addAccountDialog.deviceCodeHint')
-                    : t('addAccountDialog.browserHint')}
+                    : isKiroSocial
+                      ? 'Complete sign-in in your browser. If it does not return automatically, paste the final kiro:// callback URL below.'
+                      : t('addAccountDialog.browserHint')}
                 </p>
               </div>
 
@@ -486,13 +586,19 @@ export function AddAccountDialog({
                   {/* Callback paste field */}
                   <div className="space-y-2">
                     <Label htmlFor="callback-url" className="text-xs">
-                      {t('addAccountDialog.redirectPasteLabel')}
+                      {isKiroSocial
+                        ? 'Browser did not return? Paste the final kiro:// callback URL:'
+                        : t('addAccountDialog.redirectPasteLabel')}
                     </Label>
                     <Input
                       id="callback-url"
                       value={callbackUrl}
                       onChange={(e) => setCallbackUrl(e.target.value)}
-                      placeholder={t('addAccountDialog.callbackPlaceholder')}
+                      placeholder={
+                        isKiroSocial
+                          ? 'kiro://kiro.kiroAgent/authenticate-success?code=...&state=...'
+                          : t('addAccountDialog.callbackPlaceholder')
+                      }
                       className="font-mono text-xs"
                     />
                     <Button
@@ -516,7 +622,9 @@ export function AddAccountDialog({
 
               {!authFlow.authUrl && !authFlow.isDeviceCodeFlow && (
                 <p className="text-xs text-center text-muted-foreground">
-                  {t('addAccountDialog.preparingUrl')}
+                  {isKiroSocial
+                    ? 'Preparing the Kiro sign-in URL. If it does not open automatically, it will appear here shortly.'
+                    : t('addAccountDialog.preparingUrl')}
                 </p>
               )}
             </div>

ui/src/hooks/use-cliproxy-auth-flow.ts

@@ -26,6 +26,9 @@ interface AuthFlowState {
 interface StartAuthOptions {
   nickname?: string;
   kiroMethod?: string;
+  kiroIDCStartUrl?: string;
+  kiroIDCRegion?: string;
+  kiroIDCFlow?: 'authcode' | 'device';
   flowType?: 'authorization_code' | 'device_code';
   startEndpoint?: 'start' | 'start-url';
   riskAcknowledgement?: {
@@ -70,6 +73,7 @@ const INITIAL_STATE: AuthFlowState = {
 
 export function useCliproxyAuthFlow() {
   const [state, setState] = useState<AuthFlowState>(INITIAL_STATE);
+  const stateRef = useRef<AuthFlowState>(INITIAL_STATE);
 
   const attemptIdRef = useRef(0);
   const abortControllerRef = useRef<AbortController | null>(null);
@@ -102,6 +106,10 @@ export function useCliproxyAuthFlow() {
     };
   }, [stopPolling]);
 
+  useEffect(() => {
+    stateRef.current = state;
+  }, [state]);
+
   // Poll OAuth status
   const pollStatus = useCallback(
     async (provider: string, oauthState: string, attemptId: number) => {
@@ -253,6 +261,9 @@ export function useCliproxyAuthFlow() {
       const payload = {
         nickname: options?.nickname,
         kiroMethod: options?.kiroMethod,
+        kiroIDCStartUrl: options?.kiroIDCStartUrl,
+        kiroIDCRegion: options?.kiroIDCRegion,
+        kiroIDCFlow: options?.kiroIDCFlow,
         riskAcknowledgement: options?.riskAcknowledgement,
       };
 
@@ -368,6 +379,16 @@ export function useCliproxyAuthFlow() {
           // Start polling for completion
           if (oauthState) {
             pollStartRef.current = Date.now();
+            if (!authUrl) {
+              await pollStatus(provider, oauthState, attemptId);
+              if (!isActiveAttempt(attemptId)) {
+                return;
+              }
+              const currentState = stateRef.current;
+              if (!currentState.isAuthenticating || currentState.provider !== provider) {
+                return;
+              }
+            }
             pollIntervalRef.current = setInterval(() => {
               void pollStatus(provider, oauthState, attemptId);
             }, POLL_INTERVAL);

ui/src/lib/provider-config.ts

@@ -234,8 +234,11 @@ export function getDeviceCodeProviderInstruction(provider: unknown): string {
 }
 
 /** Kiro auth methods exposed in CCS UI (aligned with CLIProxyAPIPlus support). */
-export const KIRO_AUTH_METHODS = ['aws', 'aws-authcode', 'google', 'github'] as const;
+export const KIRO_AUTH_METHODS = ['aws', 'aws-authcode', 'google', 'github', 'idc'] as const;
 export type KiroAuthMethod = (typeof KIRO_AUTH_METHODS)[number];
+export const KIRO_IDC_FLOWS = ['authcode', 'device'] as const;
+export type KiroIDCFlow = (typeof KIRO_IDC_FLOWS)[number];
+export const DEFAULT_KIRO_IDC_FLOW: KiroIDCFlow = 'authcode';
 
 export type KiroFlowType = 'authorization_code' | 'device_code';
 export type KiroStartEndpoint = 'start' | 'start-url';
@@ -280,6 +283,13 @@ export const KIRO_AUTH_METHOD_OPTIONS: readonly KiroAuthMethodOption[] = [
     flowType: 'authorization_code',
     startEndpoint: 'start-url',
   },
+  {
+    id: 'idc',
+    label: 'AWS Identity Center (IDC)',
+    description: 'Use your organization start URL with auth code or device flow.',
+    flowType: 'authorization_code',
+    startEndpoint: 'start',
+  },
 ];
 
 export function isKiroAuthMethod(value: string): value is KiroAuthMethod {
@@ -292,7 +302,40 @@ export function normalizeKiroAuthMethod(value?: string): KiroAuthMethod {
   return isKiroAuthMethod(normalized) ? normalized : DEFAULT_KIRO_AUTH_METHOD;
 }
 
+export function isKiroIDCFlow(value: string): value is KiroIDCFlow {
+  return KIRO_IDC_FLOWS.includes(value as KiroIDCFlow);
+}
+
+export function normalizeKiroIDCFlow(value?: string): KiroIDCFlow {
+  if (!value) return DEFAULT_KIRO_IDC_FLOW;
+  const normalized = value.trim().toLowerCase();
+  return isKiroIDCFlow(normalized) ? normalized : DEFAULT_KIRO_IDC_FLOW;
+}
+
 export function getKiroAuthMethodOption(method: KiroAuthMethod): KiroAuthMethodOption {
   const option = KIRO_AUTH_METHOD_OPTIONS.find((candidate) => candidate.id === method);
   return option || KIRO_AUTH_METHOD_OPTIONS[0];
 }
+
+export function getKiroEffectiveFlowType(
+  method: KiroAuthMethod,
+  idcFlow: KiroIDCFlow = DEFAULT_KIRO_IDC_FLOW
+): KiroFlowType {
+  if (method === 'aws') {
+    return 'device_code';
+  }
+
+  if (method === 'idc') {
+    return normalizeKiroIDCFlow(idcFlow) === 'device' ? 'device_code' : 'authorization_code';
+  }
+
+  return 'authorization_code';
+}
+
+export function getKiroEffectiveStartEndpoint(method: KiroAuthMethod): KiroStartEndpoint {
+  return method === 'google' || method === 'github' ? 'start-url' : 'start';
+}
+
+export function isKiroSocialAuthMethod(method: KiroAuthMethod): boolean {
+  return method === 'google' || method === 'github';
+}

ui/tests/unit/hooks/use-cliproxy-auth-flow.test.tsx

@@ -221,6 +221,96 @@ describe('useCliproxyAuthFlow', () => {
     expect(toast.success).toHaveBeenCalledWith('codex authentication successful');
   });
 
+  it('promotes a state-first auth bootstrap into an immediate auth URL without waiting for the first interval', async () => {
+    let pollCount = 0;
+
+    vi.stubGlobal(
+      'fetch',
+      vi.fn((input: RequestInfo | URL) => {
+        const url = String(input);
+
+        if (url.includes('/start-url')) {
+          return Promise.resolve(
+            createJsonResponse({
+              success: true,
+              authUrl: null,
+              state: 'state-kiro-social',
+            })
+          );
+        }
+
+        if (url.includes('/status?state=state-kiro-social')) {
+          pollCount += 1;
+          return Promise.resolve(
+            createJsonResponse({
+              status: 'auth_url',
+              url: 'https://auth.example/kiro-social',
+            })
+          );
+        }
+
+        return Promise.reject(new Error(`Unexpected fetch: ${url}`));
+      })
+    );
+
+    const { result } = renderHook(() => useCliproxyAuthFlow(), { wrapper });
+
+    await act(async () => {
+      await result.current.startAuth('kiro', { startEndpoint: 'start-url', kiroMethod: 'google' });
+    });
+
+    expect(result.current.authUrl).toBe('https://auth.example/kiro-social');
+    expect(result.current.oauthState).toBe('state-kiro-social');
+    expect(result.current.isAuthenticating).toBe(true);
+    expect(pollCount).toBe(1);
+  });
+
+  it('forwards Kiro IDC options to the backend start endpoint payload', async () => {
+    let requestBody: Record<string, unknown> | null = null;
+
+    vi.stubGlobal(
+      'fetch',
+      vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
+        const url = String(input);
+
+        if (url.includes('/start')) {
+          requestBody = JSON.parse(String(init?.body)) as Record<string, unknown>;
+          return createJsonResponse({
+            success: true,
+            account: {
+              id: 'kiro-idc-account',
+              provider: 'kiro',
+            },
+          });
+        }
+
+        return Promise.reject(new Error(`Unexpected fetch: ${url}`));
+      })
+    );
+
+    const { result } = renderHook(() => useCliproxyAuthFlow(), { wrapper });
+
+    await act(async () => {
+      await result.current.startAuth('kiro', {
+        startEndpoint: 'start',
+        flowType: 'authorization_code',
+        kiroMethod: 'idc',
+        kiroIDCStartUrl: 'https://d-123.awsapps.com/start',
+        kiroIDCRegion: 'ca-central-1',
+        kiroIDCFlow: 'authcode',
+      });
+    });
+
+    expect(requestBody).toEqual({
+      nickname: undefined,
+      kiroMethod: 'idc',
+      kiroIDCStartUrl: 'https://d-123.awsapps.com/start',
+      kiroIDCRegion: 'ca-central-1',
+      kiroIDCFlow: 'authcode',
+      riskAcknowledgement: undefined,
+    });
+  });
+
   it('treats callback responses without an account as failures', async () => {
     vi.stubGlobal(
       'fetch',