React Router login form needs to be submitted twice before router redirects to right route

I am trying to implement an authentication/login flow with React and React Router v6.

Root router

function Router() {
  const authContextApi = useAuthContextApi();

  console.log("Root router");

  const routesConfig = [
    {
      path: "/",
      element: <LandingRoute />,
      errorElement: <ErrorRoute />,
    },
    {
      path: "/auth/login",
      element: <LoginRoute />,
      action: loginAction(authContextApi),
    },
    {
      path: "/auth/register",
      element: <RegisterRoute />,
      action: registerAction,
    },
    {
      path: "/app",
      element: (
        <AuthRoute>
          <AppRoot /> // AppRoot is a wrapper with a bunch of layout components and <Outlet />
        </AuthRoute>
      ),
      children: [
        {
          path: "profile",
          element: <ProfileRoute />,
          loader: profileLoader,
        },
      ],
    },
  ];

  const router = createBrowserRouter(routesConfig);

  return <RouterProvider router={router} />;
}

AuthContext

function AuthContextProvider({ children }: React.PropsWithChildren) {
  const [userId, setUserId] = useState<string>("");

  console.log("Value of userId in AuthContext", userId);

  const contextValue = { userId };
  const api = { setUserId };

  return (
    <AuthContextApi.Provider value={api}>
      <AuthContext.Provider value={contextValue}>
        {children}
      </AuthContext.Provider>
    </AuthContextApi.Provider>
  );
}

The root router accesses the AuthContextApi to get the state setter function for the userId state variable. Then the setter function is passed as argument to the loginAction function which returns the actual action function with access to AuthContext. Here is the loginAction function:

const loginAction =
  (authContextApi: AuthContextApi): ActionFunction =>
  async ({ request }) => {
    const url = new URL(request.url);
    const redirectTo = url.searchParams.get("redirectTo");
    const formData = Object.fromEntries(await request.formData());
    
    console.log("After input validation");

    const response = "1";
    authContextApi?.setUserId(response);

    console.log("Before redirect");

    return redirect(redirectTo || "/app/profile");
  };

Login component

function LoginRoute() {
  const [searchParams] = useSearchParams();
  const redirectTo = searchParams.get("redirectTo");
  const url = `/auth/login${
    redirectTo ? `?redirectTo=${encodeURIComponent(redirectTo)}` : ""
  }`;

  return (
    <Form method="post" action={url}>
      <input type="email" name="email" placeholder="Email address" />
      <input type="password" name="password" placeholder="Password" />
      <button type="submit">Login</button>
    </Form>
  );
}

Authentication route

function AuthRoute({ children }: React.PropsWithChildren<{}>) {
  const { userId } = useAuthContext();
  const [searchParams] = useSearchParams();
  const redirectTo = searchParams.get("redirectTo");

  if (!userId) {
    console.log("UserId is empty string");

    const redirectUrl = `/auth/login${
      redirectTo ? `?redirectTo=${encodeURIComponent(redirectTo)}` : ""
    }`;

    return <Navigate to={redirectUrl} replace />;
  }

  console.log("UserId in AuthRoute", userId);

  return children;
}

When the user submits correct login input

• The URL changes but the matching route is not rendered
• The UI stays at the login route
• The authentication route does not render on redirection even though it is the parent route of the matched route

Then when the form is submitted again

• The authentication route renders
• The matched route renders

The first form submission should redirect to the specified route in the loginAction function. What am I missing?

React state updates are not instantaneous, so the enqueued state update authContextApi?.setUserId(response); is processed after the rest of the synchronous code in the function completes.

A bit hacky but you could introduce another asynchronous "tick" that could allow for the state update to be processed by placing the rest of the async function call to the end of the event queue.

Basic Example:

const loginAction =
  (authContextApi: AuthContextApi): ActionFunction =>
  async ({ request }) => {
    const url = new URL(request.url);
    const redirectTo = url.searchParams.get("redirectTo");
    const formData = Object.fromEntries(await request.formData());
    
    const response = "1";
    authContextApi?.setUserId(response);

    await new Promise(resolve => setTimeout(resolve, 1));

    return redirect(redirectTo || "/app/profile");
  };

You can play with the delay and add more if necessary.

Another thing that would help would be to update your auth logic to start from an undefined state and handle any pending/loading auth checks in the AuthRoute component.

Example:

function AuthContextProvider({ children }: React.PropsWithChildren) {
  const [userId, setUserId] = useState<string | undefined>();

  const contextValue = { userId };
  const api = { setUserId };

  return (
    <AuthContextApi.Provider value={api}>
      <AuthContext.Provider value={contextValue}>
        {children}
      </AuthContext.Provider>
    </AuthContextApi.Provider>
  );
}

function AuthRoute({ children }: React.PropsWithChildren<{}>) {
  const { userId } = useAuthContext();
  const [searchParams] = useSearchParams();

  const redirectTo = searchParams.get("redirectTo");

  // User auth not completed yet, wait to confirm
  if (userId === undefined) {
    return null; // or loading indicator/spinner/etc
  }

  const redirectUrl = `/auth/login${
    redirectTo ? `?redirectTo=${encodeURIComponent(redirectTo)}` : ""
  }`;

  // Here we know if user authenticated or not
  return userId ? children : <Navigate to={redirectUrl} replace />;
}

Of course now you will need to explicitly update the state for authenticated and unauthenticated users.

Example Implementation:

const loginAction =
  (authContextApi: AuthContextApi): ActionFunction =>
  async ({ request }) => {
    const url = new URL(request.url);
    const redirectTo = url.searchParams.get("redirectTo");
    const formData = Object.fromEntries(await request.formData());
    
    try {
      // ... Auth logic

      // authenticated, set to response value
      authContextApi?.setUserId(response);
    } catch(error) {
      // authentication error, set to defined unauthenticated value
      authContextApi?.setUserId("");
    }

    await new Promise(resolve => setTimeout(resolve, 1));

    return redirect(redirectTo || "/app/profile");
  };