Token manager tracks instance_id in token_to_socket

Author: masenf

reflex/app.py

@@ -2033,11 +2033,13 @@ def __init__(self, namespace: str, app: App):
         self._token_manager = TokenManager.create()
 
     @property
-    def token_to_sid(self) -> dict[str, str]:
+    def token_to_sid(self) -> Mapping[str, str]:
         """Get token to SID mapping for backward compatibility.
 
+        Note: this mapping is read-only.
+
         Returns:
-            The token to SID mapping dict.
+            The token to SID mapping.
         """
         # For backward compatibility, expose the underlying dict
         return self._token_manager.token_to_sid

reflex/utils/token_manager.py

@@ -2,8 +2,11 @@
 
 from __future__ import annotations
 
+import dataclasses
+import json
 import uuid
 from abc import ABC, abstractmethod
+from types import MappingProxyType
 from typing import TYPE_CHECKING
 
 from reflex.utils import console, prerequisites
@@ -21,16 +24,37 @@ def _get_new_token() -> str:
     return str(uuid.uuid4())
 
 
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class SocketRecord:
+    """Record for a connected socket client."""
+
+    instance_id: str
+    sid: str
+
+
 class TokenManager(ABC):
     """Abstract base class for managing client token to session ID mappings."""
 
     def __init__(self):
         """Initialize the token manager with local dictionaries."""
-        # Keep a mapping between socket ID and client token.
-        self.token_to_sid: dict[str, str] = {}
+        # Each process has an instance_id to identify its own sockets.
+        self.instance_id: str = _get_new_token()
         # Keep a mapping between client token and socket ID.
+        self.token_to_socket: dict[str, SocketRecord] = {}
+        # Keep a mapping between socket ID and client token.
         self.sid_to_token: dict[str, str] = {}
 
+    @property
+    def token_to_sid(self) -> MappingProxyType[str, str]:
+        """Read-only compatibility property for token_to_socket mapping.
+
+        Returns:
+            The token to session ID mapping.
+        """
+        return MappingProxyType({
+            token: sr.sid for token, sr in self.token_to_socket.items()
+        })
+
     @abstractmethod
     async def link_token_to_sid(self, token: str, sid: str) -> str | None:
         """Link a token to a session ID.
@@ -68,7 +92,9 @@ def create(cls) -> TokenManager:
 
     async def disconnect_all(self):
         """Disconnect all tracked tokens when the server is going down."""
-        token_sid_pairs: set[tuple[str, str]] = set(self.token_to_sid.items())
+        token_sid_pairs: set[tuple[str, str]] = {
+            (token, sr.sid) for token, sr in self.token_to_socket.items()
+        }
         token_sid_pairs.update(
             ((token, sid) for sid, token in self.sid_to_token.items())
         )
@@ -95,14 +121,20 @@ async def link_token_to_sid(self, token: str, sid: str) -> str | None:
             New token if duplicate detected and new token generated, None otherwise.
         """
         # Check if token is already mapped to a different SID (duplicate tab)
-        if token in self.token_to_sid and sid != self.token_to_sid.get(token):
+        if (
+            socket_record := self.token_to_socket.get(token)
+        ) is not None and sid != socket_record.sid:
             new_token = _get_new_token()
-            self.token_to_sid[new_token] = sid
+            self.token_to_socket[new_token] = SocketRecord(
+                instance_id=self.instance_id, sid=sid
+            )
             self.sid_to_token[sid] = new_token
             return new_token
 
         # Normal case - link token to SID
-        self.token_to_sid[token] = sid
+        self.token_to_socket[token] = SocketRecord(
+            instance_id=self.instance_id, sid=sid
+        )
         self.sid_to_token[sid] = token
         return None
 
@@ -114,7 +146,7 @@ async def disconnect_token(self, token: str, sid: str) -> None:
             sid: The Socket.IO session ID.
         """
         # Clean up both mappings
-        self.token_to_sid.pop(token, None)
+        self.token_to_socket.pop(token, None)
         self.sid_to_token.pop(sid, None)
 
 
@@ -149,9 +181,9 @@ def _get_redis_key(self, token: str) -> str:
             token: The client token.
 
         Returns:
-            Redis key following Reflex conventions: {token}_sid
+            Redis key following Reflex conventions: token_manager_socket_record_{token}
         """
-        return f"{token}_sid"
+        return f"token_manager_socket_record_{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.
@@ -164,7 +196,9 @@ async def link_token_to_sid(self, token: str, sid: str) -> str | None:
             New token if duplicate detected and new token generated, None otherwise.
         """
         # Fast local check first (handles reconnections)
-        if token in self.token_to_sid and self.token_to_sid[token] == sid:
+        if (
+            socket_record := self.token_to_socket.get(token)
+        ) is not None and sid == socket_record.sid:
             return None  # Same token, same SID = reconnection, no Redis check needed
 
         # Check Redis for cross-worker duplicates
@@ -176,34 +210,29 @@ async def link_token_to_sid(self, token: str, sid: str) -> str | None:
             console.error(f"Redis error checking token existence: {e}")
             return await super().link_token_to_sid(token, sid)
 
+        new_token = None
         if token_exists_in_redis:
             # Duplicate exists somewhere - generate new token
-            new_token = _get_new_token()
-            new_redis_key = self._get_redis_key(new_token)
+            token = new_token = _get_new_token()
+            redis_key = self._get_redis_key(new_token)
 
-            try:
-                # Store in Redis
-                await self.redis.set(new_redis_key, "1", ex=self.token_expiration)
-            except Exception as e:
-                console.error(f"Redis error storing new token: {e}")
-                # Still update local dicts and continue
-
-            # Store in local dicts (always do this)
-            self.token_to_sid[new_token] = sid
-            self.sid_to_token[sid] = new_token
-            return 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
 
-        # Normal case - store in both Redis and local dicts
+        # Store in Redis if possible
         try:
-            await self.redis.set(redis_key, "1", ex=self.token_expiration)
+            await self.redis.set(
+                redis_key,
+                json.dumps(dataclasses.asdict(socket_record)),
+                ex=self.token_expiration,
+            )
         except Exception as e:
             console.error(f"Redis error storing token: {e}")
-            # Continue with local storage
-
-        # Store in local dicts (always do this)
-        self.token_to_sid[token] = sid
-        self.sid_to_token[sid] = token
-        return None
+        # Return the new token if one was generated
+        return new_token
 
     async def disconnect_token(self, token: str, sid: str) -> None:
         """Clean up token mapping when client disconnects.
@@ -213,7 +242,9 @@ async def disconnect_token(self, token: str, sid: str) -> None:
             sid: The Socket.IO session ID.
         """
         # Only clean up if we own it locally (fast ownership check)
-        if self.token_to_sid.get(token) == sid:
+        if (
+            socket_record := self.token_to_socket.get(token)
+        ) is not None and socket_record.sid == sid:
             # Clean up Redis
             redis_key = self._get_redis_key(token)
             try:

tests/units/test_state.py

@@ -53,6 +53,7 @@
     UnretrievableVarValueError,
 )
 from reflex.utils.format import json_dumps
+from reflex.utils.token_manager import SocketRecord
 from reflex.vars.base import Var, computed_var
 
 from .states import GenState
@@ -2016,7 +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_to_sid[token] = router_data.session.session_id
+    namespace._token_manager.token_to_socket[token] = SocketRecord(
+        instance_id="", 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] = (
             parent_state
@@ -2227,7 +2230,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_to_sid[token] = sid
+    namespace._token_manager.token_to_socket[token] = SocketRecord(
+        instance_id="", sid=sid
+    )
     mock_app.state_manager.state = mock_app._state = BackgroundTaskState
     async for update in rx.app.process(
         mock_app,

tests/units/utils/test_token_manager.py

@@ -1,12 +1,14 @@
 """Unit tests for TokenManager implementations."""
 
+import json
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 
 from reflex.utils.token_manager import (
     LocalTokenManager,
     RedisTokenManager,
+    SocketRecord,
     TokenManager,
 )
 
@@ -215,7 +217,7 @@ def test_get_redis_key(self, manager):
             manager: RedisTokenManager fixture instance.
         """
         token = "test_token_123"
-        expected_key = f"{token}_sid"
+        expected_key = f"token_manager_socket_record_{token}"
 
         assert manager._get_redis_key(token) == expected_key
 
@@ -232,9 +234,15 @@ async def test_link_token_to_sid_normal_case(self, manager, mock_redis):
         result = await manager.link_token_to_sid(token, sid)
 
         assert result is None
-        mock_redis.exists.assert_called_once_with(f"{token}_sid")
-        mock_redis.set.assert_called_once_with(f"{token}_sid", "1", ex=3600)
-        assert manager.token_to_sid[token] == sid
+        mock_redis.exists.assert_called_once_with(
+            f"token_manager_socket_record_{token}"
+        )
+        mock_redis.set.assert_called_once_with(
+            f"token_manager_socket_record_{token}",
+            json.dumps({"instance_id": manager.instance_id, "sid": sid}),
+            ex=3600,
+        )
+        assert manager.token_to_socket[token].sid == sid
         assert manager.sid_to_token[sid] == token
 
     async def test_link_token_to_sid_reconnection_skips_redis(
@@ -247,7 +255,9 @@ async def test_link_token_to_sid_reconnection_skips_redis(
             mock_redis: Mock Redis client fixture.
         """
         token, sid = "token1", "sid1"
-        manager.token_to_sid[token] = sid
+        manager.token_to_socket[token] = SocketRecord(
+            instance_id=manager.instance_id, sid=sid
+        )
 
         result = await manager.link_token_to_sid(token, sid)
 
@@ -271,8 +281,14 @@ async def test_link_token_to_sid_duplicate_detected(self, manager, mock_redis):
         assert result != token
         assert len(result) == 36  # UUID4 length
 
-        mock_redis.exists.assert_called_once_with(f"{token}_sid")
-        mock_redis.set.assert_called_once_with(f"{result}_sid", "1", ex=3600)
+        mock_redis.exists.assert_called_once_with(
+            f"token_manager_socket_record_{token}"
+        )
+        mock_redis.set.assert_called_once_with(
+            f"token_manager_socket_record_{result}",
+            json.dumps({"instance_id": manager.instance_id, "sid": sid}),
+            ex=3600,
+        )
         assert manager.token_to_sid[result] == sid
         assert manager.sid_to_token[sid] == result
 
@@ -323,12 +339,16 @@ async def test_disconnect_token_owned_locally(self, manager, mock_redis):
             mock_redis: Mock Redis client fixture.
         """
         token, sid = "token1", "sid1"
-        manager.token_to_sid[token] = sid
+        manager.token_to_socket[token] = SocketRecord(
+            instance_id=manager.instance_id, sid=sid
+        )
         manager.sid_to_token[sid] = token
 
         await manager.disconnect_token(token, sid)
 
-        mock_redis.delete.assert_called_once_with(f"{token}_sid")
+        mock_redis.delete.assert_called_once_with(
+            f"token_manager_socket_record_{token}"
+        )
         assert token not in manager.token_to_sid
         assert sid not in manager.sid_to_token
 
@@ -353,7 +373,9 @@ async def test_disconnect_token_redis_error(self, manager, mock_redis):
             mock_redis: Mock Redis client fixture.
         """
         token, sid = "token1", "sid1"
-        manager.token_to_sid[token] = sid
+        manager.token_to_socket[token] = SocketRecord(
+            instance_id=manager.instance_id, sid=sid
+        )
         manager.sid_to_token[sid] = token
         mock_redis.delete.side_effect = Exception("Redis delete error")