Harden Space OAuth startup

app.py

@@ -1,6 +1,7 @@
 from waterleaf.runtime import build_application
+from waterleaf.space import configure_hf_space_environment
 from waterleaf.web import create_web_app
 
+configure_hf_space_environment()
 application = build_application()
 app = create_web_app(application)
-

tests/test_space.py

@@ -0,0 +1,48 @@
+import os
+
+from waterleaf.space import (
+    OAUTH_VARIABLES,
+    configure_hf_space_environment,
+    has_hf_oauth,
+    is_hf_space,
+)
+
+
+def _clear_space_environment(monkeypatch):
+    for name in (*OAUTH_VARIABLES, "SPACE_ID", "SPACE_HOST", "SYSTEM"):
+        monkeypatch.delenv(name, raising=False)
+
+
+def test_local_environment_does_not_enable_space_oauth(monkeypatch):
+    _clear_space_environment(monkeypatch)
+
+    configure_hf_space_environment()
+
+    assert not is_hf_space()
+    assert not has_hf_oauth()
+
+
+def test_complete_oauth_environment_configures_gradio_space(monkeypatch):
+    _clear_space_environment(monkeypatch)
+    monkeypatch.setenv("SPACE_HOST", "build-small-hackathon-waterleaf.hf.space")
+    for name in OAUTH_VARIABLES:
+        monkeypatch.setenv(name, f"test-{name.casefold()}")
+
+    configure_hf_space_environment()
+
+    assert is_hf_space()
+    assert has_hf_oauth()
+    assert os.environ["SYSTEM"] == "spaces"
+    assert os.environ["SPACE_ID"] == "build-small-hackathon-waterleaf.hf.space"
+
+
+def test_incomplete_oauth_environment_is_not_enabled(monkeypatch):
+    _clear_space_environment(monkeypatch)
+    monkeypatch.setenv("SPACE_HOST", "build-small-hackathon-waterleaf.hf.space")
+    monkeypatch.setenv("OAUTH_CLIENT_ID", "client")
+
+    configure_hf_space_environment()
+
+    assert is_hf_space()
+    assert not has_hf_oauth()
+    assert os.environ.get("SYSTEM") is None

waterleaf/space.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import os
+
+OAUTH_VARIABLES = (
+    "OAUTH_CLIENT_ID",
+    "OAUTH_CLIENT_SECRET",
+    "OAUTH_SCOPES",
+    "OPENID_PROVIDER_URL",
+)
+
+
+def is_hf_space() -> bool:
+    return bool(os.getenv("SPACE_ID") or os.getenv("SPACE_HOST"))
+
+
+def has_hf_oauth() -> bool:
+    return all(os.getenv(name) for name in OAUTH_VARIABLES)
+
+
+def configure_hf_space_environment() -> None:
+    if not is_hf_space() or not has_hf_oauth():
+        return
+    space_id = os.getenv("SPACE_ID") or os.environ["SPACE_HOST"]
+    os.environ.setdefault("SYSTEM", "spaces")
+    os.environ.setdefault("SPACE_ID", space_id)

waterleaf/ui.py

@@ -19,6 +19,7 @@ from waterleaf.models import (
 )
 from waterleaf.rate_limit import SlidingWindowRateLimiter
 from waterleaf.services.demo import DemoTaxonomy, build_demo_identification
+from waterleaf.space import has_hf_oauth, is_hf_space
 
 CSS = """
 :root {
@@ -83,13 +84,14 @@ def build_ui(
 ) -> gr.Blocks:
     identification = identification or application.identification or build_demo_identification()
     taxonomy = taxonomy or application.taxonomy or DemoTaxonomy()
-    oauth_enabled = bool(os.getenv("SPACE_ID") or os.getenv("OAUTH_CLIENT_ID"))
+    space_environment = is_hf_space()
+    oauth_enabled = has_hf_oauth()
     guest_limiter = SlidingWindowRateLimiter(limit=3, window_seconds=3600)
 
     def owner_from_profile(profile: gr.OAuthProfile | None) -> str | None:
         if profile is not None:
             return profile.username
-        if not oauth_enabled:
+        if not space_environment:
             return os.getenv("WATERLEAF_LOCAL_USER", "local-gardener")
         return None