Implement lost+found for StateUpdate without websocket

masenf · masenf · commit e71eece4c2a8 · 2025-10-24T00:39:37.000-07:00
When an update is emitted for a token, but the websocket for that token is on
another instance of the app, post it to the lost+found channel where other
instances are listening for updates to send to their clients.
diff --git a/reflex/app.py b/reflex/app.py
@@ -120,7 +120,7 @@
 )
 from reflex.utils.imports import ImportVar
 from reflex.utils.misc import run_in_thread
-from reflex.utils.token_manager import TokenManager
+from reflex.utils.token_manager import RedisTokenManager, TokenManager
 from reflex.utils.types import ASGIApp, Message, Receive, Scope, Send
 
 if TYPE_CHECKING:
@@ -2061,6 +2061,9 @@ async def on_connect(self, sid: str, environ: dict):
             sid: The Socket.IO session id.
             environ: The request information, including HTTP headers.
         """
+        if isinstance(self._token_manager, RedisTokenManager):
+            # Make sure this instance is watching for updates from other instances.
+            self._token_manager.ensure_lost_and_found_task(self.emit_update)
         query_params = urllib.parse.parse_qs(environ.get("QUERY_STRING", ""))
         token_list = query_params.get("token", [])
         if token_list:
@@ -2102,16 +2105,30 @@ async def emit_update(self, update: StateUpdate, token: str) -> None:
             token: The client token (tab) associated with the event.
         """
         client_token, _ = _split_substate_key(token)
-        sid = self.token_to_sid.get(client_token)
-        if sid is None:
-            # If the sid is None, we are not connected to a client. Prevent sending
-            # updates to all clients.
-            console.warn(f"Attempting to send delta to disconnected client {token!r}")
+        socket_record = self._token_manager.token_to_socket.get(client_token)
+        if (
+            socket_record is None
+            or socket_record.instance_id != self._token_manager.instance_id
+        ):
+            if isinstance(self._token_manager, RedisTokenManager):
+                # The socket belongs to another instance of the app, send it to the lost and found.
+                if not await self._token_manager.emit_lost_and_found(
+                    client_token, update
+                ):
+                    console.warn(
+                        f"Failed to send delta to lost and found for client {token!r}"
+                    )
+            else:
+                # If the socket record is None, we are not connected to a client. Prevent sending
+                # updates to all clients.
+                console.warn(
+                    f"Attempting to send delta to disconnected client {token!r}"
+                )
             return
         # Creating a task prevents the update from being blocked behind other coroutines.
         await asyncio.create_task(
-            self.emit(str(constants.SocketEvent.EVENT), update, to=sid),
-            name=f"reflex_emit_event|{token}|{sid}|{time.time()}",
+            self.emit(str(constants.SocketEvent.EVENT), update, to=socket_record.sid),
+            name=f"reflex_emit_event|{token}|{socket_record.sid}|{time.time()}",
         )
 
     async def on_event(self, sid: str, data: Any):
diff --git a/reflex/utils/token_manager.py b/reflex/utils/token_manager.py
@@ -7,11 +7,12 @@
 import json
 import uuid
 from abc import ABC, abstractmethod
+from collections.abc import Callable, Coroutine
 from types import MappingProxyType
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from reflex.istate.manager.redis import StateManagerRedis
-from reflex.state import BaseState
+from reflex.state import BaseState, StateUpdate
 from reflex.utils import console, prerequisites
 
 if TYPE_CHECKING:
@@ -35,6 +36,14 @@ class SocketRecord:
     sid: str
 
 
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class LostAndFoundRecord:
+    """Record for a StateUpdate for a token with its socket on another instance."""
+
+    token: str
+    update: dict[str, Any]
+
+
 class TokenManager(ABC):
     """Abstract base class for managing client token to session ID mappings."""
 
@@ -176,7 +185,10 @@ def __init__(self, redis: Redis):
 
         config = get_config()
         self.token_expiration = config.redis_token_expiration
-        self._update_task = None
+
+        # Pub/sub tasks for handling sockets owned by other instances.
+        self._socket_record_task: asyncio.Task | None = None
+        self._lost_and_found_task: asyncio.Task | None = None
 
     def _get_redis_key(self, token: str) -> str:
         """Get Redis key for token mapping.
@@ -189,7 +201,53 @@ def _get_redis_key(self, token: str) -> str:
         """
         return f"token_manager_socket_record_{token}"
 
-    async def _socket_record_update_task(self) -> None:
+    def _handle_socket_record_del(self, token: str) -> None:
+        """Handle deletion of a socket record from Redis.
+
+        Args:
+            token: The client token whose record was deleted.
+        """
+        if (
+            socket_record := self.token_to_socket.pop(token, None)
+        ) is not None and socket_record.instance_id != self.instance_id:
+            self.sid_to_token.pop(socket_record.sid, None)
+
+    async def _handle_socket_record_set(self, token: str) -> None:
+        """Handle setting/updating of a socket record from Redis.
+
+        Args:
+            token: The client token whose record was set/updated.
+        """
+        # Fetch updated record from Redis
+        record_json = await self.redis.get(self._get_redis_key(token))
+        if record_json:
+            record_data = json.loads(record_json)
+            socket_record = SocketRecord(**record_data)
+            self.token_to_socket[token] = socket_record
+            self.sid_to_token[socket_record.sid] = token
+
+    async def _subscribe_socket_record_updates(self, redis_db: int) -> None:
+        """Subscribe to Redis keyspace notifications for socket record updates."""
+        async with self.redis.pubsub() as pubsub:
+            await pubsub.psubscribe(
+                f"__keyspace@{redis_db}__:token_manager_socket_record_*"
+            )
+            async for message in pubsub.listen():
+                if message["type"] == "pmessage":
+                    key = message["channel"].split(b":", 1)[1].decode()
+                    token = key.replace("token_manager_socket_record_", "")
+
+                    if token not in self.token_to_socket:
+                        # We don't know about this token, skip
+                        continue
+
+                    event = message["data"].decode()
+                    if event in ("del", "expired", "evicted"):
+                        self._handle_socket_record_del(token)
+                    elif event == "set":
+                        await self._handle_socket_record_set(token)
+
+    async def _socket_record_updates_forever(self) -> None:
         """Background task to monitor Redis keyspace notifications for socket record updates."""
         await StateManagerRedis(
             state=BaseState, redis=self.redis
@@ -203,33 +261,12 @@ async def _socket_record_update_task(self) -> None:
             except Exception as e:
                 console.error(f"RedisTokenManager socket record update task error: {e}")
 
-    async def _subscribe_socket_record_updates(self, redis_db: int) -> None:
-        """Subscribe to Redis keyspace notifications for socket record updates."""
-        pubsub = self.redis.pubsub()
-        await pubsub.psubscribe(
-            f"__keyspace@{redis_db}__:token_manager_socket_record_*"
-        )
-
-        async for message in pubsub.listen():
-            if message["type"] == "pmessage":
-                key = message["channel"].split(b":", 1)[1].decode()
-                event = message["data"].decode()
-                token = key.replace("token_manager_socket_record_", "")
-
-                if event in ("del", "expired", "evicted"):
-                    # Remove from local dicts if exists
-                    if (
-                        socket_record := self.token_to_socket.pop(token, None)
-                    ) is not None:
-                        self.sid_to_token.pop(socket_record.sid, None)
-                elif event == "set":
-                    # Fetch updated record from Redis
-                    record_json = await self.redis.get(key)
-                    if record_json:
-                        record_data = json.loads(record_json)
-                        socket_record = SocketRecord(**record_data)
-                        self.token_to_socket[token] = socket_record
-                        self.sid_to_token[socket_record.sid] = token
+    def _ensure_socket_record_task(self) -> None:
+        """Ensure the socket record updates subscriber task is running."""
+        if self._socket_record_task is None or self._socket_record_task.done():
+            self._socket_record_task = asyncio.create_task(
+                self._socket_record_updates_forever()
+            )
 
     async def link_token_to_sid(self, token: str, sid: str) -> str | None:
         """Link a token to a session ID with Redis-based duplicate detection.
@@ -248,8 +285,7 @@ async def link_token_to_sid(self, token: str, sid: str) -> str | None:
             return None  # Same token, same SID = reconnection, no Redis check needed
 
         # Make sure the update subscriber is running
-        if self._update_task is None or self._update_task.done():
-            self._update_task = asyncio.create_task(self._socket_record_update_task())
+        self._ensure_socket_record_task()
 
         # Check Redis for cross-worker duplicates
         redis_key = self._get_redis_key(token)
@@ -293,8 +329,10 @@ async def disconnect_token(self, token: str, sid: str) -> None:
         """
         # Only clean up if we own it locally (fast ownership check)
         if (
-            socket_record := self.token_to_socket.get(token)
-        ) is not None and socket_record.sid == sid:
+            (socket_record := self.token_to_socket.get(token)) is not None
+            and socket_record.sid == sid
+            and socket_record.instance_id == self.instance_id
+        ):
             # Clean up Redis
             redis_key = self._get_redis_key(token)
             try:
@@ -304,3 +342,124 @@ async def disconnect_token(self, token: str, sid: str) -> None:
 
             # Clean up local dicts (always do this)
             await super().disconnect_token(token, sid)
+
+    @staticmethod
+    def _get_lost_and_found_key(instance_id: str) -> str:
+        """Get the Redis key for lost and found deltas for an instance.
+
+        Args:
+            instance_id: The instance ID.
+
+        Returns:
+            The Redis key for lost and found deltas.
+        """
+        return f"token_manager_lost_and_found_{instance_id}"
+
+    async def _subscribe_lost_and_found_updates(
+        self,
+        emit_update: Callable[[StateUpdate, str], Coroutine[None, None, None]],
+    ) -> None:
+        """Subscribe to Redis channel notifications for lost and found deltas.
+
+        Args:
+            emit_update: The function to emit state updates.
+        """
+        async with self.redis.pubsub() as pubsub:
+            await pubsub.psubscribe(
+                f"channel:{self._get_lost_and_found_key(self.instance_id)}"
+            )
+            async for message in pubsub.listen():
+                if message["type"] == "pmessage":
+                    record = LostAndFoundRecord(**json.loads(message["data"].decode()))
+                    await emit_update(StateUpdate(**record.update), record.token)
+
+    async def _lost_and_found_updates_forever(
+        self,
+        emit_update: Callable[[StateUpdate, str], Coroutine[None, None, None]],
+    ):
+        """Background task to monitor Redis lost and found deltas.
+
+        Args:
+            emit_update: The function to emit state updates.
+        """
+        while True:
+            try:
+                await self._subscribe_lost_and_found_updates(emit_update)
+            except asyncio.CancelledError:  # noqa: PERF203
+                break
+            except Exception as e:
+                console.error(f"RedisTokenManager lost and found task error: {e}")
+
+    def ensure_lost_and_found_task(
+        self,
+        emit_update: Callable[[StateUpdate, str], Coroutine[None, None, None]],
+    ) -> None:
+        """Ensure the lost and found subscriber task is running.
+
+        Args:
+            emit_update: The function to emit state updates.
+        """
+        if self._lost_and_found_task is None or self._lost_and_found_task.done():
+            self._lost_and_found_task = asyncio.create_task(
+                self._lost_and_found_updates_forever(emit_update)
+            )
+
+    async def _get_token_owner(self, token: str, refresh: bool = False) -> str | None:
+        """Get the instance ID of the owner of a token.
+
+        Args:
+            token: The client token.
+            refresh: Whether to fetch the latest record from Redis.
+
+        Returns:
+            The instance ID of the owner, or None if not found.
+        """
+        if (
+            not refresh
+            and (socket_record := self.token_to_socket.get(token)) is not None
+        ):
+            return socket_record.instance_id
+
+        redis_key = self._get_redis_key(token)
+        try:
+            record_json = await self.redis.get(redis_key)
+            if record_json:
+                record_data = json.loads(record_json)
+                socket_record = SocketRecord(**record_data)
+                self.token_to_socket[token] = socket_record
+                self.sid_to_token[socket_record.sid] = token
+                return socket_record.instance_id
+            console.error(f"Redis token owner not found for token {token}")
+        except Exception as e:
+            console.error(f"Redis error getting token owner: {e}")
+        return None
+
+    async def emit_lost_and_found(
+        self,
+        token: str,
+        update: StateUpdate,
+    ) -> bool:
+        """Emit a lost and found delta to Redis.
+
+        Args:
+            token: The client token.
+            update: The state update.
+
+        Returns:
+            True if the delta was published, False otherwise.
+        """
+        # See where this update belongs
+        owner_instance_id = await self._get_token_owner(token)
+        if owner_instance_id is None:
+            return False
+        record = LostAndFoundRecord(token=token, update=dataclasses.asdict(update))
+        try:
+            await self.redis.publish(
+                f"channel:{self._get_lost_and_found_key(owner_instance_id)}",
+                json.dumps(dataclasses.asdict(record)),
+            )
+        except Exception as e:
+            console.error(f"Redis error publishing lost and found delta: {e}")
+        else:
+            return True
+        return False
diff --git a/tests/units/test_state.py b/tests/units/test_state.py
@@ -2017,8 +2017,9 @@ async def test_state_proxy(
     namespace = mock_app.event_namespace
     assert namespace is not None
     namespace.sid_to_token[router_data.session.session_id] = token
+    namespace._token_manager.instance_id = "mock"
     namespace._token_manager.token_to_socket[token] = SocketRecord(
-        instance_id="", sid=router_data.session.session_id
+        instance_id="mock", sid=router_data.session.session_id
     )
     if isinstance(mock_app.state_manager, (StateManagerMemory, StateManagerDisk)):
         mock_app.state_manager.states[parent_state.router.session.client_token] = (
@@ -2230,8 +2231,9 @@ async def test_background_task_no_block(mock_app: rx.App, token: str):
     namespace = mock_app.event_namespace
     assert namespace is not None
     namespace.sid_to_token[sid] = token
+    namespace._token_manager.instance_id = "mock"
     namespace._token_manager.token_to_socket[token] = SocketRecord(
-        instance_id="", sid=sid
+        instance_id="mock", sid=sid
     )
     mock_app.state_manager.state = mock_app._state = BackgroundTaskState
     async for update in rx.app.process(
diff --git a/tests/units/utils/test_token_manager.py b/tests/units/utils/test_token_manager.py
@@ -2,6 +2,7 @@
 
 import asyncio
 import json
+from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
@@ -204,9 +205,13 @@ async def listen():
                 yield
             return
 
-        psubscribe = AsyncMock()
-        psubscribe.listen = listen
-        redis.pubsub = Mock(return_value=psubscribe)
+        @asynccontextmanager
+        async def pubsub():  # noqa: RUF029
+            pubsub_mock = AsyncMock()
+            pubsub_mock.listen = listen
+            yield pubsub_mock
+
+        redis.pubsub = pubsub
         return redis
 
     @pytest.fixture
