feat: implement manual OAuth for HF Spaces Docker SDK

On HF Spaces with sdk:docker, Gradio's built-in OAuth injects the Space
owner's identity into every session. This adds custom /api/auth/* routes
that perform the full Authorization Code flow directly with HF's OAuth
provider, writing the correct visitor identity into the Starlette session
so Gradio's OAuthProfile injection works transparently.

- Add src/mosaic/ui/oauth.py with login/callback/logout routes and
server-side session store (24h TTL, mosaic_auth cookie fallback)
- Mount custom OAuth routes on Gradio app when IS_HF_SPACES
- Update login/logout links to /api/auth/login and /api/auth/logout
- Add server-session fallback to extract_user_info() and _get_username()
- Stop hashing usernames in telemetry (use raw HF usernames)
- Add pyproject.toml uv index-strategy for macOS CPU torch compatibility
- Add tests/test_oauth.py with 18 tests covering the full OAuth flow

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

pyproject.toml

@@ -42,6 +42,13 @@ disable = [
   "unspecified-encoding",
 ]
 
+[tool.uv]
+# Override PyTorch dependencies from mussel[torch-gpu] for macOS compatibility
+override-dependencies = [
+  "torch>=2.0.0; sys_platform == 'darwin'",
+  "torchvision>=0.15.0; sys_platform == 'darwin'",
+]
+
 [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/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.
@@ -169,5 +150,19 @@ 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}")
 
+    # Fallback: check server-side session store (custom OAuth flow)
+    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}")
+
     logger.debug("User not logged in: no OAuthProfile available")
     return UserInfo()

src/mosaic/ui/app.py

@@ -37,6 +37,7 @@ from mosaic.data_directory import get_tcga_cache_directory
 from mosaic.model_manager import load_all_models
 from mosaic.hardware import DEFAULT_CONCURRENCY_LIMIT, IS_HF_SPACES, IS_T4_GPU, GPU_TYPE
 from mosaic.telemetry import extract_user_info
+from mosaic.ui.oauth import mount_oauth_routes
 from mosaic.tcga import (
     compute_settings_hash,
     download_results_from_hf,
@@ -1933,7 +1934,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
                     )
                     return (
                         gr.update(
-                            value=f"Signed in as **{username}** \u00b7 [Sign out](/logout)",
+                            value=f"Signed in as **{username}** \u00b7 [Sign out](/api/auth/logout)",
                             visible=True,
                         ),  # login_status_md
                         gr.update(visible=True),  # user_storage_tabs
@@ -1943,7 +1944,7 @@ This tool is for research purposes only and not approved for clinical diagnosis.
                 else:
                     return (
                         gr.update(
-                            value="[Sign in with HuggingFace](/login/huggingface) to save slides and results",
+                            value="[Sign in with HuggingFace](/api/auth/login) to save slides and results",
                             visible=True,
                         ),  # login_status_md
                         gr.update(visible=False),  # user_storage_tabs
@@ -1967,6 +1968,10 @@ 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)
 
+    # Mount custom OAuth routes for HF Spaces (sdk:docker)
+    if IS_HF_SPACES:
+        mount_oauth_routes(demo.app)
+
     # Register cleanup handler for graceful shutdown
     import atexit
 

src/mosaic/ui/oauth.py

@@ -0,0 +1,321 @@
+"""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 and writes
+the visitor's real identity into the Starlette session so that Gradio's
+existing OAuthProfile injection works transparently.
+
+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, set session
+    GET /api/auth/logout   -> Clear session, redirect to /
+"""
+
+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
+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
+# ---------------------------------------------------------------------------
+
+# {cookie_value: {"userinfo": {...}, "created_at": float}}
+_sessions: dict[str, dict] = {}
+_SESSION_TTL_SEC = 24 * 60 * 60  # 24 hours
+_SESSION_COOKIE = "mosaic_auth"
+
+# 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 a Starlette Request; try .cookies first
+    if hasattr(request, "cookies"):
+        cookies = request.cookies
+        if isinstance(cookies, dict):
+            cookie_val = cookies.get(_SESSION_COOKIE)
+    # Fallback: parse Cookie header
+    if cookie_val is None and hasattr(request, "headers"):
+        headers = request.headers
+        cookie_header = None
+        if isinstance(headers, dict):
+            cookie_header = headers.get("cookie", "")
+        elif hasattr(headers, "get"):
+            cookie_header = headers.get("cookie", "")
+        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
+
+    if not cookie_val:
+        return None
+
+    entry = _sessions.get(cookie_val)
+    if entry is None:
+        return None
+
+    if time.time() - entry["created_at"] > _SESSION_TTL_SEC:
+        del _sessions[cookie_val]
+        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}'")
+
+    # Write into Gradio's Starlette session so OAuthProfile picks it up
+    # Gradio expects session["oauth_info"]["userinfo"] with at least
+    # "preferred_username" and optionally "name", "picture", etc.
+    oauth_info = {
+        "userinfo": userinfo,
+        "access_token": access_token,
+    }
+    request.session["oauth_info"] = oauth_info
+
+    # Also store in our server-side session as fallback
+    session_id = secrets.token_urlsafe(32)
+    _sessions[session_id] = {
+        "userinfo": userinfo,
+        "created_at": time.time(),
+    }
+
+    # Build redirect response with cookie
+    response = RedirectResponse("/", status_code=302)
+    response.set_cookie(
+        _SESSION_COOKIE,
+        session_id,
+        httponly=True,
+        secure=True,
+        samesite="lax",
+        max_age=_SESSION_TTL_SEC,
+    )
+    return response
+
+
+async def _logout(request: Request):
+    """Clear session and redirect to /."""
+    # Clear Gradio's session
+    if "oauth_info" in request.session:
+        del request.session["oauth_info"]
+
+    # Clear server-side session
+    cookie_val = request.cookies.get(_SESSION_COOKIE)
+    if cookie_val and cookie_val in _sessions:
+        del _sessions[cookie_val]
+
+    response = RedirectResponse("/", status_code=302)
+    response.delete_cookie(_SESSION_COOKIE)
+    return response
+
+
+# ---------------------------------------------------------------------------
+# 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.
+
+    Should be called after ``demo.queue()`` and before ``demo.launch()``.
+
+    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"
+    )

src/mosaic/ui/user_tabs.py

@@ -54,6 +54,19 @@ def _get_username(request: gr.Request = None, profile=None) -> tuple[str, bool]:
             except Exception:
                 pass
 
+        # Fallback: check server-side session store (custom OAuth flow)
+        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:
+                        return (username, False)
+            except Exception:
+                pass
+
         return (None, False)
     else:
         # Local mode - use debug username

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,364 @@
+"""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,
+    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
+
+
+# ---------------------------------------------------------------------------
+# 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."""
+
+    @pytest.mark.asyncio
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_ID", "test-client-id")
+    @patch("mosaic.ui.oauth.SPACE_HOST", "user-space.hf.space")
+    async def test_login_redirects_to_hf(self):
+        """Test that login redirects to HF authorize URL."""
+        from mosaic.ui.oauth import _login
+
+        request = MagicMock()
+        response = await _login(request)
+
+        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
+
+    @pytest.mark.asyncio
+    @patch("mosaic.ui.oauth.OAUTH_CLIENT_ID", "")
+    async def test_login_returns_error_when_not_configured(self):
+        """Test that login returns error when OAuth is not configured."""
+        from mosaic.ui.oauth import _login
+
+        request = MagicMock()
+        response = await _login(request)
+
+        assert response.status_code == 500
+
+
+# ---------------------------------------------------------------------------
+# Callback route tests
+# ---------------------------------------------------------------------------
+
+
+class TestCallbackRoute:
+    """Tests for the /api/auth/callback route."""
+
+    @pytest.mark.asyncio
+    async def test_callback_missing_params(self):
+        """Test callback with missing code/state parameters."""
+        from mosaic.ui.oauth import _callback
+
+        request = MagicMock()
+        request.query_params = {}
+
+        response = await _callback(request)
+        assert response.status_code == 400
+
+    @pytest.mark.asyncio
+    async def test_callback_invalid_state(self):
+        """Test callback with invalid CSRF state."""
+        from mosaic.ui.oauth import _callback
+
+        request = MagicMock()
+        request.query_params = {"code": "test-code", "state": "invalid-state"}
+
+        response = await _callback(request)
+        assert response.status_code == 400
+
+    @pytest.mark.asyncio
+    @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")
+    async def test_callback_exchanges_code_for_token(self):
+        """Test successful callback with mocked token exchange."""
+        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}
+        request.session = {}
+
+        with patch("httpx.AsyncClient", return_value=mock_client):
+            response = await _callback(request)
+
+        # Should redirect to /
+        assert response.status_code == 302
+
+        # Session should have oauth_info
+        assert "oauth_info" in request.session
+        assert (
+            request.session["oauth_info"]["userinfo"]["preferred_username"]
+            == "visitor123"
+        )
+
+        # 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."""
+
+    @pytest.mark.asyncio
+    async def test_logout_clears_session(self):
+        """Test that logout clears session data."""
+        from mosaic.ui.oauth import _logout
+
+        # Set up server-side session
+        _sessions["session_id"] = {
+            "userinfo": {"preferred_username": "user"},
+            "created_at": time.time(),
+        }
+
+        request = MagicMock()
+        request.session = {"oauth_info": {"userinfo": {"preferred_username": "user"}}}
+        request.cookies = {_SESSION_COOKIE: "session_id"}
+
+        response = await _logout(request)
+
+        assert response.status_code == 302
+        assert "oauth_info" not in request.session
+        assert "session_id" 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