fix stale values in ensemble.user (#1096)

Co-authored-by: Ehmad Saeed <justehmadsaeed@gmail.com>

Author: anserwaseem

.changeset/gentle-countries-teach.md

@@ -0,0 +1,5 @@
+---
+"@ensembleui/react-framework": patch
+---
+
+fix stale values in ensemble.user

packages/framework/jest.config.js

@@ -3,6 +3,7 @@
 module.exports = {
   preset: "ts-jest",
   testEnvironment: "jsdom",
+  setupFilesAfterEnv: ["<rootDir>/setupTests.ts"],
   moduleNameMapper: {
     "^lodash-es$": "lodash",
   },

packages/framework/setupTests.ts

@@ -0,0 +1,14 @@
+// setup i18n to avoid react-i18next warnings during tests
+import i18n from "i18next";
+import { initReactI18next } from "react-i18next";
+
+void i18n.use(initReactI18next).init({
+  lng: "en",
+  fallbackLng: "en",
+  resources: { en: { translation: {} } },
+  interpolation: { escapeValue: false },
+});
+
+// jsdom does not implement ResizeObserver; provide a polyfill for tests
+// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+global.ResizeObserver = require("resize-observer-polyfill");

packages/framework/src/hooks/__tests__/screenHydrationNoClobber.test.tsx

@@ -0,0 +1,55 @@
+import { render, screen, waitFor } from "@testing-library/react";
+import React from "react";
+import { Provider } from "jotai";
+import { ApplicationContextProvider } from "../useApplicationContext";
+import { ScreenContextProvider } from "../useScreenContext";
+import type { EnsembleAppModel } from "../../shared";
+import { useEnsembleUser } from "../useEnsembleUser";
+
+const Probe: React.FC<{ id: string }> = ({ id }) => {
+  const user = useEnsembleUser();
+  return <div data-testid={`u-${id}`}>{String(user.accessToken ?? "")}</div>;
+};
+
+const BootstrapUser: React.FC = () => {
+  const user = useEnsembleUser();
+  React.useEffect(() => {
+    user.set({ accessToken: "fresh" });
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []);
+  return null;
+};
+
+const app: EnsembleAppModel = {
+  id: "app",
+  screens: [],
+  customWidgets: [],
+  scripts: [],
+  home: {
+    id: "home",
+    name: "home",
+    body: { name: "Text", properties: { text: "home" } },
+  },
+  themes: { default: { name: "default" } },
+};
+
+describe("ScreenContextProvider does not clobber user atom", () => {
+  test("user stays latest when screen mounts", async () => {
+    render(
+      <Provider>
+        <ApplicationContextProvider app={app}>
+          <BootstrapUser />
+          <Probe id="before" />
+          <ScreenContextProvider screen={{ name: "home" }}>
+            <Probe id="after" />
+          </ScreenContextProvider>
+        </ApplicationContextProvider>
+      </Provider>,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByTestId("u-before").textContent).toBe("fresh");
+      expect(screen.getByTestId("u-after").textContent).toBe("fresh");
+    });
+  });
+});

packages/framework/src/hooks/__tests__/useEnsembleUser.test.ts

@@ -7,6 +7,11 @@ import { useEnsembleUser } from "../useEnsembleUser";
 const store = getDefaultStore();
 
 describe("useEnsembleUser", () => {
+  beforeEach(() => {
+    sessionStorage.removeItem("ensemble.user");
+    store.set(userAtom, {} as unknown as { [key: string]: unknown });
+  });
+
   test("fetches variable from ensemble.user", () => {
     store.set(userAtom, {
       accessToken: "eyJabcd",
@@ -58,4 +63,30 @@ describe("useEnsembleUser", () => {
       },
     });
   });
+
+  test("returns sessionStorage snapshot on first render when atom empty", () => {
+    sessionStorage.setItem(
+      "ensemble.user",
+      JSON.stringify({ accessToken: "seed", userDetails: { foo: "bar" } }),
+    );
+
+    const { result } = renderHook(() => useEnsembleUser());
+    const user = result.current;
+
+    expect(get(user, "accessToken")).toBe("seed");
+    expect(get(user, "userDetails.foo")).toBe("bar");
+  });
+
+  test("prefers atom value over sessionStorage when atom populated", () => {
+    sessionStorage.setItem(
+      "ensemble.user",
+      JSON.stringify({ accessToken: "seed" }),
+    );
+    store.set(userAtom, { accessToken: "live" });
+
+    const { result } = renderHook(() => useEnsembleUser());
+    const user = result.current;
+
+    expect(get(user, "accessToken")).toBe("live");
+  });
 });

packages/framework/src/hooks/__tests__/userAtomSync.test.tsx

@@ -0,0 +1,92 @@
+import { act, render, screen, waitFor } from "@testing-library/react";
+import React, { useEffect } from "react";
+import { Provider } from "jotai";
+import { useEnsembleUser } from "../useEnsembleUser";
+
+type Api = {
+  setToken: (token?: string) => void;
+};
+
+const UserTokenProbe: React.FC<{
+  id: string;
+  onReady?: (api: Api) => void;
+}> = ({ id, onReady }) => {
+  const user = useEnsembleUser();
+
+  useEffect(() => {
+    onReady?.({
+      setToken: (token?: string) => user.set({ accessToken: token }),
+    });
+  }, [onReady, user]);
+
+  return (
+    <div data-testid={`token-${id}`}>{String(user.accessToken ?? "")}</div>
+  );
+};
+
+describe("userAtom cross-store sync", () => {
+  beforeEach(() => {
+    sessionStorage.removeItem("ensemble.user");
+  });
+
+  test("initializes from sessionStorage on first mount", async () => {
+    sessionStorage.setItem(
+      "ensemble.user",
+      JSON.stringify({ accessToken: "seed" }),
+    );
+
+    render(
+      <Provider>
+        <UserTokenProbe id="one" />
+      </Provider>,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByTestId("token-one").textContent).toBe("seed");
+    });
+  });
+
+  test("propagates updates between nested Providers via storage sync", async () => {
+    let apiA: Api | undefined;
+    let apiB: Api | undefined;
+
+    render(
+      <Provider>
+        <UserTokenProbe
+          id="a"
+          onReady={(api) => {
+            apiA = api;
+          }}
+        />
+        <Provider>
+          <UserTokenProbe
+            id="b"
+            onReady={(api) => {
+              apiB = api;
+            }}
+          />
+        </Provider>
+      </Provider>,
+    );
+
+    // outer updates, inner should reflect
+    act(() => {
+      apiA?.setToken("t1");
+    });
+
+    await waitFor(() => {
+      expect(screen.getByTestId("token-a").textContent).toBe("t1");
+      expect(screen.getByTestId("token-b").textContent).toBe("t1");
+    });
+
+    // inner updates, outer should reflect
+    act(() => {
+      apiB?.setToken("t2");
+    });
+
+    await waitFor(() => {
+      expect(screen.getByTestId("token-a").textContent).toBe("t2");
+      expect(screen.getByTestId("token-b").textContent).toBe("t2");
+    });
+  });
+});

packages/framework/src/hooks/useCommandCallback.ts

@@ -1,7 +1,7 @@
 import { useAtomCallback } from "jotai/utils";
 import type { FC, ReactNode } from "react";
 import { useCallback } from "react";
-import { mapKeys } from "lodash-es";
+import { mapKeys, assign } from "lodash-es";
 import { createEvaluationContext } from "../evaluate";
 import type { EnsembleUser } from "../state";
 import { appAtom, screenAtom, themeAtom, userAtom } from "../state";
@@ -84,6 +84,8 @@ export const useCommandCallback = <
             setTheme: (name: string) => set(themeAtom, name),
             user: {
               ...user,
+              set: (userUpdate: EnsembleUser) =>
+                set(userAtom, assign({}, user, userUpdate)),
               setUser: (userUpdate: EnsembleUser) => set(userAtom, userUpdate),
             },
             storage: storageApi,

packages/framework/src/hooks/useEnsembleUser.ts

@@ -1,5 +1,5 @@
 import { useAtom } from "jotai";
-import { assign } from "lodash-es";
+import { assign, isEmpty } from "lodash-es";
 import { useMemo } from "react";
 import { userAtom, type EnsembleUser } from "../state";
 
@@ -26,5 +26,18 @@ export const useEnsembleUser = (): EnsembleUser & EnsembleUserBuffer => {
     [setUser, user],
   );
 
-  return { ...storageBuffer, ...user };
+  // ensure first render on direct loads sees latest value from sessionStorage
+  const sessionSnapshot = useMemo<EnsembleUser>(() => {
+    try {
+      const raw = sessionStorage.getItem("ensemble.user");
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+      return raw ? JSON.parse(raw) : {};
+    } catch {
+      return {};
+    }
+  }, []);
+
+  const effectiveUser = isEmpty(user) ? sessionSnapshot : user;
+
+  return { ...storageBuffer, ...effectiveUser };
 };

packages/framework/src/hooks/useScreenContext.tsx

@@ -10,7 +10,6 @@ import {
   screenDataAtom,
   screenModelAtom,
   themeModelAtom,
-  userAtom,
 } from "../state";
 import type {
   ApplicationContextDefinition,
@@ -77,16 +76,10 @@ const HydrateAtoms: React.FC<
 
   // initialising on state with prop on render here
   useHydrateAtoms([[screenAtom, screenContext]]);
-  useHydrateAtoms(
-    [
-      [appAtom, appContext],
-      [themeModelAtom, themeScope.theme],
-      [userAtom, appContext.user],
-    ],
-    {
-      dangerouslyForceHydrate: true,
-    },
-  );
+  useHydrateAtoms([
+    [appAtom, appContext],
+    [themeModelAtom, themeScope.theme],
+  ]);
 
   // initiate device resizer observer
   useDeviceObserver();

packages/framework/src/state/user.ts

@@ -17,7 +17,42 @@ const atomWithSessionStorage = <T = unknown>(
     }
     return initialValue;
   };
+  // initialise from sessionStorage so first render sees the latest value
   const baseAtom = atom<T>(getInitialValue());
+
+  // keep all jotai stores in sync by listening to storage events
+  // this is required because we create nested Providers (separate stores)
+  // and without syncing, a later Provider may see stale values
+  baseAtom.onMount = (setAtom) => {
+    // set initial value from storage on first mount for this store
+    try {
+      setAtom(getInitialValue());
+    } catch {
+      setAtom(initialValue);
+    }
+
+    const handler = (event: StorageEvent): void => {
+      if (event.key !== key) return;
+      try {
+        const next = event.newValue
+          ? (JSON.parse(event.newValue) as T)
+          : getInitialValue();
+        setAtom(next);
+      } catch {
+        setAtom(getInitialValue());
+      }
+    };
+
+    if (typeof window !== "undefined") {
+      window.addEventListener("storage", handler);
+    }
+
+    return () => {
+      if (typeof window !== "undefined") {
+        window.removeEventListener("storage", handler);
+      }
+    };
+  };
   const derivedAtom = atom<T, unknown[], unknown>(
     (get) => get(baseAtom),
     (get, set, update) => {
@@ -26,6 +61,10 @@ const atomWithSessionStorage = <T = unknown>(
       ) as T;
       set(baseAtom, nextValue);
       sessionStorage.setItem(key, JSON.stringify(nextValue));
+      // notify other stores in this tab to update immediately
+      if (typeof window !== "undefined") {
+        window.dispatchEvent(new StorageEvent("storage", { key }));
+      }
     },
   );
   return derivedAtom;