Update app.py

app.py

@@ -3,41 +3,66 @@ import json
 import os
 import random
 import time
-from typing import Dict, List, Optional, Any
+from dataclasses import dataclass, asdict
+from typing import Any, Dict, List, Optional, Tuple
 
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, Header, HTTPException, WebSocket, WebSocketDisconnect
 from fastapi.responses import FileResponse
-from fastapi.staticfiles import StaticFiles
+
 
 # ----------------------------
 # Config
 # ----------------------------
-TICK_RATE = float(os.getenv("TICK_RATE", "1.0"))  # seconds
-MARKET_LENGTH = int(os.getenv("MARKET_LENGTH", "300"))  # number of "days"/ticks
+TICK_RATE = float(os.getenv("TICK_RATE", "1.0"))          # seconds
+MARKET_LENGTH = int(os.getenv("MARKET_LENGTH", "300"))    # number of ticks/days
 START_PRICE = float(os.getenv("START_PRICE", "100.0"))
 
-# If your frontend file is named differently, adjust here.
+ADMIN_TOKEN = os.getenv("ADMIN_TOKEN", "")  # set as HF Space Secret
+ADMIN_HEADER = "X-ADMIN-TOKEN"
+
 INDEX_FILE = os.getenv("INDEX_FILE", "index.html")
+ADMIN_FILE = os.getenv("ADMIN_FILE", "admin.html")
+
+INITIAL_VOLATILITY = float(os.getenv("DEFAULT_VOLATILITY", "0.8"))
+
+
+# ----------------------------
+# Data models
+# ----------------------------
+@dataclass
+class ScenarioEvent:
+    day: int
+    shockPct: float = 0.0          # e.g. +5.0 means +5%
+    volatility: Optional[float] = None
+    news: Optional[str] = None
+
+
+@dataclass
+class Scenario:
+    name: str = "default"
+    startDay: int = 0
+    basePrice: float = START_PRICE
+    defaultVolatility: float = INITIAL_VOLATILITY
+    events: List[ScenarioEvent] = None
 
-app = FastAPI(title="Trading Game (WebSocket)")
 
 # ----------------------------
 # Market simulator
 # ----------------------------
 class MarketSimulator:
-    def __init__(self, seed: int = 42, start_price: float = 100.0):
+    def __init__(self, seed: int = 42, start_price: float = 100.0, base_vol: float = 0.8):
         self.seed = seed
         self.start_price = start_price
+        self.base_vol = base_vol
 
-    def _generate_market(self, length: int) -> List[Dict[str, float]]:
+    def generate_base_market(self, length: int) -> List[Dict[str, float]]:
         rng = random.Random(self.seed)
         price = self.start_price
         series: List[Dict[str, float]] = []
+        drift = 0.02
 
-        # Simple random walk with gaussian noise, clipped to avoid negative prices.
         for i in range(length):
-            drift = 0.02  # small upward drift per step
-            shock = rng.gauss(0.0, 0.8)
+            shock = rng.gauss(0.0, self.base_vol)
             price = max(1.0, price + drift + shock)
             series.append({"i": i, "close": round(price, 2)})
 
@@ -68,30 +93,18 @@ class ConnectionManager:
             self.leaderboard[name] = {"equity": float(equity), "roi": float(roi), "ts": now}
 
     async def _snapshot_leaderboard(self) -> List[Dict[str, Any]]:
-        # Sort by equity desc, show top N.
         async with self._lock:
             entries = [
                 {"name": n, "equity": v["equity"], "roi": v["roi"], "ts": v.get("ts", 0.0)}
                 for n, v in self.leaderboard.items()
             ]
-
         entries.sort(key=lambda x: x["equity"], reverse=True)
-        # Remove ts before sending to client
         for e in entries:
             e.pop("ts", None)
         return entries[:50]
 
-    async def broadcast_tick(self, day: int) -> None:
-        payload = {
-            "type": "TICK",
-            "payload": {
-                "day": day,
-                "leaderboard": await self._snapshot_leaderboard(),
-            },
-        }
-
-        msg = json.dumps(payload)
-        # Copy sockets to avoid holding the lock while sending
+    async def broadcast(self, obj: Dict[str, Any]) -> None:
+        msg = json.dumps(obj)
         async with self._lock:
             sockets = list(self.active.items())
 
@@ -102,40 +115,188 @@ class ConnectionManager:
             except Exception:
                 stale.append(client_id)
 
-        # Clean up dead connections
         if stale:
             async with self._lock:
                 for cid in stale:
                     self.active.pop(cid, None)
 
+    async def broadcast_tick(self, day: int) -> None:
+        payload = {
+            "type": "TICK",
+            "payload": {
+                "day": day,
+                "leaderboard": await self._snapshot_leaderboard(),
+            },
+        }
+        await self.broadcast(payload)
+
+    async def broadcast_news(self, day: int, text: str) -> None:
+        payload = {"type": "NEWS", "payload": {"day": day, "text": text}}
+        await self.broadcast(payload)
 
+
+app = FastAPI(title="Trading Game (WebSocket + Admin)")
 manager = ConnectionManager()
-market_sim = MarketSimulator(seed=42, start_price=START_PRICE)
-MARKET = market_sim._generate_market(MARKET_LENGTH)
 
-# Global game clock
-CURRENT_DAY = 0
+# ----------------------------
+# Global game state
+# ----------------------------
 DAY_LOCK = asyncio.Lock()
+STATE_LOCK = asyncio.Lock()  # protects scenario/events/volatility/market
+
+CURRENT_DAY = 0
+CURRENT_VOL = INITIAL_VOLATILITY
+
+# Market is the "timeline" clients receive on INIT; we will adjust it when events happen.
+market_sim = MarketSimulator(seed=42, start_price=START_PRICE, base_vol=INITIAL_VOLATILITY)
+MARKET: List[Dict[str, float]] = market_sim.generate_base_market(MARKET_LENGTH)
+
+# Scheduled events by day
+EVENTS: Dict[int, List[ScenarioEvent]] = {}  # day -> [ScenarioEvent,...]
+
 
 # ----------------------------
-# Static files (frontend)
+# Helpers
 # ----------------------------
-# Put your index.html + assets in the Space repo root or a ./static folder.
-# This configuration supports both patterns:
-# - If you have ./static, it will be mounted.
-# - Otherwise, root file serving will still work via "/" -> index.html
-if os.path.isdir("static"):
-    app.mount("/static", StaticFiles(directory="static"), name="static")
+def require_admin(token: Optional[str]) -> None:
+    if not ADMIN_TOKEN:
+        # If not set, keep endpoints locked down by default.
+        raise HTTPException(status_code=403, detail="Admin token not configured on server.")
+    if token != ADMIN_TOKEN:
+        raise HTTPException(status_code=401, detail="Invalid admin token.")
+
+
+def parse_event(obj: Dict[str, Any]) -> ScenarioEvent:
+    day = int(obj["day"])
+    shock = float(obj.get("shockPct", 0.0))
+    vol = obj.get("volatility", None)
+    vol_f = float(vol) if vol is not None else None
+    news = obj.get("news", None)
+    if news is not None:
+        news = str(news)
+    return ScenarioEvent(day=day, shockPct=shock, volatility=vol_f, news=news)
+
+
+def snapshot_events() -> List[Dict[str, Any]]:
+    out: List[Dict[str, Any]] = []
+    for d in sorted(EVENTS.keys()):
+        for ev in EVENTS[d]:
+            out.append(asdict(ev))
+    return out
+
+
+def apply_price_shock_from_day(day: int, shock_pct: float) -> None:
+    """
+    Applies a multiplicative factor to MARKET[day:] so the event shifts the future path.
+    This matches the "future path shift" interpretation.
+    """
+    if day < 0 or day >= len(MARKET):
+        return
+    factor = 1.0 + (shock_pct / 100.0)
+    for i in range(day, len(MARKET)):
+        MARKET[i]["close"] = round(max(1.0, MARKET[i]["close"] * factor), 2)
+
+
+def regen_market_with_volatility(seed: int, start_price: float, base_vol: float) -> List[Dict[str, float]]:
+    sim = MarketSimulator(seed=seed, start_price=start_price, base_vol=base_vol)
+    return sim.generate_base_market(MARKET_LENGTH)
 
 
+# ----------------------------
+# Static pages
+# ----------------------------
 @app.get("/")
 async def root():
-    # If you keep index.html in ./static, serve it from there; else from repo root.
-    if os.path.exists(os.path.join("static", INDEX_FILE)):
-        return FileResponse(os.path.join("static", INDEX_FILE))
     return FileResponse(INDEX_FILE)
 
 
+@app.get("/admin")
+async def admin_page():
+    return FileResponse(ADMIN_FILE)
+
+
+# ----------------------------
+# Admin REST API
+# ----------------------------
+@app.get("/admin/state")
+async def admin_state(x_admin_token: Optional[str] = Header(default=None, alias=ADMIN_HEADER)):
+    require_admin(x_admin_token)
+    async with DAY_LOCK:
+        day = CURRENT_DAY
+    async with STATE_LOCK:
+        return {
+            "day": day,
+            "tickRate": TICK_RATE,
+            "marketLength": len(MARKET),
+            "currentVolatility": CURRENT_VOL,
+            "events": snapshot_events(),
+        }
+
+
+@app.post("/admin/clear_events")
+async def admin_clear_events(x_admin_token: Optional[str] = Header(default=None, alias=ADMIN_HEADER)):
+    require_admin(x_admin_token)
+    async with STATE_LOCK:
+        EVENTS.clear()
+    return {"ok": True}
+
+
+@app.post("/admin/add_event")
+async def admin_add_event(body: Dict[str, Any], x_admin_token: Optional[str] = Header(default=None, alias=ADMIN_HEADER)):
+    require_admin(x_admin_token)
+
+    # Supports either "day" (absolute) or "offset" (relative to current day)
+    async with DAY_LOCK:
+        cur = CURRENT_DAY
+
+    if "day" in body:
+        day = int(body["day"])
+    elif "offset" in body:
+        day = cur + int(body["offset"])
+    else:
+        raise HTTPException(status_code=400, detail="Provide 'day' or 'offset'.")
+
+    if day < cur:
+        raise HTTPException(status_code=400, detail=f"Event day {day} is in the past (current day {cur}).")
+
+    ev = parse_event({**body, "day": day})
+
+    async with STATE_LOCK:
+        EVENTS.setdefault(day, []).append(ev)
+
+    return {"ok": True, "event": asdict(ev)}
+
+
+@app.post("/admin/load_scenario")
+async def admin_load_scenario(body: Dict[str, Any], x_admin_token: Optional[str] = Header(default=None, alias=ADMIN_HEADER)):
+    require_admin(x_admin_token)
+
+    name = str(body.get("name", "scenario"))
+    start_day = int(body.get("startDay", 0))
+    base_price = float(body.get("basePrice", START_PRICE))
+    default_vol = float(body.get("defaultVolatility", INITIAL_VOLATILITY))
+    evs_raw = body.get("events", [])
+    if not isinstance(evs_raw, list):
+        raise HTTPException(status_code=400, detail="'events' must be a list.")
+
+    evs = [parse_event(e) for e in evs_raw]
+
+    # Reset market + events, reset day, reset volatility
+    async with STATE_LOCK:
+        global MARKET, CURRENT_VOL
+        CURRENT_VOL = default_vol
+        MARKET = regen_market_with_volatility(seed=42, start_price=base_price, base_vol=default_vol)
+        EVENTS.clear()
+        for ev in evs:
+            EVENTS.setdefault(ev.day, []).append(ev)
+
+    async with DAY_LOCK:
+        global CURRENT_DAY
+        CURRENT_DAY = max(0, min(start_day, len(MARKET) - 1))
+
+    return {"ok": True, "name": name, "startDay": CURRENT_DAY, "eventsLoaded": len(evs)}
+
+
 # ----------------------------
 # WebSocket endpoint
 # ----------------------------
@@ -144,14 +305,9 @@ async def websocket_endpoint(websocket: WebSocket, client_id: str):
     await manager.connect(websocket, client_id)
 
     # Send INIT immediately
-    init_msg = {
-        "type": "INIT",
-        "payload": {
-            "market": MARKET,
-            "startDay": 0,
-        },
-    }
-    await websocket.send_text(json.dumps(init_msg))
+    async with STATE_LOCK:
+        init_payload = {"market": MARKET, "startDay": 0}
+    await websocket.send_text(json.dumps({"type": "INIT", "payload": init_payload}))
 
     try:
         while True:
@@ -159,7 +315,6 @@ async def websocket_endpoint(websocket: WebSocket, client_id: str):
             try:
                 data = json.loads(raw)
             except Exception:
-                # Ignore invalid JSON
                 continue
 
             msg_type = data.get("type")
@@ -167,28 +322,20 @@ async def websocket_endpoint(websocket: WebSocket, client_id: str):
 
             if msg_type == "UPDATE_EQUITY":
                 name = str(payload.get("name", client_id))
-                equity = payload.get("equity", 0.0)
-                roi = payload.get("roi", 0.0)
-
-                # Defensive parsing
                 try:
-                    equity_f = float(equity)
+                    equity_f = float(payload.get("equity", 0.0))
                 except Exception:
                     equity_f = 0.0
                 try:
-                    roi_f = float(roi)
+                    roi_f = float(payload.get("roi", 0.0))
                 except Exception:
                     roi_f = 0.0
 
                 await manager.update_equity(name=name, equity=equity_f, roi=roi_f)
-            else:
-                # Unknown message types are ignored to keep protocol forward-compatible
-                continue
 
     except WebSocketDisconnect:
         await manager.disconnect(client_id)
     except Exception:
-        # Any other error: drop connection cleanly
         await manager.disconnect(client_id)
 
 
@@ -196,16 +343,37 @@ async def websocket_endpoint(websocket: WebSocket, client_id: str):
 # Background tick loop
 # ----------------------------
 async def game_loop():
-    global CURRENT_DAY
+    global CURRENT_DAY, CURRENT_VOL
+
     while True:
         await asyncio.sleep(TICK_RATE)
+
+        # advance day
         async with DAY_LOCK:
             CURRENT_DAY = (CURRENT_DAY + 1) % len(MARKET)
             day = CURRENT_DAY
+
+        # apply scheduled events
+        news_to_broadcast: List[str] = []
+        async with STATE_LOCK:
+            if day in EVENTS and EVENTS[day]:
+                for ev in EVENTS[day]:
+                    if ev.shockPct:
+                        # shift entire future path from this day onwards
+                        apply_price_shock_from_day(day, ev.shockPct)
+                    if ev.volatility is not None:
+                        CURRENT_VOL = float(ev.volatility)
+                    if ev.news:
+                        news_to_broadcast.append(ev.news)
+
+        # broadcast optional news first
+        for text in news_to_broadcast:
+            await manager.broadcast_news(day, text)
+
+        # broadcast tick
         await manager.broadcast_tick(day)
 
 
 @app.on_event("startup")
 async def on_startup():
-    # Start ticker loop once.
     asyncio.create_task(game_loop())