From e487144ebd4ca9934d62854ab819b756df58eb88 Mon Sep 17 00:00:00 2001
From: pyr0ball <pyroballpcs@gmail.com>
Date: Sun, 22 Mar 2026 15:31:54 -0700
Subject: [PATCH] feat(tiers): add vue_ui_beta feature key and demo_tier kwarg
 to can_use

---
 app/wizard/tiers.py        | 23 +++++++++++++--
 tests/test_wizard_tiers.py | 57 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/app/wizard/tiers.py b/app/wizard/tiers.py
index 9679843..fa1310d 100644
--- a/app/wizard/tiers.py
+++ b/app/wizard/tiers.py
@@ -22,6 +22,7 @@ Features that stay gated even with BYOK:
 """
 from __future__ import annotations
 
+import os as _os
 from pathlib import Path
 
 TIERS = ["free", "paid", "premium"]
@@ -58,6 +59,9 @@ FEATURES: dict[str, str] = {
     "google_calendar_sync":         "paid",
     "apple_calendar_sync":          "paid",
     "slack_notifications":          "paid",
+
+    # Beta UI access — stays gated (access management, not compute)
+    "vue_ui_beta":                  "paid",
 }
 
 # Features that unlock when the user supplies any LLM backend (local or BYOK).
@@ -75,6 +79,10 @@ BYOK_UNLOCKABLE: frozenset[str] = frozenset({
     "survey_assistant",
 })
 
+# Demo mode flag — read from environment at module load time.
+# Allows demo toolbar to override tier without accessing st.session_state (thread-safe).
+_DEMO_MODE = _os.environ.get("DEMO_MODE", "").lower() in ("1", "true", "yes")
+
 # Free integrations (not in FEATURES):
 # google_drive_sync, dropbox_sync, onedrive_sync, mega_sync,
 # nextcloud_sync, discord_notifications, home_assistant
@@ -101,22 +109,33 @@ def has_configured_llm(config_path: Path | None = None) -> bool:
         return False
 
 
-def can_use(tier: str, feature: str, has_byok: bool = False) -> bool:
+def can_use(
+    tier: str,
+    feature: str,
+    has_byok: bool = False,
+    *,
+    demo_tier: str | None = None,
+) -> bool:
     """Return True if the given tier has access to the feature.
 
     has_byok: pass has_configured_llm() to unlock BYOK_UNLOCKABLE features
     for users who supply their own LLM backend regardless of tier.
 
+    demo_tier: when set AND _DEMO_MODE is True, substitutes for `tier`.
+               Read from st.session_state by the *caller*, not here — keeps
+               this function thread-safe for background tasks and tests.
+
     Returns True for unknown features (not gated).
     Returns False for unknown/invalid tier strings.
     """
+    effective_tier = demo_tier if (demo_tier is not None and _DEMO_MODE) else tier
     required = FEATURES.get(feature)
     if required is None:
         return True  # not gated — available to all
     if has_byok and feature in BYOK_UNLOCKABLE:
         return True
     try:
-        return TIERS.index(tier) >= TIERS.index(required)
+        return TIERS.index(effective_tier) >= TIERS.index(required)
     except ValueError:
         return False  # invalid tier string
 
diff --git a/tests/test_wizard_tiers.py b/tests/test_wizard_tiers.py
index 325f0b5..3c33ea2 100644
--- a/tests/test_wizard_tiers.py
+++ b/tests/test_wizard_tiers.py
@@ -112,3 +112,60 @@ def test_byok_false_preserves_original_gating():
     # has_byok=False (default) must not change existing behaviour
     assert can_use("free", "company_research", has_byok=False) is False
     assert can_use("paid", "company_research", has_byok=False) is True
+
+
+# ── Vue UI Beta & Demo Tier tests ──────────────────────────────────────────────
+
+def test_vue_ui_beta_free_tier():
+    assert can_use("free", "vue_ui_beta") is False
+
+
+def test_vue_ui_beta_paid_tier():
+    assert can_use("paid", "vue_ui_beta") is True
+
+
+def test_vue_ui_beta_premium_tier():
+    assert can_use("premium", "vue_ui_beta") is True
+
+
+def test_can_use_demo_tier_overrides_real_tier():
+    # demo_tier kwarg substitutes for the real tier when provided and DEMO_MODE is set
+    import os
+    os.environ["DEMO_MODE"] = "true"
+    # Need to reload the module to pick up the new DEMO_MODE value
+    import importlib
+    import app.wizard.tiers as tiers_module
+    importlib.reload(tiers_module)
+    try:
+        assert tiers_module.can_use("free", "company_research", demo_tier="paid") is True
+    finally:
+        # Cleanup: restore original state
+        os.environ.pop("DEMO_MODE", None)
+        importlib.reload(tiers_module)
+
+
+def test_can_use_demo_tier_free_restricts():
+    # demo_tier="free" should restrict access even if real tier is "paid"
+    import os
+    os.environ["DEMO_MODE"] = "true"
+    import importlib
+    import app.wizard.tiers as tiers_module
+    importlib.reload(tiers_module)
+    try:
+        assert tiers_module.can_use("paid", "model_fine_tuning", demo_tier="free") is False
+    finally:
+        os.environ.pop("DEMO_MODE", None)
+        importlib.reload(tiers_module)
+
+
+def test_can_use_demo_tier_none_falls_back_to_real():
+    # demo_tier=None means no override — real tier is used
+    assert can_use("paid", "company_research", demo_tier=None) is True
+
+
+def test_can_use_demo_tier_does_not_affect_non_demo():
+    # demo_tier is only applied when DEMO_MODE env var is set;
+    # in tests DEMO_MODE is False by default, so demo_tier is ignored
+    import os
+    os.environ.pop("DEMO_MODE", None)
+    assert can_use("free", "company_research", demo_tier="paid") is False