hf_start.py

#!/usr/bin/env python3
"""Start Potato in a Hugging Face Spaces friendly way."""

from __future__ import annotations

import os
import sys
from pathlib import Path


def running_on_hugging_face() -> bool:
    """Return true inside HF Spaces or when explicitly requested."""
    return bool(
        os.environ.get("SPACE_ID")
        or os.environ.get("SPACE_REPO_NAME")
        or os.environ.get("POTATO_ENABLE_IFRAME_COOKIES") == "1"
    )


def configure_iframe_cookies(app) -> None:
    """Allow Flask session cookies to survive inside the Hugging Face iframe."""
    app.config.update(
        SESSION_COOKIE_SAMESITE="None",
        SESSION_COOKIE_SECURE=True,
        SESSION_COOKIE_HTTPONLY=True,
    )


def patch_login_template_for_iframe() -> None:
    """Patch Potato's login form so HF iframe login navigates top-level."""
    import potato

    template_path = Path(potato.__file__).resolve().parent / "templates" / "home.html"
    template = template_path.read_text(encoding="utf-8")
    marker = "<!-- HF_DIRECT_LOGIN_PATCH -->"
    if marker in template:
        return

    patch = f"""
    {marker}
    <script>
        (function() {{
            function isInIframe() {{
                try {{
                    return window.self !== window.top;
                }} catch (error) {{
                    return true;
                }}
            }}

            document.addEventListener('DOMContentLoaded', function() {{
                if (!isInIframe()) return;
                var form = document.querySelector('form[action="/auth"]');
                var usernameInput = document.querySelector('#login-email');
                if (!form || !usernameInput) return;

                form.setAttribute('method', 'GET');
                form.setAttribute('action', '/hf-direct-login');
                form.setAttribute('target', '_top');
                usernameInput.setAttribute('name', 'username');

                var note = document.createElement('div');
                note.className = 'potato-alert';
                note.style.marginBottom = '0.75rem';
                note.style.fontSize = '0.875rem';
                note.textContent = 'Login will open the app directly so your browser can keep the session.';
                form.parentNode.insertBefore(note, form);
            }});
        }}());
    </script>
"""
    template = template.replace("</body>", patch + "\n</body>")
    template_path.write_text(template, encoding="utf-8")


def register_direct_login_route(app) -> None:
    """Register a cookie-friendly top-level login route for HF Spaces."""
    from flask import redirect, request, session, url_for
    from potato.authentication import UserAuthenticator
    from potato.flask_server import (
        UserPhase,
        get_item_state_manager,
        get_user_state_manager,
        init_user_state,
    )

    def hf_direct_login():
        username = (request.args.get("username") or "").strip()
        if not username:
            return redirect(url_for("home"))

        authenticator = UserAuthenticator.get_instance()
        if not authenticator.is_valid_username(username):
            authenticator.add_user(username, None)

        session.clear()
        session["username"] = username
        session.permanent = True

        user_state_manager = get_user_state_manager()
        if not user_state_manager.has_user(username):
            init_user_state(username)

        user_state = user_state_manager.get_user_state(username)
        if user_state and user_state.get_phase() == UserPhase.LOGIN:
            user_state_manager.advance_phase(username)
        if user_state and not user_state.has_assignments():
            get_item_state_manager().assign_instances_to_user(user_state)

        return redirect(url_for("annotate"))

    app.add_url_rule("/hf-direct-login", "hf_direct_login", hf_direct_login, methods=["GET"])


def main() -> int:
    port = os.environ.get("PORT", "7860")
    config = os.environ.get("POTATO_CONFIG", "config.yaml")

    if running_on_hugging_face():
        patch_login_template_for_iframe()

    from potato.flask_server import create_app

    app = create_app(config)
    register_direct_login_route(app)
    if running_on_hugging_face():
        configure_iframe_cookies(app)
        print("Enabled Hugging Face iframe-compatible session cookies.", flush=True)

    print(f"Starting Potato on 0.0.0.0:{port}", flush=True)
    app.run(host="0.0.0.0", port=int(port), debug=False, use_reloader=False, threaded=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())