fix: deliver linked state updates to all clients

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.

reflex/istate/shared.py

@@ -54,20 +54,23 @@ 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.
         t = asyncio.create_task(_update_client(affected_token))
         UPDATE_OTHER_CLIENT_TASKS.add(t)

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.
@@ -240,22 +251,29 @@ async def enumerate_tokens(self) -> AsyncIterator[str]:
             if not cursor:
                 break
 
-    async def _handle_socket_record_del(
-        self, token: str, expired: bool = False
-    ) -> None:
+    async def _handle_socket_record_del(self, token: str) -> None:
         """Handle deletion of a socket record from Redis.
 
+        Disconnects pop the local mapping before touching redis, so a
+        self-owned record still present locally was created by a newer link
+        and refers to a live socket; the deletion is an expiration or an
+        outdated notification, and the record is re-stored to keep it alive.
+        A cached foreign record is dropped and refetched on demand.
+
         Args:
             token: The client token whose record was deleted.
-            expired: Whether the deletion was due to expiration.
         """
-        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)
-            if expired:
-                # Keep the record alive as long as this process is alive and not deleted.
-                await self.link_token_to_sid(token, socket_record.sid)
+        if (socket_record := self.token_to_socket.get(token)) is None:
+            return
+        if socket_record.instance_id == self.instance_id:
+            # Restore only if the key is still absent: a newer record (a
+            # relink here or a claim by another instance) must win.
+            if not await self._store_socket_record(token, socket_record, nx=True):
+                # Adopt the newer record without waiting for its set
+                # notification, which may be lost across pubsub reconnects.
+                await self._get_token_owner(token, refresh=True)
+        else:
+            self.token_to_socket.pop(token, None)
 
     async def _subscribe_socket_record_updates(self) -> None:
         """Subscribe to Redis keyspace notifications for socket record updates."""
@@ -277,10 +295,7 @@ async def _subscribe_socket_record_updates(self) -> None:
 
                     event = message["data"].decode()
                     if event in ("del", "expired", "evicted"):
-                        await self._handle_socket_record_del(
-                            token,
-                            expired=(event == "expired"),
-                        )
+                        await self._handle_socket_record_del(token)
                     elif event == "set":
                         await self._get_token_owner(token, refresh=True)
 
@@ -325,25 +340,42 @@ async def link_token_to_sid(self, token: str, sid: str) -> str | None:
         if token_exists_in_redis:
             # Duplicate exists somewhere - generate new token
             token = new_token = _get_new_token()
-            redis_key = self._get_redis_key(new_token)
 
         # Store in local dicts
         socket_record = self.token_to_socket[token] = SocketRecord(
             instance_id=self.instance_id, sid=sid
         )
         self.sid_to_token[sid] = token
 
-        # Store in Redis if possible
+        await self._store_socket_record(token, socket_record)
+        # Return the new token if one was generated
+        return new_token
+
+    async def _store_socket_record(
+        self, token: str, socket_record: SocketRecord, nx: bool = False
+    ) -> bool:
+        """Store a socket record in Redis, logging errors instead of raising.
+
+        Args:
+            token: The client token.
+            socket_record: The record to store.
+            nx: Only store the record if the key does not already exist.
+
+        Returns:
+            True if the record was stored, False if rejected (nx) or on error.
+        """
         try:
-            await self.redis.set(
-                redis_key,
-                pickle.dumps(socket_record),
-                ex=self.token_expiration,
+            return bool(
+                await self.redis.set(
+                    self._get_redis_key(token),
+                    pickle.dumps(socket_record),
+                    ex=self.token_expiration,
+                    nx=nx,
+                )
             )
         except Exception as e:
             logger.error(f"Redis error storing token: {e}")
-        # Return the new token if one was generated
-        return new_token
+            return False
 
     async def disconnect_token(self, token: str, sid: str) -> None:
         """Clean up token mapping when client disconnects.
@@ -358,16 +390,17 @@ async def disconnect_token(self, token: str, sid: str) -> None:
             and socket_record.sid == sid
             and socket_record.instance_id == self.instance_id
         ):
-            # Clean up Redis
+            # Drop the local mapping before the redis round-trip so the
+            # locally-owned fast paths stop treating the token as connected
+            # while the delete is in flight.
+            await super().disconnect_token(token, sid)
+
             redis_key = self._get_redis_key(token)
             try:
                 await self.redis.delete(redis_key)
             except Exception as e:
                 logger.error(f"Redis error deleting token: {e}")
 
-            # 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.
@@ -431,17 +464,79 @@ 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.
+
+        Redis errors propagate to the caller so it can distinguish a lookup
+        failure from an absent record. This instance is authoritative for its
+        own sockets: a record claiming this instance without a live local
+        link is a stale leftover of a disconnected socket (its delete is
+        still in flight or failed) and is treated as absent.
+
+        Args:
+            token: The client token.
+
+        Returns:
+            The refreshed socket record, or None if the token has no live socket.
+        """
+        record_pkl = await self.redis.get(self._get_redis_key(token))
+        if not record_pkl:
+            return None
+        socket_record = pickle.loads(record_pkl)
+        # Stale leftover of one of this instance's own disconnected sockets.
+        if (
+            socket_record.instance_id == self.instance_id
+            and self.sid_to_token.get(socket_record.sid) != token
+        ):
+            return None
+        # 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, so it is refreshed from redis
+        (dropped if the client is gone) and trusted as-is if the refresh fails.
+
+        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,

tests/units/istate/test_shared.py

@@ -0,0 +1,111 @@
+"""Unit tests for shared state fan-out to other linked clients."""
+
+import asyncio
+import pickle
+from contextlib import asynccontextmanager
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from reflex.istate.shared import _do_update_other_tokens
+from reflex.state import State
+from reflex.utils.token_manager import (
+    LocalTokenManager,
+    RedisTokenManager,
+    SocketRecord,
+)
+
+
+@pytest.fixture
+def mock_redis():
+    """Create a mock Redis client.
+
+    Returns:
+        The mock Redis client.
+    """
+    redis = AsyncMock()
+    redis.get = AsyncMock(return_value=None)
+    redis.get_connection_kwargs = Mock(return_value={"db": 0})
+    return redis
+
+
+@pytest.fixture
+def redis_manager(mock_redis):
+    """Create a RedisTokenManager instance with mocked config.
+
+    Returns:
+        The RedisTokenManager instance.
+    """
+    with patch("reflex_base.config.get_config") as mock_get_config:
+        mock_config = Mock()
+        mock_config.redis_token_expiration = 3600
+        mock_get_config.return_value = mock_config
+
+        return RedisTokenManager(mock_redis)
+
+
+def _mock_app(token_manager) -> tuple[Mock, list[str]]:
+    """Create a mock app recording the tokens passed to modify_state.
+
+    Returns:
+        The mock app and the list collecting modified token idents.
+    """
+    modified_tokens: list[str] = []
+
+    @asynccontextmanager
+    async def modify_state(token, previous_dirty_vars=None):
+        modified_tokens.append(token.ident)
+        yield Mock()
+
+    app = Mock()
+    app.modify_state = modify_state
+    app.event_namespace = Mock()
+    app.event_namespace._token_manager = token_manager
+    return app, modified_tokens
+
+
+async def _run_update_other_tokens(app, affected_tokens: set[str]) -> None:
+    """Run _do_update_other_tokens against a mock app and await its tasks."""
+    with patch("reflex_base.registry.RegistrationContext.get") as mock_get:
+        mock_get.return_value = Mock(app=app)
+        tasks = _do_update_other_tokens(
+            affected_tokens=affected_tokens,
+            previous_dirty_vars={},
+            state_type=State,
+        )
+    await asyncio.gather(*tasks)
+
+
+async def test_update_other_tokens_local_manager():
+    """With a LocalTokenManager, only locally connected tokens are updated."""
+    manager = LocalTokenManager()
+    manager.token_to_socket["connected"] = SocketRecord(
+        instance_id=manager.instance_id, sid="sid1"
+    )
+    app, modified_tokens = _mock_app(manager)
+
+    await _run_update_other_tokens(app, {"connected", "disconnected"})
+
+    assert modified_tokens == ["connected"]
+
+
+async def test_update_other_tokens_redis_cross_instance(redis_manager, mock_redis):
+    """Tokens connected to another instance are resolved via redis and updated."""
+    redis_manager.token_to_socket["local"] = SocketRecord(
+        instance_id=redis_manager.instance_id, sid="sid1"
+    )
+    foreign_record = SocketRecord(instance_id="other-instance", sid="sid2")
+    foreign_key = redis_manager._get_redis_key("foreign")
+    mock_redis.get.side_effect = lambda key: (
+        pickle.dumps(foreign_record) if key == foreign_key else None
+    )
+    app, modified_tokens = _mock_app(redis_manager)
+
+    await _run_update_other_tokens(app, {"local", "foreign", "disconnected"})
+
+    assert sorted(modified_tokens) == ["foreign", "local"]
+    # The foreign socket record is cached locally for later emit_update routing.
+    assert redis_manager.token_to_socket["foreign"] == foreign_record
+    # Locally owned sockets are authoritative and never require a redis lookup.
+    local_key = redis_manager._get_redis_key("local")
+    assert local_key not in [call.args[0] for call in mock_redis.get.call_args_list]