From f4a0403d719b4d713e33543506c2a0afc3f497ea Mon Sep 17 00:00:00 2001 From: Khaleel Al-Adhami Date: Mon, 27 Jan 2025 13:36:35 -0800 Subject: [PATCH] use subprotocols or headers --- reflex/app.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/reflex/app.py b/reflex/app.py index 9cac205f1..18dc9f1fa 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -417,10 +417,21 @@ class App(MiddlewareMixin, LifespanMixin): original_send = send async def modified_send(message): - if message["type"] == "websocket.accept" and ( - subprotocols := scope.get("subprotocols") - ): - message["subprotocol"] = subprotocols[0] + if message["type"] == "websocket.accept": + subprotocols = scope.get("subprotocols", []) + if not subprotocols: + headers = scope.get("headers", {}) + headers_key = b"sec-websocket-protocol" + if headers_key in headers: + subprotocols = ( + headers[headers_key].decode().split(", ") + ) + else: + subprotocols = None + else: + subprotocols = subprotocols[0] + + message["subprotocol"] = subprotocols return await original_send(message)