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.

packages/reflex-base/news/+shared-state-disconnect-reap.feature.md

@@ -0,0 +1 @@
+Add `REFLEX_SHARED_STATE_DISCONNECT_GRACE` (default 30s, 0 disables), the number of seconds a disconnected client may reconnect before it is unsubscribed from its linked shared states.

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

@@ -797,6 +797,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

@@ -2312,6 +2312,8 @@ 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
+
         if otel.enabled:
             otel.record_connection(-1)
         self._client_error_counts.pop(sid, None)
@@ -2327,6 +2329,9 @@ def on_disconnect(self, sid: str) -> asyncio.Task | None:
                 and (processor := self.app._event_processor) is not None
             ):
                 processor._on_disconnect(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,11 +3,13 @@
 import asyncio
 import contextlib
 import logging
+import time
 from collections.abc import AsyncIterator
 from dataclasses import MISSING
 from typing import TYPE_CHECKING, Self, TypeVar
 
 from reflex_base.constants import ROUTER_DATA, ROUTER_VARS
+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
@@ -17,9 +19,16 @@
 from reflex.istate.manager.token import BaseStateToken
 from reflex.state import BaseState, OnLoadInternalState, State, _override_base_method
 
+if TYPE_CHECKING:
+    from reflex.app import App
+
 logger = logging.getLogger(__name__)
 
 UPDATE_OTHER_CLIENT_TASKS: set[asyncio.Task] = set()
+# The latest reap per client token, the only one allowed to unlink it.
+DISCONNECT_REAP_TASKS: dict[str, asyncio.Task] = {}
+# Reaps that started unlinking; a newer disconnect no longer cancels them.
+STARTED_REAP_TASKS: set[asyncio.Task] = set()
 LINKED_STATE = TypeVar("LINKED_STATE", bound="SharedStateBaseInternal")
 
 
@@ -73,14 +82,111 @@ async def _update_client(token: str):
             pass
 
     for affected_token in affected_tokens:
-        # 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; a cleanup already under way finishes the
+    shared state it is unlinking and leaves the rest to the new reap.
+
+    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.get(token)
+    ) is not None and previous not in STARTED_REAP_TASKS:
+        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)
+        STARTED_REAP_TASKS.discard(task)
+        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
+    token_manager = event_namespace._token_manager
+    if await 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
+                # Recheck under the lock: the client may have reconnected, or
+                # disconnected again and left this to a newer reap.
+                if (
+                    await token_manager.is_token_connected(token)
+                    or (task := asyncio.current_task()) is None
+                    or DISCONNECT_REAP_TASKS.get(token, task) is not task
+                ):
+                    return
+                # Unlinking starts; a newer disconnect must not cancel it midway.
+                STARTED_REAP_TASKS.add(task)
+                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
@@ -536,6 +642,18 @@ def _init_bookkeeping(self, parent_state: BaseState | None = None) -> None:
         super()._init_bookkeeping(parent_state)
         self._previous_dirty_vars = 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.