Added profile ARN handling in OAuth provider mapping and improved polling logic in OAuth modal for better user experience.

src/lib/oauth/providers.js

@@ -660,6 +660,7 @@ const PROVIDERS = {
             access_token: data.accessToken,
             refresh_token: data.refreshToken,
             expires_in: data.expiresIn,
+            profile_arn: data?.profileArn || null,
             // Store client credentials for refresh
             _clientId: extraData?._clientId,
             _clientSecret: extraData?._clientSecret,
@@ -675,15 +676,19 @@ const PROVIDERS = {
         },
       };
     },
-    mapTokens: (tokens) => ({
-      accessToken: tokens.access_token,
-      refreshToken: tokens.refresh_token,
-      expiresIn: tokens.expires_in,
-      providerSpecificData: {
-        clientId: tokens._clientId,
-        clientSecret: tokens._clientSecret,
-      },
-    }),
+    mapTokens: (tokens) => {
+      const mapped = {
+        accessToken: tokens.access_token,
+        refreshToken: tokens.refresh_token,
+        expiresIn: tokens.expires_in,
+        providerSpecificData: {
+          profileArn: tokens?.profile_arn || null,
+          clientId: tokens._clientId,
+          clientSecret: tokens._clientSecret,
+        },
+      };
+      return mapped;
+    },
   },
 
   cursor: {

src/shared/components/KiroOAuthWrapper.js

@@ -43,12 +43,14 @@ export default function KiroOAuthWrapper({ isOpen, providerInfo, onSuccess, onCl
     setAuthMethod(null);
     setSocialProvider(null);
     onSuccess?.();
+    onClose?.(); // Close modal after success
   };
 
   const handleDeviceSuccess = () => {
     setAuthMethod(null);
     setIdcConfig(null);
     onSuccess?.();
+    onClose?.(); // Close modal after success
   };
 
   // Show method selection first

src/shared/components/OAuthModal.js

@@ -19,6 +19,7 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess,
   const [deviceData, setDeviceData] = useState(null);
   const [polling, setPolling] = useState(false);
   const popupRef = useRef(null);
+  const pollingAbortRef = useRef(false);
   const { copied, copy } = useCopyToClipboard();
 
   // State for client-only values to avoid hydration mismatch
@@ -66,12 +67,27 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess,
 
   // Poll for device code token
   const startPolling = useCallback(async (deviceCode, codeVerifier, interval, extraData) => {
+    pollingAbortRef.current = false;
     setPolling(true);
     const maxAttempts = 60;
 
     for (let i = 0; i < maxAttempts; i++) {
+      // Check if polling should be aborted
+      if (pollingAbortRef.current) {
+        console.log("[OAuthModal] Polling aborted");
+        setPolling(false);
+        return;
+      }
+
       await new Promise((r) => setTimeout(r, interval * 1000));
 
+      // Check again after sleep
+      if (pollingAbortRef.current) {
+        console.log("[OAuthModal] Polling aborted after sleep");
+        setPolling(false);
+        return;
+      }
+
       try {
         const res = await fetch(`/api/oauth/${provider}/poll`, {
           method: "POST",
@@ -82,6 +98,7 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess,
         const data = await res.json();
 
         if (data.success) {
+          pollingAbortRef.current = true; // Stop polling immediately
           setStep("success");
           setPolling(false);
           onSuccess?.();
@@ -182,7 +199,11 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess,
       setIsDeviceCode(false);
       setDeviceData(null);
       setPolling(false);
+      pollingAbortRef.current = false;
       startOAuthFlow();
+    } else if (!isOpen) {
+      // Abort polling when modal closes
+      pollingAbortRef.current = true;
     }
   }, [isOpen, provider, startOAuthFlow]);