Render session task context for separate components

LauraBeatris · LauraBeatris · commit 0c6c23a98da1 · 2025-11-26T20:22:28.000-03:00
diff --git a/packages/clerk-js/src/ui/components/SessionTasks/index.tsx b/packages/clerk-js/src/ui/components/SessionTasks/index.tsx
@@ -1,14 +1,13 @@
 import { useClerk } from '@clerk/shared/react';
 import { eventComponentMounted } from '@clerk/shared/telemetry';
-import type { SessionResource } from '@clerk/shared/types';
 import { useEffect, useRef } from 'react';
 
 import { Flow } from '@/ui/customizables';
 import { Card } from '@/ui/elements/Card';
 import { withCardStateProvider } from '@/ui/elements/contexts';
 import { LoadingCardContainer } from '@/ui/elements/LoadingCard';
 
-import { getTaskEndpoint, INTERNAL_SESSION_TASK_ROUTE_BY_KEY } from '../../../core/sessionTasks';
+import { INTERNAL_SESSION_TASK_ROUTE_BY_KEY } from '../../../core/sessionTasks';
 import {
   SessionTasksContext,
   TaskChooseOrganizationContext,
@@ -86,7 +85,7 @@ type SessionTasksProps = {
  */
 export const SessionTasks = withCardStateProvider(({ redirectUrlComplete }: SessionTasksProps) => {
   const clerk = useClerk();
-  const { navigate, basePath, startPath } = useRouter();
+  const { navigate } = useRouter();
 
   const currentTaskContainer = useRef<HTMLDivElement>(null);
 
@@ -121,21 +120,8 @@ export const SessionTasks = withCardStateProvider(({ redirectUrlComplete }: Sess
     );
   }
 
-  const navigateOnSetActive = async ({ session }: { session: SessionResource }) => {
-    const currentTask = session.currentTask;
-    if (!currentTask) {
-      return navigate(redirectUrlComplete);
-    }
-
-    const taskEndpoint = getTaskEndpoint(currentTask);
-
-    // Base path is required for virtual routing with start path
-    // eg: to navigate from /sign-in/factor-one to /sign-in/tasks/choose-organization
-    return navigate(`/${basePath + startPath + taskEndpoint}`);
-  };
-
   return (
-    <SessionTasksContext.Provider value={{ redirectUrlComplete, currentTaskContainer, navigateOnSetActive }}>
+    <SessionTasksContext.Provider value={{ redirectUrlComplete }}>
       <SessionTasksRoutes />
     </SessionTasksContext.Provider>
   );
diff --git a/packages/clerk-js/src/ui/components/SessionTasks/tasks/TaskChooseOrganization/ChooseOrganizationScreen.tsx b/packages/clerk-js/src/ui/components/SessionTasks/tasks/TaskChooseOrganization/ChooseOrganizationScreen.tsx
@@ -16,7 +16,7 @@ import {
   sharedMainIdentifierSx,
 } from '@/ui/common/organizations/OrganizationPreview';
 import { organizationListParams, populateCacheUpdateItem } from '@/ui/components/OrganizationSwitcher/utils';
-import { useTaskChooseOrganizationContext } from '@/ui/contexts/components/SessionTasks';
+import { useSessionTasksContext, useTaskChooseOrganizationContext } from '@/ui/contexts/components/SessionTasks';
 import { Col, descriptors, localizationKeys, Text, useLocalizations } from '@/ui/customizables';
 import { Action, Actions } from '@/ui/elements/Actions';
 import { Card } from '@/ui/elements/Card';
@@ -25,7 +25,6 @@ import { Header } from '@/ui/elements/Header';
 import { OrganizationPreview } from '@/ui/elements/OrganizationPreview';
 import { useOrganizationListInView } from '@/ui/hooks/useOrganizationListInView';
 import { Add } from '@/ui/icons';
-import { useRouter } from '@/ui/router';
 import { handleError } from '@/ui/utils/errorHandler';
 
 type ChooseOrganizationScreenProps = {
@@ -107,7 +106,7 @@ export const ChooseOrganizationScreen = (props: ChooseOrganizationScreenProps) =
 const MembershipPreview = (props: { organization: OrganizationResource }) => {
   const { user } = useUser();
   const card = useCardState();
-  const { navigate } = useRouter();
+  const { navigateOnSetActive } = useSessionTasksContext();
   const { redirectUrlComplete } = useTaskChooseOrganizationContext();
   const { isLoaded, setActive } = useOrganizationList();
   const { t } = useLocalizations();
@@ -121,9 +120,8 @@ const MembershipPreview = (props: { organization: OrganizationResource }) => {
       try {
         await setActive({
           organization,
-          navigate: async () => {
-            // TODO(after-auth) ORGS-779 - Handle next tasks
-            await navigate(redirectUrlComplete);
+          navigate: async ({ session }) => {
+            await navigateOnSetActive?.({ session, redirectUrlComplete });
           },
         });
       } catch (err) {
diff --git a/packages/clerk-js/src/ui/components/SessionTasks/tasks/TaskChooseOrganization/CreateOrganizationScreen.tsx b/packages/clerk-js/src/ui/components/SessionTasks/tasks/TaskChooseOrganization/CreateOrganizationScreen.tsx
@@ -2,14 +2,13 @@ import { useOrganizationList } from '@clerk/shared/react';
 import type { CreateOrganizationParams } from '@clerk/shared/types';
 
 import { useEnvironment } from '@/ui/contexts';
-import { useTaskChooseOrganizationContext } from '@/ui/contexts/components/SessionTasks';
+import { useSessionTasksContext, useTaskChooseOrganizationContext } from '@/ui/contexts/components/SessionTasks';
 import { localizationKeys } from '@/ui/customizables';
 import { useCardState } from '@/ui/elements/contexts';
 import { Form } from '@/ui/elements/Form';
 import { FormButtonContainer } from '@/ui/elements/FormButtons';
 import { FormContainer } from '@/ui/elements/FormContainer';
 import { Header } from '@/ui/elements/Header';
-import { useRouter } from '@/ui/router';
 import { createSlug } from '@/ui/utils/createSlug';
 import { handleError } from '@/ui/utils/errorHandler';
 import { useFormControl } from '@/ui/utils/useFormControl';
@@ -22,7 +21,7 @@ type CreateOrganizationScreenProps = {
 
 export const CreateOrganizationScreen = (props: CreateOrganizationScreenProps) => {
   const card = useCardState();
-  const { navigate } = useRouter();
+  const { navigateOnSetActive } = useSessionTasksContext();
   const { redirectUrlComplete } = useTaskChooseOrganizationContext();
   const { createOrganization, isLoaded, setActive } = useOrganizationList({
     userMemberships: organizationListParams.userMemberships,
@@ -60,9 +59,8 @@ export const CreateOrganizationScreen = (props: CreateOrganizationScreenProps) =
 
       await setActive({
         organization,
-        navigate: async () => {
-          // TODO(after-auth) ORGS-779 - Handle next tasks
-          await navigate(redirectUrlComplete);
+        navigate: async ({ session }) => {
+          await navigateOnSetActive?.({ session, redirectUrlComplete });
         },
       });
     } catch (err) {
diff --git a/packages/clerk-js/src/ui/components/SessionTasks/tasks/TaskResetPassword/index.tsx b/packages/clerk-js/src/ui/components/SessionTasks/tasks/TaskResetPassword/index.tsx
@@ -25,12 +25,12 @@ const TaskResetPasswordInternal = () => {
 
   const { t, locale } = useLocalizations();
   const { redirectUrlComplete } = useTaskResetPasswordContext();
-  const { navigateOnSetActive } = useSessionTasksContext();
   const { otherSessions } = useMultipleSessions({ user: clerk.user });
   const { navigateAfterSignOut, navigateAfterMultiSessionSingleSignOutUrl } = useSignOutContext();
   const updatePasswordWithReverification = useReverification(
     (user: UserResource, opts: Parameters<UserResource['updatePassword']>) => user.updatePassword(...opts),
   );
+  const { navigateOnSetActive } = useSessionTasksContext();
 
   const handleSignOut = () => {
     if (otherSessions.length === 0) {
@@ -94,7 +94,7 @@ const TaskResetPasswordInternal = () => {
         await clerk.setActive({
           session: clerk.session,
           navigate: async ({ session }) => {
-            await navigateOnSetActive({ session, redirectUrl: redirectUrlComplete });
+            await navigateOnSetActive?.({ session, redirectUrlComplete });
           },
         });
       } catch (e) {
diff --git a/packages/clerk-js/src/ui/components/SessionTasks/tasks/shared/withTaskGuard.ts b/packages/clerk-js/src/ui/components/SessionTasks/tasks/shared/withTaskGuard.ts
@@ -6,10 +6,18 @@ import { withRedirect } from '@/ui/common';
 import { useSessionTasksContext } from '@/ui/contexts/components/SessionTasks';
 import type { AvailableComponentProps } from '@/ui/types';
 
+/**
+ * Triggers a redirect if current task is not the given task key.
+ *
+ * If there's a current session, it will redirect to the `redirectUrlComplete` prop.
+ * If there's no current session, it will redirect to the sign in URL.
+ *
+ * @internal
+ */
 export const withTaskGuard = <P extends AvailableComponentProps>(
   Component: ComponentType<P>,
   taskKey: SessionTask['key'],
-) => {
+): ((props: P) => null | JSX.Element) => {
   const displayName = Component.displayName || Component.name || 'Component';
   Component.displayName = displayName;
 
diff --git a/packages/clerk-js/src/ui/contexts/ClerkUIComponentsContext.tsx b/packages/clerk-js/src/ui/contexts/ClerkUIComponentsContext.tsx
@@ -28,7 +28,11 @@ import {
   UserVerificationContext,
   WaitlistContext,
 } from './components';
-import { TaskChooseOrganizationContext, TaskResetPasswordContext } from './components/SessionTasks';
+import {
+  SessionTasksContext,
+  TaskChooseOrganizationContext,
+  TaskResetPasswordContext,
+} from './components/SessionTasks';
 
 export function ComponentContextProvider({
   componentName,
@@ -124,15 +128,19 @@ export function ComponentContextProvider({
         <TaskChooseOrganizationContext.Provider
           value={{ componentName: 'TaskChooseOrganization', ...(props as TaskChooseOrganizationProps) }}
         >
-          {children}
+          <SessionTasksContext.Provider value={{ ...(props as TaskChooseOrganizationProps) }}>
+            {children}
+          </SessionTasksContext.Provider>
         </TaskChooseOrganizationContext.Provider>
       );
     case 'TaskResetPassword':
       return (
         <TaskResetPasswordContext.Provider
           value={{ componentName: 'TaskResetPassword', ...(props as TaskResetPasswordProps) }}
         >
-          {children}
+          <SessionTasksContext.Provider value={{ ...(props as TaskResetPasswordProps) }}>
+            {children}
+          </SessionTasksContext.Provider>
         </TaskResetPasswordContext.Provider>
       );
     default:
diff --git a/packages/clerk-js/src/ui/contexts/components/SessionTasks.ts b/packages/clerk-js/src/ui/contexts/components/SessionTasks.ts
@@ -1,17 +1,45 @@
+import type { SessionResource } from '@clerk/shared/types';
 import { createContext, useContext } from 'react';
 
+import { getTaskEndpoint } from '@/core/sessionTasks';
+import { useRouter } from '@/ui/router';
+
 import type { SessionTasksCtx, TaskChooseOrganizationCtx, TaskResetPasswordCtx } from '../../types';
 
 export const SessionTasksContext = createContext<SessionTasksCtx | null>(null);
 
-export const useSessionTasksContext = (): SessionTasksCtx => {
+type SessionTasksContextType = SessionTasksCtx & {
+  navigateOnSetActive: (opts: { session: SessionResource; redirectUrlComplete: string }) => Promise<unknown>;
+};
+
+export const useSessionTasksContext = (): SessionTasksContextType => {
   const context = useContext(SessionTasksContext);
+  const { navigate, basePath, startPath } = useRouter();
 
   if (context === null) {
     throw new Error('Clerk: useSessionTasksContext called outside of the mounted SessionTasks component.');
   }
 
-  return context;
+  const navigateOnSetActive = async ({
+    session,
+    redirectUrlComplete,
+  }: {
+    session: SessionResource;
+    redirectUrlComplete: string;
+  }) => {
+    const currentTask = session.currentTask;
+    if (!currentTask) {
+      return navigate(redirectUrlComplete);
+    }
+
+    const taskEndpoint = getTaskEndpoint(currentTask);
+
+    // Base path is required for virtual routing with start path
+    // eg: to navigate from /sign-in/factor-one to /sign-in/tasks/choose-organization
+    return navigate(`/${basePath + startPath + taskEndpoint}`);
+  };
+
+  return { ...context, navigateOnSetActive };
 };
 
 export const TaskChooseOrganizationContext = createContext<TaskChooseOrganizationCtx | null>(null);
diff --git a/packages/clerk-js/src/ui/types.ts b/packages/clerk-js/src/ui/types.ts
@@ -12,7 +12,6 @@ import type {
   OrganizationProfileProps,
   OrganizationSwitcherProps,
   PricingTableProps,
-  SessionResource,
   SignInFallbackRedirectUrl,
   SignInForceRedirectUrl,
   SignInProps,
@@ -144,8 +143,6 @@ export type CheckoutCtx = __internal_CheckoutProps & {
 
 export type SessionTasksCtx = {
   redirectUrlComplete: string;
-  currentTaskContainer?: React.RefObject<HTMLDivElement> | null;
-  navigateOnSetActive: (opts: { session: SessionResource; redirectUrl: string }) => Promise<unknown>;
 };
 
 export type TaskChooseOrganizationCtx = TaskChooseOrganizationProps & {
