File size: 5,473 Bytes
087ddd7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
"""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()