diff --git a/api/router.py b/api/router.py
index c8b7f2e..9ac4ead 100644
--- a/api/router.py
+++ b/api/router.py
@@ -150,8 +150,9 @@ async def open_file(body: dict, user: dict = Depends(current_user)):
 
 
 @router.websocket("/ws")
-async def websocket_endpoint(ws: WebSocket, token: str = ""):
+async def websocket_endpoint(ws: WebSocket):
     from auth import get_user_for_token
+    token = ws.query_params.get("token", "")
     user = get_user_for_token(token)
     if not user:
         await ws.close(code=4001)