Add Cognito groups to auth context and user profile

src/components/layout/__tests__/Header.test.tsx

@@ -20,6 +20,7 @@ vi.mock('@/context/theme-context', () => ({
 const authValue: AuthContextValue = {
   isAuthenticated: true,
   user: { name: 'Ray Liu', email: 'ray@example.com' },
+  groups: [],
   isLoading: false,
   signInWithGoogle: vi.fn(() => Promise.resolve(undefined)),
   logout: vi.fn(() => Promise.resolve(undefined)),

src/components/layout/__tests__/app-shell.test.tsx

@@ -29,6 +29,7 @@ vi.mock('@/context/theme-context', () => ({
 const authValue: AuthContextValue = {
   isAuthenticated: true,
   user: { name: 'Ray Liu', email: 'ray@example.com' },
+  groups: [],
   isLoading: false,
   signInWithGoogle: vi.fn(() => Promise.resolve(undefined)),
   logout: vi.fn(() => Promise.resolve(undefined)),

src/components/layout/header/UserMenu.tsx

@@ -16,7 +16,7 @@ import { UserProfileModal } from './UserProfileModal';
 import { UserTokenModal } from './UserTokenModal';
 
 export function UserMenu() {
-  const { user: userInformation, logout } = useAuthContext();
+  const { user: userInformation, groups, logout } = useAuthContext();
   const [settingsOpen, setSettingsOpen] = useState(false);
   const [profileOpen, setProfileOpen] = useState(false);
   const [tokenOpen, setTokenOpen] = useState(false);
@@ -129,6 +129,7 @@ export function UserMenu() {
       <ThemeSettingsModal isOpen={settingsOpen} onClose={() => setSettingsOpen(false)} />
       <UserProfileModal
         user={userInformation}
+        groups={groups}
         isOpen={profileOpen}
         onClose={() => setProfileOpen(false)}
       />

src/components/layout/header/UserProfileModal.tsx

@@ -1,4 +1,4 @@
-import { Mail, ShieldCheck, ShieldAlert, Fingerprint, Globe, User } from 'lucide-react';
+import { Mail, ShieldCheck, ShieldAlert, Fingerprint, Globe, User, Users } from 'lucide-react';
 import type { FetchUserAttributesOutput } from 'aws-amplify/auth';
 import { getUsername } from '@/utils/string';
 import { DialogFrame } from '@/components/modals/DialogFrame';
@@ -14,6 +14,7 @@ interface Identity {
 
 interface UserProfileModalProps {
   user: FetchUserAttributesOutput;
+  groups?: string[];
   isOpen: boolean;
   onClose: () => void;
 }
@@ -51,7 +52,7 @@ function IdentityRow({ label, value, mono }: { label: string; value: string; mon
   );
 }
 
-export function UserProfileModal({ user, isOpen, onClose }: UserProfileModalProps) {
+export function UserProfileModal({ user, groups = [], isOpen, onClose }: UserProfileModalProps) {
   const userName = user.name || getUsername(user.email as string);
   const userEmail = user.email ?? '';
   const isEmailVerified = user.email_verified === 'true';
@@ -127,6 +128,28 @@ export function UserProfileModal({ user, isOpen, onClose }: UserProfileModalProp
           </div>
         </div>
 
+        {/* Groups */}
+        {groups.length > 0 && (
+          <div>
+            <p className='mb-2 text-[11px] font-semibold tracking-wider text-slate-400 uppercase dark:text-neutral-500'>
+              Groups
+            </p>
+            <div className='flex items-start gap-3 rounded-lg border border-slate-200 px-4 py-3 dark:border-[#2d3540]'>
+              <Users className='mt-0.5 h-4 w-4 shrink-0 text-slate-400 dark:text-[#9dabb9]' />
+              <div className='flex min-w-0 flex-1 flex-wrap gap-1.5'>
+                {groups.map((g) => (
+                  <span
+                    key={g}
+                    className='inline-flex items-center rounded-full bg-slate-100 px-2.5 py-0.5 text-[11px] font-medium text-slate-700 dark:bg-[#2d3540] dark:text-slate-300'
+                  >
+                    {g}
+                  </span>
+                ))}
+              </div>
+            </div>
+          </div>
+        )}
+
         {/* Identity providers */}
         {identities.length > 0 && (
           <div>

src/context/AmplifyAuthProvider.tsx

@@ -1,6 +1,7 @@
 import { Suspense, use, useCallback, useEffect, useReducer, useMemo, type ReactNode } from 'react';
 import {
   fetchUserAttributes,
+  fetchAuthSession,
   signInWithRedirect,
   signOut,
   type FetchUserAttributesOutput,
@@ -36,31 +37,47 @@ function hasCognitoSession(): boolean {
   return Object.keys(localStorage).some((key) => key.startsWith(COGNITO_STORAGE_PREFIX));
 }
 
+// Cognito auto-creates a group like "ap-southeast-2_abcdefg1234_Google" for federated IDP users
+const COGNITO_IDP_GROUP_RE = /^[a-z]+-[a-z]+-\d+_[A-Za-z0-9]+_(\w+)$/;
+
+function normalizeGroupName(group: string): string {
+  const match = COGNITO_IDP_GROUP_RE.exec(group);
+  return match ? `${match[1]} (default)` : group;
+}
+
 // ---------- State management ----------
 
 interface AuthState {
   isAuthenticated: boolean;
   user: FetchUserAttributesOutput;
+  groups: string[];
   isLoading: boolean;
 }
 
 const initialState: AuthState = {
   isAuthenticated: false,
   user: {},
+  groups: [],
   isLoading: false,
 };
 
 type AuthAction =
-  | { type: 'AUTHENTICATED'; user: FetchUserAttributesOutput }
+  | { type: 'AUTHENTICATED'; user: FetchUserAttributesOutput; groups: string[] }
   | { type: 'UNAUTHENTICATED' }
   | { type: 'SET_LOADING'; isLoading: boolean };
 
 function authReducer(state: AuthState, action: AuthAction): AuthState {
   switch (action.type) {
     case 'AUTHENTICATED':
-      return { ...state, isAuthenticated: true, user: action.user, isLoading: false };
+      return {
+        ...state,
+        isAuthenticated: true,
+        user: action.user,
+        groups: action.groups,
+        isLoading: false,
+      };
     case 'UNAUTHENTICATED':
-      return { ...state, isAuthenticated: false, user: {}, isLoading: false };
+      return { ...state, isAuthenticated: false, user: {}, groups: [], isLoading: false };
     case 'SET_LOADING':
       return { ...state, isLoading: action.isLoading };
   }
@@ -72,8 +89,11 @@ async function resolveCurrentAuthState(): Promise<AuthState> {
   }
 
   try {
-    const user = await fetchUserAttributes();
-    return { ...initialState, isAuthenticated: true, user };
+    const [user, session] = await Promise.all([fetchUserAttributes(), fetchAuthSession()]);
+    const rawGroups =
+      (session.tokens?.idToken?.payload['cognito:groups'] as string[] | undefined) ?? [];
+    const groups = rawGroups.map(normalizeGroupName);
+    return { ...initialState, isAuthenticated: true, user, groups };
   } catch (error) {
     console.error('Failed to initialize auth:', error);
     toast.error('Failed to authenticate user');
@@ -124,7 +144,7 @@ function ResolvedAuthProvider({ children }: { children: ReactNode }) {
     cacheInitialAuthState(nextState);
 
     if (nextState.isAuthenticated) {
-      dispatch({ type: 'AUTHENTICATED', user: nextState.user });
+      dispatch({ type: 'AUTHENTICATED', user: nextState.user, groups: nextState.groups });
     } else {
       dispatch({ type: 'UNAUTHENTICATED' });
     }

src/context/auth-context.ts

@@ -4,6 +4,7 @@ import type { FetchUserAttributesOutput } from 'aws-amplify/auth';
 export interface AuthContextValue {
   isAuthenticated: boolean;
   user: FetchUserAttributesOutput;
+  groups: string[];
   isLoading: boolean;
   signInWithGoogle: () => Promise<void>;
   logout: () => Promise<void>;