aops02/math-annotation-demo

1#!/usr/bin/env python32"""Start Potato in a Hugging Face Spaces friendly way."""3 4from __future__ import annotations5 6import os7import sys8from pathlib import Path9 10 11def running_on_hugging_face() -> bool:12    """Return true inside HF Spaces or when explicitly requested."""13    return bool(14        os.environ.get("SPACE_ID")15        or os.environ.get("SPACE_REPO_NAME")16        or os.environ.get("POTATO_ENABLE_IFRAME_COOKIES") == "1"17    )18 19 20def configure_iframe_cookies(app) -> None:21    """Allow Flask session cookies to survive inside the Hugging Face iframe."""22    app.config.update(23        SESSION_COOKIE_SAMESITE="None",24        SESSION_COOKIE_SECURE=True,25        SESSION_COOKIE_HTTPONLY=True,26    )27 28 29def patch_login_template_for_iframe() -> None:30    """Patch Potato's login form so HF iframe login navigates top-level."""31    import potato32 33    template_path = Path(potato.__file__).resolve().parent / "templates" / "home.html"34    template = template_path.read_text(encoding="utf-8")35    marker = "<!-- HF_DIRECT_LOGIN_PATCH -->"36    if marker in template:37        return38 39    patch = f"""40    {marker}41    <script>42        (function() {{43            function isInIframe() {{44                try {{45                    return window.self !== window.top;46                }} catch (error) {{47                    return true;48                }}49            }}50 51            document.addEventListener('DOMContentLoaded', function() {{52                if (!isInIframe()) return;53                var form = document.querySelector('form[action="/auth"]');54                var usernameInput = document.querySelector('#login-email');55                if (!form || !usernameInput) return;56 57                form.setAttribute('method', 'GET');58                form.setAttribute('action', '/hf-direct-login');59                form.setAttribute('target', '_top');60                usernameInput.setAttribute('name', 'username');61 62                var note = document.createElement('div');63                note.className = 'potato-alert';64                note.style.marginBottom = '0.75rem';65                note.style.fontSize = '0.875rem';66                note.textContent = 'Login will open the app directly so your browser can keep the session.';67                form.parentNode.insertBefore(note, form);68            }});69        }}());70    </script>71"""72    template = template.replace("</body>", patch + "\n</body>")73    template_path.write_text(template, encoding="utf-8")74 75 76def register_direct_login_route(app) -> None:77    """Register a cookie-friendly top-level login route for HF Spaces."""78    from flask import redirect, request, session, url_for79    from potato.authentication import UserAuthenticator80    from potato.flask_server import (81        UserPhase,82        get_item_state_manager,83        get_user_state_manager,84        init_user_state,85    )86 87    def hf_direct_login():88        username = (request.args.get("username") or "").strip()89        if not username:90            return redirect(url_for("home"))91 92        authenticator = UserAuthenticator.get_instance()93        if not authenticator.is_valid_username(username):94            authenticator.add_user(username, None)95 96        session.clear()97        session["username"] = username98        session.permanent = True99 100        user_state_manager = get_user_state_manager()101        if not user_state_manager.has_user(username):102            init_user_state(username)103 104        user_state = user_state_manager.get_user_state(username)105        if user_state and user_state.get_phase() == UserPhase.LOGIN:106            user_state_manager.advance_phase(username)107        if user_state and not user_state.has_assignments():108            get_item_state_manager().assign_instances_to_user(user_state)109 110        return redirect(url_for("annotate"))111 112    app.add_url_rule("/hf-direct-login", "hf_direct_login", hf_direct_login, methods=["GET"])113 114 115def main() -> int:116    port = os.environ.get("PORT", "7860")117    config = os.environ.get("POTATO_CONFIG", "config.yaml")118 119    if running_on_hugging_face():120        patch_login_template_for_iframe()121 122    from potato.flask_server import create_app123 124    app = create_app(config)125    register_direct_login_route(app)126    if running_on_hugging_face():127        configure_iframe_cookies(app)128        print("Enabled Hugging Face iframe-compatible session cookies.", flush=True)129 130    print(f"Starting Potato on 0.0.0.0:{port}", flush=True)131    app.run(host="0.0.0.0", port=int(port), debug=False, use_reloader=False, threaded=True)132    return 0133 134 135if __name__ == "__main__":136    raise SystemExit(main())137