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.analysis.bar_space import fill_short_gaps from app.bars.session import bucket_duration, future_bucket_starts 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("drawings", True) and 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 session_range(runtime) -> dict | None: """High/low of the forming CME session — today's range, not prior day.""" bars = runtime.store.get(Timeframe.D1, 1) if not bars: return None bar = bars[-1] return {"t": bar.t, "high": bar.h, "low": bar.l} 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): series[tf.value] = trendline_series(runtime, tf) return {"mode": "source_tf", "series": series} def trendline_series(runtime, tf: Timeframe) -> dict: times = fill_short_gaps([bar.t for bar in runtime.store.get(tf)], tf) value = {"times": times} if tf is Timeframe.D1: value["durations"] = [bucket_duration(t, tf) for t in times] else: value["duration"] = tf.seconds future = future_bucket_starts(times[-1], tf) if times else [] value["future_times"] = future value["future_durations"] = [bucket_duration(t, tf) for t in future] return value def displayed_future_times(runtime, tf: Timeframe) -> list[int]: bars = runtime.store.get(tf, 1) return future_bucket_starts(bars[-1].t, tf) if bars else [] def bar_client_message(runtime, bar, displayed_tf: Timeframe, *, full: bool) -> dict | None: """What one socket should hear about a bar. A forming tick of an already-seen minute is just the live candle. Future calendars and HTF geometry go out only when that timeframe's timestamp advances — otherwise 1m at 4 Hz resends 180 slots and walks the store. """ if bar.tf is displayed_tf: payload = {"type": "bar", "tf": displayed_tf.value, "bar": bar.to_dict()} if not full: return payload source_tfs = trendline_timeframes(runtime) source_bars = runtime.store.get(bar.tf) source_geometry = trendline_series(runtime, bar.tf) if source_bars else {} first_source_t = source_bars[0].t if source_bars else None source_index = next( (index for index, value in enumerate(source_bars) if value.t == bar.t), None, ) previous_source_t = ( source_bars[source_index - 1].t if source_index is not None and source_index > 0 else None ) rng = session_range(runtime) extra = { "session_t": rng["t"], "session_high": rng["high"], "session_low": rng["low"], } if rng else {} payload.update( { "duration": bucket_duration(bar.t, bar.tf), "session_open": session_open(runtime), **extra, "trendline_first_t": ( first_source_t if runtime.settings.trendline_source_geometry and bar.tf in source_tfs else None ), "trendline_previous_t": previous_source_t, "trendline_future_times": source_geometry.get("future_times", []), "trendline_future_durations": source_geometry.get("future_durations", []), "future_times": source_geometry.get("future_times", []), } ) return payload if ( full and runtime.settings.trendline_source_geometry and bar.tf in trendline_timeframes(runtime) ): source_bars = runtime.store.get(bar.tf) source_geometry = trendline_series(runtime, bar.tf) if source_bars else {} first_source_t = source_bars[0].t if source_bars else None source_index = next( (index for index, value in enumerate(source_bars) if value.t == bar.t), None, ) previous_source_t = ( source_bars[source_index - 1].t if source_index is not None and source_index > 0 else None ) return { "type": "trendline_bar", "tf": bar.tf.value, "t": bar.t, "duration": bucket_duration(bar.t, bar.tf), "first_t": first_source_t, "previous_t": previous_source_t, "future_times": source_geometry.get("future_times", []), "future_durations": source_geometry.get("future_durations", []), } return None def snapshot(runtime, tf: Timeframe, prefs: dict | None = None) -> dict: events, events_more = runtime.events.recent() rng = session_range(runtime) extra = { "session_t": rng["t"], "session_high": rng["high"], "session_low": rng["low"], } if rng else {} return { "type": "snapshot", "tf": tf.value, "bars": [bar.to_dict() for bar in runtime.store.get(tf)], "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), **extra, "trendline_geometry": trendline_geometry(runtime), "future_times": displayed_future_times(runtime, tf), "events": events, "events_more": events_more, "instrument": runtime.settings.profile.payload(), } @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) last_bar_t: dict[Timeframe, int] = {} 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], } ) elif message.get("type") == "trendline_geometry": await websocket.send_json( {"type": "trendline_geometry", "geometry": trendline_geometry(runtime)} ) 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"] == "resync": # History changed behind the live edge (a backfilled outage). # Bar deltas only move the tail, so the whole series is resent. last_bar_t.clear() await websocket.send_json(snapshot(runtime, tf, prefs)) elif event["type"] == "bar": bar = event["bar"] full = last_bar_t.get(bar.tf) != bar.t if full: last_bar_t[bar.tf] = bar.t payload = bar_client_message(runtime, bar, tf, full=full) if payload is not None: await websocket.send_json(payload) 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)