From 9ef904905ba6206e86d1dc5301a9751c0f6a5aef Mon Sep 17 00:00:00 2001
From: Victor Skvortsov <vds003@gmail.com>
Date: Fri, 24 Jul 2026 10:36:56 +0500
Subject: [PATCH] Reject StaticFiles websocket requests gracefully

---
 src/dstack/_internal/server/app.py           |  4 ++--
 src/dstack/_internal/server/utils/routers.py | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/dstack/_internal/server/app.py b/src/dstack/_internal/server/app.py
index 6b1c4cff71..bf884d001a 100644
--- a/src/dstack/_internal/server/app.py
+++ b/src/dstack/_internal/server/app.py
@@ -11,7 +11,6 @@
 from fastapi import Depends, FastAPI, Request, Response, status
 from fastapi.datastructures import URL
 from fastapi.responses import HTMLResponse, RedirectResponse
-from fastapi.staticfiles import StaticFiles
 from packaging.version import Version
 
 from dstack._internal import settings as core_settings
@@ -73,6 +72,7 @@
 from dstack._internal.server.utils.logging import configure_logging
 from dstack._internal.server.utils.routers import (
     CustomORJSONResponse,
+    CustomStaticFiles,
     check_client_server_compatibility,
     error_detail,
     get_client_version,
@@ -366,7 +366,7 @@ async def healthcheck():
 
     if ui and Path(__file__).parent.joinpath("statics").exists():
         app.mount(
-            "/", StaticFiles(packages=["dstack._internal.server"], html=True), name="statics"
+            "/", CustomStaticFiles(packages=["dstack._internal.server"], html=True), name="statics"
         )
 
         @app.exception_handler(404)
diff --git a/src/dstack/_internal/server/utils/routers.py b/src/dstack/_internal/server/utils/routers.py
index 5aff751868..e56dcb7be8 100644
--- a/src/dstack/_internal/server/utils/routers.py
+++ b/src/dstack/_internal/server/utils/routers.py
@@ -3,6 +3,7 @@
 import orjson
 import packaging.version
 from fastapi import HTTPException, Request, Response, status
+from fastapi.staticfiles import StaticFiles
 
 from dstack._internal.core.errors import ServerClientError, ServerClientErrorCode
 from dstack._internal.core.models.common import CoreModel
@@ -10,6 +11,22 @@
 from dstack._internal.utils.version import parse_version
 
 
+class CustomStaticFiles(StaticFiles):
+    """
+    StaticFiles raises AssertionError on "websocket" scope type,
+    but starlette's Mount() matches both "http" and "websocket".
+    So a custom ASGI app is needed to reject WebSocket requests gracefully.
+
+    See: https://github.com/dstackai/dstack/issues/4061
+    """
+
+    async def __call__(self, scope, receive, send) -> None:
+        if scope["type"] == "websocket":
+            await send({"type": "websocket.close"})  # Reject the handshake
+            return
+        await super().__call__(scope, receive, send)
+
+
 class CustomORJSONResponse(Response):
     """
     Custom JSONResponse that uses orjson for serialization.