Merge remote-tracking branch 'hf/main'

pyproject.toml

@@ -42,6 +42,16 @@ disable = [
   "unspecified-encoding",
 ]
 
+[tool.uv]
+# Override PyTorch version constraints from mussel[torch-gpu].
+# mussel pins torch<2.6 which may be too strict; we relax to >=2.0.0.
+# The override applies on ALL platforms (not just macOS) so Docker builds
+# still get torch installed.
+override-dependencies = [
+  "torch>=2.0.0",
+  "torchvision>=0.15.0",
+]
+
 [tool.uv.sources]
 # For local dev with SSH: uv pip install -e .
 # For Docker builds with token: GH_TOKEN=<token> uv pip install -e .

src/mosaic/analysis.py

@@ -772,6 +772,7 @@ def analyze_slide(
     analysis_id=None,
     session_hash=None,
     profile=None,
+    session_id: str = "",
 ):
     """Analyze a whole slide image for cancer subtype and biomarker prediction.
 
@@ -810,7 +811,7 @@ def analyze_slide(
         raise gr.Error("Please upload a slide.")
 
     # Extract user info for telemetry (HF Spaces only)
-    user_info = extract_user_info(request, IS_HF_SPACES, profile)
+    user_info = extract_user_info(request, IS_HF_SPACES, profile, session_id)
 
     # Initialize telemetry for resource tracking
     slide_start_time = time_module.time()

src/mosaic/telemetry/__init__.py

@@ -47,7 +47,6 @@ from mosaic.telemetry.utils import (
     StageTimer,
     sanitize_error_message,
     hash_session_id,
-    hash_username,
     UserInfo,
     extract_user_info,
 )
@@ -68,7 +67,6 @@ __all__ = [
     "StageTimer",
     "sanitize_error_message",
     "hash_session_id",
-    "hash_username",
     "UserInfo",
     "extract_user_info",
 ]

src/mosaic/telemetry/tracker.py

@@ -22,7 +22,6 @@ from mosaic.telemetry.events import (
 from mosaic.telemetry.storage import TelemetryStorage
 from mosaic.telemetry.utils import (
     hash_session_id,
-    hash_username,
     sanitize_error_message,
 )
 
@@ -256,7 +255,7 @@ class TelemetryTracker:
             success=success,
             cached_slide_count=cached_slide_count,
             is_logged_in=is_logged_in,
-            hf_username=hash_username(hf_username),
+            hf_username=hf_username,
         )
         self.storage.write_usage_event(event)
 
@@ -332,7 +331,7 @@ class TelemetryTracker:
             gpu_type=gpu_type,
             peak_gpu_memory_gb=peak_gpu_memory_gb,
             is_logged_in=is_logged_in,
-            hf_username=hash_username(hf_username),
+            hf_username=hf_username,
         )
         self.storage.write_resource_event(event)
 
@@ -377,7 +376,7 @@ class TelemetryTracker:
             slide_count=slide_count,
             gpu_type=gpu_type,
             is_logged_in=is_logged_in,
-            hf_username=hash_username(hf_username),
+            hf_username=hf_username,
         )
         self.storage.write_failure_event(event)
 

src/mosaic/telemetry/utils.py

@@ -5,7 +5,7 @@ This module provides helper utilities:
 - sanitize_error_message: Remove sensitive data from error messages
 - hash_session_id: Hash session IDs for privacy
 - UserInfo: Dataclass for user information from HF Spaces
-- extract_user_info: Extract user info from Gradio request object
+- extract_user_info: Extract user info from Gradio request/OAuth profile
 """
 
 import hashlib
@@ -99,25 +99,6 @@ def hash_session_id(session_id: Optional[str]) -> Optional[str]:
     return hashlib.sha256(salted.encode()).hexdigest()[:16]
 
 
-def hash_username(username: Optional[str]) -> Optional[str]:
-    """Hash a username for privacy in telemetry.
-
-    Uses SHA-256 with a different salt than session IDs to create a one-way hash.
-    This allows distinguishing users in telemetry without storing actual usernames.
-
-    Args:
-        username: HuggingFace username (can be None for anonymous users)
-
-    Returns:
-        Hashed username or None if input is None
-    """
-    if username is None:
-        return None
-
-    salted = f"mosaic_user:{username}"
-    return hashlib.sha256(salted.encode()).hexdigest()[:16]
-
-
 @dataclass
 class UserInfo:
     """User information extracted from HF Spaces request.
@@ -131,7 +112,12 @@ class UserInfo:
     username: Optional[str] = None
 
 
-def extract_user_info(request, is_hf_spaces: bool = False, profile=None) -> UserInfo:
+def extract_user_info(
+    request,
+    is_hf_spaces: bool = False,
+    profile=None,
+    session_id: Optional[str] = None,
+) -> UserInfo:
     """Extract user info from Gradio OAuth profile.
 
     With gr.LoginButton(), Gradio provides user info via gr.OAuthProfile
@@ -145,21 +131,51 @@ def extract_user_info(request, is_hf_spaces: bool = False, profile=None) -> User
         request: Gradio request object (unused, kept for API compatibility)
         is_hf_spaces: Whether running on HuggingFace Spaces (only extract on HF)
         profile: Gradio OAuthProfile object (injected by LoginButton OAuth flow)
+        session_id: Session ID from localStorage (primary lookup on HF Spaces
+            where the HF proxy strips custom cookies)
 
     Returns:
         UserInfo with is_logged_in and username (or defaults if not available)
 
     Example:
-        user_info = extract_user_info(request, IS_HF_SPACES, profile)
+        user_info = extract_user_info(request, IS_HF_SPACES, profile, session_id)
         if user_info.is_logged_in:
             print(f"User: {user_info.username}")
     """
     if not is_hf_spaces:
         return UserInfo()
 
-    # On HF Spaces, only OAuthProfile reliably identifies the logged-in user.
-    # request.username returns the Space owner's username (not the visitor's),
-    # so we must NOT fall back to it.
+    # On HF Spaces with Docker SDK, the HF proxy strips custom cookies.
+    # The most reliable method is looking up the session ID that was stored
+    # in localStorage by the OAuth callback page.
+    if session_id:
+        try:
+            from mosaic.ui.oauth import get_user_by_session_id
+
+            userinfo = get_user_by_session_id(session_id)
+            if userinfo:
+                username = userinfo.get("preferred_username")
+                if username:
+                    logger.info(f"Extracted user from session ID: {username}")
+                    return UserInfo(is_logged_in=True, username=username)
+        except Exception as e:
+            logger.debug(f"Session ID lookup failed: {e}")
+
+    # Fallback: server-side session via mosaic_auth cookie
+    if request is not None:
+        try:
+            from mosaic.ui.oauth import get_user_from_server_session
+
+            userinfo = get_user_from_server_session(request)
+            if userinfo:
+                username = userinfo.get("preferred_username")
+                if username:
+                    logger.info(f"Extracted user from server-side session: {username}")
+                    return UserInfo(is_logged_in=True, username=username)
+        except Exception as e:
+            logger.debug(f"Server-side session lookup failed: {e}")
+
+    # Fallback: OAuthProfile (may be unreliable on Docker SDK Spaces)
     if profile is not None:
         try:
             username = getattr(profile, "username", None)
@@ -169,5 +185,5 @@ def extract_user_info(request, is_hf_spaces: bool = False, profile=None) -> User
         except Exception as e:
             logger.debug(f"Could not extract username from OAuthProfile: {e}")
 
-    logger.debug("User not logged in: no OAuthProfile available")
+    logger.debug("User not logged in: no session or OAuthProfile available")
     return UserInfo()

src/mosaic/ui/app.py

@@ -190,6 +190,7 @@ def analyze_slides(
     seg_config,
     user_dir,
     slide_ids=None,  # Mapping from filename to slide_id
+    session_id: str = "",
     progress=gr.Progress(track_tqdm=True),
     request: gr.Request = None,
     profile: Optional[gr.OAuthProfile] = None,
@@ -228,7 +229,7 @@ def analyze_slides(
 
     # Extract user info for telemetry (HF Spaces only)
     # gr.LoginButton() OAuth provides user info via OAuthProfile
-    user_info = extract_user_info(request, IS_HF_SPACES, profile)
+    user_info = extract_user_info(request, IS_HF_SPACES, profile, session_id)
 
     # Wait for core models download to complete (Paladin models can continue in background)
     if _model_download_thread is not None and not _core_models_complete:
@@ -472,6 +473,7 @@ def analyze_slides(
                     analysis_id=analysis_id,  # For telemetry tracking
                     session_hash=session_hash,  # For telemetry tracking
                     profile=profile,  # For OAuth user info
+                    session_id=session_id,  # For OAuth session lookup
                 )
 
                 # Save results to cache if this is a TCGA slide
@@ -759,8 +761,18 @@ def launch_gradio(server_name, server_port, share):
             # stays permanently hidden.
             login_button.load_event_to_attach = None
 
-            # Custom login display (OAuthProfile-based, not request.username)
-            login_status_md = gr.Markdown(value="", visible=False)
+            # Custom login display — sign-in via markdown link,
+            # sign-out via a Gradio button (stays in Gradio event system,
+            # avoiding navigation away from the SPA).
+            with gr.Row():
+                login_status_md = gr.Markdown(
+                    value="[Sign in with HuggingFace](/api/auth/login) to save slides and results",
+                    visible=True,
+                )
+                logout_button = gr.Button(
+                    "Sign out", visible=False, size="sm", scale=0,
+                    min_width=80,
+                )
 
         gr.Markdown(
             "# Mosaic: H&E Whole Slide Image Cancer Subtype and Biomarker Inference"
@@ -1211,9 +1223,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
                     gr.Group(visible=False),
                     gr.Dataframe(visible=False),
                     gr.File(visible=False),
-                    gr.Row(
-                        visible=not IS_HF_SPACES
-                    ),  # Show existing files row in upload mode
+                    gr.Row(visible=True),  # Show existing files row in upload mode
                 )
             else:
                 return (
@@ -1226,13 +1236,14 @@ This tool is for research purposes only and not approved for clinical diagnosis.
 
         # Load storage usage and show warnings
         def load_storage_usage(
+            session_id: str = "",
             request: gr.Request = None,
             profile: Optional[gr.OAuthProfile] = None,
         ):
             """Load and display storage usage with warnings."""
             from mosaic.ui.user_tabs import _get_username
 
-            username, is_local = _get_username(request, profile)
+            username, is_local = _get_username(request, profile, session_id)
 
             if not username:
                 # Not logged in on HF Spaces - hide storage info
@@ -1282,13 +1293,14 @@ This tool is for research purposes only and not approved for clinical diagnosis.
 
         # Load list of existing slides for re-analysis
         def load_existing_slides(
+            session_id: str = "",
             request: gr.Request = None,
             profile: Optional[gr.OAuthProfile] = None,
         ):
             """Load dropdown with user's existing slides."""
             from mosaic.ui.user_tabs import _get_username
 
-            username, _ = _get_username(request, profile)
+            username, _ = _get_username(request, profile, session_id)
 
             if not username:
                 return gr.Dropdown(choices=[], value=None)
@@ -1326,6 +1338,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
         def handle_slide_upload(
             files,
             slide_ids,
+            session_id: str = "",
             request: gr.Request = None,
             profile: Optional[gr.OAuthProfile] = None,
         ):
@@ -1336,7 +1349,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
             # Determine username via OAuthProfile (primary) or request (fallback)
             from mosaic.ui.user_tabs import _get_username
 
-            username, _ = _get_username(request, profile)
+            username, _ = _get_username(request, profile, session_id)
             if not username:
                 return files, slide_ids or {}
 
@@ -1500,34 +1513,47 @@ This tool is for research purposes only and not approved for clinical diagnosis.
             ],
         )
 
+        # Hidden textbox that carries the localStorage session ID into
+        # every Gradio event handler that needs the visitor's identity.
+        # On HF Spaces the js= parameter on demo.load() fills this from
+        # localStorage; locally it stays empty and _get_username() falls
+        # back to LOCAL_DEBUG_USERNAME.
+        session_id_input = gr.Textbox(visible=False, elem_id="session_id_input")
+        # Hidden flag set to "1" only by the demo.load js= parameter.
+        # On WebSocket reconnections JS doesn't re-run, so this stays
+        # at its default "".  The Python handler uses it to detect
+        # reconnections and return gr.skip() to preserve UI state.
+        _js_ran_flag = gr.Textbox(visible=False, value="")
+
         # Handle file uploads - regenerate entire settings table
         # Save uploaded slides to user storage (HF Spaces only)
         input_slides.upload(
             handle_slide_upload,
-            inputs=[input_slides, slide_ids_state],
+            inputs=[input_slides, slide_ids_state, session_id_input],
             outputs=[input_slides, slide_ids_state],
         ).then(
             # Refresh storage usage after upload
             load_storage_usage,
-            inputs=None,
+            inputs=[session_id_input],
             outputs=[storage_usage_warning, existing_file_row],
         ).then(
             # Refresh existing slides list after upload
             load_existing_slides,
-            inputs=None,
+            inputs=[session_id_input],
             outputs=[existing_slides_dropdown],
         )
 
         # Refresh existing slides dropdown
         refresh_existing_btn.click(
             load_existing_slides,
-            inputs=None,
+            inputs=[session_id_input],
             outputs=[existing_slides_dropdown],
         )
 
         # Handle selection of existing slide(s) for re-analysis
         def select_existing_slide(
             slide_ids,
+            session_id: str = "",
             request: gr.Request = None,
             profile: Optional[gr.OAuthProfile] = None,
         ):
@@ -1535,6 +1561,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
 
             Args:
                 slide_ids: Single slide_id (str) or list of slide_ids (multiselect)
+                session_id: Server-side session ID from localStorage
                 request: Gradio request object
                 profile: OAuthProfile for HF Spaces login
 
@@ -1550,7 +1577,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
 
             from mosaic.ui.user_tabs import _get_username
 
-            username, _ = _get_username(request, profile)
+            username, _ = _get_username(request, profile, session_id)
             if not username:
                 return None
 
@@ -1593,7 +1620,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
 
         existing_slides_dropdown.change(
             select_existing_slide,
-            inputs=[existing_slides_dropdown],
+            inputs=[existing_slides_dropdown, session_id_input],
             outputs=[input_slides],
         )
 
@@ -1732,6 +1759,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
             seg_config,
             user_dir,
             slide_ids,
+            session_id,
             progress=gr.Progress(track_tqdm=True),
             request: gr.Request = None,
             profile: Optional[gr.OAuthProfile] = None,
@@ -1777,6 +1805,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
                 seg_config,
                 user_dir,
                 slide_ids=slide_ids,
+                session_id=session_id,
                 progress=progress,
                 request=request,
                 profile=profile,
@@ -1797,6 +1826,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
                 seg_config_dropdown,
                 user_dir_state,
                 slide_ids_state,
+                session_id_input,
             ],
             outputs=[
                 settings_input,
@@ -1863,20 +1893,6 @@ This tool is for research purposes only and not approved for clinical diagnosis.
             outputs=[user_dir_state],
         )
 
-        # Load storage usage on page load
-        demo.load(
-            load_storage_usage,
-            inputs=None,
-            outputs=[storage_usage_warning, existing_file_row],
-        )
-
-        # Load existing slides dropdown on page load
-        demo.load(
-            load_existing_slides,
-            inputs=None,
-            outputs=[existing_slides_dropdown],
-        )
-
         # Add My Files and My Results tabs for user storage
         # On HF Spaces: visible for logged-in users
         # Locally: always visible with "local_user" for debugging
@@ -1901,81 +1917,169 @@ This tool is for research purposes only and not approved for clinical diagnosis.
             with gr.Tab("My Files"):
                 from mosaic.ui.user_tabs import create_my_files_tab
 
-                my_files_components = create_my_files_tab()
-
-                # Load files on demo load
-                demo.load(
-                    my_files_components["load_files"],
-                    inputs=None,
-                    outputs=[
-                        my_files_components["storage_usage"],
-                        my_files_components["files_table"],
-                        my_files_components["file_action_status"],
-                        my_files_components["slide_download_file"],
-                    ],
-                )
+                my_files_components = create_my_files_tab(session_id_input)
 
             with gr.Tab("My Results"):
                 from mosaic.ui.user_tabs import create_my_results_tab
 
-                my_results_components = create_my_results_tab()
+                my_results_components = create_my_results_tab(session_id_input)
 
-                # Load results on demo load
-                demo.load(
-                    my_results_components["load_results"],
-                    inputs=None,
-                    outputs=[
-                        my_results_components["results_table"],
-                        my_results_components["result_action_status"],
-                        my_results_components["result_download_file"],
-                    ],
+        # On HF Spaces: demo.load(js=...) reads the session ID from
+        # localStorage AND sets a "JS ran" flag.  Gradio re-fires
+        # demo.load() on WebSocket reconnections WITHOUT re-executing
+        # the JS — so the flag stays at its default "".  The Python
+        # handler detects this and returns gr.skip() to preserve the
+        # UI state instead of overwriting the logged-in display.
+        if IS_HF_SPACES:
+            _load_files = my_files_components["load_files"]
+            _load_results = my_results_components["load_results"]
+
+            _LS_JS = "() => [localStorage.getItem('mosaic_session_id') || '', '1']"
+
+            def on_page_load(session_id_from_browser, js_ran):
+                """Handle page load: check login, load files, load results.
+
+                On WebSocket reconnections, Gradio re-fires demo.load()
+                but does NOT re-execute the js= parameter.  The js_ran
+                flag stays at its default "" on reconnections, allowing
+                us to return gr.skip() and preserve the current UI state.
+                """
+                from mosaic.ui.oauth import get_user_by_session_id, _sessions
+
+                # Reconnection detection: JS didn't run → skip all updates
+                if js_ran != "1":
+                    logger.debug("on_page_load: reconnection detected, skipping")
+                    return gr.skip()
+
+                sid_repr = (
+                    session_id_from_browser[:8] + "..."
+                    if session_id_from_browser
+                    else "(empty)"
+                )
+                logger.info(
+                    f"on_page_load: sid={sid_repr}, "
+                    f"active_sessions={len(_sessions)}"
                 )
 
-        # Show/hide user storage features based on login state on page load.
-        # LoginButton.click() fires BEFORE the OAuth redirect so the profile
-        # is never available at that point.  Instead we rely on demo.load()
-        # which fires after the OAuth redirect completes and the page reloads.
-        #
-        # On HF Spaces with sdk:docker, LoginButton is always hidden because
-        # it reads request.username (which returns the Space owner for all
-        # visitors). Instead we use login_status_md with direct OAuth links.
-        if IS_HF_SPACES:
+                if not session_id_from_browser:
+                    return _build_logged_out_result()
 
-            def check_login_state(profile: Optional[gr.OAuthProfile] = None):
-                """Toggle login status display and user-storage UI based on OAuth state."""
-                if profile is not None:
-                    username = getattr(profile, "username", None) or getattr(
-                        profile, "name", "User"
-                    )
-                    return (
+                userinfo = get_user_by_session_id(session_id_from_browser)
+
+                if userinfo is not None:
+                    username = userinfo.get("preferred_username", "User")
+                    login_out = (
                         gr.update(
-                            value=f"Signed in as **{username}** \u00b7 [Sign out](/logout)",
+                            value=f"Signed in as **{username}**",
                             visible=True,
-                        ),  # login_status_md
-                        gr.update(visible=True),  # user_storage_tabs
-                        gr.update(visible=True),  # user_storage_header
-                        gr.update(visible=True),  # user_storage_info
+                        ),
+                        gr.update(visible=True),   # logout_button
+                        gr.update(visible=True),   # user_storage_tabs
+                        gr.update(visible=True),   # user_storage_header
+                        gr.update(visible=True),   # user_storage_info
                     )
                 else:
-                    return (
-                        gr.update(
-                            value="[Sign in with HuggingFace](/login/huggingface) to save slides and results",
-                            visible=True,
-                        ),  # login_status_md
-                        gr.update(visible=False),  # user_storage_tabs
-                        gr.update(visible=False),  # user_storage_header
-                        gr.update(visible=False),  # user_storage_info
+                    return _build_logged_out_result()
+
+                # Load storage usage + existing slides for the analysis page
+                storage_out = load_storage_usage(session_id=session_id_from_browser)
+                existing_out = load_existing_slides(session_id=session_id_from_browser)
+
+                # Load user files and results tabs
+                try:
+                    files_out = _load_files(session_id=session_id_from_browser)
+                except Exception as e:
+                    logger.error(f"on_page_load: _load_files error: {e}")
+                    files_out = (
+                        f"**Storage:** Error: {e}",
+                        [],
+                        gr.Textbox(visible=False),
+                        gr.File(visible=False),
                     )
+                try:
+                    results_out = _load_results(
+                        session_id=session_id_from_browser
+                    )
+                except Exception as e:
+                    logger.error(f"on_page_load: _load_results error: {e}")
+                    results_out = (
+                        [],
+                        gr.Textbox(visible=False),
+                        gr.File(visible=False),
+                    )
+
+                # login(5) + session_id(1) + storage(2) + existing(1)
+                # + files(4) + results(3) = 16
+                return (
+                    *login_out,
+                    session_id_from_browser,
+                    *storage_out,
+                    existing_out,
+                    *files_out,
+                    *results_out,
+                )
+
+            def _build_logged_out_result():
+                """Build the full logged-out output tuple (16 elements)."""
+                return (
+                    gr.update(
+                        value="[Sign in with HuggingFace](/api/auth/login) to save slides and results",
+                        visible=True,
+                    ),
+                    gr.update(visible=False),  # logout_button
+                    gr.update(visible=False),  # user_storage_tabs
+                    gr.update(visible=False),  # user_storage_header
+                    gr.update(visible=False),  # user_storage_info
+                    "",  # session_id_input
+                    gr.Markdown(visible=False),  # storage_usage_warning
+                    gr.Row(visible=False),  # existing_file_row
+                    gr.Dropdown(choices=[], value=None),  # existing_slides_dropdown
+                    "**Storage:** Not logged in",
+                    [],  # files_table
+                    gr.Textbox(visible=False),
+                    gr.File(visible=False),
+                    [],  # results_table
+                    gr.Textbox(visible=False),
+                    gr.File(visible=False),
+                )
 
             demo.load(
-                check_login_state,
-                inputs=None,
+                on_page_load,
+                inputs=[session_id_input, _js_ran_flag],
                 outputs=[
+                    # login state (5)
                     login_status_md,
+                    logout_button,
                     user_storage_tabs,
                     user_storage_header,
                     user_storage_info,
+                    # persist session_id for subsequent button clicks (1)
+                    session_id_input,
+                    # analysis page: storage + existing slides (3)
+                    storage_usage_warning,
+                    existing_file_row,
+                    existing_slides_dropdown,
+                    # files tab (4)
+                    my_files_components["storage_usage"],
+                    my_files_components["files_table"],
+                    my_files_components["file_action_status"],
+                    my_files_components["slide_download_file"],
+                    # results tab (3)
+                    my_results_components["results_table"],
+                    my_results_components["result_action_status"],
+                    my_results_components["result_download_file"],
                 ],
+                js=_LS_JS,
+            )
+
+            # Sign-out: clear localStorage and reload.
+            # On reload, demo.load fires; JS reads empty localStorage
+            # and returns ['', '1']; on_page_load sees empty session_id
+            # and returns the logged-out state.
+            logout_button.click(
+                fn=None,
+                inputs=None,
+                js="() => { localStorage.removeItem('mosaic_session_id'); location.reload(); }",
             )
 
     # Use hardware-specific concurrency limit
@@ -1983,6 +2087,13 @@ This tool is for research purposes only and not approved for clinical diagnosis.
     # Higher-memory GPUs and ZeroGPU can handle multiple concurrent analyses
     demo.queue(max_size=10, default_concurrency_limit=DEFAULT_CONCURRENCY_LIMIT)
 
+    # OAuth routes for HF Spaces (sdk:docker):
+    # Custom OAuth routes (/api/auth/login, /api/auth/callback, /api/auth/logout)
+    # are mounted via the lifespan protocol in create_oauth_lifespan(). This
+    # ensures routes survive launch() recreating the app. Identity is read from
+    # our server-side session (mosaic_auth cookie), not gr.OAuthProfile, which
+    # returns the Space owner's identity on Docker SDK Spaces.
+
     # Register cleanup handler for graceful shutdown
     import atexit
 
@@ -2014,10 +2125,20 @@ This tool is for research purposes only and not approved for clinical diagnosis.
     heartbeat_thread = threading.Thread(target=heartbeat_worker, daemon=True)
     heartbeat_thread.start()
 
+    # On HF Spaces, mount custom OAuth routes via the lifespan protocol.
+    # Gradio's launch() recreates the app, but the lifespan handler receives
+    # the final App instance, so routes added there survive.
+    launch_kwargs = {}
+    if IS_HF_SPACES:
+        from mosaic.ui.oauth import create_oauth_lifespan
+
+        launch_kwargs["app_kwargs"] = {"lifespan": create_oauth_lifespan}
+
     demo.launch(
         server_name=server_name,
         share=share,
         server_port=server_port,
         show_error=True,
         favicon_path=current_dir / "favicon.svg",
+        **launch_kwargs,
     )

src/mosaic/ui/oauth.py

@@ -0,0 +1,408 @@
+"""Manual OAuth flow for HF Spaces with Docker SDK.
+
+On HF Spaces with sdk:docker, Gradio's built-in OAuth (gr.LoginButton /
+gr.OAuthProfile) doesn't work because the HF reverse proxy injects the
+Space owner's identity into every session.  This module implements the
+Authorization Code flow directly against HF's OAuth provider.
+
+The HF proxy strips all custom cookies — only its own forged ``session``
+cookie survives.  To persist the visitor's session ID across page loads,
+the callback stores it in the browser's ``localStorage`` via an
+intermediate HTML page.  The Gradio ``demo.load()`` handler reads it back
+via a ``js=`` parameter and passes it to the Python handler, which looks
+up the server-side session.
+
+Environment variables (set automatically by HF Spaces):
+    OAUTH_CLIENT_ID:     OAuth application client ID
+    OAUTH_CLIENT_SECRET: OAuth application client secret
+    OAUTH_SCOPES:        Space-separated scopes (default: "openid profile")
+    SPACE_HOST:           Public hostname of the Space (e.g. "user-space.hf.space")
+
+Routes mounted on the Gradio ASGI app:
+    GET /api/auth/login    -> Redirect to HF authorize endpoint
+    GET /api/auth/callback -> Exchange code for token, store session via localStorage
+    GET /api/auth/logout   -> Clear session, clear localStorage, redirect to /
+"""
+
+import contextlib
+import html as html_module
+import json
+import os
+import secrets
+import time
+from typing import Optional
+from urllib.parse import urlencode
+
+from loguru import logger
+from starlette.requests import Request
+from starlette.responses import RedirectResponse, JSONResponse, HTMLResponse
+from starlette.routing import Route
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+OAUTH_CLIENT_ID = os.environ.get("OAUTH_CLIENT_ID", "")
+OAUTH_CLIENT_SECRET = os.environ.get("OAUTH_CLIENT_SECRET", "")
+OAUTH_SCOPES = os.environ.get("OAUTH_SCOPES", "openid profile")
+SPACE_HOST = os.environ.get("SPACE_HOST", "")
+
+HF_AUTHORIZE_URL = "https://huggingface.co/oauth/authorize"
+HF_TOKEN_URL = "https://huggingface.co/oauth/token"
+
+# ---------------------------------------------------------------------------
+# Server-side session store
+# ---------------------------------------------------------------------------
+
+# {session_id: {"userinfo": {...}, "created_at": float}}
+_sessions: dict[str, dict] = {}
+_SESSION_TTL_SEC = 24 * 60 * 60  # 24 hours
+_SESSION_COOKIE = "mosaic_auth"  # kept for backward-compat with tests
+
+# localStorage key used by the browser to persist the session ID
+_LOCALSTORAGE_KEY = "mosaic_session_id"
+
+# In-memory CSRF state tokens: {state_value: created_at_float}
+_pending_states: dict[str, float] = {}
+_STATE_TTL_SEC = 10 * 60  # 10 minutes
+
+
+def _prune_expired() -> None:
+    """Remove expired sessions and CSRF state tokens."""
+    now = time.time()
+    expired_sessions = [
+        k for k, v in _sessions.items() if now - v["created_at"] > _SESSION_TTL_SEC
+    ]
+    for k in expired_sessions:
+        del _sessions[k]
+    expired_states = [k for k, v in _pending_states.items() if now - v > _STATE_TTL_SEC]
+    for k in expired_states:
+        del _pending_states[k]
+
+
+def get_user_from_server_session(request) -> Optional[dict]:
+    """Look up user info from the server-side session store.
+
+    Checks the ``mosaic_auth`` cookie in the request and returns the
+    stored userinfo dict, or None if the cookie is missing/expired.
+
+    Args:
+        request: Starlette/Gradio request object (needs .cookies or .headers)
+
+    Returns:
+        userinfo dict with at least ``preferred_username`` key, or None
+    """
+    _prune_expired()
+
+    cookie_val = None
+    # Gradio's gr.Request wraps cookies/headers through dict_to_obj() which
+    # returns a gradio.route_utils.Obj, NOT a plain dict.  Use hasattr/get
+    # instead of isinstance(dict) to support both dict and Obj.
+    if hasattr(request, "cookies"):
+        cookies = request.cookies
+        if hasattr(cookies, "get"):
+            cookie_val = cookies.get(_SESSION_COOKIE)
+        logger.debug(
+            f"Session lookup: cookies type={type(cookies).__name__}, "
+            f"cookie_val={cookie_val!r}"
+        )
+    # Fallback: parse Cookie header
+    if cookie_val is None and hasattr(request, "headers"):
+        headers = request.headers
+        cookie_header = ""
+        if hasattr(headers, "get"):
+            cookie_header = headers.get("cookie", "") or ""
+        if cookie_header:
+            for part in cookie_header.split(";"):
+                part = part.strip()
+                if part.startswith(f"{_SESSION_COOKIE}="):
+                    cookie_val = part[len(f"{_SESSION_COOKIE}=") :]
+                    break
+        logger.debug(
+            f"Session lookup header fallback: has_cookie_header={bool(cookie_header)}, "
+            f"cookie_val={cookie_val!r}"
+        )
+
+    if not cookie_val:
+        logger.debug(
+            f"No {_SESSION_COOKIE} cookie found. "
+            f"Active sessions: {len(_sessions)}"
+        )
+        return None
+
+    entry = _sessions.get(cookie_val)
+    if entry is None:
+        logger.debug(
+            f"{_SESSION_COOKIE} cookie present ({cookie_val[:8]}...) "
+            f"but no matching server session. "
+            f"Active sessions: {list(_sessions.keys())[:3]}"
+        )
+        return None
+
+    if time.time() - entry["created_at"] > _SESSION_TTL_SEC:
+        del _sessions[cookie_val]
+        return None
+
+    return entry.get("userinfo")
+
+
+def get_user_by_session_id(session_id: Optional[str]) -> Optional[dict]:
+    """Look up user info by a raw session ID (from localStorage).
+
+    This is the primary lookup method on HF Spaces where the HF proxy
+    strips custom cookies.  The session ID is stored in the browser's
+    localStorage by the callback page and passed to Gradio handlers via
+    the ``js=`` parameter on ``demo.load()``.
+
+    Args:
+        session_id: The session ID string (from localStorage), or empty/None.
+
+    Returns:
+        userinfo dict with at least ``preferred_username`` key, or None.
+    """
+    _prune_expired()
+
+    if not session_id:
+        return None
+
+    entry = _sessions.get(session_id)
+    if entry is None:
+        logger.debug(
+            f"Session ID ({session_id[:8]}...) not found. "
+            f"Active sessions: {len(_sessions)}"
+        )
+        return None
+
+    if time.time() - entry["created_at"] > _SESSION_TTL_SEC:
+        del _sessions[session_id]
+        return None
+
+    return entry.get("userinfo")
+
+
+# ---------------------------------------------------------------------------
+# Route handlers
+# ---------------------------------------------------------------------------
+
+
+async def _login(request: Request):
+    """Redirect to HF OAuth authorize endpoint."""
+    if not OAUTH_CLIENT_ID or not SPACE_HOST:
+        return JSONResponse(
+            {"error": "OAuth not configured (missing OAUTH_CLIENT_ID or SPACE_HOST)"},
+            status_code=500,
+        )
+
+    _prune_expired()
+    state = secrets.token_urlsafe(32)
+    _pending_states[state] = time.time()
+
+    # Build redirect URI pointing back to our callback
+    redirect_uri = f"https://{SPACE_HOST}/api/auth/callback"
+
+    params = {
+        "client_id": OAUTH_CLIENT_ID,
+        "redirect_uri": redirect_uri,
+        "response_type": "code",
+        "scope": OAUTH_SCOPES,
+        "state": state,
+    }
+    authorize_url = f"{HF_AUTHORIZE_URL}?{urlencode(params)}"
+    logger.info(f"OAuth login: redirecting to HF authorize endpoint")
+    return RedirectResponse(authorize_url, status_code=302)
+
+
+async def _callback(request: Request):
+    """Handle OAuth callback: exchange code for token, set session."""
+    import httpx
+
+    code = request.query_params.get("code")
+    state = request.query_params.get("state")
+
+    if not code or not state:
+        return JSONResponse(
+            {"error": "Missing code or state parameter"}, status_code=400
+        )
+
+    # Validate CSRF state
+    _prune_expired()
+    created_at = _pending_states.pop(state, None)
+    if created_at is None:
+        return JSONResponse(
+            {"error": "Invalid or expired state parameter"}, status_code=400
+        )
+
+    # Exchange authorization code for access token
+    redirect_uri = f"https://{SPACE_HOST}/api/auth/callback"
+    token_data = {
+        "grant_type": "authorization_code",
+        "code": code,
+        "redirect_uri": redirect_uri,
+        "client_id": OAUTH_CLIENT_ID,
+        "client_secret": OAUTH_CLIENT_SECRET,
+    }
+
+    try:
+        async with httpx.AsyncClient() as client:
+            resp = await client.post(
+                HF_TOKEN_URL,
+                data=token_data,
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+            )
+            resp.raise_for_status()
+            token_response = resp.json()
+    except httpx.HTTPStatusError as e:
+        logger.error(
+            f"OAuth token exchange failed: {e.response.status_code} {e.response.text}"
+        )
+        return JSONResponse({"error": "Token exchange failed"}, status_code=502)
+    except Exception as e:
+        logger.error(f"OAuth token exchange error: {e}")
+        return JSONResponse({"error": "Token exchange failed"}, status_code=502)
+
+    # Extract userinfo from the id_token (JWT) or use the userinfo endpoint
+    access_token = token_response.get("access_token")
+    userinfo = None
+
+    # Try to decode the id_token (JWT) for userinfo
+    id_token = token_response.get("id_token")
+    if id_token:
+        try:
+            # JWT is base64url-encoded: header.payload.signature
+            # We only need the payload (claims) — no signature verification
+            # since we just received this directly from HF's token endpoint
+            import base64
+
+            payload_b64 = id_token.split(".")[1]
+            # Add padding if needed
+            padding = 4 - len(payload_b64) % 4
+            if padding != 4:
+                payload_b64 += "=" * padding
+            payload_bytes = base64.urlsafe_b64decode(payload_b64)
+            userinfo = json.loads(payload_bytes)
+        except Exception as e:
+            logger.warning(f"Failed to decode id_token: {e}")
+
+    # Fallback: call userinfo endpoint
+    if userinfo is None and access_token:
+        try:
+            async with httpx.AsyncClient() as client:
+                resp = await client.get(
+                    "https://huggingface.co/oauth/userinfo",
+                    headers={"Authorization": f"Bearer {access_token}"},
+                )
+                resp.raise_for_status()
+                userinfo = resp.json()
+        except Exception as e:
+            logger.error(f"Failed to fetch userinfo: {e}")
+            return JSONResponse({"error": "Failed to get user info"}, status_code=502)
+
+    if not userinfo:
+        return JSONResponse({"error": "No user info received"}, status_code=502)
+
+    # Extract username
+    username = userinfo.get("preferred_username") or userinfo.get("sub", "unknown")
+    logger.info(f"OAuth callback: authenticated user '{username}'")
+
+    # Store in server-side session dict (keyed by random session ID)
+    session_id = secrets.token_urlsafe(32)
+    _sessions[session_id] = {
+        "userinfo": userinfo,
+        "created_at": time.time(),
+    }
+
+    # The HF proxy strips custom cookies, so we can't use Set-Cookie to
+    # persist the session ID.  Instead, return a small HTML page that
+    # stores the session ID in localStorage and then redirects to /.
+    safe_sid = html_module.escape(session_id)
+    page = f"""<!DOCTYPE html>
+<html><head><title>Signing in...</title></head>
+<body><p>Signing in&hellip;</p>
+<script>
+try {{ localStorage.setItem('{_LOCALSTORAGE_KEY}', '{safe_sid}'); }}
+catch(e) {{ console.warn('Could not store session:', e); }}
+window.location.replace('/');
+</script>
+<noscript><a href="/">Click here to continue</a></noscript>
+</body></html>"""
+    return HTMLResponse(page)
+
+
+async def _logout(request: Request):
+    """Clear server-side session and browser localStorage, then redirect."""
+    # Try to clear server-side session from query param (set by JS)
+    sid = request.query_params.get("sid", "")
+    if sid and sid in _sessions:
+        del _sessions[sid]
+        logger.info(f"Cleared server-side session for logout ({sid[:8]}...)")
+
+    # Return HTML that clears localStorage and redirects
+    page = f"""<!DOCTYPE html>
+<html><head><title>Signing out...</title></head>
+<body><p>Signing out&hellip;</p>
+<script>
+try {{ localStorage.removeItem('{_LOCALSTORAGE_KEY}'); }}
+catch(e) {{ console.warn('Could not clear session:', e); }}
+window.location.replace('/');
+</script>
+<noscript><a href="/">Click here to continue</a></noscript>
+</body></html>"""
+    return HTMLResponse(page)
+
+
+# ---------------------------------------------------------------------------
+# Mount helper
+# ---------------------------------------------------------------------------
+
+_oauth_routes = [
+    Route("/api/auth/login", _login, methods=["GET"]),
+    Route("/api/auth/callback", _callback, methods=["GET"]),
+    Route("/api/auth/logout", _logout, methods=["GET"]),
+]
+
+def mount_oauth_routes(app) -> None:
+    """Mount custom OAuth routes on the Gradio ASGI app.
+
+    .. deprecated::
+        Gradio's ``launch()`` recreates ``demo.app`` from scratch, discarding
+        any routes prepended beforehand.  Use :func:`create_oauth_lifespan`
+        instead, which mounts routes on the **final** app instance via the
+        lifespan protocol.
+
+    Args:
+        app: The Starlette/FastAPI app (``demo.app``)
+    """
+    if not OAUTH_CLIENT_ID:
+        logger.warning(
+            "OAuth routes not mounted: OAUTH_CLIENT_ID not set. "
+            "Custom login will not work."
+        )
+        return
+
+    # Insert our routes at the beginning so they take priority
+    app.routes[:0] = _oauth_routes
+    logger.info(
+        f"Mounted custom OAuth routes: /api/auth/login, /api/auth/callback, /api/auth/logout"
+    )
+
+
+@contextlib.asynccontextmanager
+async def create_oauth_lifespan(app):
+    """ASGI lifespan handler that mounts custom OAuth routes.
+
+    Gradio's ``launch()`` recreates ``demo.app`` via ``App.create_app()``,
+    discarding any routes prepended before launch.  However,
+    ``create_lifespan_handler()`` (in Gradio's ``route_utils.py``) wraps
+    user-supplied lifespan and calls it with the **final** App instance.
+    Routes added here are on the real app with the full middleware stack
+    (including ``SessionMiddleware``).
+
+    Pass this to ``demo.launch(app_kwargs={"lifespan": create_oauth_lifespan})``.
+    """
+    if OAUTH_CLIENT_ID:
+        app.routes[:0] = _oauth_routes
+        logger.info(
+            "Custom OAuth routes mounted via lifespan: "
+            "/api/auth/login, /api/auth/callback, /api/auth/logout"
+        )
+    yield

src/mosaic/ui/user_tabs.py

@@ -29,40 +29,57 @@ from mosaic.user_results import (
 LOCAL_DEBUG_USERNAME = "local_user"
 
 
-def _get_username(request: gr.Request = None, profile=None) -> tuple[str, bool]:
+def _get_username(
+    request: gr.Request = None, profile=None, session_id: str = ""
+) -> tuple[str, bool]:
     """Get username for storage operations.
 
-    On HF Spaces, uses OAuthProfile from gr.LoginButton() to identify
-    the logged-in user. Does NOT use request.username, which returns
-    the Space owner's username rather than the visitor's.
+    On HF Spaces with Docker SDK the HF proxy strips custom cookies and
+    forges the Starlette session so ``gr.OAuthProfile`` always returns the
+    Space owner's identity.  The only trustworthy channel is the session ID
+    stored in the browser's ``localStorage`` and passed via Gradio inputs.
+
+    Args:
+        request: Gradio request (auto-injected).
+        profile: OAuthProfile (auto-injected, **unreliable on HF Spaces**).
+        session_id: Server-side session ID from localStorage (reliable).
 
     Returns:
         Tuple of (username, is_local_mode)
-        - In HF Spaces: (profile.username, False) if logged in
+        - In HF Spaces: (username, False) if session_id is valid
         - Locally: (LOCAL_DEBUG_USERNAME, True)
         - HF Spaces not logged in: (None, False)
     """
     if IS_HF_SPACES:
-        # On HF Spaces, only OAuthProfile reliably identifies the logged-in user.
-        # request.username returns the Space owner's username (not the visitor's),
-        # so we must NOT fall back to it.
-        if profile is not None:
+        # Look up user by session ID passed from localStorage via Gradio js=
+        if session_id:
             try:
-                username = getattr(profile, "username", None)
-                if username:
-                    return (username, False)
+                from mosaic.ui.oauth import get_user_by_session_id
+
+                userinfo = get_user_by_session_id(session_id)
+                if userinfo:
+                    username = userinfo.get("preferred_username")
+                    if username:
+                        return (username, False)
             except Exception:
                 pass
 
+        # Do NOT fall back to OAuthProfile — it always returns the Space
+        # owner's identity on Docker SDK Spaces.
         return (None, False)
     else:
         # Local mode - use debug username
         return (LOCAL_DEBUG_USERNAME, True)
 
 
-def create_my_files_tab():
+def create_my_files_tab(session_id_component=None):
     """Create the My Files tab UI for managing uploaded slides.
 
+    Args:
+        session_id_component: Hidden gr.Textbox carrying the localStorage
+            session ID.  Passed as input to event handlers so that
+            ``_get_username`` can resolve the visitor's identity on HF Spaces.
+
     Returns:
         Dictionary of UI components for event handlers
     """
@@ -107,9 +124,13 @@ def create_my_files_tab():
         # Download output
         slide_download_file = gr.File(label="Download", visible=False)
 
-        def load_files(request: gr.Request = None, profile: gr.OAuthProfile = None):
+        def load_files(
+            session_id: str = "",
+            request: gr.Request = None,
+            profile: gr.OAuthProfile = None,
+        ):
             """Load user's uploaded files."""
-            username, is_local = _get_username(request, profile)
+            username, is_local = _get_username(request, profile, session_id)
 
             if not username:
                 return (
@@ -188,11 +209,12 @@ def create_my_files_tab():
 
         def download_slide(
             slide_id_or_name,
+            session_id: str = "",
             request: gr.Request = None,
             profile: gr.OAuthProfile = None,
         ):
             """Download the original slide file."""
-            username, _ = _get_username(request, profile)
+            username, _ = _get_username(request, profile, session_id)
 
             if not username or not slide_id_or_name:
                 return gr.File(visible=False), gr.Textbox(
@@ -226,11 +248,12 @@ def create_my_files_tab():
 
         def delete_slide(
             slide_id_or_name,
+            session_id: str = "",
             request: gr.Request = None,
             profile: gr.OAuthProfile = None,
         ):
             """Delete a slide and all associated results."""
-            username, _ = _get_username(request, profile)
+            username, _ = _get_username(request, profile, session_id)
 
             if not username or not slide_id_or_name:
                 return gr.Textbox(
@@ -257,10 +280,14 @@ def create_my_files_tab():
                 logger.error(f"Failed to delete slide {slide_id_or_name}: {e}")
                 return gr.Textbox(value=f"Error: {str(e)}", visible=True)
 
+        # Build inputs lists — include session_id_component when provided
+        # (HF Spaces) so _get_username can resolve the visitor's identity.
+        _sid = [session_id_component] if session_id_component is not None else []
+
         # Wire up events
         refresh_files_btn.click(
             load_files,
-            inputs=None,
+            inputs=_sid or None,
             outputs=[
                 storage_usage_md,
                 files_table,
@@ -271,17 +298,17 @@ def create_my_files_tab():
 
         download_slide_btn.click(
             download_slide,
-            inputs=[selected_slide_id],
+            inputs=[selected_slide_id] + _sid,
             outputs=[slide_download_file, file_action_status],
         )
 
         delete_slide_btn.click(
             delete_slide,
-            inputs=[selected_slide_id],
+            inputs=[selected_slide_id] + _sid,
             outputs=[file_action_status],
         ).then(
             load_files,  # Refresh after delete
-            inputs=None,
+            inputs=_sid or None,
             outputs=[
                 storage_usage_md,
                 files_table,
@@ -313,9 +340,13 @@ def create_my_files_tab():
     }
 
 
-def create_my_results_tab():
+def create_my_results_tab(session_id_component=None):
     """Create the My Results tab UI for browsing analysis results.
 
+    Args:
+        session_id_component: Hidden gr.Textbox carrying the localStorage
+            session ID (see ``create_my_files_tab`` docstring).
+
     Returns:
         Dictionary of UI components for event handlers
     """
@@ -364,9 +395,13 @@ def create_my_results_tab():
         )
         result_download_file = gr.File(label="Download", visible=False)
 
-        def load_results(request: gr.Request = None, profile: gr.OAuthProfile = None):
+        def load_results(
+            session_id: str = "",
+            request: gr.Request = None,
+            profile: gr.OAuthProfile = None,
+        ):
             """Load user's analysis results."""
-            username, _ = _get_username(request, profile)
+            username, _ = _get_username(request, profile, session_id)
 
             if not username:
                 return [], gr.Textbox(visible=False), gr.File(visible=False)
@@ -403,10 +438,13 @@ def create_my_results_tab():
                 )
 
         def view_result_details(
-            analysis_id, request: gr.Request = None, profile: gr.OAuthProfile = None
+            analysis_id,
+            session_id: str = "",
+            request: gr.Request = None,
+            profile: gr.OAuthProfile = None,
         ):
             """Load and display result details."""
-            username, _ = _get_username(request, profile)
+            username, _ = _get_username(request, profile, session_id)
 
             if not username or not analysis_id:
                 return (
@@ -515,10 +553,13 @@ def create_my_results_tab():
                 )
 
         def download_result_zip(
-            analysis_id, request: gr.Request = None, profile: gr.OAuthProfile = None
+            analysis_id,
+            session_id: str = "",
+            request: gr.Request = None,
+            profile: gr.OAuthProfile = None,
         ):
             """Create and download ZIP of result files."""
-            username, _ = _get_username(request, profile)
+            username, _ = _get_username(request, profile, session_id)
 
             if not username or not analysis_id:
                 return gr.File(visible=False), gr.Textbox(
@@ -544,10 +585,13 @@ def create_my_results_tab():
                 )
 
         def delete_result(
-            analysis_id, request: gr.Request = None, profile: gr.OAuthProfile = None
+            analysis_id,
+            session_id: str = "",
+            request: gr.Request = None,
+            profile: gr.OAuthProfile = None,
         ):
             """Delete an analysis result."""
-            username, _ = _get_username(request, profile)
+            username, _ = _get_username(request, profile, session_id)
 
             if not username or not analysis_id:
                 return gr.Textbox(
@@ -568,16 +612,19 @@ def create_my_results_tab():
                 logger.error(f"Failed to delete result {analysis_id}: {e}")
                 return gr.Textbox(value=f"Error: {str(e)}", visible=True)
 
+        # Build inputs lists — include session_id_component when provided
+        _sid = [session_id_component] if session_id_component is not None else []
+
         # Wire up events
         refresh_results_btn.click(
             load_results,
-            inputs=None,
+            inputs=_sid or None,
             outputs=[results_table, result_action_status, result_download_file],
         )
 
         view_result_btn.click(
             view_result_details,
-            inputs=[selected_analysis_id],
+            inputs=[selected_analysis_id] + _sid,
             outputs=[
                 result_details,
                 result_metadata_md,
@@ -590,17 +637,17 @@ def create_my_results_tab():
 
         download_zip_btn.click(
             download_result_zip,
-            inputs=[selected_analysis_id],
+            inputs=[selected_analysis_id] + _sid,
             outputs=[result_download_file, result_action_status],
         )
 
         delete_result_btn.click(
             delete_result,
-            inputs=[selected_analysis_id],
+            inputs=[selected_analysis_id] + _sid,
             outputs=[result_action_status],
         ).then(
             load_results,  # Refresh after delete
-            inputs=None,
+            inputs=_sid or None,
             outputs=[results_table, result_action_status, result_download_file],
         )
 

tests/telemetry/test_tracker.py

@@ -8,7 +8,6 @@ from pathlib import Path
 import pytest
 
 from mosaic.telemetry import TelemetryTracker, TelemetryConfig
-from mosaic.telemetry.utils import hash_username
 
 
 @pytest.fixture
@@ -153,7 +152,7 @@ class TestUsageEvents:
             event = json.loads(f.read().strip())
 
         assert event["is_logged_in"] is True
-        assert event["hf_username"] == hash_username("testuser")
+        assert event["hf_username"] == "testuser"
 
     def test_log_analysis_complete(self, tracker, temp_dir):
         """Test logging analysis complete event."""
@@ -287,7 +286,7 @@ class TestResourceEvents:
             event = json.loads(f.read().strip())
 
         assert event["is_logged_in"] is True
-        assert event["hf_username"] == hash_username("testuser")
+        assert event["hf_username"] == "testuser"
 
 
 class TestFailureEvents:

tests/telemetry/test_utils.py

@@ -8,7 +8,6 @@ from mosaic.telemetry.utils import (
     StageTimer,
     sanitize_error_message,
     hash_session_id,
-    hash_username,
     UserInfo,
     extract_user_info,
 )
@@ -151,40 +150,6 @@ class TestHashSessionId:
         assert len(set(hashes)) == 1  # All hashes should be identical
 
 
-class TestHashUsername:
-    """Tests for username hashing."""
-
-    def test_hash_username(self):
-        """Test basic username hashing."""
-        hashed = hash_username("testuser")
-        assert hashed is not None
-        assert hashed != "testuser"
-        assert len(hashed) == 16
-
-    def test_hash_none_returns_none(self):
-        """Test that None input returns None."""
-        assert hash_username(None) is None
-
-    def test_hash_is_deterministic(self):
-        """Test that same input produces same hash."""
-        hash1 = hash_username("alice")
-        hash2 = hash_username("alice")
-        assert hash1 == hash2
-
-    def test_different_inputs_different_hashes(self):
-        """Test that different usernames produce different hashes."""
-        hash1 = hash_username("alice")
-        hash2 = hash_username("bob")
-        assert hash1 != hash2
-
-    def test_different_salt_from_session_id(self):
-        """Test that username hash uses different salt than session hash."""
-        value = "same_value"
-        username_hash = hash_username(value)
-        session_hash = hash_session_id(value)
-        assert username_hash != session_hash
-
-
 class TestUserInfo:
     """Tests for UserInfo dataclass."""
 

tests/test_oauth.py

@@ -0,0 +1,479 @@
+"""Tests for the custom OAuth module (mosaic.ui.oauth)."""
+
+import time
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from mosaic.ui.oauth import (
+    _prune_expired,
+    _sessions,
+    _pending_states,
+    _SESSION_COOKIE,
+    _SESSION_TTL_SEC,
+    _STATE_TTL_SEC,
+    create_oauth_lifespan,
+    get_user_by_session_id,
+    get_user_from_server_session,
+    mount_oauth_routes,
+)
+
+
+@pytest.fixture(autouse=True)
+def clean_sessions():
+    """Clear session and state stores before each test."""
+    _sessions.clear()
+    _pending_states.clear()
+    yield
+    _sessions.clear()
+    _pending_states.clear()
+
+
+# ---------------------------------------------------------------------------
+# Session store tests
+# ---------------------------------------------------------------------------
+
+
+class TestSessionStore:
+    """Tests for the server-side session store."""
+
+    def test_store_and_retrieve_session(self):
+        """Test basic session CRUD."""
+        _sessions["abc123"] = {
+            "userinfo": {"preferred_username": "alice"},
+            "created_at": time.time(),
+        }
+
+        request = MagicMock()
+        request.cookies = {_SESSION_COOKIE: "abc123"}
+        request.headers = {}
+
+        userinfo = get_user_from_server_session(request)
+        assert userinfo is not None
+        assert userinfo["preferred_username"] == "alice"
+
+    def test_missing_cookie_returns_none(self):
+        """Test that missing cookie returns None."""
+        request = MagicMock()
+        request.cookies = {}
+        request.headers = {}
+
+        assert get_user_from_server_session(request) is None
+
+    def test_unknown_session_id_returns_none(self):
+        """Test that unknown session ID returns None."""
+        request = MagicMock()
+        request.cookies = {_SESSION_COOKIE: "unknown"}
+        request.headers = {}
+
+        assert get_user_from_server_session(request) is None
+
+    def test_expired_session_returns_none(self):
+        """Test that expired sessions are pruned."""
+        _sessions["expired"] = {
+            "userinfo": {"preferred_username": "bob"},
+            "created_at": time.time() - _SESSION_TTL_SEC - 1,
+        }
+
+        request = MagicMock()
+        request.cookies = {_SESSION_COOKIE: "expired"}
+        request.headers = {}
+
+        assert get_user_from_server_session(request) is None
+        assert "expired" not in _sessions
+
+    def test_cookie_from_header_fallback(self):
+        """Test parsing cookie from raw Cookie header."""
+        _sessions["header_val"] = {
+            "userinfo": {"preferred_username": "carol"},
+            "created_at": time.time(),
+        }
+
+        request = MagicMock()
+        request.cookies = {}  # Empty dict, fallback to header
+        request.headers = {"cookie": f"other=x; {_SESSION_COOKIE}=header_val; foo=bar"}
+
+        userinfo = get_user_from_server_session(request)
+        assert userinfo is not None
+        assert userinfo["preferred_username"] == "carol"
+
+    def test_none_request_returns_none(self):
+        """Test that None request returns None gracefully."""
+        # get_user_from_server_session checks hasattr, so pass an object
+        # without cookies or headers
+        assert get_user_from_server_session(None) is None
+
+
+# ---------------------------------------------------------------------------
+# get_user_by_session_id tests
+# ---------------------------------------------------------------------------
+
+
+class TestGetUserBySessionId:
+    """Tests for get_user_by_session_id (localStorage-based lookup)."""
+
+    def test_valid_session_id(self):
+        """Test lookup with a valid session ID."""
+        _sessions["sid_abc"] = {
+            "userinfo": {"preferred_username": "alice"},
+            "created_at": time.time(),
+        }
+        result = get_user_by_session_id("sid_abc")
+        assert result is not None
+        assert result["preferred_username"] == "alice"
+
+    def test_empty_session_id(self):
+        """Test that empty string returns None."""
+        assert get_user_by_session_id("") is None
+
+    def test_none_session_id(self):
+        """Test that None returns None."""
+        assert get_user_by_session_id(None) is None
+
+    def test_unknown_session_id(self):
+        """Test that unknown session ID returns None."""
+        assert get_user_by_session_id("nonexistent") is None
+
+    def test_expired_session_id(self):
+        """Test that expired session returns None and is cleaned up."""
+        _sessions["expired_sid"] = {
+            "userinfo": {"preferred_username": "bob"},
+            "created_at": time.time() - _SESSION_TTL_SEC - 1,
+        }
+        assert get_user_by_session_id("expired_sid") is None
+        assert "expired_sid" not in _sessions
+
+
+# ---------------------------------------------------------------------------
+# Prune tests
+# ---------------------------------------------------------------------------
+
+
+class TestPruneExpired:
+    """Tests for _prune_expired."""
+
+    def test_prune_expired_sessions(self):
+        """Test that expired sessions are removed."""
+        _sessions["old"] = {
+            "userinfo": {"preferred_username": "old_user"},
+            "created_at": time.time() - _SESSION_TTL_SEC - 100,
+        }
+        _sessions["fresh"] = {
+            "userinfo": {"preferred_username": "new_user"},
+            "created_at": time.time(),
+        }
+
+        _prune_expired()
+
+        assert "old" not in _sessions
+        assert "fresh" in _sessions
+
+    def test_prune_expired_states(self):
+        """Test that expired CSRF states are removed."""
+        _pending_states["old_state"] = time.time() - _STATE_TTL_SEC - 100
+        _pending_states["fresh_state"] = time.time()
+
+        _prune_expired()
+
+        assert "old_state" not in _pending_states
+        assert "fresh_state" in _pending_states
+
+
+# ---------------------------------------------------------------------------
+# mount_oauth_routes tests
+# ---------------------------------------------------------------------------
+
+
+class TestMountOAuthRoutes:
+    """Tests for mount_oauth_routes."""
+
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_ID", "")
+    def test_no_routes_when_client_id_missing(self):
+        """Test that routes are not mounted when OAUTH_CLIENT_ID is empty."""
+        app = MagicMock()
+        app.routes = []
+
+        mount_oauth_routes(app)
+
+        assert len(app.routes) == 0
+
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_ID", "test-client-id")
+    def test_routes_mounted_when_configured(self):
+        """Test that routes are mounted when OAuth is configured."""
+        app = MagicMock()
+        app.routes = []
+
+        mount_oauth_routes(app)
+
+        assert len(app.routes) == 3
+        paths = [r.path for r in app.routes]
+        assert "/api/auth/login" in paths
+        assert "/api/auth/callback" in paths
+        assert "/api/auth/logout" in paths
+
+
+# ---------------------------------------------------------------------------
+# Login route tests
+# ---------------------------------------------------------------------------
+
+
+class TestLoginRoute:
+    """Tests for the /api/auth/login route."""
+
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_ID", "test-client-id")
+    @patch("mosaic.ui.oauth.SPACE_HOST", "user-space.hf.space")
+    def test_login_redirects_to_hf(self):
+        """Test that login redirects to HF authorize URL."""
+        import asyncio
+
+        from mosaic.ui.oauth import _login
+
+        async def _run():
+            request = MagicMock()
+            return await _login(request)
+
+        response = asyncio.run(_run())
+
+        assert response.status_code == 302
+        location = response.headers["location"]
+        assert "huggingface.co/oauth/authorize" in location
+        assert "client_id=test-client-id" in location
+        assert "redirect_uri=" in location
+        assert "state=" in location
+        # Should have stored a CSRF state
+        assert len(_pending_states) == 1
+
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_ID", "")
+    def test_login_returns_error_when_not_configured(self):
+        """Test that login returns error when OAuth is not configured."""
+        import asyncio
+
+        from mosaic.ui.oauth import _login
+
+        async def _run():
+            request = MagicMock()
+            return await _login(request)
+
+        response = asyncio.run(_run())
+
+        assert response.status_code == 500
+
+
+# ---------------------------------------------------------------------------
+# Callback route tests
+# ---------------------------------------------------------------------------
+
+
+class TestCallbackRoute:
+    """Tests for the /api/auth/callback route."""
+
+    def test_callback_missing_params(self):
+        """Test callback with missing code/state parameters."""
+        import asyncio
+
+        from mosaic.ui.oauth import _callback
+
+        async def _run():
+            request = MagicMock()
+            request.query_params = {}
+            return await _callback(request)
+
+        response = asyncio.run(_run())
+        assert response.status_code == 400
+
+    def test_callback_invalid_state(self):
+        """Test callback with invalid CSRF state."""
+        import asyncio
+
+        from mosaic.ui.oauth import _callback
+
+        async def _run():
+            request = MagicMock()
+            request.query_params = {"code": "test-code", "state": "invalid-state"}
+            return await _callback(request)
+
+        response = asyncio.run(_run())
+        assert response.status_code == 400
+
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_ID", "test-client-id")
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_SECRET", "test-secret")
+    @patch("mosaic.ui.oauth.SPACE_HOST", "user-space.hf.space")
+    def test_callback_exchanges_code_for_token(self):
+        """Test successful callback with mocked token exchange."""
+        import asyncio
+        import json
+        import base64
+
+        from mosaic.ui.oauth import _callback
+
+        # Set up a valid CSRF state
+        valid_state = "valid-state-token"
+        _pending_states[valid_state] = time.time()
+
+        # Create a mock id_token JWT
+        header = base64.urlsafe_b64encode(b'{"alg":"RS256"}').rstrip(b"=").decode()
+        payload_data = json.dumps(
+            {"preferred_username": "visitor123", "sub": "visitor123"}
+        )
+        payload = base64.urlsafe_b64encode(payload_data.encode()).rstrip(b"=").decode()
+        mock_id_token = f"{header}.{payload}.fake_signature"
+
+        # Mock httpx — json() is a regular method (not async) on httpx.Response
+        mock_response = AsyncMock()
+        mock_response.json = MagicMock(
+            return_value={
+                "access_token": "test-access-token",
+                "id_token": mock_id_token,
+            }
+        )
+        mock_response.raise_for_status = MagicMock()
+
+        mock_client = AsyncMock()
+        mock_client.post.return_value = mock_response
+        mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+        mock_client.__aexit__ = AsyncMock(return_value=False)
+
+        request = MagicMock()
+        request.query_params = {"code": "auth-code", "state": valid_state}
+
+        async def _run():
+            with patch("httpx.AsyncClient", return_value=mock_client):
+                return await _callback(request)
+
+        response = asyncio.run(_run())
+
+        # Callback returns an HTML page that stores session in localStorage
+        assert response.status_code == 200
+        assert "localStorage.setItem" in response.body.decode()
+        assert "mosaic_session_id" in response.body.decode()
+
+        # Server-side session should be stored
+        assert len(_sessions) == 1
+
+        # CSRF state should be consumed
+        assert valid_state not in _pending_states
+
+
+# ---------------------------------------------------------------------------
+# Logout route tests
+# ---------------------------------------------------------------------------
+
+
+class TestLogoutRoute:
+    """Tests for the /api/auth/logout route."""
+
+    def test_logout_clears_session(self):
+        """Test that logout clears session data and returns localStorage-clearing HTML."""
+        import asyncio
+
+        from mosaic.ui.oauth import _logout
+
+        # Set up server-side session
+        _sessions["session_id_abc"] = {
+            "userinfo": {"preferred_username": "user"},
+            "created_at": time.time(),
+        }
+
+        request = MagicMock()
+        request.query_params = {"sid": "session_id_abc"}
+
+        async def _run():
+            return await _logout(request)
+
+        response = asyncio.run(_run())
+
+        # Returns HTML page that clears localStorage
+        assert response.status_code == 200
+        assert "localStorage.removeItem" in response.body.decode()
+        assert "mosaic_session_id" in response.body.decode()
+        assert "session_id_abc" not in _sessions
+
+
+# ---------------------------------------------------------------------------
+# get_user_from_server_session integration
+# ---------------------------------------------------------------------------
+
+
+class TestGetUserFromServerSession:
+    """Integration tests for the server-session fallback."""
+
+    def test_works_with_gradio_request_like_object(self):
+        """Test with an object mimicking gr.Request."""
+        _sessions["gr_session"] = {
+            "userinfo": {"preferred_username": "gradio_user", "name": "Test User"},
+            "created_at": time.time(),
+        }
+
+        class MockGradioRequest:
+            cookies = {_SESSION_COOKIE: "gr_session"}
+            headers = {}
+
+        result = get_user_from_server_session(MockGradioRequest())
+        assert result is not None
+        assert result["preferred_username"] == "gradio_user"
+        assert result["name"] == "Test User"
+
+    def test_returns_none_for_object_without_cookies(self):
+        """Test graceful handling of objects without cookies."""
+
+        class Bare:
+            pass
+
+        assert get_user_from_server_session(Bare()) is None
+
+
+# ---------------------------------------------------------------------------
+# create_oauth_lifespan tests
+# ---------------------------------------------------------------------------
+
+
+class TestCreateOAuthLifespan:
+    """Tests for the create_oauth_lifespan async context manager."""
+
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_ID", "test-client-id")
+    def test_routes_mounted_when_configured(self):
+        """Test that OAuth routes are mounted on app when OAUTH_CLIENT_ID is set."""
+        import asyncio
+
+        app = MagicMock()
+        app.routes = []
+
+        async def _run():
+            async with create_oauth_lifespan(app):
+                assert len(app.routes) == 3
+                paths = [r.path for r in app.routes]
+                assert "/api/auth/login" in paths
+                assert "/api/auth/callback" in paths
+                assert "/api/auth/logout" in paths
+
+        asyncio.run(_run())
+
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_ID", "")
+    def test_routes_skipped_when_not_configured(self):
+        """Test that no routes are mounted when OAUTH_CLIENT_ID is empty."""
+        import asyncio
+
+        app = MagicMock()
+        app.routes = []
+
+        async def _run():
+            async with create_oauth_lifespan(app):
+                assert len(app.routes) == 0
+
+        asyncio.run(_run())
+
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_ID", "test-client-id")
+    def test_routes_persist_after_yield(self):
+        """Test that routes remain on the app after the lifespan exits."""
+        import asyncio
+
+        app = MagicMock()
+        app.routes = []
+
+        async def _run():
+            async with create_oauth_lifespan(app):
+                pass
+            # Routes should still be present after context manager exits
+            assert len(app.routes) == 3
+
+        asyncio.run(_run())

tests/test_ui_user_storage.py

@@ -86,16 +86,23 @@ class TestUsernameExtraction:
 
     @patch("mosaic.ui.user_tabs.IS_HF_SPACES", True)
     def test_hf_logged_in_returns_username(self, mock_request_hf_logged_in):
-        """HF Spaces logged in should return OAuthProfile username, not request.username."""
-
-        class MockProfile:
-            username = "test_user"
-
-        username, is_local = _get_username(
-            mock_request_hf_logged_in, profile=MockProfile()
-        )
-        assert username == "test_user"
-        assert is_local is False
+        """HF Spaces logged in should return username from session_id lookup."""
+        import time
+        from mosaic.ui.oauth import _sessions
+
+        # Set up a server-side session
+        _sessions["test_session_123"] = {
+            "userinfo": {"preferred_username": "test_user"},
+            "created_at": time.time(),
+        }
+        try:
+            username, is_local = _get_username(
+                mock_request_hf_logged_in, session_id="test_session_123"
+            )
+            assert username == "test_user"
+            assert is_local is False
+        finally:
+            _sessions.pop("test_session_123", None)
 
     @patch("mosaic.ui.user_tabs.IS_HF_SPACES", True)
     def test_hf_request_username_ignored(self, mock_request_hf_logged_in):