@imtbl/auth-next-client 2.12.7-alpha.0 → 2.12.7-alpha.10

package/src/hooks.tsx

@@ -1,6 +1,8 @@
 'use client';
 
-import { useCallback, useRef, useState } from 'react';
+import {
+  useCallback, useEffect, useRef, useState,
+} from 'react';
 import { useSession, signIn, signOut } from 'next-auth/react';
 import type { Session } from 'next-auth';
 import type {
@@ -16,12 +18,35 @@ import {
   loginWithRedirect as rawLoginWithRedirect,
   logoutWithRedirect as rawLogoutWithRedirect,
 } from '@imtbl/auth';
-import { IMMUTABLE_PROVIDER_ID } from './constants';
+import { IMMUTABLE_PROVIDER_ID, TOKEN_EXPIRY_BUFFER_MS } from './constants';
+import { storeIdToken, getStoredIdToken, clearStoredIdToken } from './idTokenStorage';
+
+// ---------------------------------------------------------------------------
+// Module-level deduplication for session refresh
+// ---------------------------------------------------------------------------
+
+/**
+ * Deduplicates concurrent session refresh calls.
+ * Multiple components may mount useImmutableSession simultaneously; without
+ * deduplication each would trigger its own update() call, which could fail
+ * if the auth server rotates refresh tokens.
+ */
+let pendingRefresh: Promise<Session | null | undefined> | null = null;
+
+function deduplicatedUpdate(
+  update: () => Promise<Session | null | undefined>,
+): Promise<Session | null | undefined> {
+  if (!pendingRefresh) {
+    pendingRefresh = update().finally(() => { pendingRefresh = null; });
+  }
+  return pendingRefresh;
+}
 
 /**
- * Extended session type with Immutable token data
+ * Internal session type with full token data (not exported).
+ * Used internally by the hook for token validation, refresh logic, and getUser/getAccessToken.
  */
-export interface ImmutableSession extends Session {
+interface ImmutableSessionInternal extends Session {
   accessToken: string;
   refreshToken?: string;
   idToken?: string;
@@ -30,6 +55,15 @@ export interface ImmutableSession extends Session {
   error?: string;
 }
 
+/**
+ * Public session type exposed to consumers.
+ *
+ * Does **not** include `accessToken` -- consumers must use the `getAccessToken()`
+ * function returned by `useImmutableSession()` to obtain a guaranteed-fresh token.
+ * This prevents accidental use of stale/expired tokens.
+ */
+export type ImmutableSession = Omit<ImmutableSessionInternal, 'accessToken'>;
+
 /**
  * Return type for useImmutableSession hook
  */
@@ -53,6 +87,13 @@ export interface UseImmutableSessionReturn {
    *   The refreshed session will include updated zkEvm data if available.
    */
   getUser: (forceRefresh?: boolean) => Promise<User | null>;
+  /**
+   * Get a guaranteed-fresh access token.
+   * Returns immediately if the current token is valid.
+   * If expired, triggers a refresh and blocks (awaits) until the fresh token is available.
+   * Throws if the user is not authenticated or if refresh fails.
+   */
+  getAccessToken: () => Promise<string>;
 }
 
 /**
@@ -92,25 +133,29 @@ export function useImmutableSession(): UseImmutableSessionReturn {
   // Track when a manual refresh is in progress (via getUser(true))
   const [isRefreshing, setIsRefreshing] = useState(false);
 
-  // Cast session to our extended type
-  const session = sessionData as ImmutableSession | null;
+  // Cast session to our internal type (includes accessToken for internal logic)
+  const session = sessionData as ImmutableSessionInternal | null;
 
   const isLoading = status === 'loading';
 
-  // Core authentication check from NextAuth
-  const hasSession = status === 'authenticated' && !!session;
+  // Core authentication check - user has a valid session with usable access token.
+  // A session can exist but be unusable if the access token is missing or refresh failed.
+  const hasValidSession = status === 'authenticated'
+    && !!session
+    && !!session.accessToken
+    && !session.error;
 
-  // During loading/refreshing, keep showing authenticated if we had a session (avoids UI flicker
+  // During loading/refreshing, keep showing authenticated if we had a valid session (avoids UI flicker
   // when NextAuth refetches on window focus or after getUser(forceRefresh)).
   const hadSessionRef = useRef(false);
-  if (hasSession) hadSessionRef.current = true;
-  if (!hasSession && !isLoading && !isRefreshing) hadSessionRef.current = false;
-  const isAuthenticated = hasSession || ((isLoading || isRefreshing) && hadSessionRef.current);
+  if (hasValidSession) hadSessionRef.current = true;
+  if (!hasValidSession && !isLoading && !isRefreshing) hadSessionRef.current = false;
+  const isAuthenticated = hasValidSession || ((isLoading || isRefreshing) && hadSessionRef.current);
 
   // Use a ref to always have access to the latest session.
   // This avoids stale closure issues when the wallet stores the getUser function
   // and calls it later - the ref always points to the current session.
-  const sessionRef = useRef<ImmutableSession | null>(session);
+  const sessionRef = useRef<ImmutableSessionInternal | null>(session);
   sessionRef.current = session;
 
   // Also store update in a ref so the callback is stable
@@ -121,6 +166,44 @@ export function useImmutableSession(): UseImmutableSessionReturn {
   const setIsRefreshingRef = useRef(setIsRefreshing);
   setIsRefreshingRef.current = setIsRefreshing;
 
+  // ---------------------------------------------------------------------------
+  // Proactive token refresh
+  // ---------------------------------------------------------------------------
+
+  // Reactive refresh: when the effect runs and the token is already expired
+  // (e.g., after tab regains focus), trigger an immediate silent refresh.
+  // For tokens that are still valid, getAccessToken() handles refresh on demand.
+  //
+  // NOTE: This intentionally does NOT set isRefreshing. isRefreshing is reserved
+  // for explicit user-triggered refreshes (e.g., getUser(true) after wallet
+  // registration). Background token refreshes must be invisible to consumers --
+  // setting isRefreshing would cause downstream hooks that gate SWR keys on
+  // `!isRefreshing` to briefly lose their cached data, resulting in UI flicker.
+  useEffect(() => {
+    if (!session?.accessTokenExpires) return;
+
+    const timeUntilExpiry = session.accessTokenExpires - Date.now() - TOKEN_EXPIRY_BUFFER_MS;
+
+    if (timeUntilExpiry <= 0) {
+      // Already expired -- refresh silently
+      deduplicatedUpdate(() => updateRef.current());
+    }
+  }, [session?.accessTokenExpires]);
+
+  // ---------------------------------------------------------------------------
+  // Sync idToken to localStorage
+  // ---------------------------------------------------------------------------
+
+  // The idToken is stripped from the cookie by jwt.encode on the server to avoid
+  // CloudFront 413 errors. It is only present in the session response transiently
+  // after sign-in or token refresh. When present, persist it in localStorage so
+  // that getUser() can always return it (used by wallet's MagicTEESigner).
+  useEffect(() => {
+    if (session?.idToken) {
+      storeIdToken(session.idToken);
+    }
+  }, [session?.idToken]);
+
   /**
    * Get user function for wallet integration.
    * Returns a User object compatible with @imtbl/wallet's getUser option.
@@ -131,7 +214,7 @@ export function useImmutableSession(): UseImmutableSessionReturn {
    * @param forceRefresh - When true, triggers a server-side token refresh
    */
   const getUser = useCallback(async (forceRefresh?: boolean): Promise<User | null> => {
-    let currentSession: ImmutableSession | null;
+    let currentSession: ImmutableSessionInternal | null;
 
     // If forceRefresh is requested, trigger server-side refresh via NextAuth
     // This calls the jwt callback with trigger='update' and sessionUpdate.forceRefresh=true
@@ -141,10 +224,14 @@ export function useImmutableSession(): UseImmutableSessionReturn {
       try {
         // update() returns the refreshed session
         const updatedSession = await updateRef.current({ forceRefresh: true });
-        currentSession = updatedSession as ImmutableSession | null;
+        currentSession = updatedSession as ImmutableSessionInternal | null;
         // Also update the ref so subsequent calls get the fresh data
         if (currentSession) {
           sessionRef.current = currentSession;
+          // Immediately persist fresh idToken to localStorage (avoids race with useEffect)
+          if (currentSession.idToken) {
+            storeIdToken(currentSession.idToken);
+          }
         }
       } catch (error) {
         // eslint-disable-next-line no-console
@@ -154,6 +241,20 @@ export function useImmutableSession(): UseImmutableSessionReturn {
       } finally {
         setIsRefreshingRef.current(false);
       }
+    } else if (pendingRefresh) {
+      // If a refresh is in-flight (proactive timer or another getAccessToken call),
+      // wait for it and use the refreshed session rather than returning a stale token.
+      const refreshed = await pendingRefresh;
+      if (refreshed) {
+        currentSession = refreshed as ImmutableSessionInternal;
+        sessionRef.current = currentSession;
+        // Persist fresh idToken to localStorage immediately
+        if (currentSession.idToken) {
+          storeIdToken(currentSession.idToken);
+        }
+      } else {
+        currentSession = sessionRef.current;
+      }
     } else {
       // Read from ref - instant, no network call
       // The ref is always updated on each render with the latest session
@@ -174,7 +275,9 @@ export function useImmutableSession(): UseImmutableSessionReturn {
     return {
       accessToken: currentSession.accessToken,
       refreshToken: currentSession.refreshToken,
-      idToken: currentSession.idToken,
+      // Prefer session idToken (fresh after sign-in or refresh, before useEffect
+      // stores it), fall back to localStorage for normal reads (cookie has no idToken).
+      idToken: currentSession.idToken || getStoredIdToken(),
       profile: {
         sub: currentSession.user?.sub ?? '',
         email: currentSession.user?.email ?? undefined,
@@ -184,13 +287,55 @@ export function useImmutableSession(): UseImmutableSessionReturn {
     };
   }, []); // Empty deps - uses refs for latest values
 
+  /**
+   * Get a guaranteed-fresh access token.
+   * Returns immediately if the current token is valid (fast path, no network call).
+   * If expired, triggers a server-side refresh and blocks (awaits) until the fresh
+   * token is available. Piggybacks on any in-flight refresh to avoid duplicate calls.
+   *
+   * @throws Error if the user is not authenticated or if the refresh fails.
+   */
+  const getAccessToken = useCallback(async (): Promise<string> => {
+    const currentSession = sessionRef.current;
+
+    // Fast path: token is valid -- return immediately
+    if (
+      currentSession?.accessToken
+      && currentSession.accessTokenExpires
+      && Date.now() < currentSession.accessTokenExpires - TOKEN_EXPIRY_BUFFER_MS
+      && !currentSession.error
+    ) {
+      return currentSession.accessToken;
+    }
+
+    // Token is expired or missing -- wait for in-flight refresh or trigger one
+    const refreshed = await deduplicatedUpdate(
+      () => updateRef.current(),
+    ) as ImmutableSessionInternal | null;
+
+    if (!refreshed?.accessToken || refreshed.error) {
+      throw new Error(
+        `[auth-next-client] Failed to get access token: ${refreshed?.error || 'no session'}`,
+      );
+    }
+
+    // Update ref so subsequent sync reads get the fresh data
+    sessionRef.current = refreshed;
+    return refreshed.accessToken;
+  }, []); // Empty deps -- uses refs for latest values
+
+  // Cast to public type (omits accessToken) to prevent consumers from
+  // accidentally using a potentially stale token. Use getAccessToken() instead.
+  const publicSession = session as ImmutableSession | null;
+
   return {
-    session,
+    session: publicSession,
     status,
     isLoading,
     isAuthenticated,
     isRefreshing,
     getUser,
+    getAccessToken,
   };
 }
 
@@ -267,6 +412,12 @@ export function useLogin(): UseLoginReturn {
     profile: { sub: string; email?: string; nickname?: string };
     zkEvm?: ZkEvmInfo;
   }) => {
+    // Persist idToken to localStorage before signIn so it's available immediately.
+    // The cookie won't contain idToken (stripped by jwt.encode on the server).
+    if (tokens.idToken) {
+      storeIdToken(tokens.idToken);
+    }
+
     const result = await signIn(IMMUTABLE_PROVIDER_ID, {
       tokens: JSON.stringify(tokens),
       redirect: false,
@@ -430,6 +581,9 @@ export function useLogout(): UseLogoutReturn {
     setError(null);
 
     try {
+      // Clear idToken from localStorage before clearing session
+      clearStoredIdToken();
+
       // First, clear the NextAuth session (this clears the JWT cookie)
       // We use redirect: false to handle the redirect ourselves for federated logout
       await signOut({ redirect: false });

package/src/idTokenStorage.ts

@@ -0,0 +1,56 @@
+/**
+ * Utility for persisting idToken in localStorage.
+ *
+ * The idToken is stripped from the NextAuth session cookie (via a custom
+ * jwt.encode in @imtbl/auth-next-server) to keep cookie size under CDN header
+ * limits (CloudFront 20 KB). Instead, the client stores idToken in
+ * localStorage so that wallet operations (e.g., MagicTEESigner) can still
+ * access it via getUser().
+ *
+ * All functions are safe to call during SSR or in restricted environments
+ * (e.g., incognito mode with localStorage disabled) -- they silently no-op.
+ */
+
+const ID_TOKEN_STORAGE_KEY = 'imtbl_id_token';
+
+/**
+ * Store the idToken in localStorage.
+ * @param idToken - The raw ID token JWT string
+ */
+export function storeIdToken(idToken: string): void {
+  try {
+    if (typeof window !== 'undefined' && window.localStorage) {
+      window.localStorage.setItem(ID_TOKEN_STORAGE_KEY, idToken);
+    }
+  } catch {
+    // Silently ignore -- localStorage may be unavailable (SSR, incognito, etc.)
+  }
+}
+
+/**
+ * Retrieve the idToken from localStorage.
+ * @returns The stored idToken, or undefined if not available.
+ */
+export function getStoredIdToken(): string | undefined {
+  try {
+    if (typeof window !== 'undefined' && window.localStorage) {
+      return window.localStorage.getItem(ID_TOKEN_STORAGE_KEY) ?? undefined;
+    }
+  } catch {
+    // Silently ignore
+  }
+  return undefined;
+}
+
+/**
+ * Remove the idToken from localStorage (e.g., on logout).
+ */
+export function clearStoredIdToken(): void {
+  try {
+    if (typeof window !== 'undefined' && window.localStorage) {
+      window.localStorage.removeItem(ID_TOKEN_STORAGE_KEY);
+    }
+  } catch {
+    // Silently ignore
+  }
+}