Use sync authToken in AuthKit React

react/src/lib/auth-context.ts

@@ -0,0 +1,10 @@
+import * as React from "react";
+
+interface AuthContextType {
+  authToken: string | null;
+}
+
+const AuthContext = React.createContext<AuthContextType | null>(null);
+AuthContext.displayName = "AuthContext";
+
+export { AuthContext, type AuthContextType };

react/src/lib/auth-provider.tsx

@@ -0,0 +1,22 @@
+import * as React from "react";
+import { AuthKitProvider } from "@workos-inc/authkit-react";
+import { AuthContext } from "./auth-context";
+
+export const AuthProvider: React.FC<
+  React.ComponentProps<typeof AuthKitProvider>
+> = ({ children, ...props }) => {
+  const [authToken, setAuthToken] = React.useState<string | null>(null);
+  return (
+    <AuthKitProvider
+      {...props}
+      onRefresh={(payload) => {
+        props.onRefresh?.(payload);
+        setAuthToken(payload.accessToken);
+      }}
+    >
+      <AuthContext.Provider value={{ authToken }}>
+        {children}
+      </AuthContext.Provider>
+    </AuthKitProvider>
+  );
+};

react/src/lib/use-auth.ts

@@ -0,0 +1,83 @@
+import * as React from "react";
+import { useAuth as useAuthKit } from "@workos-inc/authkit-react";
+import { AuthContext } from "./auth-context";
+
+type AuthKit = ReturnType<typeof useAuthKit>;
+type User = Exclude<AuthKit["user"], null>;
+
+type UseAuthReturnType = Omit<AuthKit, "user" | "isLoading"> & State;
+
+export function useAuth(): UseAuthReturnType {
+  const authKit = useAuthKit();
+  const context = React.useContext(AuthContext);
+  if (!context) {
+    throw new Error("useAuth must be used within an AuthProvider");
+  }
+
+  if (authKit.isLoading) {
+    return {
+      ...authKit,
+      ...context,
+      status: "loading",
+      isLoading: true,
+      user: null,
+      authToken: null,
+    };
+  }
+
+  if (!authKit.user) {
+    return {
+      ...authKit,
+      ...context,
+      status: "unauthenticated",
+      isLoading: false,
+      user: null,
+      authToken: null,
+    };
+  }
+
+  return {
+    ...authKit,
+    ...context,
+    status: "authenticated",
+    isLoading: false,
+    user: authKit.user,
+    authToken: context.authToken,
+  };
+}
+
+type Status =
+  | "loading"
+  | "authenticated"
+  | "authenticated-with-token"
+  | "unauthenticated";
+
+interface SharedState {
+  status: Status;
+  isLoading: boolean;
+  user: User | null;
+  authToken: string | null;
+}
+
+interface LoadingState extends SharedState {
+  status: "loading";
+  isLoading: true;
+  user: null;
+  authToken: null;
+}
+
+interface AuthenticatedState extends SharedState {
+  status: "authenticated";
+  isLoading: false;
+  user: User;
+  authToken: string | null;
+}
+
+interface UnauthenticatedState extends SharedState {
+  status: "unauthenticated";
+  isLoading: false;
+  user: null;
+  authToken: null;
+}
+
+export type State = LoadingState | AuthenticatedState | UnauthenticatedState;

react/src/root.tsx

@@ -1,8 +1,6 @@
 import { Flex, Heading, Link, Text, Theme } from "@radix-ui/themes";
 import { ErrorBoundary as ReactErrorBoundary } from "react-error-boundary";
 import {
-  Links,
-  Meta,
   Outlet,
   Link as RouterLink,
   ScrollRestoration,
@@ -12,7 +10,7 @@ import {
 import * as React from "react";
 import { preconnect } from "react-dom";
 import { getApiProps, getDevtools } from "~/lib/env";
-import { AuthKitProvider } from "@workos-inc/authkit-react";
+import { AuthProvider } from "~/lib/auth-provider";
 import { WorkOsWidgets } from "@workos-inc/widgets";
 import { ReactQueryDevtools } from "@tanstack/react-query-devtools";
 
@@ -40,13 +38,13 @@ function Providers({ children }: { children: React.ReactNode }) {
         </Theme>
       )}
     >
-      <AuthKitProvider {...getApiProps()}>
+      <AuthProvider {...getApiProps()}>
         <WorkOsWidgets {...getApiProps()}>
           <Theme data-is-root-theme="false" style={{ height: "100%" }}>
             {children}
           </Theme>
         </WorkOsWidgets>
-      </AuthKitProvider>
+      </AuthProvider>
     </ReactErrorBoundary>
   );
 }

react/src/routes/_auth.signin.tsx

@@ -1,6 +1,6 @@
 import { Button, Flex, Heading } from "@radix-ui/themes";
 import { Link } from "react-router";
-import { useAuth } from "@workos-inc/authkit-react";
+import { useAuth } from "~/lib/use-auth";
 
 export default function HomePage() {
   const auth = useAuth();

react/src/routes/_dashboard.tsx

@@ -11,10 +11,10 @@ import {
   Text,
   VisuallyHidden,
 } from "@radix-ui/themes";
-import { useAuth } from "@workos-inc/authkit-react";
+import { useAuth } from "~/lib/use-auth";
 import { Outlet, useNavigation } from "react-router";
-import { PrimaryNav, PrimaryNavItem } from "src/ui/dashboard";
-import { SuperAppMark } from "src/ui/svgs";
+import { PrimaryNav, PrimaryNavItem } from "~/ui/dashboard";
+import { SuperAppMark } from "~/ui/svgs";
 
 export default function DashboardLayout() {
   const navigation = useNavigation();

react/src/routes/_dashboard.users.tsx

@@ -1,16 +1,19 @@
-import { UsersManagement } from "@workos-inc/widgets";
-import { Flex, Heading, Separator, Text } from "@radix-ui/themes";
+import { UsersManagement, UsersManagementLoading } from "@workos-inc/widgets";
+import { Button, Flex, Heading, Separator, Text } from "@radix-ui/themes";
 import { Navigate } from "react-router";
-import { useAuth } from "@workos-inc/authkit-react";
+import { useAuth } from "~/lib/use-auth";
 
 export default function Users() {
   const auth = useAuth();
 
   if (auth.isLoading) {
     return (
-      <Flex justify="center" align="center" height="100%">
-        Authenticating...
-      </Flex>
+      <>
+        <title>SuperApp | Users</title>
+        <Layout subtitle="Manage and invite users for the SuperApp team.">
+          <UsersManagementLoading />
+        </Layout>
+      </>
     );
   }
 
@@ -22,41 +25,48 @@ export default function Users() {
     return (
       <>
         <title>SuperApp | No organization</title>
-        <Flex direction="column" align="center" px="9">
-          <Flex direction="column" gap="5" maxWidth="640px" width="100%" my="9">
-            <main>
-              <Heading size="8" mb="2">
-                User does not belong to an organization
-              </Heading>
-              <Text asChild>
-                <button type="button" onClick={() => void auth.signOut()}>
-                  Sign out
-                </button>
-              </Text>
-            </main>
-          </Flex>
-        </Flex>
+        <Layout subtitle="User does not belong to an organization.">
+          <Button type="button" onClick={() => void auth.signOut()}>
+            Sign out
+          </Button>
+        </Layout>
       </>
     );
   }
 
   return (
     <>
       <title>SuperApp | Users</title>
-      <Flex direction="column" align="center" px="9">
-        <Flex direction="column" gap="5" maxWidth="640px" width="100%" my="9">
-          <header>
-            <Heading size="8" mb="2">
-              Users
-            </Heading>
-            <Text>Manage and invite users for the SuperApp team</Text>
-          </header>
-          <Separator size="4" />
-          <main>
-            <UsersManagement authToken={auth.getAccessToken} />
-          </main>
-        </Flex>
-      </Flex>
+      <Layout subtitle="Manage and invite users for the SuperApp team.">
+        {auth.authToken ? (
+          <UsersManagement authToken={auth.authToken} />
+        ) : (
+          <UsersManagementLoading />
+        )}
+      </Layout>
     </>
   );
 }
+
+function Layout({
+  children,
+  subtitle,
+}: {
+  children: React.ReactNode;
+  subtitle: string;
+}) {
+  return (
+    <Flex direction="column" align="center" px="9">
+      <Flex direction="column" gap="5" maxWidth="640px" width="100%" my="9">
+        <header>
+          <Heading size="8" mb="2">
+            Users
+          </Heading>
+          <Text>{subtitle}</Text>
+        </header>
+        <Separator size="4" />
+        <main>{children}</main>
+      </Flex>
+    </Flex>
+  );
+}