initial commit

backend/websocket_manager.py

@@ -0,0 +1,182 @@
+
+import asyncio
+import websockets
+import json
+import logging
+import datetime
+from typing import List, Set
+from database import SessionLocal, Stock, StockPrice
+from kis_api import kis
+
+logger = logging.getLogger("WEBSOCKET")
+
+class KisWebSocketManager:
+    def __init__(self):
+        self.active_frontend_connections: List[any] = []
+        self.subscribed_codes: Set[str] = set()
+        self.running = False
+        self.approval_key = None
+        self.msg_queue = asyncio.Queue() # For outgoing subscription requests
+        
+        # KIS Environment
+        if kis.is_paper:
+            self.url = "ws://ops.koreainvestment.com:31000"
+        else:
+            self.url = "ws://ops.koreainvestment.com:21000"
+
+    # ... (connect/disconnect/broadcast remains same)
+
+    # ... (connect/disconnect/broadcast remains same)
+
+    # ... (_handle_realtime_data remains same)
+
+    async def subscribe_stock(self, code):
+        if code in self.subscribed_codes:
+            return
+        self.subscribed_codes.add(code)
+        await self.msg_queue.put(code)
+        logger.info(f"Queued Subscription for {code}")
+            
+    async def connect_frontend(self, websocket):
+        await websocket.accept()
+        self.active_frontend_connections.append(websocket)
+        logger.info(f"Frontend Client Connected. Total: {len(self.active_frontend_connections)}")
+
+    def disconnect_frontend(self, websocket):
+        if websocket in self.active_frontend_connections:
+            self.active_frontend_connections.remove(websocket)
+            logger.info("Frontend Client Disconnected")
+
+    async def broadcast_to_frontend(self, message: dict):
+        # Broadcast to all connected frontend clients
+        for connection in self.active_frontend_connections:
+            try:
+                await connection.send_json(message)
+            except Exception as e:
+                logger.error(f"Broadcast error: {e}")
+                self.disconnect_frontend(connection)
+
+    async def start_kis_socket(self):
+        self.running = True
+        logger.info(f"Starting KIS WebSocket Service... Target: {self.url}")
+        
+        while self.running:
+            # 1. Ensure Approval Key
+            if not self.approval_key:
+                self.approval_key = kis.get_websocket_key()
+                if not self.approval_key:
+                    logger.error("Failed to get WebSocket Approval Key. Retrying in 10s...")
+                    await asyncio.sleep(10)
+                    continue
+                logger.info(f"Got WS Key: {self.approval_key[:10]}...")
+
+            # 2. Connect
+            try:
+                # KIS doesn't use standard ping frames often, handle manually or disable auto-ping
+                async with websockets.connect(self.url, ping_interval=None, open_timeout=20) as ws:
+                    logger.info("Connected to KIS WebSocket Server")
+                    
+                    # Process initial subscriptions
+                    for code in self.subscribed_codes:
+                        await self._send_subscription(ws, code)
+                        
+                    while self.running:
+                        try:
+                            # 1. Check for incoming data
+                            msg = await asyncio.wait_for(ws.recv(), timeout=0.1)
+                            
+                            # PING/PONG (String starting with 0 or 1 usually means data)
+                            if msg[0] in ['0', '1']: 
+                                await self._handle_realtime_data(msg)
+                            else:
+                                # JSON Message (System, PINGPONG)
+                                try:
+                                    data = json.loads(msg)
+                                    if data.get('header', {}).get('tr_id') == 'PINGPONG':
+                                        await ws.send(msg) # Echo back
+                                        continue
+                                except:
+                                    pass 
+                                    
+                        except asyncio.TimeoutError:
+                            pass
+                            
+                        except websockets.ConnectionClosed:
+                            logger.warning("KIS WS Closed. Reconnecting...")
+                            break
+                            
+            except Exception as e:
+                logger.error(f"WS Connection Error: {e}")
+                # If auth failed (maybe expired key?), clear key to force refresh
+                # simplified check: if "Approval key" error in exception message?
+                # For now just retry.
+                await asyncio.sleep(5)
+
+    async def _handle_realtime_data(self, msg: str):
+        # Format: 0|TR_ID|DATA_CNT|Code^Time^Price...
+        try:
+            parts = msg.split('|')
+            if len(parts) < 4: return
+            
+            tr_id = parts[1]
+            data_part = parts[3]
+            
+            if tr_id == "H0STCNT0": # Domestic Stock Price
+                # Data format: Code^Time^CurrentPrice^Sign^Change...
+                # Actually, data_part is delimiter separated.
+                values = data_part.split('^')
+                code = values[0]
+                price = values[2]
+                change = values[4]
+                rate = values[5]
+                
+                # Broadcast
+                payload = {
+                    "type": "PRICE",
+                    "code": code,
+                    "price": price,
+                    "change": change,
+                    "rate": rate,
+                    "timestamp": datetime.datetime.now().isoformat()
+                }
+                await self.broadcast_to_frontend(payload)
+                
+                # Update DB (Optional? Too frequent writes maybe bad)
+                # Let's save only significant updates or throttle?
+                # For now just log/broadcast.
+                
+        except Exception as e:
+            logger.error(f"Data Parse Error: {e} | Msg: {msg[:50]}")
+
+    async def subscribe_stock(self, code):
+        if code in self.subscribed_codes:
+            return
+        self.subscribed_codes.add(code)
+        # If socket is active, send subscription (Implementation complexity: need access to active 'ws' object)
+        # Will handle by restarting connection or using a queue? 
+        # Better: just set it in set, and the main loop will pick it up on reconnect, 
+        # BUT for immediate sub, we need a way to signal the running loop. 
+        # For MVP, let's assume we subscribe on startup or bulk. 
+        # Real-time dynamic sub needs a queue.
+        logger.info(f"Subscribed to {code} (Pending next reconnect/sweep)")
+        
+    async def _send_subscription(self, ws, code):
+        # Domestic Stock Realtime Price: H0STCNT0
+        body = {
+            "header": {
+                "approval_key": self.approval_key,
+                "custtype": "P",
+                "tr_type": "1", # 1: Register, 2: Unregister
+                "content-type": "utf-8"
+            },
+            "body": {
+                "input": {
+                    "tr_id": "H0STCNT0", 
+                    "tr_key": code 
+                }
+            }
+        }
+        await ws.send(json.dumps(body))
+        logger.info(f"Sent Subscription Request for {code}")
+
+ws_manager = KisWebSocketManager()