From 33f71c6eefcb21ac371f30ceec0191039bf53476 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Thomas=20Brand=C3=A9ho?= Date: Tue, 28 May 2024 19:53:26 +0200 Subject: [PATCH] add mapping between client_token and socket id (#3388) --- reflex/app.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/reflex/app.py b/reflex/app.py index 9481204cb9..98ad1eacf6 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -1302,6 +1302,12 @@ class EventNamespace(AsyncNamespace): # The application object. app: App + # Keep a mapping between socket ID and client token. + token_to_sid: dict[str, str] = {} + + # Keep a mapping between client token and socket ID. + sid_to_token: dict[str, str] = {} + def __init__(self, namespace: str, app: App): """Initialize the event namespace. @@ -1327,7 +1333,9 @@ def on_disconnect(self, sid): Args: sid: The Socket.IO session id. """ - pass + disconnect_token = self.sid_to_token.pop(sid, None) + if disconnect_token: + self.token_to_sid.pop(disconnect_token, None) async def emit_update(self, update: StateUpdate, sid: str) -> None: """Emit an update to the client. @@ -1351,6 +1359,9 @@ async def on_event(self, sid, data): # Get the event. event = Event.parse_raw(data) + self.token_to_sid[event.token] = sid + self.sid_to_token[sid] = event.token + # Get the event environment. assert self.app.sio is not None environ = self.app.sio.get_environ(sid, self.namespace)