Remove token/sid associations when server is exiting

This allows existing tokens to reconnect to redis after a hot or cold reload of
the app. Otherwise, the old associations for the token remain in place and when
the same client reconnects, it is given a new_token, since the requested token
is already "taken" in redis.

Author: masenf

reflex/app_mixins/lifespan.py

@@ -60,6 +60,17 @@ async def _run_lifespan_tasks(self, app: Starlette):
             for task in running_tasks:
                 console.debug(f"Canceling lifespan task: {task}")
                 task.cancel(msg="lifespan_cleanup")
+        # Disassociate sid / token pairings so they can be reconnected properly.
+        try:
+            event_namespace = self.event_namespace  # pyright: ignore[reportAttributeAccessIssue]
+        except AttributeError:
+            pass
+        else:
+            try:
+                if event_namespace:
+                    await event_namespace._token_manager.disconnect_all()
+            except Exception as e:
+                console.error(f"Error during lifespan cleanup: {e}")
 
     def register_lifespan_task(self, task: Callable | asyncio.Task, **task_kwargs):
         """Register a task to run during the lifespan of the app.

reflex/utils/token_manager.py

@@ -66,6 +66,16 @@ def create(cls) -> TokenManager:
 
         return LocalTokenManager()
 
+    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.update(
+            ((token, sid) for sid, token in self.sid_to_token.items())
+        )
+        # Perform the disconnection logic here
+        for token, sid in token_sid_pairs:
+            await self.disconnect_token(token, sid)
+
 
 class LocalTokenManager(TokenManager):
     """Token manager using local in-memory dictionaries (single worker)."""