fix(fxa-settings): Update hasPassword in localStorage after createPassword and fix loading state flash

After creating a password via createPassword/createPasswordWithJwt, localStorage
was not updated with hasPassword: true, causing the app to still show the
"set password" flow and resulting in errno 206 on retry.

The loading/error state in useAccountData was routed through localStorage where
transient fields are stripped, so isLoading was always false. This caused a flash
of default values (security features shown as disabled) before the data fetch
completed. Switched to local React state for loading/error tracking.

Author: vbudhram

packages/fxa-settings/src/components/Settings/index.test.tsx

@@ -19,6 +19,7 @@ import AppLocalizationProvider from 'fxa-react/lib/AppLocalizationProvider';
 import { Subject } from './mocks';
 
 const mockSessionStatus = jest.fn();
+const mockUseAccountData = jest.fn();
 const mockAccountState = {
   isLoading: false,
   error: null,
@@ -61,6 +62,11 @@ jest.mock('../../models/contexts/AccountStateContext', () => ({
   useAccountState: jest.fn(),
 }));
 
+jest.mock('../../lib/hooks/useAccountData', () => ({
+  ...jest.requireActual('../../lib/hooks/useAccountData'),
+  useAccountData: (...args: any[]) => mockUseAccountData(...args),
+}));
+
 jest.mock('../../lib/totp-utils', () => {
   const mockBackupCodes = ['0123456789'];
   return {
@@ -115,6 +121,13 @@ describe('Settings App', () => {
     jest.spyOn(console, 'error').mockImplementation(() => {});
     // Reset account state mock to default values
     (useAccountState as jest.Mock).mockReturnValue({ ...mockAccountState, isLoading: false, error: null });
+    mockUseAccountData.mockReturnValue({
+      isLoading: false,
+      error: null,
+      data: mockAccountState,
+      refetch: jest.fn(),
+      refetchField: jest.fn(),
+    });
     mockNavigate.mockReset();
     mockSessionStatus.mockResolvedValue({
       details: {
@@ -131,9 +144,12 @@ describe('Settings App', () => {
   });
 
   it('renders `LoadingSpinner` component when loading initial state is true', () => {
-    (useAccountState as jest.Mock).mockReturnValueOnce({
-      ...mockAccountState,
+    mockUseAccountData.mockReturnValue({
       isLoading: true,
+      error: null,
+      data: mockAccountState,
+      refetch: jest.fn(),
+      refetchField: jest.fn(),
     });
     const { getByLabelText } = renderWithRouter(
       <AppContext.Provider value={mockAppContext()}>
@@ -145,11 +161,13 @@ describe('Settings App', () => {
   });
 
   it('renders `AppErrorDialog` component when settings query errors', async () => {
-    (useAccountState as jest.Mock).mockImplementation(() => ({
-      ...mockAccountState,
+    mockUseAccountData.mockReturnValue({
       isLoading: false,
       error: new Error('Error'),
-    }));
+      data: mockAccountState,
+      refetch: jest.fn(),
+      refetchField: jest.fn(),
+    });
     const { getByTestId } = renderWithRouter(
       <AppContext.Provider value={mockAppContext()}>
         <Subject />

packages/fxa-settings/src/lib/hooks/useAccountData.test.ts

@@ -0,0 +1,49 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+import { renderHook, act } from '@testing-library/react';
+import { useAccountData } from './useAccountData';
+import { sessionToken as getSessionToken } from '../cache';
+
+const mockSetAccountData = jest.fn();
+
+jest.mock('../cache', () => ({ sessionToken: jest.fn() }));
+jest.mock('../../models/contexts/AccountStateContext', () => ({
+  useAccountState: () => ({ setAccountData: mockSetAccountData }),
+}));
+jest.mock('../config', () => ({
+  oauth: { clientId: 'test' },
+  servers: { profile: { url: 'http://localhost' } },
+}));
+jest.mock('@sentry/browser', () => ({ captureMessage: jest.fn() }));
+
+const authClient = {
+  account: jest.fn(),
+  attachedClients: jest.fn(),
+  createOAuthToken: jest.fn(),
+} as any;
+
+describe('useAccountData', () => {
+  beforeEach(() => jest.clearAllMocks());
+
+  it('starts with isLoading=true', () => {
+    (getSessionToken as jest.Mock).mockReturnValue('tok');
+    authClient.account.mockReturnValue(new Promise(() => {}));
+    authClient.attachedClients.mockReturnValue(new Promise(() => {}));
+    authClient.createOAuthToken.mockReturnValue(new Promise(() => {}));
+
+    const { result } = renderHook(() => useAccountData({ authClient }));
+    expect(result.current.isLoading).toBe(true);
+  });
+
+  it('sets error when no session token', async () => {
+    (getSessionToken as jest.Mock).mockReturnValue(null);
+    let result: any;
+    await act(async () => {
+      ({ result } = renderHook(() => useAccountData({ authClient })));
+    });
+    expect(result.current.isLoading).toBe(false);
+    expect(result.current.error?.message).toBe('No session token available');
+  });
+});

packages/fxa-settings/src/lib/hooks/useAccountData.ts

@@ -2,7 +2,7 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
-import { useCallback, useEffect, useRef } from 'react';
+import { useCallback, useEffect, useRef, useState } from 'react';
 import AuthClient from 'fxa-auth-client/browser';
 import { sessionToken as getSessionToken } from '../cache';
 import {
@@ -185,37 +185,40 @@ export function useAccountData({
   authClient,
   onError,
 }: UseAccountDataOptions): AccountDataResult {
-  // Ref prevents infinite re-render from unstable onError callback
+  // Refs prevent infinite re-renders from unstable callback/object references
   const onErrorRef = useRef(onError);
   onErrorRef.current = onError;
+  const authClientRef = useRef(authClient);
+  authClientRef.current = authClient;
 
   const accountState = useAccountState();
   const {
     setAccountData,
-    setLoading,
-    setError,
-    isLoading,
-    error,
     ...stateData
   } = accountState;
 
+  const [localLoading, setLocalLoading] = useState(true);
+  const [localError, setLocalError] = useState<Error | null>(null);
+
   const fetchAccountData = useCallback(async () => {
+    const client = authClientRef.current;
     const token = getSessionToken();
     if (!token) {
-      setError(new Error('No session token available'));
+      setLocalError(new Error('No session token available'));
+      setLocalLoading(false);
       return;
     }
 
-    setLoading(true);
-    setError(null);
+    setLocalLoading(true);
+    setLocalError(null);
 
     try {
       // allSettled (not .all) so a single failure doesn't discard other results
       const [accountResult, profileResult, attachedClientsResult] =
         await Promise.allSettled([
-          authClient.account(token),
-          fetchProfileData(authClient, token),
-          authClient.attachedClients(token),
+          client.account(token),
+          fetchProfileData(client, token),
+          client.attachedClients(token),
         ]);
 
       // Check for invalid token errors - user needs to sign in again
@@ -238,28 +241,29 @@ export function useAccountData({
         if (displayName !== null) accountData.displayName = displayName;
         if (avatar !== null) accountData.avatar = avatar;
       } else {
-        Sentry.captureMessage(`Failed to fetch profile: ${profileResult.reason}`);  
+        Sentry.captureMessage(`Failed to fetch profile: ${profileResult.reason}`);
       }
 
       if (attachedClientsResult.status === 'fulfilled') {
         accountData.attachedClients = attachedClientsResult.value.map(mapAttachedClient);
       } else {
-        Sentry.captureMessage(`Failed to fetch attached clients: ${attachedClientsResult.reason}`);  
+        Sentry.captureMessage(`Failed to fetch attached clients: ${attachedClientsResult.reason}`);
         accountData.attachedClients = [];
       }
 
       setAccountData(accountData);
     } catch (err) {
       const error = err instanceof Error ? err : new Error('Unknown error');
-      setError(error);
+      setLocalError(error);
       onErrorRef.current?.(error);
     } finally {
-      setLoading(false);
+      setLocalLoading(false);
     }
-  }, [authClient, setAccountData, setLoading, setError]);
+  }, [setAccountData]);
 
   const refetchField = useCallback(
     async (field: keyof AccountState) => {
+      const client = authClientRef.current;
       const token = getSessionToken();
       if (!token) return;
 
@@ -268,19 +272,19 @@ export function useAccountData({
 
         switch (field) {
           case 'attachedClients': {
-            const clients = await authClient.attachedClients(token);
+            const clients = await client.attachedClients(token);
             fieldData.attachedClients = clients.map(mapAttachedClient);
             break;
           }
           case 'displayName':
           case 'avatar': {
-            const { displayName, avatar } = await fetchProfileData(authClient, token);
+            const { displayName, avatar } = await fetchProfileData(client, token);
             if (displayName !== null) fieldData.displayName = displayName;
             if (avatar !== null) fieldData.avatar = avatar;
             break;
           }
           default: {
-            const account = await authClient.account(token);
+            const account = await client.account(token);
             fieldData = { ...fieldData, ...transformAccountResponse(account) };
             break;
           }
@@ -291,17 +295,17 @@ export function useAccountData({
         console.error(`Failed to refetch ${field}:`, err);
       }
     },
-    [authClient, setAccountData]
+    [setAccountData]
   );
 
   useEffect(() => {
     fetchAccountData();
   }, [fetchAccountData]);
 
   return {
-    data: { ...stateData, isLoading, error } as AccountState,
-    isLoading,
-    error,
+    data: { ...stateData, isLoading: localLoading, error: localError } as AccountState,
+    isLoading: localLoading,
+    error: localError,
     refetch: fetchAccountData,
     refetchField,
   };

packages/fxa-settings/src/models/Account.test.ts

@@ -2,14 +2,23 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
-import { getNextAvatar } from './Account';
+import { Account, getNextAvatar } from './Account';
+import { sessionToken, JwtTokenCache } from '../lib/cache';
+import { getFullAccountData, updateExtendedAccountState } from '../lib/account-storage';
 
 jest.mock('../lib/config', () => ({
-  servers: {
-    profile: {
-      url: 'default-url',
-    },
-  },
+  servers: { profile: { url: 'default-url' } },
+  oauth: { clientId: 'test' },
+}));
+jest.mock('../lib/cache', () => ({
+  sessionToken: jest.fn(),
+  JwtTokenCache: { hasToken: jest.fn(() => true), getToken: jest.fn() },
+  isSigningOut: jest.fn(() => false),
+}));
+jest.mock('../lib/account-storage', () => ({
+  getFullAccountData: jest.fn(),
+  updateExtendedAccountState: jest.fn(),
+  updateBasicAccountData: jest.fn(),
 }));
 
 describe('Account', () => {
@@ -54,4 +63,38 @@ describe('Account', () => {
       expect(avatar.url).toBe(null);
     });
   });
+
+  describe('createPassword', () => {
+    let account: Account;
+    const authClient: any = {
+      createPassword: jest.fn().mockResolvedValue({ passwordCreated: 123 }),
+      createPasswordWithJwt: jest.fn().mockResolvedValue({ passwordCreated: 123 }),
+    };
+
+    beforeEach(() => {
+      jest.clearAllMocks();
+      (getFullAccountData as jest.Mock).mockReturnValue({
+        uid: 'abc',
+        emails: [{ email: '[email protected]', isPrimary: true, verified: true }],
+        primaryEmail: { email: '[email protected]', isPrimary: true, verified: true },
+      });
+      (sessionToken as jest.Mock).mockReturnValue('tok');
+      (JwtTokenCache.getToken as jest.Mock).mockReturnValue('jwt');
+      account = new Account(authClient);
+    });
+
+    it('sets hasPassword after createPassword', async () => {
+      await account.createPassword('pw');
+      expect(updateExtendedAccountState).toHaveBeenCalledWith(
+        expect.objectContaining({ hasPassword: true })
+      );
+    });
+
+    it('sets hasPassword after createPasswordWithJwt', async () => {
+      await account.createPasswordWithJwt('pw');
+      expect(updateExtendedAccountState).toHaveBeenCalledWith(
+        expect.objectContaining({ hasPassword: true })
+      );
+    });
+  });
 });

packages/fxa-settings/src/models/Account.ts

@@ -662,6 +662,7 @@ export class Account implements AccountData {
 
     updateExtendedAccountState({
       passwordCreated: passwordCreatedResult.passwordCreated,
+      hasPassword: true,
     });
   }
 
@@ -677,6 +678,7 @@ export class Account implements AccountData {
 
     updateExtendedAccountState({
       passwordCreated: passwordCreatedResult.passwordCreated,
+      hasPassword: true,
     });
   }