fix: flatten providers for large state trees (#7441)

* fix: reduce client state provider nesting

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* fix: avoid copying all client substates per delta

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* fix: flatten large state provider trees

* fix: address state provider review feedback

* docs: add PR 7441 performance news fragment

* docs: add reflex package performance news fragment

* docs: use valid orphan news fragment name

* fix: preserve legacy state context reads

* fix: keep legacy state context reads compatible

* Fix compiler pre-commit checks for legacy context rewrite

* Install greenlet in legacy playground check

* Address PR 7441 review feedback

* Fix remaining PR 7441 review findings

* Fix legacy state read scanning and rewriting

* fix: remove legacy StateContexts rewriter

Drop the compile-time JavaScript rewriter that turned
useContext(StateContexts.x) into a shared-store read. Reflex now
generates useStateContext(...) for every state read on each compile, so
the rewriter only served hand-written hooks. It missed dynamic
components and asset JS, where legacy reads silently became null.

Remove the StateContexts export and getStateContext as well, so any
remaining legacy read fails the build with a missing-export error
instead of reading null. Add a breaking news fragment with the
migration to useStateContext.

* chore: drop breaking fragment for internal StateContexts removal

* fix: address review findings on the shared state store

- Raise reflex's reflex-base floor to 0.10.1.dev0: code_uses_state_contexts
  looks for useStateContext(, which only the new reflex-base generates.
- Drop the unrelated greenlet change in examples.yml.
- Compare applyDelta's result with the current state instead of
  pre-checking the delta, so a deep-equal delta no longer notifies.
- Skip syncInitialState when the store already holds that initial state.
- Serve the server render from a read-only store.
- Remove the duplicate listener set creation and use one useMemo in
  useStateContext.
- Mark the news fragments as bugfix with per-package text.
- Make the JS store test use a structural-sharing applyDelta stub,
  subscribe by state name, and cover an equal list delta.

* test: subscribe through useStateContext in the store test

Attach the test listeners through the subscribe callback the hook passes
to useSyncExternalStore, so the test fails if the hook subscribes to the
wrong substate.

* refactor: seed the state store from its argument

createStateStore takes the initial state and registers substates through
one addSubstates path, which hot reload calls for the substates of new
state classes. Fields added to existing substates are not merged here:
they arrive with the hydrate delta after the backend restarts.

* fix: merge new default fields into live substates on hot reload

mergeInitialState adds new substates and gives existing substates the
fields of the recompiled initial state at their defaults, keeping live
values. A substate is replaced and notified only when it gains a field,
so the mount-time call changes nothing.

---------

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Masen Furer <m_github@0x26.net>
Co-authored-by: Farhan <www.mfarvirus@gmail.com>

Author: 4 people

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:
@@ -343,13 +343,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}"
@@ -416,11 +409,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()
@@ -449,18 +437,17 @@ 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 stateSchema = {json.dumps(state_schema)}
 
 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}
@@ -543,78 +530,80 @@ 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(initialState) {{
+  const state = {{}};
+  const listeners = new Map();
+  const dispatchers = {{}};
+  const notify = (substateName) =>
+    listeners.get(substateName)?.forEach((listener) => listener());
+  // Hot reload passes the recompiled initial state: new substates are added,
+  // and existing substates gain new fields at their defaults while keeping
+  // their live values.
+  const mergeInitialState = (nextInitialState) => {{
+    for (const [substateName, value] of Object.entries(nextInitialState)) {{
+      const liveState = state[substateName];
+      if (liveState !== undefined) {{
+        if (Object.keys(value).some((key) => !Object.hasOwn(liveState, key))) {{
+          state[substateName] = {{ ...value, ...liveState }};
+          notify(substateName);
+        }}
+        continue;
+      }}
+      state[substateName] = value;
+      dispatchers[substateName] = (delta) => {{
+        const currentState = state[substateName];
+        const nextState = applyDelta(currentState, delta);
+        if (nextState === currentState) return;
+        state[substateName] = nextState;
+        notify(substateName);
+      }};
+      notify(substateName);
+    }}
+  }};
+  mergeInitialState(initialState);
 
-// 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,
+    mergeInitialState,
+    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(initialState));
+  useEffect(() => store.mergeInitialState(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(),
         )

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

@@ -3024,9 +3024,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(