from fastapi import APIRouter, Query, WebSocket
from yuxi.channel.gateway.server import gateway_ws_server

router = APIRouter(tags=["Gateway WebSocket"])


@router.websocket("/ws/gateway")
async def gateway_ws_endpoint(
    ws: WebSocket,
    token: str | None = Query(None),
):
    await gateway_ws_server.handle_connection(ws, token)