Skip to content

Replace useState calls with discriminated union #352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sundaray wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Replace useState calls with discriminated union #352

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
77ebb89
replace useState calls with discriminated union
sundaray Dec 28, 2025
5093a36
refactor: replace useState with useReducer for auth state
sundaray Jan 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
150 changes: 72 additions & 78 deletions src/components/authkit-provider.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,17 @@ import {
refreshAuthAction,
switchToOrganizationAction,
} from '../actions.js';
import type { Impersonator, User } from '@workos-inc/node';
import type { User } from '@workos-inc/node';
import type { UserInfo, SwitchToOrganizationOptions, NoUserInfo } from '../interfaces.js';

type AuthContextType = {
user: User | null;
sessionId: string | undefined;
organizationId: string | undefined;
role: string | undefined;
roles: string[] | undefined;
permissions: string[] | undefined;
entitlements: string[] | undefined;
featureFlags: string[] | undefined;
impersonator: Impersonator | undefined;
loading: boolean;
type AuthState =
| {
status: 'loading';
}
| { status: 'unauthenticated'; data: Omit<NoUserInfo, 'accessToken'> }
| { status: 'authenticated'; data: Omit<UserInfo, 'accessToken'> };

type AuthContextType = Omit<UserInfo | NoUserInfo, 'accessToken'> & { loading: boolean } & {
getAuth: (options?: { ensureSignedIn?: boolean }) => Promise<void>;
refreshAuth: (options?: { ensureSignedIn?: boolean; organizationId?: string }) => Promise<void | { error: string }>;
signOut: (options?: { returnTo?: string }) => Promise<void>;
Expand All @@ -46,43 +43,63 @@ interface AuthKitProviderProps {
initialAuth?: Omit<UserInfo | NoUserInfo, 'accessToken'>;
}

const unauthenticatedState: { status: 'unauthenticated'; data: Omit<NoUserInfo, 'accessToken'> } = {
status: 'unauthenticated',
data: {
user: null,
sessionId: undefined,
organizationId: undefined,
role: undefined,
roles: undefined,
permissions: undefined,
entitlements: undefined,
featureFlags: undefined,
impersonator: undefined,
},
};

function createAuthState(auth: Omit<UserInfo | NoUserInfo, 'accessToken'> | undefined): AuthState {
if (!auth) {
return { status: 'loading' };
}

if (!auth.user) {
return unauthenticatedState;
}

return {
status: 'authenticated',
data: {
user: auth.user,
sessionId: auth.sessionId,
organizationId: auth.organizationId,
role: auth.role,
roles: auth.roles,
permissions: auth.permissions,
entitlements: auth.entitlements,
featureFlags: auth.featureFlags,
impersonator: auth.impersonator,
} as Omit<UserInfo, 'accessToken'>,
};
}

function getAuthStateData(authState: AuthState): Omit<UserInfo | NoUserInfo, 'accessToken'> {
if (authState.status === 'loading') {
return unauthenticatedState.data;
}
return authState.data;
}

export const AuthKitProvider = ({ children, onSessionExpired, initialAuth }: AuthKitProviderProps) => {
const [user, setUser] = useState<User | null>(initialAuth?.user ?? null);
const [sessionId, setSessionId] = useState<string | undefined>(initialAuth?.sessionId);
const [organizationId, setOrganizationId] = useState<string | undefined>(initialAuth?.organizationId);
const [role, setRole] = useState<string | undefined>(initialAuth?.role);
const [roles, setRoles] = useState<string[] | undefined>(initialAuth?.roles);
const [permissions, setPermissions] = useState<string[] | undefined>(initialAuth?.permissions);
const [entitlements, setEntitlements] = useState<string[] | undefined>(initialAuth?.entitlements);
const [featureFlags, setFeatureFlags] = useState<string[] | undefined>(initialAuth?.featureFlags);
const [impersonator, setImpersonator] = useState<Impersonator | undefined>(initialAuth?.impersonator);
const [loading, setLoading] = useState(!initialAuth);
const [authState, setAuthState] = useState<AuthState>(() => createAuthState(initialAuth));
Copy link
Copy Markdown
Member

@nicknisi nicknisi Jan 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if this should be a useReducer call instead? That way the auth state and loading boolean can be updated independently?

Copy link
Copy Markdown
Contributor Author

@sundaray sundaray Jan 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right. I have replaced useState with useReducer.


const getAuth = useCallback(async ({ ensureSignedIn = false }: { ensureSignedIn?: boolean } = {}) => {
setLoading(true);
setAuthState({ status: 'loading' });
try {
const auth = await getAuthAction({ ensureSignedIn });
setUser(auth.user);
setSessionId(auth.sessionId);
setOrganizationId(auth.organizationId);
setRole(auth.role);
setRoles(auth.roles);
setPermissions(auth.permissions);
setEntitlements(auth.entitlements);
setFeatureFlags(auth.featureFlags);
setImpersonator(auth.impersonator);
setAuthState(createAuthState(auth));
} catch (error) {
setUser(null);
setSessionId(undefined);
setOrganizationId(undefined);
setRole(undefined);
setRoles(undefined);
setPermissions(undefined);
setEntitlements(undefined);
setFeatureFlags(undefined);
setImpersonator(undefined);
} finally {
setLoading(false);
setAuthState(unauthenticatedState);
}
}, []);

Expand All @@ -105,23 +122,12 @@ export const AuthKitProvider = ({ children, onSessionExpired, initialAuth }: Aut

const refreshAuth = useCallback(
async ({ ensureSignedIn = false, organizationId }: { ensureSignedIn?: boolean; organizationId?: string } = {}) => {
setAuthState({ status: 'loading' });
try {
setLoading(true);
const auth = await refreshAuthAction({ ensureSignedIn, organizationId });

setUser(auth.user);
setSessionId(auth.sessionId);
setOrganizationId(auth.organizationId);
setRole(auth.role);
setRoles(auth.roles);
setPermissions(auth.permissions);
setEntitlements(auth.entitlements);
setFeatureFlags(auth.featureFlags);
setImpersonator(auth.impersonator);
setAuthState(createAuthState(auth));
} catch (error) {
return error instanceof Error ? { error: error.message } : { error: String(error) };
} finally {
setLoading(false);
Comment on lines -123 to -124
Copy link
Copy Markdown
Member

@nicknisi nicknisi Jan 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If this action call fails above then without this finally the loading state stays loading forever. We probably need to keep this in some form to reset the loading state.

Copy link
Copy Markdown
Contributor Author

@sundaray sundaray Jan 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed.

}
},
[],
Expand Down Expand Up @@ -184,34 +190,22 @@ export const AuthKitProvider = ({ children, onSessionExpired, initialAuth }: Aut
};
}, [onSessionExpired]);

return (
<AuthContext.Provider
value={{
user,
sessionId,
organizationId,
role,
roles,
permissions,
entitlements,
featureFlags,
impersonator,
loading,
getAuth,
refreshAuth,
signOut,
switchToOrganization,
}}
>
{children}
</AuthContext.Provider>
);
const contextValue: AuthContextType = {
loading: authState.status === 'loading',
...getAuthStateData(authState),
getAuth,
refreshAuth,
signOut,
switchToOrganization,
};
return <AuthContext.Provider value={contextValue}>{children}</AuthContext.Provider>;
};

export function useAuth(options: {
ensureSignedIn: true;
}): AuthContextType & ({ loading: true; user: User | null } | { loading: false; user: User });
export function useAuth(options?: { ensureSignedIn?: false }): AuthContextType;

export function useAuth({ ensureSignedIn = false }: { ensureSignedIn?: boolean } = {}) {
const context = useContext(AuthContext);

Expand Down