web_development/backend/app/api/websocket.py

import asyncio
import json
import logging
from datetime import datetime

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from app.services.market_data import market_data_service
from app.services.order_manager import order_manager

logger = logging.getLogger(__name__)
ws_router = APIRouter()


class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        dead = []
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except Exception:
                dead.append(connection)
        for conn in dead:
            self.disconnect(conn)


manager = ConnectionManager()


def json_serializer(obj):
    if isinstance(obj, datetime):
        return obj.isoformat()
    raise TypeError(f"Type {type(obj)} not serializable")


@ws_router.websocket("/ws/market")
async def market_ws(websocket: WebSocket):
    await manager.connect(websocket)
    queue = market_data_service.subscribe()
    try:
        while True:
            data = await queue.get()
            msg = json.dumps({
                "type": "market_data",
                "data": data.model_dump(),
            }, default=json_serializer)
            await websocket.send_text(msg)
    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error("WebSocket error: %s", e)
    finally:
        market_data_service.unsubscribe(queue)
        manager.disconnect(websocket)


@ws_router.websocket("/ws/account")
async def account_ws(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            account = order_manager.account
            msg = json.dumps({
                "type": "account_update",
                "data": account.model_dump(),
            }, default=json_serializer)
            await websocket.send_text(msg)
            await asyncio.sleep(2)
    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error("Account WS error: %s", e)
    finally:
        manager.disconnect(websocket)