From 827c41925137570599e876435ca2f65b3a6007b2 Mon Sep 17 00:00:00 2001
From: Jun Siang Cheah <git@jscheah.me>
Date: Mon, 9 Sep 2024 23:17:17 +0100
Subject: [PATCH] feat: add ENABLE_WEBSOCKET_SUPPORT to force socket.io to
 ignore websocket upgrades

---
 backend/open_webui/apps/socket/main.py | 9 ++++-----
 backend/open_webui/config.py           | 6 ------
 backend/open_webui/env.py              | 4 ++++
 3 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/backend/open_webui/apps/socket/main.py b/backend/open_webui/apps/socket/main.py
index 777d877bf..d6f1a4999 100644
--- a/backend/open_webui/apps/socket/main.py
+++ b/backend/open_webui/apps/socket/main.py
@@ -2,16 +2,15 @@ import asyncio
 
 import socketio
 from open_webui.apps.webui.models.users import Users
-from open_webui.config import ENABLE_WEBSOCKET_SUPPORT
+from open_webui.env import ENABLE_WEBSOCKET_SUPPORT
 from open_webui.utils.utils import decode_token
 
 sio = socketio.AsyncServer(
     cors_allowed_origins=[],
     async_mode="asgi",
-    transports=(
-        ["polling", "websocket"] if ENABLE_WEBSOCKET_SUPPORT.value else ["polling"]
-    ),
-    allow_upgrades=ENABLE_WEBSOCKET_SUPPORT.value,
+    transports=(["polling", "websocket"] if ENABLE_WEBSOCKET_SUPPORT else ["polling"]),
+    allow_upgrades=ENABLE_WEBSOCKET_SUPPORT,
+    always_connect=True,
 )
 app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io")
 
diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py
index ac34001d6..5ccb40d47 100644
--- a/backend/open_webui/config.py
+++ b/backend/open_webui/config.py
@@ -810,12 +810,6 @@ ENABLE_MESSAGE_RATING = PersistentConfig(
     os.environ.get("ENABLE_MESSAGE_RATING", "True").lower() == "true",
 )
 
-ENABLE_WEBSOCKET_SUPPORT = PersistentConfig(
-    "ENABLE_WEBSOCKET_SUPPORT",
-    "ui.enable_websocket_support",
-    os.environ.get("ENABLE_WEBSOCKET_SUPPORT", "True").lower() == "true",
-)
-
 
 def validate_cors_origins(origins):
     for origin in origins:
diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py
index b716769c2..8683bb370 100644
--- a/backend/open_webui/env.py
+++ b/backend/open_webui/env.py
@@ -273,3 +273,7 @@ WEBUI_SESSION_COOKIE_SECURE = os.environ.get(
 
 if WEBUI_AUTH and WEBUI_SECRET_KEY == "":
     raise ValueError(ERROR_MESSAGES.ENV_VAR_NOT_FOUND)
+
+ENABLE_WEBSOCKET_SUPPORT = (
+    os.environ.get("ENABLE_WEBSOCKET_SUPPORT", "True").lower() == "true"
+)