""" ========================================= CV6 AI Trading OS WebSocket Router Routes: /ws/live /ws/market /ws/ticks ========================================= """ import asyncio import json import uuid from typing import Optional from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query from loguru import logger from app.websocket.websocket_manager import manager from app.websocket.market_stream import stream_engine router = APIRouter(tags=["WebSocket"]) def _ensure_broker_subscribed(symbol: str) -> None: """ PHASE-2.5 FIX (Goal 3): stream_engine.start(symbol) only starts the internal poll/broadcast loop — it never told the actual Angel One WebSocket feed to subscribe to a NEW symbol outside the fixed startup watchlist. Any client requesting a symbol outside that ~48-symbol list would silently get simulated data forever, even with the broker fully connected. live_feed.subscribe() is incremental/idempotent — safe to call on every client subscribe. """ try: from app.market.live_feed import live_feed live_feed.subscribe([symbol]) except Exception as e: logger.debug(f"[WS] Could not subscribe {symbol} to live feed: {e}") # ─── /ws/live ───────────────────────────────────────────────────── # Protocol (JSON commands from client): # {"action":"subscribe","symbol":"NIFTY"} # {"action":"unsubscribe","symbol":"NIFTY"} # {"action":"status"} # {"action":"ping"} @router.websocket("/ws/live") async def ws_live( ws: WebSocket, client_id: Optional[str] = Query(default=None), ): """ General-purpose live feed. Subscribe/unsubscribe symbols via JSON commands. """ cid = client_id or str(uuid.uuid4()) await manager.connect(ws, cid) try: while True: raw = await ws.receive_text() try: msg = json.loads(raw) except json.JSONDecodeError: await manager.send(cid, {"type": "error", "msg": "Invalid JSON"}) continue action = msg.get("action", "").lower() if action == "subscribe": sym = msg.get("symbol", "").upper() if not sym: await manager.send(cid, {"type": "error", "msg": "symbol required"}) continue await manager.subscribe(cid, sym) _ensure_broker_subscribed(sym) await stream_engine.start(sym) last = stream_engine.last_tick(sym) if last: await manager.send(cid, last) elif action == "unsubscribe": sym = msg.get("symbol", "").upper() await manager.unsubscribe(cid, sym) elif action == "status": await manager.send(cid, { "type": "status", **manager.status(), "streaming": stream_engine.streaming(), **stream_engine.status(), }) elif action == "ping": await manager.send(cid, {"type": "pong"}) else: await manager.send(cid, {"type": "error", "msg": f"Unknown: {action}"}) except WebSocketDisconnect: logger.info(f"[/ws/live] Disconnected: {cid}") finally: await manager.disconnect(cid) # ─── /ws/market ─────────────────────────────────────────────────── # Auto-subscribes to ?symbol=NIFTY on connect. @router.websocket("/ws/market") async def ws_market( ws: WebSocket, symbol: str = Query(default="NIFTY"), client_id: Optional[str] = Query(default=None), ): """ Single-symbol auto-feed. Connect with ?symbol=BANKNIFTY to receive that symbol's ticks. """ cid = client_id or str(uuid.uuid4()) sym = symbol.upper() await manager.connect(ws, cid) await manager.subscribe(cid, sym) _ensure_broker_subscribed(sym) await stream_engine.start(sym) last = stream_engine.last_tick(sym) if last: await manager.send(cid, last) try: while True: raw = await ws.receive_text() try: msg = json.loads(raw) if msg.get("action") == "ping": await manager.send(cid, {"type": "pong"}) elif msg.get("action") == "change_symbol": new_sym = msg.get("symbol", "").upper() if new_sym: await manager.unsubscribe(cid, sym) sym = new_sym await manager.subscribe(cid, sym) _ensure_broker_subscribed(sym) await stream_engine.start(sym) except json.JSONDecodeError as e: logger.debug("[/ws/market] JSON decode error: {}", e) except WebSocketDisconnect: logger.info(f"[/ws/market] Disconnected: {cid}") finally: await manager.disconnect(cid) # ─── /ws/ticks ──────────────────────────────────────────────────── # Multi-symbol: ?symbols=NIFTY,BANKNIFTY,RELIANCE @router.websocket("/ws/ticks") async def ws_ticks( ws: WebSocket, symbols: str = Query(default="NIFTY,BANKNIFTY"), client_id: Optional[str] = Query(default=None), ): """ Multi-symbol raw tick stream. Connect with ?symbols=NIFTY,BANKNIFTY,RELIANCE """ cid = client_id or str(uuid.uuid4()) sym_list = [s.strip().upper() for s in symbols.split(",") if s.strip()] await manager.connect(ws, cid) for sym in sym_list: await manager.subscribe(cid, sym) _ensure_broker_subscribed(sym) await stream_engine.start(sym) last = stream_engine.last_tick(sym) if last: await manager.send(cid, last) try: while True: raw = await ws.receive_text() try: msg = json.loads(raw) action = msg.get("action", "").lower() if action == "ping": await manager.send(cid, {"type": "pong"}) elif action == "add": s = msg.get("symbol", "").upper() if s: await manager.subscribe(cid, s) _ensure_broker_subscribed(s) await stream_engine.start(s) elif action == "remove": s = msg.get("symbol", "").upper() if s: await manager.unsubscribe(cid, s) elif action == "status": await manager.send(cid, { "type": "status", "streaming": stream_engine.streaming(), **stream_engine.status(), }) except json.JSONDecodeError as e: logger.debug("[/ws/ticks] JSON decode error: {}", e) except WebSocketDisconnect: logger.info(f"[/ws/ticks] Disconnected: {cid}") finally: await manager.disconnect(cid) # ─── /ws/state — real-time source-of-truth push ──────────────────── # Streams the TradingStateManager snapshot (positions, capital, PnL) # whenever it changes. Version-gated: only re-sends when TSM.version() # advances, so an idle account produces no traffic. @router.websocket("/ws/state") async def ws_state( ws: WebSocket, client_id: Optional[str] = Query(default=None), ): from app.state.trading_state_manager import trading_state_manager cid = client_id or str(uuid.uuid4()) await manager.connect(ws, cid) last_version = -1 try: # Send the current snapshot immediately on connect snap = trading_state_manager.snapshot() last_version = snap["version"] await manager.send(cid, {"type": "state", **snap}) while True: await asyncio.sleep(1.0) v = trading_state_manager.version() if v != last_version: last_version = v await manager.send(cid, {"type": "state", **trading_state_manager.snapshot()}) except WebSocketDisconnect: logger.info(f"[/ws/state] Disconnected: {cid}") finally: await manager.disconnect(cid) # ─── REST status endpoint ────────────────────────────────────────── @router.get("/ws/status", tags=["WebSocket"]) def ws_status(): """ WebSocket connection status and active streams. PHASE-2.5 (Goal 3): includes live_provider_registered + a per-symbol LIVE/SIMULATED breakdown so a caller never has to guess whether it's getting real Angel One ticks or the simulated fallback. """ return { "success": True, "websocket": manager.status(), "streaming_symbols": stream_engine.streaming(), **stream_engine.status(), }