[REF-1885] Shard Substates when serializing to Redis

integration/test_client_storage.py

@@ -449,7 +449,7 @@ def set_sub_sub(var: str, value: str):
     assert l1s.text == "l1s value"
 
     # reset the backend state to force refresh from client storage
-    async with client_side.modify_state(token) as state:
+    async with client_side.modify_state(f"{token}_state.client_side_state") as state:
         state.reset()
     driver.refresh()
 

integration/test_dynamic_routes.py

@@ -85,6 +85,7 @@ def dynamic_route(
     """
     with app_harness_env.create(
         root=tmp_path_factory.mktemp(f"dynamic_route"),
+        app_name=f"dynamicroute_{app_harness_env.__name__.lower()}",
         app_source=DynamicRoute,  # type: ignore
     ) as harness:
         yield harness
@@ -146,7 +147,7 @@ def poll_for_order(
 
     async def _poll_for_order(exp_order: list[str]):
         async def _backend_state():
-            return await dynamic_route.get_state(token)
+            return await dynamic_route.get_state(f"{token}_state.dynamic_state")
 
         async def _check():
             return (await _backend_state()).substates[
@@ -194,7 +195,9 @@ async def test_on_load_navigate(
         assert link
         assert page_id_input
 
-        assert dynamic_route.poll_for_value(page_id_input) == str(ix)
+        assert dynamic_route.poll_for_value(
+            page_id_input, exp_not_equal=str(ix - 1)
+        ) == str(ix)
         assert dynamic_route.poll_for_value(raw_path_input) == f"/page/{ix}/"
     await poll_for_order(exp_order)
 
@@ -220,7 +223,9 @@ async def test_on_load_navigate(
     with poll_for_navigation(driver):
         driver.get(f"{driver.current_url}?foo=bar")
     await poll_for_order(exp_order)
-    assert (await dynamic_route.get_state(token)).router.page.params["foo"] == "bar"
+    assert (
+        await dynamic_route.get_state(f"{token}_state.dynamic_state")
+    ).router.page.params["foo"] == "bar"
 
     # hit a 404 and ensure we still hydrate
     exp_order += ["/404-no page id"]

integration/test_event_actions.py

@@ -207,7 +207,7 @@ def poll_for_order(
 
     async def _poll_for_order(exp_order: list[str]):
         async def _backend_state():
-            return await event_action.get_state(token)
+            return await event_action.get_state(f"{token}_state.event_action_state")
 
         async def _check():
             return (await _backend_state()).substates[

integration/test_event_chain.py

@@ -298,7 +298,7 @@ def assert_token(event_chain: AppHarness, driver: WebDriver) -> str:
     token = event_chain.poll_for_value(token_input)
     assert token is not None
 
-    return token
+    return f"{token}_state.state"
 
 
 @pytest.mark.parametrize(

integration/test_form_submit.py

@@ -221,7 +221,11 @@ async def test_submit(driver, form_submit: AppHarness):
     submit_input.click()
 
     async def get_form_data():
-        return (await form_submit.get_state(token)).substates["form_state"].form_data
+        return (
+            (await form_submit.get_state(f"{token}_state.form_state"))
+            .substates["form_state"]
+            .form_data
+        )
 
     # wait for the form data to arrive at the backend
     form_data = await AppHarness._poll_for_async(get_form_data)

integration/test_input.py

@@ -76,6 +76,10 @@ async def test_fully_controlled_input(fully_controlled_input: AppHarness):
     token = fully_controlled_input.poll_for_value(token_input)
     assert token
 
+    async def get_state_text():
+        state = await fully_controlled_input.get_state(f"{token}_state.state")
+        return state.substates["state"].text
+
     # find the input and wait for it to have the initial state value
     debounce_input = driver.find_element(By.ID, "debounce_input_input")
     value_input = driver.find_element(By.ID, "value_input")
@@ -95,16 +99,14 @@ async def test_fully_controlled_input(fully_controlled_input: AppHarness):
     debounce_input.send_keys("foo")
     time.sleep(0.5)
     assert debounce_input.get_attribute("value") == "ifoonitial"
-    assert (await fully_controlled_input.get_state(token)).substates[
-        "state"
-    ].text == "ifoonitial"
+    assert await get_state_text() == "ifoonitial"
     assert fully_controlled_input.poll_for_value(value_input) == "ifoonitial"
     assert fully_controlled_input.poll_for_value(plain_value_input) == "ifoonitial"
 
     # clear the input on the backend
-    async with fully_controlled_input.modify_state(token) as state:
+    async with fully_controlled_input.modify_state(f"{token}_state.state") as state:
         state.substates["state"].text = ""
-    assert (await fully_controlled_input.get_state(token)).substates["state"].text == ""
+    assert await get_state_text() == ""
     assert (
         fully_controlled_input.poll_for_value(
             debounce_input, exp_not_equal="ifoonitial"
@@ -116,9 +118,7 @@ async def test_fully_controlled_input(fully_controlled_input: AppHarness):
     debounce_input.send_keys("getting testing done")
     time.sleep(0.5)
     assert debounce_input.get_attribute("value") == "getting testing done"
-    assert (await fully_controlled_input.get_state(token)).substates[
-        "state"
-    ].text == "getting testing done"
+    assert await get_state_text() == "getting testing done"
     assert fully_controlled_input.poll_for_value(value_input) == "getting testing done"
     assert (
         fully_controlled_input.poll_for_value(plain_value_input)
@@ -130,9 +130,7 @@ async def test_fully_controlled_input(fully_controlled_input: AppHarness):
     time.sleep(0.5)
     assert debounce_input.get_attribute("value") == "overwrite the state"
     assert on_change_input.get_attribute("value") == "overwrite the state"
-    assert (await fully_controlled_input.get_state(token)).substates[
-        "state"
-    ].text == "overwrite the state"
+    assert await get_state_text() == "overwrite the state"
     assert fully_controlled_input.poll_for_value(value_input) == "overwrite the state"
     assert (
         fully_controlled_input.poll_for_value(plain_value_input)

integration/test_upload.py

@@ -171,6 +171,7 @@ async def test_upload_file(
     # wait for the backend connection to send the token
     token = upload_file.poll_for_value(token_input)
     assert token is not None
+    substate_token = f"{token}_state.upload_state"
 
     suffix = "_secondary" if secondary else ""
 
@@ -191,7 +192,11 @@ async def test_upload_file(
 
     # look up the backend state and assert on uploaded contents
     async def get_file_data():
-        return (await upload_file.get_state(token)).substates["upload_state"]._file_data
+        return (
+            (await upload_file.get_state(substate_token))
+            .substates["upload_state"]
+            ._file_data
+        )
 
     file_data = await AppHarness._poll_for_async(get_file_data)
     assert isinstance(file_data, dict)
@@ -201,7 +206,7 @@ async def get_file_data():
     selected_files = driver.find_element(By.ID, f"selected_files{suffix}")
     assert selected_files.text == exp_name
 
-    state = await upload_file.get_state(token)
+    state = await upload_file.get_state(substate_token)
     if secondary:
         # only the secondary form tracks progress and chain events
         assert state.substates["upload_state"].event_order.count("upload_progress") == 1
@@ -223,6 +228,7 @@ async def test_upload_file_multiple(tmp_path, upload_file: AppHarness, driver):
     # wait for the backend connection to send the token
     token = upload_file.poll_for_value(token_input)
     assert token is not None
+    substate_token = f"{token}_state.upload_state"
 
     upload_box = driver.find_element(By.XPATH, "//input[@type='file']")
     assert upload_box
@@ -250,7 +256,11 @@ async def test_upload_file_multiple(tmp_path, upload_file: AppHarness, driver):
 
     # look up the backend state and assert on uploaded contents
     async def get_file_data():
-        return (await upload_file.get_state(token)).substates["upload_state"]._file_data
+        return (
+            (await upload_file.get_state(substate_token))
+            .substates["upload_state"]
+            ._file_data
+        )
 
     file_data = await AppHarness._poll_for_async(get_file_data)
     assert isinstance(file_data, dict)
@@ -330,6 +340,7 @@ async def test_cancel_upload(tmp_path, upload_file: AppHarness, driver: WebDrive
     # wait for the backend connection to send the token
     token = upload_file.poll_for_value(token_input)
     assert token is not None
+    substate_token = f"{token}_state.upload_state"
 
     upload_box = driver.find_elements(By.XPATH, "//input[@type='file']")[1]
     upload_button = driver.find_element(By.ID, f"upload_button_secondary")
@@ -347,7 +358,7 @@ async def test_cancel_upload(tmp_path, upload_file: AppHarness, driver: WebDrive
     cancel_button.click()
 
     # look up the backend state and assert on progress
-    state = await upload_file.get_state(token)
+    state = await upload_file.get_state(substate_token)
     assert state.substates["upload_state"].progress_dicts
     assert exp_name not in state.substates["upload_state"]._file_data
 

reflex/app.py

@@ -926,7 +926,7 @@ async def process(
         }
     )
     # Get the state for the session exclusively.
-    async with app.state_manager.modify_state(event.token) as state:
+    async with app.state_manager.modify_state(event.substate_token) as state:
         # re-assign only when the value is different
         if state.router_data != router_data:
             # assignment will recurse into substates and force recalculation of
@@ -1002,7 +1002,8 @@ async def upload_file(request: Request, files: List[UploadFile]):
             )
 
         # Get the state for the session.
-        state = await app.state_manager.get_state(token)
+        substate_token = token + "_" + handler.rpartition(".")[0]
+        state = await app.state_manager.get_state(substate_token)
 
         # get the current session ID
         # get the current state(parent state/substate)
@@ -1049,7 +1050,7 @@ async def _ndjson_updates():
                 Each state update as JSON followed by a new line.
             """
             # Process the event.
-            async with app.state_manager.modify_state(token) as state:
+            async with app.state_manager.modify_state(event.substate_token) as state:
                 async for update in state._process(event):
                     # Postprocess the event.
                     update = await app.postprocess(state, event, update)

reflex/event.py

@@ -41,6 +41,16 @@ class Event(Base):
     # The event payload.
     payload: Dict[str, Any] = {}
 
+    @property
+    def substate_token(self) -> str:
+        """Get the substate token for the event.
+
+        Returns:
+            The substate token.
+        """
+        substate = self.name.rpartition(".")[0]
+        return f"{self.token}_{substate}"
+
 
 BACKGROUND_TASK_MARKER = "_reflex_background_task"