fix: websocket connection when submounting chainlit (#1337)

* fix: websocket connection when submounting chainlit
* fix: starting chainlit with --root-path

backend/chainlit/config.py

@@ -285,9 +285,9 @@ class CodeSettings:
     password_auth_callback: Optional[
         Callable[[str, str], Awaitable[Optional["User"]]]
     ] = None
-    header_auth_callback: Optional[Callable[[Headers], Awaitable[Optional["User"]]]] = (
-        None
-    )
+    header_auth_callback: Optional[
+        Callable[[Headers], Awaitable[Optional["User"]]]
+    ] = None
     oauth_callback: Optional[
         Callable[[str, str, Dict[str, str], "User"], Awaitable[Optional["User"]]]
     ] = None
@@ -305,9 +305,9 @@ class CodeSettings:
     set_chat_profiles: Optional[
         Callable[[Optional["User"]], Awaitable[List["ChatProfile"]]]
     ] = None
-    set_starters: Optional[Callable[[Optional["User"]], Awaitable[List["Starter"]]]] = (
-        None
-    )
+    set_starters: Optional[
+        Callable[[Optional["User"]], Awaitable[List["Starter"]]]
+    ] = None
 
 
 @dataclass()

backend/chainlit/langchain/callbacks.py

@@ -590,10 +590,14 @@ def _on_run_update(self, run: Run) -> None:
 
         if output_keys:
             output = outputs.get(output_keys[0], outputs)
-            
+
         if current_step:
             current_step.output = (
-                output[0] if isinstance(output, Sequence) and not isinstance(output, str) and len(output) else output
+                output[0]
+                if isinstance(output, Sequence)
+                and not isinstance(output, str)
+                and len(output)
+                else output
             )
             current_step.end = utc_now()
             self._run_sync(current_step.update())

backend/chainlit/server.py

@@ -64,6 +64,8 @@
 
 ROOT_PATH = os.environ.get("CHAINLIT_ROOT_PATH", "")
 IS_SUBMOUNT = os.environ.get("CHAINLIT_SUBMOUNT", "") == "true"
+# If the app is a submount, no need to set the prefix
+PREFIX = ROOT_PATH if ROOT_PATH and not IS_SUBMOUNT else ""
 
 
 @asynccontextmanager
@@ -185,14 +187,12 @@ def get_build_dir(local_target: str, packaged_target: str) -> str:
 
 sio = socketio.AsyncServer(cors_allowed_origins=[], async_mode="asgi")
 
-sio_mount_location = f"{ROOT_PATH}/ws" if ROOT_PATH else "ws"
-
 asgi_app = socketio.ASGIApp(
     socketio_server=sio,
-    socketio_path=f"{sio_mount_location}/socket.io",
+    socketio_path="",
 )
 
-app.mount(f"/{sio_mount_location}", asgi_app)
+app.mount(f"{PREFIX}/ws/socket.io", asgi_app)
 
 app.add_middleware(
     CORSMiddleware,
@@ -202,16 +202,16 @@ def get_build_dir(local_target: str, packaged_target: str) -> str:
     allow_headers=["*"],
 )
 
-router = APIRouter(prefix=ROOT_PATH)
+router = APIRouter(prefix=PREFIX)
 
 app.mount(
-    f"{ROOT_PATH}/public",
+    f"{PREFIX}/public",
     StaticFiles(directory="public", check_dir=False),
     name="public",
 )
 
 app.mount(
-    f"{ROOT_PATH}/assets",
+    f"{PREFIX}/assets",
     StaticFiles(
         packages=[("chainlit", os.path.join(build_dir, "assets"))],
         follow_symlink=config.project.follow_symlink,
@@ -220,7 +220,7 @@ def get_build_dir(local_target: str, packaged_target: str) -> str:
 )
 
 app.mount(
-    f"{ROOT_PATH}/copilot",
+    f"{PREFIX}/copilot",
     StaticFiles(
         packages=[("chainlit", copilot_build_dir)],
         follow_symlink=config.project.follow_symlink,

backend/chainlit/utils.py

@@ -130,4 +130,4 @@ def mount_chainlit(app: FastAPI, target: str, path="/chainlit"):
 
     ensure_jwt_secret()
 
-    app.mount("/", chainlit_app)
+    app.mount(path, chainlit_app)

backend/tests/conftest.py

@@ -3,7 +3,6 @@
 
 import pytest
 import pytest_asyncio
-
 from chainlit.context import ChainlitContext, context_var
 from chainlit.session import HTTPSession, WebsocketSession
 from chainlit.user_session import UserSession

backend/tests/test_callbacks.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
 import pytest
-
 from chainlit.callbacks import password_auth_callback
-from chainlit import config
 from chainlit.user import User
 
+from chainlit import config
+
 
 @pytest.fixture
 def test_config(monkeypatch: pytest.MonkeyPatch):