233 lines
9.5 KiB
Python
233 lines
9.5 KiB
Python
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, 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 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 = [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
|
|
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 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)],
|
|
"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],
|
|
}
|
|
)
|
|
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"] == "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
|
|
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
|
|
)
|
|
source_geometry = trendline_series(runtime, bar.tf) if source_bars else {}
|
|
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
|
|
),
|
|
"trendline_previous_t": previous_source_t,
|
|
"trendline_future_times": source_geometry.get("future_times", []),
|
|
"trendline_future_durations": source_geometry.get("future_durations", []),
|
|
}
|
|
)
|
|
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,
|
|
"previous_t": previous_source_t,
|
|
"future_times": source_geometry.get("future_times", []),
|
|
"future_durations": source_geometry.get("future_durations", []),
|
|
}
|
|
)
|
|
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)
|