patch_fiftyone_frontend.py

"""Patches the packaged FiftyOne frontend for a stateless hosted demo."""

import importlib.util
import json
import logging
from pathlib import Path
import re


ENV_ASSIGNMENT_PATTERN = re.compile(
    r"(__vite_import_meta_env__\$\w+=\{)([^}]*)(\})"
)
HEARTBEAT_INTERVAL_MS = 60_000
HEARTBEAT_MARKER = "fiftyone-hf-session-heartbeat"
MANIFEST_PATH = Path("/app/datasets.json")
LOG_FORMAT = "%(asctime)s | %(levelname)s | %(name)s | %(message)s"

logging.basicConfig(level=logging.INFO, format=LOG_FORMAT)
logger = logging.getLogger(__name__)


def main():
    """Enables stateless state and browser session heartbeats."""
    static_dir = _get_static_dir()
    dataset_names = _get_dataset_names()
    _enable_stateless_frontend(static_dir)
    _add_session_heartbeat(static_dir, dataset_names)


def _get_static_dir():
    """Returns the static frontend directory in the installed package."""
    spec = importlib.util.find_spec("fiftyone")
    if spec is None or spec.origin is None:
        raise RuntimeError("Could not locate the installed FiftyOne package")

    static_dir = Path(spec.origin).parent / "server" / "static"
    if not static_dir.is_dir():
        raise RuntimeError(
            f"Could not locate FiftyOne static assets at {static_dir}"
        )

    logger.info("Located FiftyOne frontend at %s", static_dir)
    return static_dir


def _get_dataset_names():
    """Gets the configured base dataset names from the manifest."""
    with open(MANIFEST_PATH) as file:
        manifest = json.load(file)

    names = [entry["name"] for entry in manifest["datasets"]]
    if not names:
        raise ValueError("datasets.json must contain at least one dataset")

    return names


def _enable_stateless_frontend(static_dir):
    """Enables URL-local state in every compiled Vite environment object."""
    javascript_files = sorted((static_dir / "assets").glob("*.js"))
    reference_count = 0
    environment_assignments = 0
    patched_assignments = 0

    def _inject_flag(match):
        nonlocal environment_assignments, patched_assignments
        environment_assignments += 1
        properties = match.group(2)
        if "VITE_NO_STATE:" in properties:
            return match.group(0)

        patched_assignments += 1
        return (
            match.group(1)
            + "VITE_NO_STATE:!0,"
            + properties
            + match.group(3)
        )

    for path in javascript_files:
        content = path.read_text()
        reference_count += content.count(".VITE_NO_STATE")
        patched = ENV_ASSIGNMENT_PATTERN.sub(_inject_flag, content)
        if patched != content:
            path.write_text(patched)
            logger.info("Enabled stateless mode in %s", path.name)

    if reference_count == 0:
        raise RuntimeError(
            "The packaged frontend contains no VITE_NO_STATE references"
        )

    if environment_assignments == 0:
        raise RuntimeError("No compiled Vite environment objects were found")

    logger.info(
        "Stateless patch complete: references=%d environments=%d patched=%d",
        reference_count,
        environment_assignments,
        patched_assignments,
    )


def _add_session_heartbeat(static_dir, dataset_names):
    """Adds a heartbeat that keeps the browser's dataset clone alive."""
    index_path = static_dir / "index.html"
    content = index_path.read_text()
    if HEARTBEAT_MARKER in content:
        logger.info("Browser session heartbeat is already installed")
        return

    session_script = _get_session_script(dataset_names)
    if "</head>" not in content:
        raise RuntimeError(f"Could not patch missing </head> in {index_path}")

    index_path.write_text(
        content.replace("</head>", session_script + "</head>", 1)
    )
    logger.info(
        "Installed browser heartbeat with interval=%dms",
        HEARTBEAT_INTERVAL_MS,
    )


def _get_session_script(dataset_names):
    """Returns the inline browser identity and heartbeat script."""
    base_paths = json.dumps(
        [f"/datasets/{name}" for name in dataset_names],
        separators=(",", ":"),
    )
    return (
        f'<script id="{HEARTBEAT_MARKER}">'
        f"if({base_paths}."
        'includes(location.pathname))location.replace("/");'
        'const foSessionKey="fiftyone_demo_session";'
        'const foFetch=window.fetch.bind(window);'
        'window.fetch=(input,init={})=>{'
        'const token=localStorage.getItem(foSessionKey);'
        'const target=typeof input==="string"?input:'
        'input instanceof Request?input.url:String(input);'
        'const url=new URL(target,location.href);'
        'if(token&&url.origin===location.origin){'
        'const sourceHeaders=init.headers||'
        '(input instanceof Request?input.headers:void 0);'
        'const headers=new Headers(sourceHeaders);'
        'headers.set("X-FiftyOne-Session",token);'
        'init={...init,headers};'
        '}return foFetch(input,init);'
        '};'
        'setInterval(()=>{'
        'const dataset=location.pathname.split("/").pop();'
        'fetch("/__session/heartbeat?dataset="+encodeURIComponent(dataset),'
        '{method:"POST",credentials:"same-origin"}).then(response=>{'
        'if(response.status===401){'
        'localStorage.removeItem(foSessionKey);'
        'location.replace("/");'
        '}'
        '});'
        '},'
        f"{HEARTBEAT_INTERVAL_MS});"
        "</script>"
    )


if __name__ == "__main__":
    main()