Blame for app.py

"""Lantern web server: runs Beam episodes and streams them to the browser over SSE.



Live runs spend the Space owner's inference credits, so they are limited to one at a

time and a daily cap. Every finished episode is saved and can be replayed for free.

"""

from __future__ import annotations

import asyncio
import json
import hashlib
import logging
import hmac
import os
import secrets
import threading
import time
import uuid
from datetime import date
from pathlib import Path

import httpx
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from lantern.agent import EpisodeConfig, run_episode
from lantern.coach import Hint, HintQueue
from lantern.guard import check_hint
from lantern.env import TASKS, DungeonEnv
from lantern.journal import Journal
from lantern.llm import ChatClient, OutOfCredits
from lantern.store import EpisodeStore

DATA_DIR = Path(os.environ.get("LANTERN_DATA", "data"))
DAILY_LIVE_RUNS = int(os.environ.get("LANTERN_DAILY_RUNS", "20"))
VISITOR_DAILY_RUNS = int(os.environ.get("LANTERN_VISITOR_DAILY_RUNS", "3"))
MAX_CALLS_PER_RUN = int(os.environ.get("LANTERN_MAX_CALLS", "60"))
WEB_DIR = Path(__file__).parent / "web"

app = FastAPI(title="Lantern")
log = logging.getLogger("lantern.app")
# Episodes and the journal are mirrored to a Hub dataset so they survive restarts.
store = EpisodeStore(DATA_DIR, os.environ.get("LANTERN_DATASET"), os.environ.get("LANTERN_DATASET_TOKEN"))
store.pull()
journal = Journal(store.journal_path)


class Run:
    """One live episode running in a background thread, buffering its events."""

    def __init__(self, task_id: str, seed: int, use_lessons: bool):
        self.id = uuid.uuid4().hex[:12]
        self.task_id, self.seed, self.use_lessons = task_id, seed, use_lessons
        self.events: list[dict] = []
        self.done = False
        self.stop = threading.Event()
        self._lock = threading.Lock()
        self.visitor: str | None = None   # keyed hash of who started it; only they may coach
        self.hints = HintQueue()
        self.hints_sent = 0
        self.coached = False

    def start(self) -> None:
        threading.Thread(target=self._play, daemon=True).start()

    def request_stop(self) -> None:
        """End the run for viewers right away.



        A Beam call already in flight cannot be cancelled cleanly, so the worker thread

        lets it finish and discards the result; viewers and the run slot don't wait for it.

        """
        self.stop.set()
        with self._lock:
            if self.done:
                return
            turns = [e for e in self.events if e["type"] == "turn"]
            # Usage up to the last finished call; a call still in flight isn't counted.
            usage = dict(turns[-1].get("usage", {})) if turns else {}
            usage.setdefault("calls", len(turns))
            self.events.append({"type": "end", "task": self.task_id, "outcome": "stopped",
                                "steps": turns[-1]["steps"] if turns else 0, "reward": 0.0,
                                "usage": usage, "new_lessons": []})
            self._finish()

    def add_hint(self, hint: Hint) -> None:
        """Queue advice from the viewer; it interrupts Beam's current plan between moves."""
        with self._lock:
            self.hints.add(hint)
            self.hints_sent += 1
            self.coached = True

    def _emit(self, event: dict) -> bool:
        """Record an event unless the run was already stopped; returns False once stopped."""
        with self._lock:
            if self.done:
                return False
            self.events.append(event)
            return True

    def _play(self) -> None:
        env = None
        try:
            env = DungeonEnv(TASKS[self.task_id], seed=self.seed)
            lessons = journal.lessons(self.task_id)
            # Journal v2: Beam curates the journal from each lesson's win/loss record.
            # With "Use journal" off Beam never sees the lessons, so the run must neither credit
            # them with its outcome nor rewrite the journal (same as the benchmark's "off").
            config = EpisodeConfig(use_lessons=self.use_lessons, max_llm_calls=MAX_CALLS_PER_RUN,
                                   reflect=self.use_lessons, curate=True,
                                   lesson_records=journal.records(self.task_id) if self.use_lessons else None)
            for event in run_episode(env, ChatClient(), lessons, config, self.stop, self.hints):
                if self.stop.is_set():
                    break
                # A coached run's lessons could carry a viewer's text into the journal that
                # every later run reads in its system prompt, so only uncoached runs write to it.
                if event["type"] == "end" and event.get("journal_records") is not None and not self.coached:
                    event["journal"] = journal.set_records(self.task_id, event["journal_records"])
                    store.push_journal()
                if not self._emit(event):
                    break
        except OutOfCredits as exc:
            log.warning("Live runs paused, provider said: %s", exc)
            credits.mark_out()
            self._emit({"type": "error", "kind": "credits", "message": CREDITS_MESSAGE})
        except httpx.HTTPStatusError as exc:  # refused even after retries; usually temporary
            log.warning("Provider refused a call: %s", exc)
            self._emit({"type": "error", "kind": "busy", "message": BUSY_MESSAGE})
        except Exception as exc:  # surface the failure to the viewer instead of hanging
            log.exception("Run %s failed", self.id)
            # Details go to the server log only; viewers and the saved episode get a plain message.
            self._emit({"type": "error", "message": "Something went wrong on the server, so this run stopped."})
        finally:
            try:
                if env is not None:
                    env.close()
            finally:  # a failing cleanup must never leave the Space stuck "busy"
                with self._lock:
                    if not self.done:
                        self._finish()

    def _finish(self) -> None:
        """Mark done and save; callers hold the lock."""
        self.done = True
        end = next((e for e in reversed(self.events) if e["type"] == "end"), {})
        # Coached runs hold text typed by visitors (also quoted in Beam's reasoning), so they
        # stay on the Space and never go to the public dataset.
        store.save({"id": self.id, "task": self.task_id, "seed": self.seed, "source": "space",
                    "use_lessons": self.use_lessons, "finished_at": time.time(),
                    "outcome": end.get("outcome", "error"), "steps": end.get("steps", 0),
                    "coached": self.coached,
                    "events": list(self.events)}, upload=not self.coached)


runs: dict[str, Run] = {}
_quota = {"day": date.today(), "used": 0}


def _roll_quota() -> None:
    """Start a fresh daily count once the date changes (called on every read and write)."""
    if _quota["day"] != date.today():
        _quota.update(day=date.today(), used=0)
_lock = threading.Lock()

CREDITS_MESSAGE = ("Live runs are paused: this demo has reached its inference spending limit "
                   "for now. Replays still work.")
BUSY_MESSAGE = ("The model is busy right now and didn't answer. Please try again in a minute; "
                "replays work in the meantime.")


class CreditsState:
    """Remembers an out-of-credits error; allows a fresh attempt after `retry_s` in case of a top-up."""

    def __init__(self, retry_s: float = 600):
        self.retry_s = retry_s
        self.out_since: float | None = None

    def mark_out(self) -> None:
        self.out_since = time.time()

    @property
    def available(self) -> bool:
        if self.out_since is not None and time.time() - self.out_since > self.retry_s:
            self.out_since = None
        return self.out_since is None


credits = CreditsState()


class VisitorQuota:
    """Per-visitor daily live-run allowance.



    Visitors are identified by a keyed hash of their IP address. The raw IP is never

    stored or logged, the key lives only in memory (new on every restart), and counts

    reset each day, so nothing here can identify a person.

    """

    def __init__(self, limit: int):
        self.limit = limit
        self._key = secrets.token_bytes(32)
        self._day = date.today()
        self._counts: dict[str, int] = {}

    def visitor(self, request: Request) -> str:
        # HF's proxy appends the real client address to any X-Forwarded-For the client sent,
        # so the left-most entry is client-controlled (verified: a fake one dodged the limit).
        # Use the right-most entry, the one the proxy added.
        forwarded = [part.strip() for part in request.headers.get("x-forwarded-for", "").split(",") if part.strip()]
        ip = forwarded[-1] if forwarded else (request.client.host if request.client else "unknown")
        return hmac.new(self._key, ip.encode(), hashlib.sha256).hexdigest()

    def _roll(self) -> None:
        if self._day != date.today():
            self._day, self._counts = date.today(), {}

    def left(self, visitor: str) -> int:
        self._roll()
        return max(0, self.limit - self._counts.get(visitor, 0))

    def use(self, visitor: str) -> None:
        self._roll()
        self._counts[visitor] = self._counts.get(visitor, 0) + 1


visitors = VisitorQuota(VISITOR_DAILY_RUNS)
VISITOR_LIMIT_MESSAGE = (f"You've used your {VISITOR_DAILY_RUNS} live runs for today. Thanks for exploring! "
                         "Replays are free and unlimited, and live runs reset at midnight UTC.")


MAX_HINT_CHARS = 120
MAX_HINTS_PER_RUN = 30


class HintRequest(BaseModel):
    text: str = Field(min_length=1, max_length=MAX_HINT_CHARS)
    # Where the viewer was looking when they sent it; their screen can lag the game.
    seen_step: int | None = Field(default=None, ge=0, le=100_000)
    seen_x: int | None = Field(default=None, ge=0, le=200)
    seen_y: int | None = Field(default=None, ge=0, le=200)


def clean_hint(text: str) -> str:
    """Printable characters only, whitespace collapsed; hints are shown to Beam as data."""
    kept = "".join(" " if ch.isspace() else ch for ch in text if ch.isspace() or ch.isprintable())
    kept = kept.replace('"', "'")  # hints are shown to Beam inside double quotes; don't let one close them
    return " ".join(kept.split())[:MAX_HINT_CHARS]


class RunRequest(BaseModel):
    task: str
    seed: int = Field(default=0, ge=0, le=1_000_000)
    use_lessons: bool = True


@app.get("/api/tasks")
def list_tasks() -> list[dict]:
    return [{"id": t.id, "title": t.title, "goal": t.goal, "max_steps": t.max_steps,
             "lessons": len(journal.lessons(t.id))} for t in TASKS.values()]


@app.get("/api/status")
def status(request: Request) -> dict:
    _roll_quota()
    busy = any(not r.done for r in runs.values())
    return {"busy": busy, "runs_left_today": max(0, DAILY_LIVE_RUNS - _quota["used"]),
            "your_runs_left": visitors.left(visitors.visitor(request)),
            "live_available": credits.available,
            "live_message": None if credits.available else CREDITS_MESSAGE}


@app.post("/api/runs")
def start_run(req: RunRequest, request: Request) -> dict:
    if req.task not in TASKS:
        raise HTTPException(404, f"Unknown level '{req.task}'.")
    with _lock:
        _roll_quota()
        if any(not r.done for r in runs.values()):
            raise HTTPException(409, "Beam is already exploring. Watch that run or open a replay.")
        if not credits.available:
            raise HTTPException(402, CREDITS_MESSAGE)
        visitor = visitors.visitor(request)
        if visitors.left(visitor) <= 0:
            raise HTTPException(429, VISITOR_LIMIT_MESSAGE)
        if _quota["used"] >= DAILY_LIVE_RUNS:
            raise HTTPException(429, "This demo's live runs for today are used up. Replays are free "
                                     "and unlimited, and live runs reset at midnight UTC.")
        _quota["used"] += 1
        visitors.use(visitor)
        run = Run(req.task, req.seed, req.use_lessons)
        run.visitor = visitor
        runs[run.id] = run
    run.start()
    return {"id": run.id}


@app.post("/api/runs/{run_id}/hint")
def send_hint(run_id: str, req: HintRequest, request: Request) -> dict:
    run = runs.get(run_id)
    if run is None or run.done:
        raise HTTPException(404, "That run has finished.")
    if run.visitor != visitors.visitor(request):
        raise HTTPException(403, "Only the person who started this run can coach it.")
    if run.hints_sent >= MAX_HINTS_PER_RUN:
        raise HTTPException(429, f"That's the limit of {MAX_HINTS_PER_RUN} hints for one run.")
    text = clean_hint(req.text)
    if not text:
        raise HTTPException(400, "Type a hint first.")
    problem = check_hint(text)   # PII, health information, profanity; the text itself is never logged
    if problem:
        raise HTTPException(400, f"That hint looks like it contains {problem}. Please keep hints about the game.")
    run.add_hint(Hint(text, req.seen_step, req.seen_x, req.seen_y))
    return {"queued": text}


@app.post("/api/runs/{run_id}/stop")
def stop_run(run_id: str, request: Request) -> dict:
    run = runs.get(run_id)
    if run is None:
        raise HTTPException(404, "No such run.")
    if run.visitor is not None and run.visitor != visitors.visitor(request):
        raise HTTPException(403, "Only the person who started this run can stop it.")
    run.request_stop()
    return {"stopped": True}


@app.get("/api/runs/{run_id}/events")
async def run_events(run_id: str) -> StreamingResponse:
    run = runs.get(run_id)
    if run is None:
        raise HTTPException(404, "No such run.")

    async def stream():
        sent = 0
        while True:
            while sent < len(run.events):
                yield f"data: {json.dumps(run.events[sent])}\n\n"
                sent += 1
            if run.done and sent >= len(run.events):
                yield "event: close\ndata: {}\n\n"
                return
            await asyncio.sleep(0.25)

    return StreamingResponse(stream(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


@app.get("/api/episodes")
def list_episodes() -> list[dict]:
    return store.summaries(50)


@app.get("/api/episodes/{episode_id}")
def get_episode(episode_id: str) -> dict:
    if not episode_id.isalnum():
        raise HTTPException(400, "Bad episode id.")
    record = store.load(episode_id)
    if record is None:
        raise HTTPException(404, "No such episode.")
    return record


def _asset_version() -> str:
    """Short hash of the page's own scripts and styles; changes whenever they change."""
    digest = hashlib.sha256()
    for name in ("app.js", "audio.js", "scene3d.js", "style.css"):
        digest.update((WEB_DIR / name).read_bytes())
    return digest.hexdigest()[:10]


ASSET_VERSION = _asset_version()


@app.get("/")
def index() -> HTMLResponse:
    # Stamp asset URLs so a new deploy is never served stale scripts from a browser cache
    # (the Space runs in an iframe, where a hard refresh doesn't reach the app's files).
    html = (WEB_DIR / "index.html").read_text("utf-8")
    for name in ("app.js", "audio.js", "scene3d.js", "style.css"):
        html = html.replace(f"/static/{name}", f"/static/{name}?v={ASSET_VERSION}")
    return HTMLResponse(html, headers={"Cache-Control": "no-cache"})


@app.middleware("http")
async def revalidate_static(request: Request, call_next):
    """Browsers must check static files with the server (cheap, via ETag) before reusing them."""
    response = await call_next(request)
    if request.url.path.startswith("/static/"):
        response.headers["Cache-Control"] = "no-cache"
    return response


app.mount("/static", StaticFiles(directory=WEB_DIR), name="static")