RedisTokenManager: keep local dicts globally updated via pub/sub

Author: masenf

reflex/istate/manager/redis.py

@@ -67,6 +67,7 @@ class StateManagerRedis(StateManager):
     # The keyspace subscription string when redis is waiting for lock to be released.
     _redis_notify_keyspace_events: str = dataclasses.field(
         default="K"  # Enable keyspace notifications (target a particular key)
+        "$"  # For String commands (like setting keys)
         "g"  # For generic commands (DEL, EXPIRE, etc)
         "x"  # For expired events
         "e"  # For evicted events (i.e. maxmemory exceeded)
@@ -76,7 +77,6 @@ class StateManagerRedis(StateManager):
     _redis_keyspace_lock_release_events: set[bytes] = dataclasses.field(
         default_factory=lambda: {
             b"del",
-            b"expire",
             b"expired",
             b"evicted",
         }

reflex/utils/token_manager.py

@@ -2,13 +2,16 @@
 
 from __future__ import annotations
 
+import asyncio
 import dataclasses
 import json
 import uuid
 from abc import ABC, abstractmethod
 from types import MappingProxyType
 from typing import TYPE_CHECKING
 
+from reflex.istate.manager.redis import StateManagerRedis
+from reflex.state import BaseState
 from reflex.utils import console, prerequisites
 
 if TYPE_CHECKING:
@@ -173,6 +176,7 @@ def __init__(self, redis: Redis):
 
         config = get_config()
         self.token_expiration = config.redis_token_expiration
+        self._update_task = None
 
     def _get_redis_key(self, token: str) -> str:
         """Get Redis key for token mapping.
@@ -185,6 +189,48 @@ def _get_redis_key(self, token: str) -> str:
         """
         return f"token_manager_socket_record_{token}"
 
+    async def _socket_record_update_task(self) -> None:
+        """Background task to monitor Redis keyspace notifications for socket record updates."""
+        await StateManagerRedis(
+            state=BaseState, redis=self.redis
+        )._enable_keyspace_notifications()
+        redis_db = self.redis.get_connection_kwargs().get("db", 0)
+        while True:
+            try:
+                await self._subscribe_socket_record_updates(redis_db)
+            except asyncio.CancelledError:  # noqa: PERF203
+                break
+            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
+
     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.
 
@@ -201,6 +247,10 @@ async def link_token_to_sid(self, token: str, sid: str) -> str | None:
         ) is not None and sid == socket_record.sid:
             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())
+
         # Check Redis for cross-worker duplicates
         redis_key = self._get_redis_key(token)
 

tests/units/utils/test_token_manager.py

@@ -1,5 +1,6 @@
 """Unit tests for TokenManager implementations."""
 
+import asyncio
 import json
 from unittest.mock import AsyncMock, Mock, patch
 
@@ -63,6 +64,7 @@ def test_create_redis_when_redis_available(
         """
         mock_check_redis_used.return_value = True
         mock_redis_client = Mock()
+        mock_redis_client.get_connection_kwargs.return_value = {"db": 0}
         mock_get_redis.return_value = mock_redis_client
 
         manager = TokenManager.create()
@@ -191,6 +193,20 @@ def mock_redis(self):
         redis.exists = AsyncMock()
         redis.set = AsyncMock()
         redis.delete = AsyncMock()
+
+        # Non-async call
+        redis.get_connection_kwargs = Mock(return_value={"db": 0})
+
+        # Mock out pubsub
+        async def listen():
+            await asyncio.sleep(1)
+            if False:
+                yield
+            return
+
+        psubscribe = AsyncMock()
+        psubscribe.listen = listen
+        redis.pubsub = Mock(return_value=psubscribe)
         return redis
 
     @pytest.fixture