Fix guardComponent resolution failure

Author: IzumiSy

packages/core/src/contexts/auth-context.test.tsx

@@ -78,7 +78,7 @@ describe("AuthProvider", () => {
       expect(screen.getByText("Test Content")).toBeDefined();
     });
 
-    it("should show guard component when not ready", () => {
+    it("should always render children even with guardComponent when not ready", () => {
       const state = {
         isAuthenticated: false,
         error: null,
@@ -92,11 +92,11 @@ describe("AuthProvider", () => {
         </AuthProvider>,
       );
 
-      expect(screen.getByText("Loading...")).toBeDefined();
-      expect(screen.queryByText("Protected Content")).toBeNull();
+      // AuthProvider always renders children; guard rendering is handled by the router
+      expect(screen.getByText("Protected Content")).toBeDefined();
     });
 
-    it("should show guard component when not authenticated", async () => {
+    it("should always render children even with guardComponent when not authenticated", async () => {
       const state = {
         isAuthenticated: false,
         error: null,
@@ -110,10 +110,8 @@ describe("AuthProvider", () => {
         </AuthProvider>,
       );
 
-      await waitFor(() => {
-        expect(screen.getByText("Please log in")).toBeDefined();
-      });
-      expect(screen.queryByText("Protected Content")).toBeNull();
+      // AuthProvider always renders children; guard rendering is handled by the router
+      expect(screen.getByText("Protected Content")).toBeDefined();
     });
 
     it("should show children when authenticated", async () => {
@@ -165,7 +163,7 @@ describe("AuthProvider", () => {
       });
 
       expect(result.current).not.toBeNull();
-      const response = await result.current!(new URL("http://localhost/"));
+      const response = await result.current!.loader(new URL("http://localhost/"));
       expect(mockCheckAuthStatus).toHaveBeenCalled();
       expect(response).toBeNull();
     });
@@ -187,7 +185,9 @@ describe("AuthProvider", () => {
       });
 
       expect(result.current).not.toBeNull();
-      const response = await result.current!(new URL("http://localhost/?code=auth-code-123"));
+      const response = await result.current!.loader(
+        new URL("http://localhost/?code=auth-code-123"),
+      );
       expect(mockHandleCallback).toHaveBeenCalled();
       expect(response).toBeNull();
     });
@@ -492,7 +492,7 @@ describe("AuthProvider", () => {
       });
 
       expect(result.current).not.toBeNull();
-      await result.current!(new URL("http://localhost/"));
+      await result.current!.loader(new URL("http://localhost/"));
       expect(mockCheckAuthStatus).toHaveBeenCalled();
       expect(mockLogin).toHaveBeenCalled();
     });

packages/core/src/contexts/auth-context.tsx

@@ -182,67 +182,78 @@ const AuthContext = createContext<AuthContextType | null>(null);
 type AuthClientContextType = {
   client: EnhancedAuthClient;
   autoLogin?: boolean;
+  guardComponent?: () => React.ReactNode;
 };
 
 const AuthClientContext = createContext<AuthClientContextType | null>(null);
 
 /** @internal */
 export type AuthLoader = (requestUrl: URL) => Promise<Response | null>;
 
+/** @internal */
+export type AuthLoaderContext = {
+  loader: AuthLoader;
+  guardComponent?: () => React.ReactNode;
+};
+
 /**
  * Internal hook that returns a loader function for handling OAuth callbacks
- * and auth status checks, or null if AuthProvider is not present.
+ * and auth status checks, along with the guard component configuration,
+ * or null if AuthProvider is not present.
  *
  * The returned loader is intended to be called from a react-router loader,
  * which runs outside React's lifecycle and is therefore unaffected by strict
  * mode double-invocation.
  *
- * @returns A loader function, or null if not inside an AuthProvider.
+ * @returns An object with loader and guardComponent, or null if not inside an AuthProvider.
  * @internal
  */
-export const useAuthLoader = (): AuthLoader | null => {
+export const useAuthLoader = (): AuthLoaderContext | null => {
   const authClientCtx = useContext(AuthClientContext);
   if (!authClientCtx) return null;
 
-  const { client, autoLogin } = authClientCtx;
-  return async (requestUrl: URL): Promise<Response | null> => {
-    // The "code" query parameter indicates a redirect back from the OAuth provider.
-    // handleCallback() internally cleans up the OAuth-related query parameters
-    // from the URL, so no additional URL cleanup is needed here.
-    if (requestUrl.searchParams.has("code")) {
-      try {
-        await client.handleCallback();
-      } catch (error) {
-        console.error("Failed to handle callback:", error);
+  const { client, autoLogin, guardComponent } = authClientCtx;
+  return {
+    guardComponent,
+    loader: async (requestUrl: URL): Promise<Response | null> => {
+      // The "code" query parameter indicates a redirect back from the OAuth provider.
+      // handleCallback() internally cleans up the OAuth-related query parameters
+      // from the URL, so no additional URL cleanup is needed here.
+      if (requestUrl.searchParams.has("code")) {
+        try {
+          await client.handleCallback();
+        } catch (error) {
+          console.error("Failed to handle callback:", error);
+        }
       }
-    }
-
-    // Only check auth status on first load (when isReady is false).
-    // Subsequent navigations skip this because the client already holds
-    // the cached auth state via useSyncExternalStore.
-    if (!client.getState().isReady) {
-      try {
-        await client.checkAuthStatus();
-      } catch (error) {
-        // Intentionally swallow errors to avoid rendering the error boundary
-        // on transient failures (e.g. network timeouts). The next navigation
-        // will re-run this loader and retry automatically.
-        console.error("Failed to check auth status:", error);
+
+      // Only check auth status on first load (when isReady is false).
+      // Subsequent navigations skip this because the client already holds
+      // the cached auth state via useSyncExternalStore.
+      if (!client.getState().isReady) {
+        try {
+          await client.checkAuthStatus();
+        } catch (error) {
+          // Intentionally swallow errors to avoid rendering the error boundary
+          // on transient failures (e.g. network timeouts). The next navigation
+          // will re-run this loader and retry automatically.
+          console.error("Failed to check auth status:", error);
+        }
       }
-    }
-
-    // autoLogin is evaluated separately from checkAuthStatus so that it
-    // still fires after handleCallback updates the internal state via
-    // getState() (e.g. setting isAuthenticated to true on success).
-    if (autoLogin && !client.getState().isAuthenticated) {
-      try {
-        await client.login();
-      } catch (error) {
-        console.error("Failed to login:", error);
+
+      // autoLogin is evaluated separately from checkAuthStatus so that it
+      // still fires after handleCallback updates the internal state via
+      // getState() (e.g. setting isAuthenticated to true on success).
+      if (autoLogin && !client.getState().isAuthenticated) {
+        try {
+          await client.login();
+        } catch (error) {
+          console.error("Failed to login:", error);
+        }
       }
-    }
 
-    return null;
+      return null;
+    },
   };
 };
 
@@ -310,17 +321,15 @@ export const AuthProvider = (props: React.PropsWithChildren<AuthProviderProps>)
 
   // Use useSyncExternalStore for state management
   const authState = useSyncExternalStore(subscribe, getSnapshot);
-  const isAuthenticated = authState.isAuthenticated;
-
-  const contents =
-    props.guardComponent && (!authState.isReady || !isAuthenticated) ? (
-      <props.guardComponent />
-    ) : (
-      props.children
-    );
 
   return (
-    <AuthClientContext.Provider value={{ client, autoLogin: props.autoLogin }}>
+    <AuthClientContext.Provider
+      value={{
+        client,
+        autoLogin: props.autoLogin,
+        guardComponent: props.guardComponent,
+      }}
+    >
       <AuthContext.Provider
         value={{
           authState,
@@ -330,7 +339,7 @@ export const AuthProvider = (props: React.PropsWithChildren<AuthProviderProps>)
           ready: () => client.ready(),
         }}
       >
-        {contents}
+        {props.children}
       </AuthContext.Provider>
     </AuthClientContext.Provider>
   );

packages/core/src/routing/router.test.tsx

@@ -1,5 +1,5 @@
 import { afterEach, describe, expect, it, vi } from "vitest";
-import { cleanup, fireEvent, render, screen, waitFor } from "@testing-library/react";
+import { act, cleanup, fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { RouterContainer } from "./router";
 import { AppShellConfigContext, AppShellDataContext } from "@/contexts/appshell-context";
 import { Link, Outlet, useNavigate } from "react-router";
@@ -27,6 +27,7 @@ const renderWithConfig = ({
   contextData = {},
   authClient,
   autoLogin,
+  guardComponent,
 }: {
   modules?: Array<Module>;
   basePath?: string;
@@ -35,6 +36,7 @@ const renderWithConfig = ({
   contextData?: ContextData;
   authClient?: EnhancedAuthClient;
   autoLogin?: boolean;
+  guardComponent?: () => React.ReactNode;
 }) => {
   const configurations = {
     modules,
@@ -58,7 +60,7 @@ const renderWithConfig = ({
 
   render(
     authClient ? (
-      <AuthProvider client={authClient} autoLogin={autoLogin}>
+      <AuthProvider client={authClient} autoLogin={autoLogin} guardComponent={guardComponent}>
         {tree}
       </AuthProvider>
     ) : (
@@ -492,4 +494,129 @@ describe("RouterContainer with AuthProvider", () => {
     await screen.findByText("Home");
     expect(mockLogin).not.toHaveBeenCalled();
   });
+
+  it("shows guard component when not ready", async () => {
+    const authClient = createMockAuthClient({
+      isAuthenticated: false,
+      error: null,
+      isReady: false,
+    });
+
+    renderWithConfig({
+      modules: [],
+      rootComponent: () => <div>Home</div>,
+      initialEntries: ["/"],
+      authClient,
+      guardComponent: () => <div>Loading...</div>,
+    });
+
+    expect(await screen.findByText("Loading...")).toBeDefined();
+    expect(screen.queryByText("Home")).toBeNull();
+  });
+
+  it("shows guard component when not authenticated", async () => {
+    const authClient = createMockAuthClient({
+      isAuthenticated: false,
+      error: null,
+      isReady: true,
+    });
+
+    renderWithConfig({
+      modules: [],
+      rootComponent: () => <div>Home</div>,
+      initialEntries: ["/"],
+      authClient,
+      guardComponent: () => <div>Please log in</div>,
+    });
+
+    expect(await screen.findByText("Please log in")).toBeDefined();
+    expect(screen.queryByText("Home")).toBeNull();
+  });
+
+  it("shows children when authenticated with guardComponent", async () => {
+    const mockCheckAuthStatus = vi.fn().mockResolvedValue({
+      isAuthenticated: true,
+      error: null,
+      isReady: true,
+    });
+    const authClient = createMockAuthClient(
+      { isAuthenticated: true, error: null, isReady: true },
+      { checkAuthStatus: mockCheckAuthStatus },
+    );
+
+    renderWithConfig({
+      modules: [],
+      rootComponent: () => <div>Home</div>,
+      initialEntries: ["/"],
+      authClient,
+      guardComponent: () => <div>Please log in</div>,
+    });
+
+    expect(await screen.findByText("Home")).toBeDefined();
+    expect(screen.queryByText("Please log in")).toBeNull();
+  });
+
+  it("transitions from guard to children when auth state changes", async () => {
+    // Mutable snapshot; initially not ready, not authenticated.
+    // The loader calls checkAuthStatus, but the mock does NOT update the
+    // snapshot during the loader — so the guard is shown on first render.
+    let snapshot = {
+      isAuthenticated: false,
+      error: null as string | null,
+      isReady: false,
+    };
+
+    // Collect listeners so we can trigger state change notifications
+    const listeners: Array<(event: { type: string }) => void> = [];
+
+    // checkAuthStatus resolves without updating snapshot, mimicking a
+    // still-pending auth check from the guard's perspective.
+    const mockCheckAuthStatus = vi.fn().mockResolvedValue({
+      isAuthenticated: false,
+      error: null,
+      isReady: true,
+    });
+
+    const authClient = createMockAuthClient(snapshot, {
+      checkAuthStatus: mockCheckAuthStatus,
+      // Return the same reference until we reassign snapshot (required by useSyncExternalStore)
+      getState: vi.fn(() => snapshot),
+      addEventListener: vi.fn((listener) => {
+        listeners.push(listener);
+        return () => {
+          const idx = listeners.indexOf(listener);
+          if (idx >= 0) listeners.splice(idx, 1);
+        };
+      }),
+    });
+
+    renderWithConfig({
+      modules: [],
+      rootComponent: () => <div>Home</div>,
+      initialEntries: ["/"],
+      authClient,
+      guardComponent: () => <div>Loading...</div>,
+    });
+
+    // Initially the guard should be shown
+    expect(await screen.findByText("Loading...")).toBeDefined();
+    expect(screen.queryByText("Home")).toBeNull();
+
+    // Simulate auth state becoming ready and authenticated (e.g. token refresh
+    // or login flow completing outside the loader).
+    act(() => {
+      snapshot = {
+        isAuthenticated: true,
+        error: null,
+        isReady: true,
+      };
+      for (const listener of listeners) {
+        listener({ type: "auth_state_changed" });
+      }
+    });
+
+    // After auth state transitions, children should replace the guard
+    expect(await screen.findByText("Home")).toBeDefined();
+    expect(screen.queryByText("Loading...")).toBeNull();
+  });
 });