Skip to content

Events against SharedState should not mark root state dirty all the time #6392

@masenf

Description

@masenf

Describe the bug

(added to tests/units/test_state.py)

class _LinkedStatePatchRoot(BaseState):
    """Root state for testing _patch_state dirty propagation regression."""


class _LinkedStatePatchShared(_LinkedStatePatchRoot):
    """Substate used to exercise _patch_state without full SharedState infrastructure."""

    counter: int = 0


@pytest.mark.asyncio
async def test_linked_state_event_does_not_dirty_root_state():
    """Linked-state event must not leave root state dirty.

    Regression: _patch_state unconditionally added "router" and ROUTER_DATA to
    root_state.dirty_vars even for ordinary linked-state events (full_delta=False).
    Because _get_resolved_delta() does not call _clean(), those dirty vars
    survived the yield and were included in the final get_delta() call, producing
    a spurious root-state delta on every linked-state event and causing extra
    frontend re-renders.
    """
    from reflex.istate.shared import _patch_state
    from reflex_base.constants import ROUTER_DATA

    # Two separate state trees: "original" (private token) and "linked" (shared token).
    tree1 = _LinkedStatePatchRoot()
    tree2 = _LinkedStatePatchRoot()

    shared_state_name = _LinkedStatePatchShared.get_name()
    original_state = tree1.substates[shared_state_name]
    linked_state = tree2.substates[shared_state_name]

    assert isinstance(original_state, _LinkedStatePatchShared)
    assert isinstance(linked_state, _LinkedStatePatchShared)

    tree1._clean()

    # _patch_state with full_delta=False simulates a regular event on a linked state.
    async with _patch_state(original_state, linked_state, full_delta=False):
        linked_state.counter = 1  # simulate event handler modifying the shared state

    # After the event, the root state (tree1) must not carry "router" or ROUTER_DATA
    # in its dirty_vars.  Those vars were set inside _patch_state to help compute an
    # initial delta, but they must be cleaned before yielding so they do not leak
    # into the final delta computation.
    assert "router" not in tree1.dirty_vars, (
        "root state was left dirty with 'router' after a linked-state event; "
        "_patch_state must clean root dirty vars before yielding"
    )
    assert ROUTER_DATA not in tree1.dirty_vars, (
        f"root state was left dirty with '{ROUTER_DATA}' after a linked-state event; "
        "_patch_state must clean root dirty vars before yielding"
    )
    # The root state itself must not appear in the delta — only the shared substate should.
    delta = tree1.get_delta()
    assert tree1.get_full_name() not in delta, (
        f"root state appeared in delta after a linked-state event: {delta}"
    )

It appears that this is by design however...

# Apply the updates into the existing state tree for rehydrate.
root_state = original_state._get_root_state()
root_state.dirty_vars.add("router")
root_state.dirty_vars.add(ROUTER_DATA)
root_state._mark_dirty()

The code is explicitly dirtying the root state and causing it to be sent in each event delta after linked states are set up. The purpose for this is to force recomputation of vars that depend on router_data, like the client_token. If a shared state had a computed var that depended on the client_token, then it would need to be explicitly marked dirty in order to recompute the var for each token it is linked from.

The problem comes in that the root state router_data has not actually changed, so it's wasteful to send it to the client, and in some cases can cause excess re-renders on the frontend.

A nice fix for this might look through the computed vars in the patched state for ones that depend on other states and then mark those dirty whenever they get patched in. It's not great for performance, but it's really the only way to prevent data leakage between private states when a shared state depends on private data.

The proper pattern for implementing shared state with computed views using private data would put the computed var in the private state that depends on data from the shared state. Then the computed data is cached with the private state and only needs to be recomputed if the shared data or private data changes, rather than every time the shared state is patched to a different private state tree.

  • Python Version: 3.10+
  • Reflex Version: 0.8.23+
  • OS: any

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions