ENG-8540: avoid dataclasses.asdict in Lost+Found path

Use the reflex serializers registry to serialize StateUpdate objects for
Lost+Found usage.

Author: masenf

reflex/utils/token_manager.py

@@ -9,11 +9,12 @@
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Callable, Coroutine
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
 from reflex.istate.manager.redis import StateManagerRedis
 from reflex.state import BaseState, StateUpdate
 from reflex.utils import console, prerequisites
+from reflex.utils.format import json_dumps
 from reflex.utils.tasks import ensure_task
 
 if TYPE_CHECKING:
@@ -42,7 +43,7 @@ class LostAndFoundRecord:
     """Record for a StateUpdate for a token with its socket on another instance."""
 
     token: str
-    update: dict[str, Any]
+    update: StateUpdate
 
 
 class TokenManager(ABC):
@@ -386,8 +387,12 @@ async def _subscribe_lost_and_found_updates(
             )
             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)
+                    record_dict = json.loads(message["data"].decode())
+                    record = LostAndFoundRecord(
+                        token=record_dict["token"],
+                        update=StateUpdate(**record_dict["update"]),
+                    )
+                    await emit_update(record.update, record.token)
 
     def ensure_lost_and_found_task(
         self,
@@ -454,11 +459,11 @@ async def emit_lost_and_found(
         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))
+        record = LostAndFoundRecord(token=token, update=update)
         try:
             await self.redis.publish(
                 f"channel:{self._get_lost_and_found_key(owner_instance_id)}",
-                json.dumps(dataclasses.asdict(record)),
+                json_dumps(record),
             )
         except Exception as e:
             console.error(f"Redis error publishing lost and found delta: {e}")

tests/units/utils/test_token_manager.py

@@ -11,7 +11,9 @@
 
 from reflex import config
 from reflex.app import EventNamespace
+from reflex.istate.data import RouterData
 from reflex.state import StateUpdate
+from reflex.utils.format import json_dumps
 from reflex.utils.token_manager import (
     LocalTokenManager,
     RedisTokenManager,
@@ -670,3 +672,38 @@ async def test_redis_token_manager_lost_and_found(
     emit2_mock.assert_not_called()
     emit1_mock.assert_called_once()
     emit1_mock.reset_mock()
+
+
+@pytest.mark.usefixtures("redis_url")
+@pytest.mark.asyncio
+async def test_redis_token_manager_lost_and_found_router_data(
+    event_namespace_factory: Callable[[], EventNamespace],
+):
+    """Updates emitted for lost and found tokens should serialize properly.
+
+    Args:
+        event_namespace_factory: Factory fixture for EventNamespace instances.
+    """
+    event_namespace1 = event_namespace_factory()
+    emit1_mock: Mock = event_namespace1.emit  # pyright: ignore[reportAssignmentType]
+    event_namespace2 = event_namespace_factory()
+    emit2_mock: Mock = event_namespace2.emit  # pyright: ignore[reportAssignmentType]
+
+    await event_namespace1.on_connect(sid="sid1", environ=query_string_for("token1"))
+    await event_namespace2.on_connect(sid="sid2", environ=query_string_for("token2"))
+
+    router = RouterData.from_router_data(
+        {"headers": {"x-test": "value"}},
+    )
+
+    await event_namespace2.emit_update(
+        StateUpdate(delta={"state": {"router": router}}), token="token1"
+    )
+    await _wait_for_call_count_positive(emit1_mock)
+    emit2_mock.assert_not_called()
+    emit1_mock.assert_called_once()
+    assert isinstance(emit1_mock.call_args[0][1], StateUpdate)
+    assert emit1_mock.call_args[0][1].delta["state"]["router"] == json.loads(
+        json_dumps(router)
+    )
+    emit1_mock.reset_mock()