Add top-level direct login fallback for HF iframe

hf_start.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import os
 import sys
+from pathlib import Path
 
 
 def running_on_hugging_face() -> bool:
@@ -25,13 +26,103 @@ def configure_iframe_cookies(app) -> None:
     )
 
 
+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)