Prechádzať zdrojové kódy

add mapping between client_token and socket id (#3388)

Thomas Brandého 1 rok pred
rodič
commit
33f71c6eef
1 zmenil súbory, kde vykonal 12 pridanie a 1 odobranie
  1. 12 1
      reflex/app.py

+ 12 - 1
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 @@ class EventNamespace(AsyncNamespace):
         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 @@ class EventNamespace(AsyncNamespace):
         # 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)