import asyncio from urllib.parse import urlsplit from fastapi import APIRouter, WebSocket, WebSocketDisconnect from app.api.deps import SESSION_COOKIE, session_matches, token_matches from app.analysis.levels import LevelKind from app.bars.models import Timeframe from app.bars.session import bucket_duration from app.analysis.confluence import cluster_levels router = APIRouter() def same_origin(websocket: WebSocket) -> bool: origin = websocket.headers.get("origin", "") host = websocket.headers.get("host", "") return bool(origin and host) and urlsplit(origin).netloc == host def level_enabled(level, enabled: dict) -> bool: kind = level.kind.value if kind == "ma": return level.period in enabled.get("ma", {}).get(level.tf.value, []) if kind == "manual": return enabled.get("manual", True) if kind == "trendline": return enabled.get("auto", False) if kind == "horizontal": return enabled.get("horizontal", True) if kind == "vwap": return enabled.get("vwap", True) return False def enabled_levels(runtime, prefs: dict | None): if not prefs or prefs.get("hidden_levels_score"): return runtime.levels enabled = prefs.get("enabled", {}) return [level for level in runtime.levels if level_enabled(level, enabled)] def connection_clusters(runtime, prefs: dict | None): if runtime.price is None or runtime.stream.last_bar_t is None: return [] return cluster_levels( enabled_levels(runtime, prefs), runtime.stream.last_bar_t, runtime.price, runtime.atr15 ) def session_open(runtime) -> float | None: bars = runtime.store.get(Timeframe.D1, 1) return bars[-1].o if bars else None def trendline_timeframes(runtime) -> set[Timeframe]: return { level.tf for level in runtime.levels if level.kind is LevelKind.MANUAL and level.slope } def trendline_geometry(runtime) -> dict: if not runtime.settings.trendline_source_geometry: return {"mode": "legacy", "series": {}} series = {} for tf in sorted(trendline_timeframes(runtime), key=lambda value: value.value): times = [bar.t for bar in runtime.store.get(tf)] value = {"times": times} if tf is Timeframe.D1: value["durations"] = [bucket_duration(t, tf) for t in times] else: value["duration"] = tf.seconds series[tf.value] = value return {"mode": "source_tf", "series": series} def snapshot(runtime, tf: Timeframe, prefs: dict | None = None) -> dict: return { "type": "snapshot", "tf": tf.value, "bars": [bar.to_dict() for bar in runtime.store.get(tf, 1000)], "levels": [level.to_dict() for level in runtime.levels], "clusters": [cluster.to_dict() for cluster in connection_clusters(runtime, prefs)], "price": runtime.store.get(Timeframe.M1, 1)[-1].c if runtime.store.get(Timeframe.M1, 1) else None, "session_open": session_open(runtime), "trendline_geometry": trendline_geometry(runtime), } @router.websocket("/ws") async def websocket_endpoint(websocket: WebSocket): # Browsers automatically include the HttpOnly session cookie in the # handshake. Query-token support remains for non-browser clients and for # tabs migrating from the previous localStorage-based login. token_ok = token_matches(websocket.app, websocket.query_params.get("token", "")) session_ok = same_origin(websocket) and session_matches( websocket.app, websocket.cookies.get(SESSION_COOKIE, "") ) if not token_ok and not session_ok: await websocket.close(code=1008, reason="Missing or invalid chart token") return await websocket.accept() runtime = websocket.app.state.runtime queue: asyncio.Queue = asyncio.Queue(maxsize=100) runtime.subscribers.add(queue) tf = Timeframe.M1 prefs = None await websocket.send_json(snapshot(runtime, tf, prefs)) geometry_tfs = trendline_timeframes(runtime) async def receive(): nonlocal tf, prefs try: while True: message = await websocket.receive_json() if message.get("type") == "subscribe": tf = Timeframe(message.get("tf", "1m")) await websocket.send_json(snapshot(runtime, tf, prefs)) elif message.get("type") == "prefs": prefs = message clusters = connection_clusters(runtime, prefs) await websocket.send_json( { "type": "clusters", "price": runtime.price, "session_open": session_open(runtime), "clusters": [cluster.to_dict() for cluster in clusters], } ) except WebSocketDisconnect: queue.put_nowait({"type": "disconnect"}) receiver = asyncio.create_task(receive()) try: while True: event = await queue.get() if event["type"] == "disconnect": break if event["type"] == "bar": bar = event["bar"] source_tfs = trendline_timeframes(runtime) source_bars = runtime.store.get(bar.tf) first_source_t = source_bars[0].t if source_bars else None if bar.tf is tf: await websocket.send_json( { "type": "bar", "tf": tf.value, "bar": bar.to_dict(), "session_open": session_open(runtime), "trendline_first_t": ( first_source_t if runtime.settings.trendline_source_geometry and bar.tf in source_tfs else None ), } ) elif ( runtime.settings.trendline_source_geometry and bar.tf in source_tfs ): await websocket.send_json( { "type": "trendline_bar", "tf": bar.tf.value, "t": bar.t, "duration": bucket_duration(bar.t, bar.tf), "first_t": first_source_t, } ) elif event["type"] == "levels": message = { "type": "levels", "changed": event["changed"], "removed": event["removed"], } current_tfs = trendline_timeframes(runtime) if current_tfs != geometry_tfs: message["trendline_geometry"] = trendline_geometry(runtime) geometry_tfs = current_tfs await websocket.send_json(message) elif event["type"] == "clusters": clusters = connection_clusters(runtime, prefs) await websocket.send_json( { "type": "clusters", "price": runtime.price, "session_open": session_open(runtime), "clusters": [cluster.to_dict() for cluster in clusters], } ) elif event["type"] == "alert": # Alerts are produced once, server-side. This socket only relays # them, so opening a second tab cannot double-notify. await websocket.send_json( { "type": "alert", "number": event.get("number", 0), "at": event.get("at", 0), "cluster": event["cluster"].to_dict(), "message": event["message"], } ) except (WebSocketDisconnect, asyncio.CancelledError): pass finally: receiver.cancel() runtime.subscribers.discard(queue)