Merge pull request #213 from kaitranntt/kai/fix/kiro-callback-fallback

fix(kiro): add fallback import from Kiro IDE when OAuth callback redirects

src/cliproxy/auth/auth-types.ts

@@ -173,4 +173,6 @@ export interface OAuthOptions {
   fromUI?: boolean;
   /** If true, use --no-incognito flag (Kiro only - use normal browser instead of incognito) */
   noIncognito?: boolean;
+  /** If true, skip OAuth and import token from Kiro IDE directly (Kiro only) */
+  import?: boolean;
 }

src/cliproxy/auth/kiro-import.ts

@@ -0,0 +1,159 @@
+/**
+ * Kiro Import Helper
+ *
+ * Imports Kiro token from Kiro IDE when OAuth callback redirects to IDE instead of CLI.
+ * Spawns cli-proxy-api-plus --kiro-import to import token from Kiro IDE's storage.
+ */
+
+import { spawn } from 'child_process';
+import { info, ok, fail } from '../../utils/ui';
+import { ensureCLIProxyBinary } from '../binary-manager';
+import { generateConfig } from '../config-generator';
+import { getProviderTokenDir } from './token-manager';
+
+export interface KiroImportResult {
+  success: boolean;
+  provider?: string;
+  email?: string;
+  error?: string;
+}
+
+/**
+ * Try to import Kiro token from Kiro IDE
+ * Uses cli-proxy-api-plus --kiro-import flag
+ */
+export async function tryKiroImport(tokenDir: string, verbose = false): Promise<KiroImportResult> {
+  const log = (msg: string) => {
+    if (verbose) console.error(`[kiro-import] ${msg}`);
+  };
+
+  try {
+    log('Ensuring CLIProxy binary is available...');
+    const binaryPath = await ensureCLIProxyBinary(verbose);
+    const configPath = generateConfig('kiro');
+
+    log(`Binary: ${binaryPath}`);
+    log(`Config: ${configPath}`);
+    log(`Token dir: ${tokenDir}`);
+
+    return new Promise<KiroImportResult>((resolve) => {
+      const args = ['--config', configPath, '--kiro-import'];
+
+      log(`Running: ${binaryPath} ${args.join(' ')}`);
+
+      const proc = spawn(binaryPath, args, {
+        stdio: ['ignore', 'pipe', 'pipe'],
+        env: { ...process.env, CLI_PROXY_AUTH_DIR: tokenDir },
+      });
+
+      let stdout = '';
+      let stderr = '';
+      let resolved = false;
+
+      const safeResolve = (result: KiroImportResult) => {
+        if (resolved) return;
+        resolved = true;
+        clearTimeout(timeoutId);
+        resolve(result);
+      };
+
+      proc.stdout?.on('data', (data: Buffer) => {
+        const output = data.toString();
+        stdout += output;
+        log(`stdout: ${output.trim()}`);
+      });
+
+      proc.stderr?.on('data', (data: Buffer) => {
+        const output = data.toString();
+        stderr += output;
+        log(`stderr: ${output.trim()}`);
+      });
+
+      proc.on('exit', (code) => {
+        log(`Exit code: ${code}`);
+
+        if (code === 0) {
+          // Parse output for provider info
+          const providerMatch = stdout.match(/Provider:\s*(\w+)/i);
+          const emailMatch = stdout.match(/email[:\s]+([^\s,)]+)/i);
+          const successMatch =
+            stdout.includes('Kiro token import successful') ||
+            stdout.includes('Imported Kiro token') ||
+            stdout.includes('Authentication saved');
+
+          if (successMatch) {
+            safeResolve({
+              success: true,
+              provider: providerMatch?.[1],
+              email: emailMatch?.[1],
+            });
+          } else {
+            safeResolve({
+              success: false,
+              error: 'Import completed but token not confirmed',
+            });
+          }
+        } else {
+          const errorLine = stderr.trim().split('\n')[0] || stdout.trim().split('\n')[0];
+          safeResolve({
+            success: false,
+            error: errorLine || `Exit code ${code}`,
+          });
+        }
+      });
+
+      proc.on('error', (error) => {
+        log(`Process error: ${error.message}`);
+        safeResolve({
+          success: false,
+          error: error.message,
+        });
+      });
+
+      // Timeout after 30 seconds
+      const timeoutId = setTimeout(() => {
+        if (!resolved && !proc.killed) {
+          proc.kill();
+          safeResolve({
+            success: false,
+            error: 'Import timed out after 30 seconds',
+          });
+        }
+      }, 30000);
+    });
+  } catch (error) {
+    log(`Error: ${(error as Error).message}`);
+    return {
+      success: false,
+      error: (error as Error).message,
+    };
+  }
+}
+
+/**
+ * Import Kiro token with user-facing output
+ * Shows progress and result to user
+ */
+export async function importKiroToken(verbose = false): Promise<boolean> {
+  const tokenDir = getProviderTokenDir('kiro');
+
+  console.log('');
+  console.log(info('Importing token from Kiro IDE...'));
+
+  const result = await tryKiroImport(tokenDir, verbose);
+
+  if (result.success) {
+    const providerInfo = result.provider ? ` (Provider: ${result.provider})` : '';
+    console.log(ok(`Imported Kiro token from IDE${providerInfo}`));
+    return true;
+  }
+
+  console.log(fail(`Import failed: ${result.error}`));
+  console.log('');
+  console.log('Make sure you are logged into Kiro IDE first:');
+  console.log('  1. Open Kiro IDE');
+  console.log('  2. Sign in with your AWS/Google account');
+  console.log('  3. Run: ccs kiro --import');
+
+  return false;
+}

src/cliproxy/auth/oauth-handler.ts

@@ -27,8 +27,9 @@ import {
 } from '../../management/oauth-port-diagnostics';
 import { OAuthOptions, OAUTH_CALLBACK_PORTS, getOAuthConfig } from './auth-types';
 import { isHeadlessEnvironment, killProcessOnPort, showStep } from './environment-detector';
-import { getProviderTokenDir, isAuthenticated } from './token-manager';
+import { getProviderTokenDir, isAuthenticated, registerAccountFromToken } from './token-manager';
 import { executeOAuthProcess } from './oauth-process';
+import { importKiroToken } from './kiro-import';
 
 /**
  * Prompt user to add another account
@@ -127,6 +128,17 @@ export async function triggerOAuth(
 ): Promise<AccountInfo | null> {
   const oauthConfig = getOAuthConfig(provider);
   const { verbose = false, add = false, nickname, fromUI = false, noIncognito = true } = options;
+
+  // Handle --import flag: skip OAuth and import from Kiro IDE directly
+  if (options.import && provider === 'kiro') {
+    const tokenDir = getProviderTokenDir(provider);
+    const success = await importKiroToken(verbose);
+    if (success) {
+      return registerAccountFromToken(provider, tokenDir, nickname);
+    }
+    return null;
+  }
+
   const callbackPort = OAUTH_PORTS[provider];
   const isCLI = !fromUI;
   const headless = options.headless ?? isHeadlessEnvironment();

src/cliproxy/auth/oauth-process.ts

@@ -6,7 +6,8 @@
  */
 
 import { spawn, ChildProcess } from 'child_process';
-import { ok, fail, info } from '../../utils/ui';
+import { ok, fail, info, warn } from '../../utils/ui';
+import { tryKiroImport } from './kiro-import';
 import { CLIProxyProvider } from '../types';
 import { AccountInfo } from '../account-manager';
 import {
@@ -205,7 +206,35 @@ function displayUrlFromStderr(
 }
 
 /** Handle token not found after successful process exit */
-function handleTokenNotFound(provider: CLIProxyProvider, callbackPort: number | null): void {
+async function handleTokenNotFound(
+  provider: CLIProxyProvider,
+  callbackPort: number | null,
+  tokenDir: string,
+  nickname: string | undefined,
+  verbose: boolean
+): Promise<AccountInfo | null> {
+  // Kiro-specific: Try auto-import from Kiro IDE
+  if (provider === 'kiro') {
+    console.log('');
+    console.log(warn('Callback redirected to Kiro IDE. Attempting to import token...'));
+
+    const result = await tryKiroImport(tokenDir, verbose);
+
+    if (result.success) {
+      const providerInfo = result.provider ? ` (Provider: ${result.provider})` : '';
+      console.log(ok(`Imported Kiro token from IDE${providerInfo}`));
+      return registerAccountFromToken(provider, tokenDir, nickname);
+    }
+
+    console.log(fail(`Auto-import failed: ${result.error}`));
+    console.log('');
+    console.log('To manually import from Kiro IDE:');
+    console.log('  1. Ensure you are logged into Kiro IDE');
+    console.log('  2. Run: ccs kiro --import');
+    return null;
+  }
+
+  // Default behavior for other providers
   console.log('');
   console.log(fail('Token not found after authentication'));
   console.log('');
@@ -225,6 +254,7 @@ function handleTokenNotFound(provider: CLIProxyProvider, callbackPort: number |
 
   console.log('');
   console.log(`Try: ccs ${provider} --auth --verbose`);
+  return null;
 }
 
 /** Handle process exit with error */
@@ -356,7 +386,7 @@ export function executeOAuthProcess(options: OAuthProcessOptions): Promise<Accou
       resolve(null);
     }, timeoutMs);
 
-    authProcess.on('exit', (code) => {
+    authProcess.on('exit', async (code) => {
       clearTimeout(timeout);
       // H5: Remove signal handlers to prevent memory leaks
       process.removeListener('SIGINT', cleanup);
@@ -383,8 +413,15 @@ export function executeOAuthProcess(options: OAuthProcessOptions): Promise<Accou
             });
           }
 
-          handleTokenNotFound(provider, callbackPort);
-          resolve(null);
+          // Try auto-import for Kiro, show error for others
+          const account = await handleTokenNotFound(
+            provider,
+            callbackPort,
+            tokenDir,
+            nickname,
+            verbose
+          );
+          resolve(account);
         }
       } else {
         // Emit device code failure event for UI

src/cliproxy/cliproxy-executor.ts

@@ -255,6 +255,8 @@ export async function execClaudeWithCLIProxy(
   const forceConfig = argsWithoutProxy.includes('--config');
   const addAccount = argsWithoutProxy.includes('--add');
   const showAccounts = argsWithoutProxy.includes('--accounts');
+  // Kiro-specific: --import to import token from Kiro IDE directly
+  const forceImport = argsWithoutProxy.includes('--import');
   // Kiro-specific: browser mode for OAuth
   // Default to normal browser (noIncognito=true) for reliability - incognito often fails on Linux
   // --incognito flag opts into incognito mode, --no-incognito is legacy (now default)
@@ -367,6 +369,38 @@ export async function execClaudeWithCLIProxy(
     process.exit(0);
   }
 
+  // Handle --import: import token from Kiro IDE directly (Kiro only)
+  if (forceImport) {
+    if (provider !== 'kiro') {
+      console.error(fail('--import is only available for Kiro'));
+      console.error(`    Run "ccs ${provider} --auth" to authenticate`);
+      process.exit(1);
+    }
+    // Validate flag conflicts
+    if (forceAuth) {
+      console.error(fail('Cannot use --import with --auth'));
+      console.error('    --import: Import existing token from Kiro IDE');
+      console.error('    --auth: Trigger new OAuth flow in browser');
+      process.exit(1);
+    }
+    if (forceLogout) {
+      console.error(fail('Cannot use --import with --logout'));
+      process.exit(1);
+    }
+    const { triggerOAuth } = await import('./auth-handler');
+    const authSuccess = await triggerOAuth(provider, {
+      verbose,
+      import: true,
+      ...(setNickname ? { nickname: setNickname } : {}),
+    });
+    if (!authSuccess) {
+      console.error(fail('Failed to import Kiro token from IDE'));
+      console.error('    Make sure you are logged into Kiro IDE first');
+      process.exit(1);
+    }
+    process.exit(0);
+  }
+
   // 3. Ensure OAuth completed (if provider requires it)
   if (providerConfig.requiresOAuth) {
     log(`Checking authentication for ${provider}`);
@@ -636,6 +670,7 @@ export async function execClaudeWithCLIProxy(
     '--nickname',
     '--incognito',
     '--no-incognito',
+    '--import',
     // Proxy flags are handled by resolveProxyConfig, but list for documentation
     ...PROXY_CLI_FLAGS,
   ];

src/commands/help-command.ts

@@ -173,6 +173,7 @@ Run ${color('ccs config', 'command')} for web dashboard`.trim();
       ['ccs <provider> --config', 'Change model (agy, gemini)'],
       ['ccs <provider> --logout', 'Clear authentication'],
       ['ccs <provider> --headless', 'Headless auth (for SSH)'],
+      ['ccs kiro --import', 'Import token from Kiro IDE'],
       ['ccs kiro --incognito', 'Use incognito browser (default: normal)'],
       ['ccs codex "explain code"', 'Use with prompt'],
     ]

src/web-server/routes/cliproxy-auth-routes.ts

@@ -24,6 +24,8 @@ import {
 import { getProxyTarget } from '../../cliproxy/proxy-target-resolver';
 import { fetchRemoteAuthStatus } from '../../cliproxy/remote-auth-fetcher';
 import { loadOrCreateUnifiedConfig } from '../../config/unified-config-loader';
+import { tryKiroImport } from '../../cliproxy/auth/kiro-import';
+import { getProviderTokenDir } from '../../cliproxy/auth/token-manager';
 import type { CLIProxyProvider } from '../../cliproxy/types';
 
 const router = Router();
@@ -350,4 +352,52 @@ router.post('/project-selection/:sessionId', (req: Request, res: Response): void
   }
 });
 
+/**
+ * POST /api/cliproxy/auth/kiro/import - Import Kiro token from Kiro IDE
+ * Alternative auth path when OAuth callback fails to redirect properly
+ */
+router.post('/kiro/import', async (_req: Request, res: Response): Promise<void> => {
+  // Check if remote mode is enabled - import not available remotely
+  const target = getProxyTarget();
+  if (target.isRemote) {
+    res.status(501).json({
+      error: 'Kiro import not available in remote mode',
+    });
+    return;
+  }
+
+  try {
+    const tokenDir = getProviderTokenDir('kiro');
+    const result = await tryKiroImport(tokenDir, false);
+
+    if (result.success) {
+      // Re-initialize accounts to pick up new token
+      initializeAccounts();
+
+      // Get the newly added account
+      const accounts = getProviderAccounts('kiro');
+      const newAccount = accounts.find((a) => a.isDefault) || accounts[0];
+
+      res.json({
+        success: true,
+        account: newAccount
+          ? {
+              id: newAccount.id,
+              email: newAccount.email,
+              provider: 'kiro',
+              isDefault: newAccount.isDefault,
+            }
+          : null,
+      });
+    } else {
+      res.status(400).json({
+        success: false,
+        error: result.error || 'Failed to import Kiro token',
+      });
+    }
+  } catch (error) {
+    res.status(500).json({ error: (error as Error).message });
+  }
+});
+
 export default router;

ui/package.json

@@ -6,7 +6,7 @@
   "scripts": {
     "dev": "vite",
     "build": "tsc -b && vite build",
-    "typecheck": "tsc -b --noEmit",
+    "typecheck": "tsc --noEmit",
     "lint": "eslint .",
     "lint:fix": "eslint . --fix",
     "format": "prettier --write src/ tests/",

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

@@ -2,6 +2,7 @@
  * Add Account Dialog Component
  * Triggers OAuth flow server-side to add another account to a provider
  * Applies default preset when adding first account
+ * For Kiro: Also shows "Import from IDE" option as fallback
  */
 
 import { useState } from 'react';
@@ -15,8 +16,8 @@ import {
 import { Button } from '@/components/ui/button';
 import { Input } from '@/components/ui/input';
 import { Label } from '@/components/ui/label';
-import { Loader2, ExternalLink, User } from 'lucide-react';
-import { useStartAuth } from '@/hooks/use-cliproxy';
+import { Loader2, ExternalLink, User, Download } from 'lucide-react';
+import { useStartAuth, useKiroImport } from '@/hooks/use-cliproxy';
 import { applyDefaultPreset } from '@/lib/preset-utils';
 import { toast } from 'sonner';
 
@@ -38,6 +39,10 @@ export function AddAccountDialog({
 }: AddAccountDialogProps) {
   const [nickname, setNickname] = useState('');
   const startAuthMutation = useStartAuth();
+  const kiroImportMutation = useKiroImport();
+
+  const isKiro = provider === 'kiro';
+  const isPending = startAuthMutation.isPending || kiroImportMutation.isPending;
 
   const handleStartAuth = () => {
     startAuthMutation.mutate(
@@ -60,8 +65,24 @@ export function AddAccountDialog({
     );
   };
 
+  const handleKiroImport = () => {
+    kiroImportMutation.mutate(undefined, {
+      onSuccess: async () => {
+        // Apply default preset if this is the first account
+        if (isFirstAccount) {
+          const result = await applyDefaultPreset('kiro');
+          if (result.success && result.presetName) {
+            toast.success(`Applied "${result.presetName}" preset`);
+          }
+        }
+        setNickname('');
+        onClose();
+      },
+    });
+  };
+
   const handleOpenChange = (isOpen: boolean) => {
-    if (!isOpen && !startAuthMutation.isPending) {
+    if (!isOpen && !isPending) {
       setNickname('');
       onClose();
     }
@@ -73,8 +94,9 @@ export function AddAccountDialog({
         <DialogHeader>
           <DialogTitle>Add {displayName} Account</DialogTitle>
           <DialogDescription>
-            Click the button below to authenticate a new account. A browser window will open for
-            OAuth.
+            {isKiro
+              ? 'Authenticate via browser or import an existing token from Kiro IDE.'
+              : 'Click the button below to authenticate a new account. A browser window will open for OAuth.'}
           </DialogDescription>
         </DialogHeader>
 
@@ -88,7 +110,7 @@ export function AddAccountDialog({
                 value={nickname}
                 onChange={(e) => setNickname(e.target.value)}
                 placeholder="e.g., work, personal"
-                disabled={startAuthMutation.isPending}
+                disabled={isPending}
                 className="flex-1"
               />
             </div>
@@ -98,10 +120,25 @@ export function AddAccountDialog({
           </div>
 
           <div className="flex items-center justify-end gap-2 pt-2">
-            <Button variant="ghost" onClick={onClose} disabled={startAuthMutation.isPending}>
+            <Button variant="ghost" onClick={onClose} disabled={isPending}>
               Cancel
             </Button>
-            <Button onClick={handleStartAuth} disabled={startAuthMutation.isPending}>
+            {isKiro && (
+              <Button variant="outline" onClick={handleKiroImport} disabled={isPending}>
+                {kiroImportMutation.isPending ? (
+                  <>
+                    <Loader2 className="w-4 h-4 mr-2 animate-spin" />
+                    Importing...
+                  </>
+                ) : (
+                  <>
+                    <Download className="w-4 h-4 mr-2" />
+                    Import from IDE
+                  </>
+                )}
+              </Button>
+            )}
+            <Button onClick={handleStartAuth} disabled={isPending}>
               {startAuthMutation.isPending ? (
                 <>
                   <Loader2 className="w-4 h-4 mr-2 animate-spin" />
@@ -121,6 +158,11 @@ export function AddAccountDialog({
               Complete the OAuth flow in your browser...
             </p>
           )}
+          {kiroImportMutation.isPending && (
+            <p className="text-sm text-center text-muted-foreground">
+              Importing token from Kiro IDE...
+            </p>
+          )}
         </div>
       </DialogContent>
     </Dialog>

ui/src/hooks/use-cliproxy.ts

@@ -136,6 +136,27 @@ export function useStartAuth() {
   });
 }
 
+// Kiro IDE import hook (alternative auth path when OAuth callback fails)
+export function useKiroImport() {
+  const queryClient = useQueryClient();
+
+  return useMutation({
+    mutationFn: () => api.cliproxy.auth.kiroImport(),
+    onSuccess: (data) => {
+      queryClient.invalidateQueries({ queryKey: ['cliproxy-accounts'] });
+      queryClient.invalidateQueries({ queryKey: ['cliproxy-auth'] });
+      if (data.account) {
+        toast.success(`Imported Kiro account: ${data.account.email || data.account.id}`);
+      } else {
+        toast.success('Kiro token imported');
+      }
+    },
+    onError: (error: Error) => {
+      toast.error(error.message);
+    },
+  });
+}
+
 // Stats and models hooks for Overview tab
 export function useCliproxyStats() {
   return useQuery({

ui/src/lib/api-client.ts

@@ -331,6 +331,12 @@ export const api = {
           method: 'POST',
           body: JSON.stringify({ nickname }),
         }),
+      /** Import Kiro token from Kiro IDE (Kiro only) */
+      kiroImport: () =>
+        request<{ success: boolean; account: OAuthAccount | null; error?: string }>(
+          '/cliproxy/auth/kiro/import',
+          { method: 'POST' }
+        ),
     },
     // Error logs
     errorLogs: {