fix: flatten providers for large state trees

docs/app/tests/frontend_quality.test.mjs

@@ -147,7 +147,7 @@ with RegistrationContext(), patch('reflex.compiler.compiler.get_config', return_
     bundle_library('quality-lazy-fixture')
     print(compile_app_root(rx.el.div('Quality'))[1])
 `], {cwd: root, encoding: "utf8"});
-  const setup = new Function("window", "React", "emotion_react", "utils_context", "utils_state", "loadFixture",
+  const setup = new Function("window", "React", "emotion_react", "utils_context", "utils_context_registry", "utils_state", "loadFixture",
     rootCode.replace(/^import .*$/gm, "")
       .replace(/^export default function/gm, "function")
       .replace(/^export /gm, "")
@@ -165,7 +165,7 @@ with RegistrationContext(), patch('reflex.compiler.compiler.get_config', return_
   try {
     let calls = 0;
     globalThis.window = {};
-    setup(window, react, {}, {}, {}, async () => {calls++; return {value: 42};});
+    setup(window, react, {}, {}, {}, {}, async () => {calls++; return {value: 42};});
     assert.equal(calls, 0);
     assert.equal(window.__reflex.react, react);
     assert.equal(window.__reflex["quality-lazy-fixture"], undefined);
@@ -179,7 +179,7 @@ with RegistrationContext(), patch('reflex.compiler.compiler.get_config', return_
 
     let attempts = 0;
     globalThis.window = {};
-    setup(window, react, {}, {}, {}, async () => {
+    setup(window, react, {}, {}, {}, {}, async () => {
       if (++attempts === 1) throw new Error("Temporary download failure");
       return {value: 43};
     });

news/7441.bugfix.md

@@ -0,0 +1 @@
+Detect state reads from the new shared state store when checking that an app that reads state defines an `rx.State` subclass.

packages/reflex-base/news/+large-state-provider-tree.bugfix.md

@@ -0,0 +1 @@
+Fix a blank page with `Maximum call stack size exceeded` in apps with thousands of state classes. State is now read from one shared store, so the React provider tree no longer grows with the number of states.

packages/reflex-base/src/reflex_base/.templates/web/utils/context-registry.js

@@ -4,7 +4,12 @@
 // a module the compiler rewrites: Vite re-executes a rewritten module on hot
 // update, and every module it imports along with it, but this module imports
 // nothing that can change and is never re-executed.
-import { createContext } from "react";
+import {
+  createContext,
+  useContext,
+  useMemo,
+  useSyncExternalStore,
+} from "react";
 
 export const ColorModeContext = createContext({
   rawColorMode: "system",
@@ -15,31 +20,30 @@ export const ColorModeContext = createContext({
 export const UploadFilesContext = createContext(null);
 export const DispatchContext = createContext(null);
 export const EventLoopContext = createContext(null);
+export const StateStoreContext = createContext(null);
 
 ColorModeContext.displayName = "ColorModeContext";
 UploadFilesContext.displayName = "UploadFilesContext";
 DispatchContext.displayName = "DispatchContext";
 EventLoopContext.displayName = "EventLoopContext";
-
-const stateContexts = new Map();
+StateStoreContext.displayName = "StateStoreContext";
 
 /**
- * Return the context object carrying the named Python state.
- *
- * The first call for a name creates the context; later calls, including ones
- * made from a re-executed generated module, return the same object.
+ * Subscribe a component to one Python substate in the shared store.
  *
  * @param {string} name - The dotted Python state name.
- * @returns {React.Context} The state's context object.
+ * @returns {object} The current state value for the named substate.
  */
-export function getStateContext(name) {
-  let context = stateContexts.get(name);
-  if (context === undefined) {
-    context = createContext(null);
-    context.displayName = `StateContext(${name})`;
-    stateContexts.set(name, context);
-  }
-  return context;
+export function useStateContext(name) {
+  const store = useContext(StateStoreContext);
+  const [subscribe, getSnapshot] = useMemo(
+    () => [
+      (listener) => store.subscribe(name, listener),
+      () => store.getSnapshot(name) ?? app.initialState?.[name],
+    ],
+    [store, name],
+  );
+  return useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
 }
 
 // Values the generated context module registers each time it executes. The

packages/reflex-base/src/reflex_base/compiler/templates.py

@@ -10,7 +10,7 @@
 from reflex_base import constants
 from reflex_base.constants import Hooks
 from reflex_base.utils import memo_paths
-from reflex_base.utils.format import format_state_name, json_dumps
+from reflex_base.utils.format import json_dumps
 from reflex_base.vars.base import VarData
 
 if TYPE_CHECKING:
@@ -338,13 +338,6 @@ def context_template(
     initial_state = initial_state or {}
     if initial_state_json is None:
         initial_state_json = json_dumps(initial_state)
-    # Context objects come from the static registry so they survive hot
-    # updates of this module; only the lookup table is generated here.
-    state_contexts_str = "".join([
-        f'{format_state_name(state_name)}: getStateContext("{state_name}"),'
-        for state_name in initial_state
-    ])
-
     state_str = (
         rf"""
 export const state_name = "{state_name}"
@@ -411,11 +404,6 @@ def context_template(
 """
     )
 
-    substates_str = "".join(
-        f"\n  ['{state_name}', '{format_state_name(state_name)}'],"
-        for state_name in initial_state
-    )
-
     disable_owner_stacks_str = (
         r"""
 // Disable React dev-build owner-stack capture: the per-element Error()
@@ -444,17 +432,16 @@ def context_template(
         else ""
     )
 
-    return rf"""import {"React, " if disable_react_owner_stacks else ""}{{ useContext, useMemo, useReducer, useRef, useState, createElement, useEffect, useLayoutEffect }} from "react"
+    return rf"""import {"React, " if disable_react_owner_stacks else ""}{{ useContext, useMemo, useState, createElement, useEffect }} from "react"
 import {{ applyDelta, ReflexEvent, hydrateClientStorage, useEventLoop, refs }} from "$/utils/state"
-import {{ ColorModeContext, UploadFilesContext, DispatchContext, EventLoopContext, getStateContext, registerApp, eventLoop }} from "$/utils/context-registry"
+import {{ ColorModeContext, UploadFilesContext, DispatchContext, EventLoopContext, StateStoreContext, registerApp, eventLoop }} from "$/utils/context-registry"
 import {{ jsx }} from "@emotion/react";
 {disable_owner_stacks_str}
-export {{ ColorModeContext, UploadFilesContext, DispatchContext, EventLoopContext }};
+export {{ ColorModeContext, UploadFilesContext, DispatchContext, EventLoopContext, StateStoreContext }};
 export const initialState = {initial_state_json}
 export const initialStateHashes = {"[]" if not initial_state_hashes else json_dumps(initial_state_hashes)}
 
 export const defaultColorMode = {default_color_mode}
-export const StateContexts = {{{state_contexts_str}}};
 export const clientStorage = {"{}" if client_storage is None else json.dumps(client_storage)}
 
 {state_str}
@@ -536,78 +523,88 @@ def context_template(
   );
 }}
 
-// ``useLayoutEffect`` warns when rendered on the server, where no effect runs
-// at all, so fall back to ``useEffect`` there.
-const useIsomorphicLayoutEffect =
-  typeof document !== "undefined" ? useLayoutEffect : useEffect;
-
-// Holds the mutable substate -> dispatch registry that ``SubstateProvider``
-// writes into and ``EventLoopProvider`` reads. The registry object identity is
-// stable for the lifetime of the tree, so neither adding a dispatcher nor
-// updating a substate re-renders the consumers of ``DispatchContext``.
-const DispatchProvider = ({{ children }}) => {{
-  const dispatchers = useRef({{}});
-  return useMemo(
-    () =>
-      createElement(DispatchContext, {{ value: dispatchers.current }}, children),
-    [children],
+// Dispatch updates happen outside render; useSyncExternalStore keeps reads
+// consistent while notifying only consumers of the changed substate.
+function createStateStore() {{
+  const state = {{ ...initialState }};
+  const listeners = new Map();
+  const dispatchers = {{}};
+  let syncedInitialState = initialState;
+  const registerSubstate = (substateName, initialValue) => {{
+    if (!Object.hasOwn(state, substateName)) {{
+      state[substateName] = initialValue;
+    }}
+    if (!Object.hasOwn(dispatchers, substateName)) {{
+      dispatchers[substateName] = (delta) => {{
+        const currentState = state[substateName];
+        const nextState = applyDelta(currentState, delta);
+        if (nextState === currentState) return;
+        state[substateName] = nextState;
+        listeners.get(substateName)?.forEach((listener) => listener());
+      }};
+    }}
+  }};
+  Object.entries(state).forEach(([substateName, value]) =>
+    registerSubstate(substateName, value),
   );
-}};
 
-// One provider per substate: each owns its own reducer, so a delta for one
-// substate only re-renders its provider instead of recreating every provider.
-const SubstateProvider = ({{ children, substateName, contextName }}) => {{
-  const dispatchers = useContext(DispatchContext);
-  const [state, dispatchSubstate] = useReducer(
-    applyDelta,
-    initialState[substateName],
-  );
-  // A layout effect, not a passive one: layout effects for the whole commit
-  // run before any passive effect, so every dispatcher is registered before
-  // ``EventLoopProvider`` (mounted below this provider) connects the socket.
-  // A delta naming an unregistered substate is skipped with a warning.
-  useIsomorphicLayoutEffect(() => {{
-    dispatchers[substateName] = dispatchSubstate;
-    return () => {{
-      delete dispatchers[substateName];
-    }};
-  }}, [dispatchers, dispatchSubstate, substateName]);
-  return useMemo(
-    () => createElement(StateContexts[contextName], {{ value: state }}, children),
-    [children, state, contextName],
-  );
-}};
+  return {{
+    dispatchers,
+    syncInitialState: (nextInitialState) => {{
+      if (nextInitialState === syncedInitialState) return;
+      syncedInitialState = nextInitialState;
+      Object.entries(nextInitialState).forEach(([substateName, value]) => {{
+        const isNewSubstate = !Object.hasOwn(state, substateName);
+        registerSubstate(substateName, value);
+        const currentState = state[substateName];
+        const addedFields = Object.fromEntries(
+          Object.entries(value).filter(([key]) => !Object.hasOwn(currentState, key)),
+        );
+        if (Object.keys(addedFields).length > 0) {{
+          state[substateName] = {{ ...currentState, ...addedFields }};
+        }}
+        if (isNewSubstate || Object.keys(addedFields).length > 0) {{
+          listeners.get(substateName)?.forEach((listener) => listener());
+        }}
+      }});
+    }},
+    getSnapshot: (substateName) => state[substateName],
+    subscribe: (substateName, listener) => {{
+      let substateListeners = listeners.get(substateName);
+      if (!substateListeners) {{
+        substateListeners = new Set();
+        listeners.set(substateName, substateListeners);
+      }}
+      substateListeners.add(listener);
+      return () => substateListeners.delete(listener);
+    }},
+  }};
+}}
 
-// ``[substateName, contextName]`` for every substate, outermost first.
-const SUBSTATES = [{substates_str}
-];
+function StateStoreProviders({{ children, store }}) {{
+  return createElement(
+    StateStoreContext,
+    {{ value: store }},
+    createElement(DispatchContext, {{ value: store.dispatchers }}, children),
+  );
+}}
 
 function ClientStateProvider({{ children }}) {{
-  return useMemo(() => {{
-    let tree = children;
-    for (let i = SUBSTATES.length - 1; i >= 0; i--) {{
-      const [substateName, contextName] = SUBSTATES[i];
-      tree = createElement(SubstateProvider, {{ substateName, contextName }}, tree);
-    }}
-    return createElement(DispatchProvider, {{}}, tree);
-  }}, [children]);
+  const [store] = useState(createStateStore);
+  useEffect(() => store.syncInitialState(initialState), [store, initialState]);
+  return createElement(StateStoreProviders, {{ children, store }});
 }}
 
-// The server renders once and never applies a delta, so it provides the
-// initial state through bare context providers. ``SubstateProvider`` would add
-// a second render level per substate, and the server renderer recurses once per
-// level, so with many substates rendering a page can exhaust the stack.
+// The server renders once and never applies a delta, so it only reads the
+// initial state.
+const serverStateStore = {{
+  dispatchers: {{}},
+  getSnapshot: (substateName) => initialState[substateName],
+  subscribe: () => () => {{}},
+}};
+
 function ServerStateProvider({{ children }}) {{
-  let tree = children;
-  for (let i = SUBSTATES.length - 1; i >= 0; i--) {{
-    const [substateName, contextName] = SUBSTATES[i];
-    tree = createElement(
-      StateContexts[contextName],
-      {{ value: initialState[substateName] }},
-      tree,
-    );
-  }}
-  return createElement(DispatchContext, {{ value: {{}} }}, tree);
+  return createElement(StateStoreProviders, {{ children, store: serverStateStore }});
 }}
 
 export const StateProvider =

packages/reflex-base/src/reflex_base/registry.py

@@ -32,6 +32,7 @@ def _default_bundled_libraries() -> list[str]:
         "react",
         "@emotion/react",
         f"$/{constants.Dirs.CONTEXTS_PATH}",
+        f"$/{constants.Dirs.UTILS}/context-registry",
         f"$/{constants.Dirs.STATE_PATH}",
     ]
 

packages/reflex-base/src/reflex_base/vars/base.py

@@ -636,19 +636,19 @@ def from_state(cls, state: type[BaseState] | str, field_name: str = "") -> VarDa
         from reflex_base.utils import format
 
         state_name = state if isinstance(state, str) else state.get_full_name()
+        state_var_name = format.format_state_name(state_name)
         return VarData(
             state=state_name,
             field_name=field_name,
             hooks={
-                "const {0} = useContext(StateContexts.{0})".format(
-                    format.format_state_name(state_name)
-                ): None
+                f"const {state_var_name} = useStateContext({json.dumps(state_name)})": None
             },
             imports={
-                f"$/{constants.Dirs.CONTEXTS_PATH}": [ImportVar(tag="StateContexts")],
-                "react": [ImportVar(tag="useContext")],
+                f"$/{constants.Dirs.UTILS}/context-registry": [
+                    ImportVar(tag="useStateContext")
+                ],
             },
-            # State Vars read ``StateContexts``/``EventLoopContext``, so the
+            # State Vars read the shared store/``EventLoopContext``, so the
             # providers must enclose every component that uses them.
             app_wraps=get_event_app_wraps(),
         )

pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
   "starlette >=1.3.1",
   "typing_extensions >=4.13.0",
   "wrapt >=1.17.0,<2.6",
-  "reflex-base >= 0.10.0a2.dev0",
+  "reflex-base >= 0.10.1.dev0",
   "reflex-components-code >= 0.10.0",
   "reflex-components-core >= 0.10.0",
   "reflex-components-dataeditor >= 0.10.0",

reflex/compiler/utils.py

@@ -528,7 +528,7 @@ def compile_experimental_component_memo(
             )
             render._imports_cache = analysis.imports
         # Strings returned by the root's ``add_hooks`` can reference symbols
-        # (``refs``, ``StateContexts``, etc.) that normally reach this module
+        # (``refs``, ``EventLoopContext``, etc.) that normally reach this module
         # through descendants' ``_get_hooks_imports`` / ``_get_imports``. JS
         # imports are side-effect-free and dedup cleanly, so pulling the
         # whole subtree's imports here is safe even when some go unused.

reflex/state.py

@@ -2951,9 +2951,9 @@ def code_uses_state_contexts(javascript_code: str) -> bool:
         javascript_code: The Javascript code to check.
 
     Returns:
-        True if the code attempts to access a member of StateContexts.
+        True if the code subscribes to a state context.
     """
-    return bool("useContext(StateContexts" in javascript_code)
+    return "useStateContext(" in javascript_code
 
 
 def reload_state_module(