fix(account): wait for remote config before the first balances fetch

extension/e2e-tests/helpers/stubs.ts

@@ -3388,6 +3388,7 @@ export const stubAllExternalApis = async (
 
   // Feature flags
   await stubFeatureFlags(context);
+  await stubExperimentDefaults(context);
 
   // RPC health is fetched by the background service worker
   await stubRpcHealth(context);
@@ -3577,6 +3578,17 @@ export const stubReportTransactionWarning = async (
 const AMPLITUDE_EXPERIMENT_VARDATA_ROUTE = "**/sdk/v2/vardata**";
 const MAINTENANCE_VARIANT_VALUE = "on";
 
+/**
+ * Answers the Experiment flag request with no variants, so every flag keeps
+ * its default and balances don't wait on a real Amplitude request. Specs that
+ * need a variant register a page route, which takes precedence over this one.
+ */
+export const stubExperimentDefaults = async (context: BrowserContext) => {
+  await context.route(AMPLITUDE_EXPERIMENT_VARDATA_ROUTE, async (route) => {
+    await route.fulfill({ json: {} });
+  });
+};
+
 const stubExperimentVariants = async (
   page: Page,
   variants: Record<string, { payload?: unknown }>,

extension/src/helpers/__tests__/useGetBalances.test.tsx

@@ -8,6 +8,10 @@ import {
 } from "../hooks/useGetBalances";
 import { getAccountBalances } from "@shared/api/internal";
 import { makeDummyStore, TEST_PUBLIC_KEY } from "popup/__testHelpers__";
+import {
+  fetchFeatureFlags,
+  reducer as remoteConfigReducer,
+} from "popup/ducks/remoteConfig";
 import { TESTNET_NETWORK_DETAILS } from "@shared/constants/stellar";
 import { getIconUrlFromIssuer } from "@shared/api/helpers/getIconUrlFromIssuer";
 import { getCombinedAssetListData } from "@shared/api/helpers/token-list";
@@ -260,6 +264,54 @@ describe("useGetBalances (flag routing)", () => {
     );
   });
 
+  it("waits for the remote config flags before choosing the balances endpoint", async () => {
+    const store = makeDummyStore({
+      cache: { balanceData: {}, icons: {}, tokenLists: [] },
+      settings: { assetsLists: [] },
+      remoteConfig: {
+        ...remoteConfigReducer(undefined, { type: "@@init" }),
+        isInitialized: false,
+      },
+    });
+    const { result } = renderHook(
+      () => useGetBalances({ showHidden: false, includeIcons: false }),
+      { wrapper: Wrapper(store) },
+    );
+
+    let fetchPromise: Promise<unknown> = Promise.resolve();
+    await act(async () => {
+      fetchPromise = result.current.fetchData(
+        publicKey,
+        true,
+        TESTNET_NETWORK_DETAILS,
+        false,
+      );
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+    expect(getAccountBalances).not.toHaveBeenCalled();
+
+    await act(async () => {
+      store.dispatch({
+        type: fetchFeatureFlags.fulfilled.type,
+        payload: {
+          ...remoteConfigReducer(undefined, { type: "@@init" }),
+          isInitialized: true,
+          use_balances_v2: true,
+        },
+      });
+      await fetchPromise;
+    });
+
+    expect(getAccountBalances).toHaveBeenCalledTimes(1);
+    expect(getAccountBalances).toHaveBeenCalledWith(
+      publicKey,
+      TESTNET_NETWORK_DETAILS,
+      true,
+      false,
+      true,
+    );
+  });
+
   it("passes localOnlyTokenIds through to the payload", async () => {
     const localOnlyTokenIds = ["CLOCALTOKENCONTRACTID"];
     (getAccountBalances as jest.Mock).mockResolvedValue({

extension/src/helpers/__tests__/useGetTokenPrices.test.tsx

@@ -4,7 +4,10 @@ import { Provider } from "react-redux";
 import { renderHook, act } from "@testing-library/react";
 import { useGetTokenPrices } from "../hooks/useGetTokenPrices";
 import { makeDummyStore } from "popup/__testHelpers__";
-import { reducer as remoteConfigReducer } from "popup/ducks/remoteConfig";
+import {
+  fetchFeatureFlags,
+  reducer as remoteConfigReducer,
+} from "popup/ducks/remoteConfig";
 
 // Flags settled, as in a running popup once `fetchFeatureFlags` has resolved.
 const SETTLED_REMOTE_CONFIG = {
@@ -80,6 +83,63 @@ describe("useGetTokenPrices", () => {
       },
     });
   });
+  it("waits for the remote config flags before choosing the prices endpoint", async () => {
+    const getTokenPricesSpy = jest
+      .spyOn(ApiInternal, "getTokenPrices")
+      .mockResolvedValue({});
+    const store = makeDummyStore({
+      remoteConfig: {
+        ...remoteConfigReducer(undefined, { type: "@@init" }),
+        isInitialized: false,
+      },
+      cache: { tokenPrices: {} },
+    });
+    const Wrapper =
+      (store: ReturnType<typeof makeDummyStore>) =>
+      ({ children }: { children: React.ReactNode }) => (
+        <Provider store={store}>{children}</Provider>
+      );
+    const { result } = renderHook(() => useGetTokenPrices(), {
+      wrapper: Wrapper(store),
+    });
+
+    let fetchPromise: Promise<unknown> = Promise.resolve();
+    await act(async () => {
+      fetchPromise = result.current.fetchData({
+        publicKey: "G123",
+        balances: [
+          {
+            token: { type: "native", code: "XLM" },
+            total: new BigNumber("50"),
+            available: new BigNumber("50"),
+            blockaidData: defaultBlockaidScanAssetResult,
+          },
+        ],
+        networkDetails: MAINNET_NETWORK_DETAILS,
+      } as any);
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+    expect(getTokenPricesSpy).not.toHaveBeenCalled();
+
+    await act(async () => {
+      store.dispatch({
+        type: fetchFeatureFlags.fulfilled.type,
+        payload: {
+          ...SETTLED_REMOTE_CONFIG,
+          use_token_prices_v2: false,
+        },
+      });
+      await fetchPromise;
+    });
+
+    expect(getTokenPricesSpy).toHaveBeenCalledTimes(1);
+    expect(getTokenPricesSpy).toHaveBeenCalledWith(
+      ["native"],
+      MAINNET_NETWORK_DETAILS,
+      false,
+    );
+  });
+
   it("should return token prices from API with expired cache", async () => {
     const getTokenPricesSpy = jest
       .spyOn(ApiInternal, "getTokenPrices")

extension/src/helpers/experimentClient.ts

@@ -1,10 +1,18 @@
-import { Experiment, ExperimentClient } from "@amplitude/experiment-js-client";
+import {
+  Experiment,
+  ExperimentClient,
+  ExperimentConfig,
+} from "@amplitude/experiment-js-client";
 
 import { AMPLITUDE_EXPERIMENT_DEPLOYMENT_KEY } from "constants/env";
 import { isDev } from "@shared/helpers/dev";
 
 const TEST_EXPERIMENT_DEPLOYMENT_KEY = "playwright-test-deployment-key";
 
+// Balance and price fetches wait on the flags, so a slow Amplitude host holds
+// the spinner for this long before falling back to the default flags.
+const CONFIG: ExperimentConfig = { fetchTimeoutMillis: 1500 };
+
 // Console log message constants
 const LOG_MESSAGES = {
   EXPERIMENT_PREFIX: "[Experiment]",
@@ -48,8 +56,8 @@ export const initExperimentClient = (): void => {
 
   try {
     client = AMPLITUDE_EXPERIMENT_DEPLOYMENT_KEY
-      ? Experiment.initializeWithAmplitudeAnalytics(deploymentKey)
-      : Experiment.initialize(deploymentKey);
+      ? Experiment.initializeWithAmplitudeAnalytics(deploymentKey, CONFIG)
+      : Experiment.initialize(deploymentKey, CONFIG);
   } catch (e) {
     console.error(
       `${LOG_MESSAGES.EXPERIMENT_PREFIX} ${LOG_MESSAGES.INIT_FAILED}`,

extension/src/helpers/hooks/useGetBalances.tsx

@@ -21,7 +21,10 @@ import { AssetType } from "@shared/api/types/account-balance";
 import { settingsSelector } from "popup/ducks/settings";
 import { getCombinedAssetListData } from "@shared/api/helpers/token-list";
 import { AppDispatch, AppState } from "popup/App";
-import { balancesV2Selector } from "popup/ducks/remoteConfig";
+import {
+  balancesV2Selector,
+  waitForRemoteConfig,
+} from "popup/ducks/remoteConfig";
 import {
   balancesSelector,
   iconsSelector,
@@ -106,15 +109,17 @@ function useGetBalances(options: {
       const accountBalances =
         useCache && isBalancesCacheValid
           ? cachedBalanceData
-          : await getAccountBalances(
-              publicKey,
-              networkDetails,
-              isMainnet,
-              shouldSkipScan,
-              // Read the flag from the store at call time (not a
-              // render-captured value) so a freshly resolved Amplitude flag
-              // isn't missed — mirrors useGetTokenPrices.
-              balancesV2Selector(store.getState()),
+          : await waitForRemoteConfig(store).then(() =>
+              getAccountBalances(
+                publicKey,
+                networkDetails,
+                isMainnet,
+                shouldSkipScan,
+                // Read the flag from the store at call time (not a
+                // render-captured value) so a freshly resolved Amplitude flag
+                // isn't missed — mirrors useGetTokenPrices.
+                balancesV2Selector(store.getState()),
+              ),
             );
 
       const { balances, unfilteredBalances } = await formatBalances({

extension/src/helpers/hooks/useGetTokenPrices.tsx

@@ -2,7 +2,10 @@ import { useReducer } from "react";
 import { captureException } from "@sentry/browser";
 import { useDispatch, useSelector, useStore } from "react-redux";
 import { tokenPricesSelector } from "popup/ducks/cache";
-import { tokenPricesV2Selector } from "popup/ducks/remoteConfig";
+import {
+  tokenPricesV2Selector,
+  waitForRemoteConfig,
+} from "popup/ducks/remoteConfig";
 import { getTokenPrices } from "@shared/api/internal";
 import { initialState, isCacheValid, reducer } from "helpers/request";
 import { AppDataType } from "helpers/hooks/useGetAppData";
@@ -80,6 +83,7 @@ export function useGetTokenPrices() {
       );
       if (unseenAssetIds.length) {
         try {
+          await waitForRemoteConfig(store);
           const useV2 = tokenPricesV2Selector(store.getState());
           const unseenPrices = await getTokenPrices(
             unseenAssetIds,
@@ -117,6 +121,7 @@ export function useGetTokenPrices() {
           // value) so a freshly resolved Amplitude flag isn't missed when this
           // fetch runs inside a long-lived async flow that closed over a stale
           // default.
+          await waitForRemoteConfig(store);
           const useV2 = tokenPricesV2Selector(store.getState());
           const fetchedTokenPrices = await getTokenPrices(
             assetIds,
@@ -148,6 +153,7 @@ export function useGetTokenPrices() {
       const missingExtra = additionalAssetIds.filter((id) => !(id in resolved));
       if (missingExtra.length) {
         try {
+          await waitForRemoteConfig(store);
           const useV2 = tokenPricesV2Selector(store.getState());
           const extraPrices = await getTokenPrices(
             missingExtra,

extension/src/popup/__testHelpers__/index.tsx

@@ -69,6 +69,15 @@ export const Wrapper: React.FunctionComponent<any> = ({
       applicationState: APPLICATION_STATE.MNEMONIC_PHRASE_CONFIRMED,
     },
     transactionSubmission: transactionSubmissionInitialState,
+    // Loaded so balance fetches don't wait forever; kill switches stay off.
+    remoteConfig: {
+      isInitialized: true,
+      use_token_prices_v2: true,
+      use_balances_v2: false,
+      use_token_price_history: false,
+      maintenance_banner: { enabled: false, payload: undefined },
+      maintenance_screen: { enabled: false, payload: undefined },
+    },
     ...state,
   });
   testStore = store;

extension/src/popup/ducks/remoteConfig.ts

@@ -2,6 +2,7 @@ import {
   createAsyncThunk,
   createSelector,
   createSlice,
+  Store,
 } from "@reduxjs/toolkit";
 import type { Variant } from "@amplitude/experiment-js-client";
 
@@ -278,3 +279,24 @@ export const isRemoteConfigInitializedSelector = createSelector(
   remoteConfigSelector,
   (rc) => rc.isInitialized,
 );
+
+/**
+ * Resolves once the Experiment flags have been fetched (or failed). Await it
+ * before reading a flag that picks a backend, since every popup, tab, and
+ * sidebar starts with the flags unloaded.
+ */
+export const waitForRemoteConfig = (
+  store: Store<{ remoteConfig: RemoteConfigState }>,
+): Promise<void> =>
+  new Promise((resolve) => {
+    if (isRemoteConfigInitializedSelector(store.getState())) {
+      resolve();
+      return;
+    }
+    const unsubscribe = store.subscribe(() => {
+      if (isRemoteConfigInitializedSelector(store.getState())) {
+        unsubscribe();
+        resolve();
+      }
+    });
+  });

extension/src/popup/views/AddToken/index.tsx

@@ -24,7 +24,10 @@ import { newTabHref, parsedSearchParam, TokenToAdd } from "helpers/urls";
 
 import { AppState } from "popup/App";
 import { rejectToken, addToken } from "popup/ducks/access";
-import { balancesV2Selector } from "popup/ducks/remoteConfig";
+import {
+  balancesV2Selector,
+  waitForRemoteConfig,
+} from "popup/ducks/remoteConfig";
 import { isNonSSLEnabledSelector } from "popup/ducks/settings";
 import { useSetupAddTokenFlow } from "popup/helpers/useSetupAddTokenFlow";
 import {
@@ -396,6 +399,7 @@ export const AddToken = () => {
 
     const fetchClassicTrustline = async () => {
       try {
+        await waitForRemoteConfig(store);
         const accountBalances = await getAccountBalances(
           hydratedPublicKey,
           hydratedNetworkDetails,