DO NOT MERGE - Disconnect shared state

news/+shared-state-disconnect-reap.feature.md

@@ -0,0 +1 @@
+Disconnected clients are unsubscribed from their linked shared states after a reconnect grace period (`REFLEX_SHARED_STATE_DISCONNECT_GRACE`, default 30s, 0 disables). The new `SharedState._on_subscriber_disconnected` hook lets shared states clean up per-client data (e.g. presence bookkeeping); a client reconnecting later re-subscribes automatically with its next event.

news/6934.bugfix.md

@@ -0,0 +1 @@
+Shared state updates now reach linked clients connected to other backend instances — the fan-out previously skipped any client whose websocket was not connected to the instance processing the event, so with redis and multiple workers only same-instance clients received live updates.

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

@@ -688,6 +688,10 @@ class EnvironmentVariables:
     # The address to bind the HTTP client to. You can set this to "::" to enable IPv6.
     REFLEX_HTTP_CLIENT_BIND_ADDRESS: EnvVar[str | None] = env_var(None)
 
+    # Seconds a disconnected client may reconnect before it is unsubscribed
+    # from the shared states it was linked to. 0 disables the unsubscription.
+    REFLEX_SHARED_STATE_DISCONNECT_GRACE: EnvVar[int] = env_var(30)
+
     # Maximum size of the message in the websocket server in bytes.
     REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE: EnvVar[int] = env_var(
         constants.POLLING_MAX_HTTP_BUFFER_SIZE

reflex/app.py

@@ -2010,10 +2010,15 @@ def on_disconnect(self, sid: str) -> asyncio.Task | None:
         Returns:
             An asyncio Task for cleaning up the token, or None.
         """
+        from reflex.istate.shared import schedule_disconnect_reap
+
         self._client_error_counts.pop(sid, None)
         # Get token before cleaning up
         disconnect_token = self.sid_to_token.get(sid)
         if disconnect_token:
+            # Unsubscribe the client from its linked shared states unless it
+            # reconnects within the grace period.
+            schedule_disconnect_reap(self.app, disconnect_token)
             # Use async cleanup through token manager
             task = asyncio.create_task(
                 self._token_manager.disconnect_token(disconnect_token, sid),

reflex/istate/shared.py

@@ -3,10 +3,12 @@
 import asyncio
 import contextlib
 import logging
+import time
 from collections.abc import AsyncIterator
-from typing import TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
 from reflex_base.constants import ROUTER_DATA
+from reflex_base.environment import environment
 from reflex_base.event import Event, get_hydrate_event
 from reflex_base.registry import RegistrationContext
 from reflex_base.utils.exceptions import ReflexRuntimeError
@@ -15,9 +17,13 @@
 from reflex.istate.manager.token import BaseStateToken
 from reflex.state import BaseState, State, _override_base_method
 
+if TYPE_CHECKING:
+    from reflex.app import App
+
 logger = logging.getLogger(__name__)
 
 UPDATE_OTHER_CLIENT_TASKS: set[asyncio.Task] = set()
+DISCONNECT_REAP_TASKS: dict[str, asyncio.Task] = {}
 LINKED_STATE = TypeVar("LINKED_STATE", bound="SharedStateBaseInternal")
 
 
@@ -54,28 +60,113 @@ def _do_update_other_tokens(
     """
     app = RegistrationContext.get().app
 
+    tasks = []
+    if (event_namespace := app.event_namespace) is None:
+        return tasks
+    token_manager = event_namespace._token_manager
+
     async def _update_client(token: str):
+        # Don't send updates for disconnected clients; emit_update relays the
+        # delta to the owning instance if the socket lives elsewhere.
+        if not await token_manager.is_token_connected(token):
+            return
         async with app.modify_state(
             BaseStateToken(ident=token, cls=state_type),
             previous_dirty_vars=previous_dirty_vars,
         ):
             pass
 
-    tasks = []
-    if (event_namespace := app.event_namespace) is None:
-        return tasks
     for affected_token in affected_tokens:
-        # Don't send updates for disconnected clients.
-        if affected_token not in event_namespace._token_manager.token_to_socket:
-            continue
-        # TODO: remove disconnected clients after some time.
+        # Disconnected clients are removed by schedule_disconnect_reap after
+        # the reconnect grace; until then the connectivity check above skips
+        # them.
         t = asyncio.create_task(_update_client(affected_token))
         UPDATE_OTHER_CLIENT_TASKS.add(t)
         t.add_done_callback(_log_update_client_errors)
         tasks.append(t)
     return tasks
 
 
+def schedule_disconnect_reap(app: "App", token: str) -> asyncio.Task | None:
+    """Schedule unsubscribing a disconnected client from its linked shared states.
+
+    Called on client disconnect. The reap runs after a grace period
+    (REFLEX_SHARED_STATE_DISCONNECT_GRACE) so brief reconnects (page reloads,
+    network blips) are no-ops; a client reaped too eagerly re-subscribes on its
+    next event through _internal_patch_linked_state. A new disconnect for the
+    same token restarts the grace.
+
+    Args:
+        app: The application object.
+        token: The client token that disconnected.
+
+    Returns:
+        The scheduled reap task, or None when reaping is disabled.
+    """
+    grace = environment.REFLEX_SHARED_STATE_DISCONNECT_GRACE.get()
+    if grace <= 0 or app._state is None:
+        return None
+    if (previous := DISCONNECT_REAP_TASKS.pop(token, None)) is not None:
+        previous.cancel()
+
+    task = asyncio.create_task(
+        _reap_disconnected_client(app, token, grace),
+        name=f"reflex_shared_state_reap|{token}|{time.time()}",
+    )
+    DISCONNECT_REAP_TASKS[token] = task
+
+    def _on_done(task: asyncio.Task) -> None:
+        if DISCONNECT_REAP_TASKS.get(token) is task:
+            DISCONNECT_REAP_TASKS.pop(token, None)
+        if not task.cancelled() and (exc := task.exception()) is not None:
+            logger.warning(f"Error reaping disconnected shared state client: {exc}")
+
+    task.add_done_callback(_on_done)
+    return task
+
+
+async def _reap_disconnected_client(app: "App", token: str, grace: float) -> None:
+    """Unsubscribe a client from its linked shared states unless it reconnected.
+
+    Args:
+        app: The application object.
+        token: The client token that disconnected.
+        grace: Seconds to wait for a reconnect before unsubscribing.
+    """
+    await asyncio.sleep(grace)
+    if (event_namespace := app.event_namespace) is None:
+        return
+    if await event_namespace._token_manager.is_token_connected(token):
+        # The client reconnected, possibly to another instance.
+        return
+    if app._state is None:
+        return
+    # Read-only peek at the client's links; racing a concurrent event is fine,
+    # since any later event through the link re-adds the client anyway.
+    root_state = await app.state_manager.get_state(
+        BaseStateToken(ident=token, cls=app._state)
+    )
+    if not isinstance(root_state, State):
+        return
+    links = dict(root_state._reflex_internal_links or {})
+    for state_name, linked_token in links.items():
+        try:
+            state_cls = app._state.get_class_substate(state_name)
+            async with app.modify_state(
+                BaseStateToken(ident=linked_token, cls=state_cls)
+            ) as shared_root:
+                shared = await shared_root.get_state(state_cls)
+                if (
+                    not isinstance(shared, SharedState)
+                    or token not in shared._linked_from
+                ):
+                    continue
+                shared._linked_from.discard(token)
+                await shared._on_subscriber_disconnected(token)
+        except Exception as e:
+            logger.warning(f"Error unsubscribing disconnected shared state client: {e}")
+
+
 @contextlib.asynccontextmanager
 async def _patch_state(
     original_state: BaseState, linked_state: BaseState, full_delta: bool = False
@@ -503,6 +594,18 @@ class SharedState(SharedStateBaseInternal, mixin=True):
     _linked_to: str = ""
     _previous_dirty_vars: set[str] = set()
 
+    async def _on_subscriber_disconnected(self, client_token: str) -> None:
+        """Hook called when a subscribed client is unsubscribed after disconnecting.
+
+        Override to clean up per-client data kept on the shared state (e.g.
+        presence bookkeeping). Runs with the shared token's state locked, after
+        the client token was removed from the subscriber set; mutations
+        propagate to the remaining linked clients.
+
+        Args:
+            client_token: The client token that was unsubscribed.
+        """
+
     @classmethod
     def __init_subclass__(cls, **kwargs):
         """Initialize subclass and set up shared state fields.

reflex/utils/token_manager.py

@@ -81,6 +81,17 @@ async def enumerate_tokens(self) -> AsyncIterator[str]:
         for token in self.token_to_socket:
             yield token
 
+    async def is_token_connected(self, token: str) -> bool:
+        """Whether the token has a connected client socket on any instance.
+
+        Args:
+            token: The client token.
+
+        Returns:
+            True if the token has a connected socket.
+        """
+        return token in self.token_to_socket
+
     @abstractmethod
     async def link_token_to_sid(self, token: str, sid: str) -> str | None:
         """Link a token to a session ID.
@@ -431,17 +442,72 @@ async def _get_token_owner(self, token: str, refresh: bool = False) -> str | Non
         ):
             return socket_record.instance_id
 
-        redis_key = self._get_redis_key(token)
         try:
-            record_pkl = await self.redis.get(redis_key)
-            if record_pkl:
-                socket_record = pickle.loads(record_pkl)
-                self.token_to_socket[token] = socket_record
-                self.sid_to_token[socket_record.sid] = token
-                return socket_record.instance_id
+            socket_record = await self._fetch_socket_record(token)
         except Exception as e:
             logger.error(f"Redis error getting token owner: {e}")
-        return None
+            return None
+        return socket_record.instance_id if socket_record is not None else None
+
+    async def _fetch_socket_record(self, token: str) -> SocketRecord | None:
+        """Fetch the socket record for a token from redis and cache it.
+
+        Unlike _get_token_owner, redis errors propagate to the caller so it
+        can distinguish a lookup failure from an absent record.
+
+        Args:
+            token: The client token.
+
+        Returns:
+            The refreshed socket record, or None if the token has none.
+        """
+        record_pkl = await self.redis.get(self._get_redis_key(token))
+        if not record_pkl:
+            return None
+        socket_record = pickle.loads(record_pkl)
+        # Drop the reverse mapping of a superseded record (client moved sids).
+        if (
+            (previous := self.token_to_socket.get(token)) is not None
+            and previous.sid != socket_record.sid
+            and self.sid_to_token.get(previous.sid) == token
+        ):
+            self.sid_to_token.pop(previous.sid, None)
+        self.token_to_socket[token] = socket_record
+        self.sid_to_token[socket_record.sid] = token
+        return socket_record
+
+    async def is_token_connected(self, token: str) -> bool:
+        """Whether the token has a connected client socket on any instance.
+
+        A record owned by this instance is authoritative. A cached record
+        from another instance may be stale (the client may have reconnected
+        elsewhere), so the socket record is refreshed from redis instead,
+        and dropped from the local cache if the client is gone. If the
+        refresh fails, the cached record is preserved and trusted.
+
+        Args:
+            token: The client token.
+
+        Returns:
+            True if the token has a connected socket on any instance.
+        """
+        if (
+            socket_record := self.token_to_socket.get(token)
+        ) is not None and socket_record.instance_id == self.instance_id:
+            return True
+        try:
+            if await self._fetch_socket_record(token) is not None:
+                return True
+        except Exception as e:
+            logger.warning(f"Redis error checking token connection: {e}")
+            return socket_record is not None
+        if (
+            socket_record is not None
+            and self.token_to_socket.get(token) is socket_record
+        ):
+            self.token_to_socket.pop(token, None)
+            self.sid_to_token.pop(socket_record.sid, None)
+        return False
 
     async def emit_lost_and_found(
         self,