From e39c26d3212a7b656bd5629c6feaafcb259f71e3 Mon Sep 17 00:00:00 2001
From: markov <markov@uix.su>
Date: Mon, 23 Feb 2026 23:22:09 +0100
Subject: [PATCH] feat: JWT auth in WebSocket handler

---
 src/tracker/ws/handler.py | 54 ++++++++++++++++++++++++---------------
 1 file changed, 33 insertions(+), 21 deletions(-)

diff --git a/src/tracker/ws/handler.py b/src/tracker/ws/handler.py
index d7bce2e..03339c5 100644
--- a/src/tracker/ws/handler.py
+++ b/src/tracker/ws/handler.py
@@ -16,23 +16,29 @@ router = APIRouter()
 
 
 @router.websocket("/ws")
-async def websocket_endpoint(ws: WebSocket):
+async def websocket_endpoint(ws: WebSocket, token: str = ""):
     await ws.accept()
     session_id = None
 
     try:
-        # Wait for auth message
-        auth_msg = await ws.receive_json()
-        if auth_msg.get("type") != "auth":
-            await ws.send_json({"type": "auth.error", "message": "First message must be auth"})
-            await ws.close()
-            return
+        # Try query param token first (for direct JWT auth)
+        if token:
+            session_id = await _authenticate(ws, token)
+            if not session_id:
+                return
+        else:
+            # Wait for auth message (backward compatibility with agents)
+            auth_msg = await ws.receive_json()
+            if auth_msg.get("type") != "auth":
+                await ws.send_json({"type": "auth.error", "message": "First message must be auth"})
+                await ws.close()
+                return
 
-        token = auth_msg.get("token", "")
-        on_behalf_of = auth_msg.get("on_behalf_of")
-        session_id = await _authenticate(ws, token, on_behalf_of=on_behalf_of)
-        if not session_id:
-            return
+            token = auth_msg.get("token", "")
+            on_behalf_of = auth_msg.get("on_behalf_of")
+            session_id = await _authenticate(ws, token, on_behalf_of=on_behalf_of)
+            if not session_id:
+                return
 
         client = manager.sessions.get(session_id)
         slug = client.member_slug if client else None
@@ -86,22 +92,28 @@ async def websocket_endpoint(ws: WebSocket):
 async def _authenticate(ws: WebSocket, token: str, on_behalf_of: str | None = None) -> str | None:
     """Authenticate and register session. Returns session_id or None."""
     async with async_session() as db:
-        result = await db.execute(
-            select(Member).where(Member.token == token).options(selectinload(Member.agent_config))
-        )
-        member = result.scalar_one_or_none()
-
-        if not member:
+        member = None
+        
+        # Check if it's an agent token (starts with 'tb-')
+        if token.startswith("tb-"):
+            result = await db.execute(
+                select(Member).where(Member.token == token).options(selectinload(Member.agent_config))
+            )
+            member = result.scalar_one_or_none()
+        else:
+            # Try JWT decode
             from tracker.api.auth import decode_jwt
             try:
                 payload = decode_jwt(token)
+                member_id = payload["sub"]
                 result = await db.execute(
-                    select(Member).where(Member.id == payload["sub"])
+                    select(Member).where(Member.id == member_id)
                     .options(selectinload(Member.agent_config))
                 )
                 member = result.scalar_one_or_none()
-            except Exception:
-                pass
+                logger.info("JWT auth successful for member_id=%s", member_id)
+            except Exception as e:
+                logger.warning("JWT decode failed: %s", e)
 
         if not member:
             await ws.send_json({"type": "auth.error", "message": "Invalid token"})