Fix: Show loading state during initial token fetch to prevent flash

__tests__/useAccessToken.spec.tsx

@@ -57,15 +57,15 @@ describe('useAccessToken', () => {
     );
   };
 
-  it('should fetch an access token on mount without showing loading state', async () => {
+  it('should fetch an access token on mount and show loading state initially', async () => {
     const mockToken =
       'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwic2lkIjoic2Vzc2lvbl8xMjMiLCJleHAiOjk5OTk5OTk5OTl9.mock-signature';
     (getAccessTokenAction as jest.Mock).mockResolvedValueOnce(mockToken);
 
     const { getByTestId } = render(<TestComponent />);
 
-    // Loading should remain false for background fetches
-    expect(getByTestId('loading')).toHaveTextContent('false');
+    // Loading should be true during initial fetch
+    expect(getByTestId('loading')).toHaveTextContent('true');
 
     await waitFor(() => {
       expect(getAccessTokenAction).toHaveBeenCalledTimes(1);
@@ -77,7 +77,7 @@ describe('useAccessToken', () => {
     });
   });
 
-  it('should handle token refresh when an expiring token is received without showing loading', async () => {
+  it('should handle token refresh when an expiring token is received', async () => {
     // Create a token that's about to expire (exp is very close to current time)
     const currentTimeInSeconds = Math.floor(Date.now() / 1000);
     // Use 25 seconds to ensure it's within the 30-second buffer for short-lived tokens
@@ -97,8 +97,8 @@ describe('useAccessToken', () => {
 
     const { getByTestId } = render(<TestComponent />);
 
-    // Loading should remain false throughout
-    expect(getByTestId('loading')).toHaveTextContent('false');
+    // Loading should be true initially during token fetch
+    expect(getByTestId('loading')).toHaveTextContent('true');
 
     await waitFor(() => {
       expect(getByTestId('loading')).toHaveTextContent('false');
@@ -155,14 +155,14 @@ describe('useAccessToken', () => {
     });
   });
 
-  it('should handle errors during token fetch without showing loading', async () => {
+  it('should handle errors during token fetch', async () => {
     const error = new Error('Failed to fetch token');
     (getAccessTokenAction as jest.Mock).mockRejectedValueOnce(error);
 
     const { getByTestId } = render(<TestComponent />);
 
-    // Loading should remain false even when there's an error
-    expect(getByTestId('loading')).toHaveTextContent('false');
+    // Loading should be true initially
+    expect(getByTestId('loading')).toHaveTextContent('true');
 
     await waitFor(() => {
       expect(getByTestId('loading')).toHaveTextContent('false');
@@ -381,8 +381,8 @@ describe('useAccessToken', () => {
 
     const { getByTestId } = render(<TestComponent />);
 
-    // Loading should remain false for background operations
-    expect(getByTestId('loading')).toHaveTextContent('false');
+    // Loading should be true initially during fetch
+    expect(getByTestId('loading')).toHaveTextContent('true');
 
     await waitFor(() => {
       expect(fetchCalls).toBe(1);
@@ -459,6 +459,44 @@ describe('useAccessToken', () => {
     });
   });
 
+  it('should show loading state immediately on first render when user exists but no token', () => {
+    // Mock user with no token initially
+    (useAuth as jest.Mock).mockImplementation(() => ({
+      user: { id: 'user_123' },
+      sessionId: 'session_123',
+      refreshAuth: jest.fn().mockResolvedValue({}),
+    }));
+
+    (getAccessTokenAction as jest.Mock).mockImplementation(
+      () => new Promise((resolve) => setTimeout(() => resolve('token'), 100)),
+    );
+
+    const { getByTestId } = render(<TestComponent />);
+
+    expect(getByTestId('loading')).toHaveTextContent('true');
+    expect(getByTestId('token')).toHaveTextContent('no-token');
+  });
+
+  it('should not show loading when a valid token already exists', async () => {
+    const existingToken =
+      'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJleGlzdGluZyIsInNpZCI6InNlc3Npb24xMjMiLCJleHAiOjk5OTk5OTk5OTl9.existing';
+
+    await act(async () => {
+      (getAccessTokenAction as jest.Mock).mockResolvedValueOnce(existingToken);
+      await tokenStore.getAccessTokenSilently();
+    });
+
+    // Reset the mock to track new calls
+    (getAccessTokenAction as jest.Mock).mockClear();
+
+    const { getByTestId } = render(<TestComponent />);
+
+    expect(getByTestId('loading')).toHaveTextContent('false');
+    expect(getByTestId('token')).toHaveTextContent(existingToken);
+
+    expect(getAccessTokenAction).not.toHaveBeenCalled();
+  });
+
   // Additional test cases to increase coverage
   it('should handle concurrent manual refresh attempts', async () => {
     const initialToken =

src/components/useAccessToken.ts

@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useRef, useSyncExternalStore } from 'react';
+import { useCallback, useEffect, useRef, useState, useSyncExternalStore } from 'react';
 import { useAuth } from './authkit-provider.js';
 import { tokenStore } from './tokenStore.js';
 
@@ -42,9 +42,17 @@ export function useAccessToken(): UseAccessTokenReturn {
 
   const tokenState = useSyncExternalStore(tokenStore.subscribe, tokenStore.getSnapshot, tokenStore.getServerSnapshot);
 
+  // Track if we're waiting for the initial token fetch for the current user
+  // Initialize synchronously to prevent first-paint flash
+  const [isInitialTokenLoading, setIsInitialTokenLoading] = useState(() => {
+    // Only show loading if we have a user but no token yet
+    return Boolean(user && !tokenState.token && !tokenState.error);
+  });
+
   useEffect(() => {
     if (!user) {
       tokenStore.clearToken();
+      setIsInitialTokenLoading(false);
       return;
     }
 
@@ -59,10 +67,28 @@ export function useAccessToken(): UseAccessTokenReturn {
     prevSessionRef.current = sessionId;
     prevUserIdRef.current = userId;
 
+    // Check if getAccessTokenSilently will actually fetch (not just return cached)
+    const currentToken = tokenStore.getSnapshot().token;
+    const tokenData = currentToken ? tokenStore.parseToken(currentToken) : null;
+    const willActuallyFetch = !currentToken || (tokenData && tokenData.isExpiring);
+
+    // Only show loading if we're actually going to fetch
+    if (willActuallyFetch) {
+      setIsInitialTokenLoading(true);
+    }
+
     /* istanbul ignore next */
-    tokenStore.getAccessTokenSilently().catch(() => {
-      // Error is handled in the store
-    });
+    tokenStore
+      .getAccessTokenSilently()
+      .catch(() => {
+        // Error is handled in the store
+      })
+      .finally(() => {
+        // Only clear loading if we were actually loading
+        if (willActuallyFetch) {
+          setIsInitialTokenLoading(false);
+        }
+      });
   }, [userId, sessionId]);
 
   useEffect(() => {
@@ -112,9 +138,12 @@ export function useAccessToken(): UseAccessTokenReturn {
     return tokenStore.refreshToken();
   }, []);
 
+  // Combine loading states: initial token fetch OR token store is loading
+  const isLoading = isInitialTokenLoading || tokenState.loading;
+
   return {
     accessToken: tokenState.token,
-    loading: tokenState.loading,
+    loading: isLoading,
     error: tokenState.error,
     refresh,
     getAccessToken,