Merge pull request #561 from kaitranntt/fix/552-review-followups

fix(cliproxy): harden Kiro device-code auth flow consistency

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

@@ -1,7 +1,7 @@
 /**
  * Add Account Dialog Component
  * Uses /start-url to get OAuth URL + polls for completion via management API.
- * For Device Code flows (ghcp, qwen): Uses /start endpoint which spawns CLIProxy
+ * For Device Code flows (ghcp, qwen, kiro): Uses /start endpoint which spawns CLIProxy
  * binary and emits WebSocket events. DeviceCodeDialog handles user code display.
  * Shows auth URL + callback paste field. Polling auto-closes on success.
  * For Kiro: Also shows "Import from IDE" option.
@@ -18,11 +18,25 @@ import {
 import { Button } from '@/components/ui/button';
 import { Input } from '@/components/ui/input';
 import { Label } from '@/components/ui/label';
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from '@/components/ui/select';
 import { Loader2, ExternalLink, User, Download, Copy, Check } from 'lucide-react';
 import { useKiroImport } from '@/hooks/use-cliproxy';
 import { useCliproxyAuthFlow } from '@/hooks/use-cliproxy-auth-flow';
 import { applyDefaultPreset } from '@/lib/preset-utils';
-import { isDeviceCodeProvider } from '@/lib/provider-config';
+import {
+  DEFAULT_KIRO_AUTH_METHOD,
+  getKiroAuthMethodOption,
+  isDeviceCodeProvider,
+  isNicknameRequiredProvider,
+  KIRO_AUTH_METHOD_OPTIONS,
+} from '@/lib/provider-config';
+import type { KiroAuthMethod } from '@/lib/provider-config';
 import { toast } from 'sonner';
 
 interface AddAccountDialogProps {
@@ -44,18 +58,27 @@ export function AddAccountDialog({
   const [nickname, setNickname] = useState('');
   const [callbackUrl, setCallbackUrl] = useState('');
   const [copied, setCopied] = useState(false);
+  const [localError, setLocalError] = useState<string | null>(null);
+  const [kiroAuthMethod, setKiroAuthMethod] = useState<KiroAuthMethod>(DEFAULT_KIRO_AUTH_METHOD);
   const wasAuthenticatingRef = useRef(false);
   const authFlow = useCliproxyAuthFlow();
   const kiroImportMutation = useKiroImport();
 
   const isKiro = provider === 'kiro';
-  const isDeviceCode = isDeviceCodeProvider(provider);
+  const defaultDeviceCode = isDeviceCodeProvider(provider);
+  const requiresNickname = isNicknameRequiredProvider(provider);
+  const kiroMethodOption = getKiroAuthMethodOption(kiroAuthMethod);
+  const isDeviceCode = isKiro ? kiroMethodOption.flowType === 'device_code' : defaultDeviceCode;
   const isPending = authFlow.isAuthenticating || kiroImportMutation.isPending;
+  const nicknameTrimmed = nickname.trim();
+  const errorMessage = localError || authFlow.error;
 
   const resetAndClose = () => {
     setNickname('');
     setCallbackUrl('');
     setCopied(false);
+    setLocalError(null);
+    setKiroAuthMethod(DEFAULT_KIRO_AUTH_METHOD);
     wasAuthenticatingRef.current = false;
     onClose();
   };
@@ -103,13 +126,23 @@ export function AddAccountDialog({
   };
 
   /**
-   * Authenticate via /start-url + polling only.
-   * Does NOT call /start (which spawns a local CLIProxy binary that kills running instances).
-   * /start-url uses the management API to get auth URL, then polls for completion.
+   * Start auth flow using provider capabilities.
+   * - Device code providers use /start and rely on WebSocket events for code display.
+   * - Authorization code providers use /start-url and polling.
    */
   const handleAuthenticate = () => {
+    if (requiresNickname && !nicknameTrimmed) {
+      setLocalError(`Nickname is required for ${displayName} accounts.`);
+      return;
+    }
+    setLocalError(null);
     wasAuthenticatingRef.current = true;
-    authFlow.startAuth(provider, { nickname: nickname.trim() || undefined });
+    authFlow.startAuth(provider, {
+      nickname: nicknameTrimmed || undefined,
+      kiroMethod: isKiro ? kiroAuthMethod : undefined,
+      flowType: isKiro ? kiroMethodOption.flowType : undefined,
+      startEndpoint: isKiro ? kiroMethodOption.startEndpoint : undefined,
+    });
   };
 
   const handleKiroImport = () => {
@@ -146,7 +179,7 @@ export function AddAccountDialog({
           <DialogTitle>Add {displayName} Account</DialogTitle>
           <DialogDescription>
             {isKiro
-              ? 'Authenticate via browser or import an existing token from Kiro IDE.'
+              ? 'Choose a Kiro auth method, then authenticate via browser or import from Kiro IDE.'
               : isDeviceCode
                 ? 'Click Authenticate. A verification code will appear for you to enter on the provider website.'
                 : 'Click Authenticate to get an OAuth URL. Open it in any browser to sign in.'}
@@ -154,23 +187,56 @@ export function AddAccountDialog({
         </DialogHeader>
 
         <div className="space-y-4 py-4">
+          {/* Kiro auth method */}
+          {isKiro && !showAuthUI && (
+            <div className="space-y-2">
+              <Label htmlFor="kiro-auth-method">Auth Method</Label>
+              <Select
+                value={kiroAuthMethod}
+                onValueChange={(value) => {
+                  setKiroAuthMethod(value as KiroAuthMethod);
+                  setLocalError(null);
+                }}
+              >
+                <SelectTrigger id="kiro-auth-method">
+                  <SelectValue placeholder="Select Kiro auth method" />
+                </SelectTrigger>
+                <SelectContent>
+                  {KIRO_AUTH_METHOD_OPTIONS.map((option) => (
+                    <SelectItem key={option.id} value={option.id}>
+                      {option.label}
+                    </SelectItem>
+                  ))}
+                </SelectContent>
+              </Select>
+              <p className="text-xs text-muted-foreground">{kiroMethodOption.description}</p>
+            </div>
+          )}
+
           {/* Nickname input - only show before auth starts */}
           {!showAuthUI && (
             <div className="space-y-2">
-              <Label htmlFor="nickname">Nickname (optional)</Label>
+              <Label htmlFor="nickname">
+                {requiresNickname ? 'Nickname (required)' : 'Nickname (optional)'}
+              </Label>
               <div className="flex items-center gap-2">
                 <User className="w-4 h-4 text-muted-foreground" />
                 <Input
                   id="nickname"
                   value={nickname}
-                  onChange={(e) => setNickname(e.target.value)}
+                  onChange={(e) => {
+                    setNickname(e.target.value);
+                    setLocalError(null);
+                  }}
                   placeholder="e.g., work, personal"
                   disabled={isPending}
                   className="flex-1"
                 />
               </div>
               <p className="text-xs text-muted-foreground">
-                A friendly name to identify this account. Auto-generated from email if left empty.
+                {requiresNickname
+                  ? 'Required for this provider. Use a unique friendly name (e.g., work, personal).'
+                  : 'A friendly name to identify this account. Auto-generated from email if left empty.'}
               </p>
             </div>
           )}
@@ -191,11 +257,6 @@ export function AddAccountDialog({
                 </p>
               </div>
 
-              {/* Error display */}
-              {authFlow.error && !authFlow.authUrl && (
-                <p className="text-xs text-center text-destructive">{authFlow.error}</p>
-              )}
-
               {/* Auth URL section - only for Authorization Code flows, NOT Device Code */}
               {authFlow.authUrl && !authFlow.isDeviceCodeFlow && (
                 <div className="space-y-3">
@@ -270,9 +331,18 @@ export function AddAccountDialog({
                   </div>
                 </div>
               )}
+
+              {!authFlow.authUrl && !authFlow.isDeviceCodeFlow && (
+                <p className="text-xs text-center text-muted-foreground">
+                  Preparing sign-in URL...
+                </p>
+              )}
             </div>
           )}
 
+          {/* Persist error visibility outside auth-only UI states */}
+          {errorMessage && <p className="text-xs text-center text-destructive">{errorMessage}</p>}
+
           {/* Kiro import loading */}
           {kiroImportMutation.isPending && (
             <p className="text-sm text-center text-muted-foreground">
@@ -302,7 +372,10 @@ export function AddAccountDialog({
               </Button>
             )}
             {!showAuthUI && (
-              <Button onClick={handleAuthenticate} disabled={isPending}>
+              <Button
+                onClick={handleAuthenticate}
+                disabled={isPending || (requiresNickname && !nicknameTrimmed)}
+              >
                 <ExternalLink className="w-4 h-4 mr-2" />
                 Authenticate
               </Button>

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

@@ -19,12 +19,15 @@ interface AuthFlowState {
   oauthState: string | null;
   /** Whether callback is being submitted */
   isSubmittingCallback: boolean;
-  /** Whether this is a device code flow (ghcp, qwen) - dialog handled separately via WebSocket */
+  /** Whether this is a device code flow (ghcp, qwen, kiro) - dialog handled separately via WebSocket */
   isDeviceCodeFlow: boolean;
 }
 
 interface StartAuthOptions {
   nickname?: string;
+  kiroMethod?: string;
+  flowType?: 'authorization_code' | 'device_code';
+  startEndpoint?: 'start' | 'start-url';
 }
 
 /** Polling interval for OAuth status check (3 seconds) */
@@ -49,6 +52,7 @@ export function useCliproxyAuthFlow() {
   const abortControllerRef = useRef<AbortController | null>(null);
   const pollIntervalRef = useRef<ReturnType<typeof setInterval> | null>(null);
   const pollStartRef = useRef<number>(0);
+  const openedAuthUrlRef = useRef(false);
   const queryClient = useQueryClient();
 
   // Clear polling
@@ -64,6 +68,7 @@ export function useCliproxyAuthFlow() {
     return () => {
       abortControllerRef.current?.abort();
       stopPolling();
+      openedAuthUrlRef.current = false;
     };
   }, [stopPolling]);
 
@@ -85,14 +90,46 @@ export function useCliproxyAuthFlow() {
         const response = await fetch(
           `/api/cliproxy/auth/${provider}/status?state=${encodeURIComponent(oauthState)}`
         );
-        const data = await response.json();
+        const data = (await response.json()) as {
+          status?: string;
+          error?: string;
+          url?: string;
+          auth_url?: string;
+          verification_url?: string;
+          user_code?: string;
+        };
 
         if (data.status === 'ok') {
           stopPolling();
           queryClient.invalidateQueries({ queryKey: ['cliproxy-auth'] });
           queryClient.invalidateQueries({ queryKey: ['account-quota'] });
           toast.success(`${provider} authentication successful`);
+          openedAuthUrlRef.current = false;
           setState(INITIAL_STATE);
+        } else if (data.status === 'auth_url') {
+          const authUrl = data.url || data.auth_url;
+          if (authUrl) {
+            setState((prev) => ({
+              ...prev,
+              authUrl,
+            }));
+            if (!openedAuthUrlRef.current) {
+              openedAuthUrlRef.current = true;
+              window.open(authUrl, '_blank');
+            }
+          }
+        } else if (data.status === 'device_code') {
+          stopPolling();
+          const details =
+            data.user_code && data.verification_url
+              ? `Open ${data.verification_url} and enter code: ${data.user_code}`
+              : 'Switch to Device Code method and try again.';
+          toast.error('Provider returned Device Code flow in callback mode');
+          setState((prev) => ({
+            ...prev,
+            isAuthenticating: false,
+            error: details,
+          }));
         } else if (data.status === 'error') {
           stopPolling();
           const errorMsg = data.error || 'Authentication failed';
@@ -103,7 +140,7 @@ export function useCliproxyAuthFlow() {
             error: errorMsg,
           }));
         }
-        // status === 'pending' means continue polling
+        // status === 'wait' (or pending) means continue polling
       } catch {
         // Network error - continue polling
       }
@@ -124,12 +161,21 @@ export function useCliproxyAuthFlow() {
       // Abort any in-progress auth
       abortControllerRef.current?.abort();
       stopPolling();
+      openedAuthUrlRef.current = false;
 
       // Create fresh controller and capture locally to avoid race with cancelAuth
       const controller = new AbortController();
       abortControllerRef.current = controller;
 
-      const deviceCodeFlow = isDeviceCodeProvider(provider);
+      const flowType =
+        options?.flowType ||
+        (isDeviceCodeProvider(provider) ? 'device_code' : 'authorization_code');
+      const deviceCodeFlow = flowType === 'device_code';
+      const startEndpoint = options?.startEndpoint || (deviceCodeFlow ? 'start' : 'start-url');
+      const payload = {
+        nickname: options?.nickname,
+        kiroMethod: options?.kiroMethod,
+      };
 
       setState({
         provider,
@@ -142,14 +188,13 @@ export function useCliproxyAuthFlow() {
       });
 
       try {
-        if (deviceCodeFlow) {
-          // Device Code Flow: Call /start endpoint which spawns CLIProxyAPI binary.
-          // This emits WebSocket events with userCode that DeviceCodeDialog will display.
-          // The /start endpoint blocks until completion, so we don't await it here.
+        if (startEndpoint === 'start') {
+          // /start spawns CLIProxy binary and blocks until completion.
+          // For Device Code flows, userCode is delivered via WebSocket.
           fetch(`/api/cliproxy/auth/${provider}/start`, {
             method: 'POST',
             headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify({ nickname: options?.nickname }),
+            body: JSON.stringify(payload),
             signal: controller.signal,
           })
             .then(async (response) => {
@@ -159,6 +204,7 @@ export function useCliproxyAuthFlow() {
                 queryClient.invalidateQueries({ queryKey: ['account-quota'] });
                 // Note: No toast here - DeviceCodeDialog's useDeviceCode hook handles success toast
                 // via deviceCodeCompleted WebSocket event to avoid duplicate toasts
+                openedAuthUrlRef.current = false;
                 setState(INITIAL_STATE);
               } else {
                 const errorMsg = data.error || 'Authentication failed';
@@ -183,13 +229,13 @@ export function useCliproxyAuthFlow() {
                 error: message,
               }));
             });
-          // Don't await - let the request run in background while DeviceCodeDialog handles UI
+          // Don't await - keeps UI responsive while backend auth is in progress
         } else {
-          // Authorization Code Flow: Call /start-url to get auth URL immediately (non-blocking)
+          // /start-url uses management API to bootstrap callback/social flows.
           const response = await fetch(`/api/cliproxy/auth/${provider}/start-url`, {
             method: 'POST',
             headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify({ nickname: options?.nickname }),
+            body: JSON.stringify(payload),
             signal: controller.signal,
           });
 
@@ -202,12 +248,13 @@ export function useCliproxyAuthFlow() {
           // Update state with auth URL
           setState((prev) => ({
             ...prev,
-            authUrl: data.authUrl,
+            authUrl: data.authUrl || null,
             oauthState: data.state,
           }));
 
           // Auto-open auth URL in new browser tab (fallback URL still shown in dialog)
           if (data.authUrl) {
+            openedAuthUrlRef.current = true;
             window.open(data.authUrl, '_blank');
           }
 
@@ -221,6 +268,7 @@ export function useCliproxyAuthFlow() {
         }
       } catch (error) {
         if (error instanceof Error && error.name === 'AbortError') {
+          openedAuthUrlRef.current = false;
           setState(INITIAL_STATE);
           return;
         }
@@ -240,6 +288,7 @@ export function useCliproxyAuthFlow() {
     const currentProvider = state.provider;
     abortControllerRef.current?.abort();
     stopPolling();
+    openedAuthUrlRef.current = false;
     setState(INITIAL_STATE);
     // Also cancel on backend
     if (currentProvider) {

ui/src/hooks/use-device-code.ts

@@ -3,7 +3,7 @@
  *
  * Listens for WebSocket device code events and manages dialog state.
  * Similar to useProjectSelection but for Device Code OAuth flows
- * (GitHub Copilot, Qwen, etc.)
+ * (GitHub Copilot, Qwen, Kiro, etc.)
  */
 
 import { useState, useEffect, useCallback, useMemo } from 'react';
@@ -26,6 +26,7 @@ interface DeviceCodeState {
 /** Provider display names for user-friendly messages */
 const PROVIDER_DISPLAY_NAMES: Record<string, string> = {
   ghcp: 'GitHub Copilot',
+  kiro: 'Kiro (AWS)',
   qwen: 'Qwen Code',
 };
 

ui/src/lib/provider-config.ts

@@ -82,3 +82,75 @@ export const DEVICE_CODE_PROVIDERS: CLIProxyProvider[] = ['ghcp', 'kiro', 'qwen'
 export function isDeviceCodeProvider(provider: string): boolean {
   return DEVICE_CODE_PROVIDERS.includes(provider as CLIProxyProvider);
 }
+
+/** Providers that require nickname because token payload may not include email. */
+export const NICKNAME_REQUIRED_PROVIDERS: CLIProxyProvider[] = ['ghcp', 'kiro'];
+
+/** Check if provider requires user-supplied nickname in auth flow */
+export function isNicknameRequiredProvider(provider: string): boolean {
+  return NICKNAME_REQUIRED_PROVIDERS.includes(provider as CLIProxyProvider);
+}
+
+/** Kiro auth methods exposed in CCS UI (aligned with CLIProxyAPIPlus support). */
+export const KIRO_AUTH_METHODS = ['aws', 'aws-authcode', 'google', 'github'] as const;
+export type KiroAuthMethod = (typeof KIRO_AUTH_METHODS)[number];
+
+export type KiroFlowType = 'authorization_code' | 'device_code';
+export type KiroStartEndpoint = 'start' | 'start-url';
+
+export interface KiroAuthMethodOption {
+  id: KiroAuthMethod;
+  label: string;
+  description: string;
+  flowType: KiroFlowType;
+  startEndpoint: KiroStartEndpoint;
+}
+
+/** UX-first default for issue #233: AWS Builder ID device flow. */
+export const DEFAULT_KIRO_AUTH_METHOD: KiroAuthMethod = 'aws';
+
+export const KIRO_AUTH_METHOD_OPTIONS: readonly KiroAuthMethodOption[] = [
+  {
+    id: 'aws',
+    label: 'AWS Builder ID (Recommended)',
+    description: 'Device code flow for AWS organizations and Builder ID accounts.',
+    flowType: 'device_code',
+    startEndpoint: 'start',
+  },
+  {
+    id: 'aws-authcode',
+    label: 'AWS Builder ID (Auth Code)',
+    description: 'Authorization code flow via CLI binary.',
+    flowType: 'authorization_code',
+    startEndpoint: 'start',
+  },
+  {
+    id: 'google',
+    label: 'Google OAuth',
+    description: 'Social OAuth flow with callback URL support.',
+    flowType: 'authorization_code',
+    startEndpoint: 'start-url',
+  },
+  {
+    id: 'github',
+    label: 'GitHub OAuth',
+    description: 'Social OAuth flow via management API callback.',
+    flowType: 'authorization_code',
+    startEndpoint: 'start-url',
+  },
+];
+
+export function isKiroAuthMethod(value: string): value is KiroAuthMethod {
+  return KIRO_AUTH_METHODS.includes(value as KiroAuthMethod);
+}
+
+export function normalizeKiroAuthMethod(value?: string): KiroAuthMethod {
+  if (!value) return DEFAULT_KIRO_AUTH_METHOD;
+  const normalized = value.trim().toLowerCase();
+  return isKiroAuthMethod(normalized) ? normalized : DEFAULT_KIRO_AUTH_METHOD;
+}
+
+export function getKiroAuthMethodOption(method: KiroAuthMethod): KiroAuthMethodOption {
+  const option = KIRO_AUTH_METHOD_OPTIONS.find((candidate) => candidate.id === method);
+  return option || KIRO_AUTH_METHOD_OPTIONS[0];
+}