80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
from typing import Dict, List
|
|
from fastapi import WebSocket
|
|
import json
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class WebSocketManager:
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, WebSocket] = {}
|
|
|
|
async def connect(self, websocket: WebSocket, client_id: str):
|
|
"""클라이언트 연결"""
|
|
await websocket.accept()
|
|
self.active_connections[client_id] = websocket
|
|
logger.info(f"Client {client_id} connected. Total connections: {len(self.active_connections)}")
|
|
|
|
def disconnect(self, client_id: str):
|
|
"""클라이언트 연결 해제"""
|
|
if client_id in self.active_connections:
|
|
del self.active_connections[client_id]
|
|
logger.info(f"Client {client_id} disconnected. Total connections: {len(self.active_connections)}")
|
|
|
|
async def send_personal_message(self, message: str, client_id: str):
|
|
"""특정 클라이언트에게 메시지 전송"""
|
|
if client_id in self.active_connections:
|
|
try:
|
|
await self.active_connections[client_id].send_text(message)
|
|
except Exception as e:
|
|
logger.error(f"Error sending message to {client_id}: {e}")
|
|
self.disconnect(client_id)
|
|
|
|
async def broadcast(self, message: str):
|
|
"""모든 연결된 클라이언트에게 메시지 브로드캐스트"""
|
|
disconnected_clients = []
|
|
for client_id, connection in self.active_connections.items():
|
|
try:
|
|
await connection.send_text(message)
|
|
except Exception as e:
|
|
logger.error(f"Error broadcasting to {client_id}: {e}")
|
|
disconnected_clients.append(client_id)
|
|
|
|
# 연결 실패한 클라이언트 제거
|
|
for client_id in disconnected_clients:
|
|
self.disconnect(client_id)
|
|
|
|
async def broadcast_json(self, data: dict):
|
|
"""JSON 데이터를 모든 클라이언트에게 브로드캐스트"""
|
|
message = json.dumps(data)
|
|
await self.broadcast(message)
|
|
|
|
async def send_market_update(self, stock_code: str, data: dict):
|
|
"""시장 데이터 업데이트 전송"""
|
|
message = {
|
|
"type": "market_update",
|
|
"stock_code": stock_code,
|
|
"data": data
|
|
}
|
|
await self.broadcast_json(message)
|
|
|
|
async def send_trade_notification(self, trade_data: dict):
|
|
"""거래 알림 전송"""
|
|
message = {
|
|
"type": "trade_notification",
|
|
"data": trade_data
|
|
}
|
|
await self.broadcast_json(message)
|
|
|
|
async def send_mcp_status(self, status: dict):
|
|
"""MCP 상태 업데이트 전송"""
|
|
message = {
|
|
"type": "mcp_status",
|
|
"data": status
|
|
}
|
|
await self.broadcast_json(message)
|
|
|
|
|
|
websocket_manager = WebSocketManager()
|