From dfceedc4e30ce864e7afedc710d9410491ae6030 Mon Sep 17 00:00:00 2001 From: Khaleel Al-Adhami Date: Mon, 27 Jan 2025 13:47:37 -0800 Subject: [PATCH] separate the logic --- reflex/app.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/reflex/app.py b/reflex/app.py index 18dc9f1fa..f86d348dc 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -418,20 +418,17 @@ class App(MiddlewareMixin, LifespanMixin): async def modified_send(message): 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] + if scope.get("subprotocols"): + # The following *does* say "subprotocol" instead of "subprotocols", intentionally. + message["subprotocol"] = scope["subprotocols"][0] - message["subprotocol"] = subprotocols + headers = dict(message.get("headers", [])) + header_key = b"sec-websocket-protocol" + if subprotocol := headers.get(header_key): + message["headers"] = [ + *message.get("headers", []), + (header_key, subprotocol), + ] return await original_send(message)