feat: KIS 실시간 스트림 서비스, WS 구독 프로토콜, MCP 리스크 승인 연동

2026-02-24 18:16:20 +09:00
parent 8c9b4abc72
commit 24f3082693
3 changed files with 510 additions and 53 deletions
--- a/backend/app/services/websocket_manager.py
+++ b/backend/app/services/websocket_manager.py
@@ -1,4 +1,4 @@
-from typing import Dict, List
+from typing import Dict, Set
 from fastapi import WebSocket
 from app.core.json_helpers import _dumps
 import logging
@@ -9,71 +9,125 @@ logger = logging.getLogger(__name__)
 class WebSocketManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
-    
+        # 클라이언트별 구독 종목 세트
+        self._subscriptions: Dict[str, Set[str]] = {}
+
+    # ------------------------------------------------------------------
+    # 연결 관리
+    # ------------------------------------------------------------------
    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)}")
-    
+        self._subscriptions.setdefault(client_id, set())
+        logger.info(f"Client {client_id} connected. Total: {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)}")
-    
+        self.active_connections.pop(client_id, None)
+        self._subscriptions.pop(client_id, None)
+        logger.info(f"Client {client_id} disconnected. Total: {len(self.active_connections)}")
+
+    # ------------------------------------------------------------------
+    # 구독 관리
+    # ------------------------------------------------------------------
+    def subscribe(self, client_id: str, stock_codes: list[str]):
+        """클라이언트에 종목 추가 구독."""
+        bucket = self._subscriptions.setdefault(client_id, set())
+        bucket.update(stock_codes)
+        logger.debug(f"[sub+] {client_id}: {stock_codes} → 현재 {sorted(bucket)}")
+
+    def unsubscribe(self, client_id: str, stock_codes: list[str]):
+        """클라이언트의 특정 종목 구독 해제."""
+        bucket = self._subscriptions.get(client_id, set())
+        bucket.difference_update(stock_codes)
+        logger.debug(f"[sub-] {client_id}: {stock_codes} → 현재 {sorted(bucket)}")
+
+    def set_subscriptions(self, client_id: str, stock_codes: list[str]):
+        """클라이언트의 구독 목록을 전체 교체."""
+        self._subscriptions[client_id] = set(stock_codes)
+        logger.debug(f"[sub=] {client_id}: 구독 목록 재설정 → {sorted(stock_codes)}")
+
+    def get_subscriptions(self, client_id: str) -> list[str]:
+        """클라이언트의 현재 구독 목록 반환."""
+        return sorted(self._subscriptions.get(client_id, set()))
+
+    def get_all_subscribed_codes(self) -> Set[str]:
+        """모든 활성 클라이언트에 걸쳐 구독 중인 종목 합집합 반환.
+        KisRealtimeStream이 이 목록을 KIS 측에 구독한다.
+        """
+        result: Set[str] = set()
+        for client_id in self.active_connections:
+            result.update(self._subscriptions.get(client_id, set()))
+        return result
+
+    # ------------------------------------------------------------------
+    # 메시지 전송
+    # ------------------------------------------------------------------
    async def send_personal_message(self, message: str, client_id: str):
-        """특정 클라이언트에게 메시지 전송"""
-        if client_id in self.active_connections:
+        """특정 클라이언트에게 텍스트 메시지 전송"""
+        ws = self.active_connections.get(client_id)
+        if ws:
            try:
-                await self.active_connections[client_id].send_text(message)
+                await ws.send_text(message)
            except Exception as e:
-                logger.error(f"Error sending message to {client_id}: {e}")
+                logger.error(f"Error sending 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():
+        disconnected = []
+        for client_id, ws in list(self.active_connections.items()):
            try:
-                await connection.send_text(message)
+                await ws.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:
+                disconnected.append(client_id)
+        for client_id in disconnected:
            self.disconnect(client_id)
-    
+
    async def broadcast_json(self, data: dict):
        """JSON 데이터를 모든 클라이언트에게 브로드캐스트"""
-        message = _dumps(data)
-        await self.broadcast(message)
-    
+        await self.broadcast(_dumps(data))
+
    async def send_market_update(self, stock_code: str, data: dict):
-        """시장 데이터 업데이트 전송"""
-        message = {
+        """시장 데이터를 해당 종목을 구독한 클라이언트에게만 전송.
+        구독자가 없으면 전체 브로드캐스트(기존 동작 유지).
+        """
+        message = _dumps({
            "type": "market_update",
            "stock_code": stock_code,
-            "data": data
-        }
-        await self.broadcast_json(message)
-    
+            "data": data,
+        })
+        # 해당 종목 구독 클라이언트 필터링
+        targets = [
+            cid for cid, codes in self._subscriptions.items()
+            if stock_code in codes and cid in self.active_connections
+        ]
+        if not targets:
+            # 구독자 없음 — 전체 브로드캐스트(fallback)
+            await self.broadcast(message)
+            return
+        disconnected = []
+        for cid in targets:
+            ws = self.active_connections.get(cid)
+            if not ws:
+                continue
+            try:
+                await ws.send_text(message)
+            except Exception as e:
+                logger.error(f"send_market_update error → {cid}: {e}")
+                disconnected.append(cid)
+        for cid in disconnected:
+            self.disconnect(cid)
+
    async def send_trade_notification(self, trade_data: dict):
        """거래 알림 전송"""
-        message = {
-            "type": "trade_notification",
-            "data": trade_data
-        }
-        await self.broadcast_json(message)
-    
+        await self.broadcast_json({"type": "trade_notification", "data": trade_data})
+
    async def send_mcp_status(self, status: dict):
        """MCP 상태 업데이트 전송"""
-        message = {
-            "type": "mcp_status",
-            "data": status
-        }
-        await self.broadcast_json(message)
+        await self.broadcast_json({"type": "mcp_status", "data": status})


 websocket_manager = WebSocketManager()