diff --git a/backend/onyx/server/features/build/rate_limit.py b/backend/onyx/server/features/build/rate_limit.py index 5750ac408d..299536ef9a 100644 --- a/backend/onyx/server/features/build/rate_limit.py +++ b/backend/onyx/server/features/build/rate_limit.py @@ -8,6 +8,8 @@ from typing import Literal from pydantic import BaseModel from sqlalchemy.orm import Session +from onyx.configs.app_configs import AUTH_TYPE +from onyx.configs.constants import AuthType from onyx.db.models import User from onyx.feature_flags.factory import get_default_feature_flag_provider from onyx.server.features.build.configs import CRAFT_PAID_USER_RATE_LIMIT @@ -17,6 +19,7 @@ from onyx.server.features.build.db.rate_limit import get_oldest_message_timestam from onyx.server.features.build.subscription_check import is_user_subscribed from onyx.server.features.build.utils import CRAFT_HAS_USAGE_LIMITS from shared_configs.configs import MULTI_TENANT +from shared_configs.contextvars import get_current_tenant_id # Default limit for free/non-subscribed users (not configurable) FREE_USER_RATE_LIMIT = 5 @@ -36,8 +39,9 @@ def _should_skip_rate_limiting(user: User) -> bool: """ Check if rate limiting should be skipped for this user. - Currently grants unlimited usage to dev tenant users (tenant_dev). - Controlled via PostHog feature flag. + Grants unlimited usage to tenants/users whose `craft-has-usage-limits` + PostHog flag is off (e.g. tenant_dev). Only reached on Onyx Cloud, where + PostHog is configured. Returns: True to skip rate limiting (unlimited), False to apply normal limits @@ -49,9 +53,10 @@ def _should_skip_rate_limiting(user: User) -> bool: feature_flag_provider = get_default_feature_flag_provider() # Flag returns True for users who SHOULD be rate limited # We negate to get: True = skip rate limiting - has_rate_limit = feature_flag_provider.feature_enabled( + has_rate_limit = feature_flag_provider.feature_enabled_for_user_tenant( CRAFT_HAS_USAGE_LIMITS, - user.id, + user, + get_current_tenant_id(), ) return not has_rate_limit @@ -63,14 +68,13 @@ def get_user_rate_limit_status( """ Get the rate limit status for a user. - Rate limits: - - Cloud (MULTI_TENANT=true): - - Subscribed users: CRAFT_PAID_USER_RATE_LIMIT messages per week - (configurable, default 25) - - Non-subscribed users: 5 messages (lifetime total) - - Per-user overrides via PostHog feature flag - - Self-hosted (MULTI_TENANT=false): - - Unlimited (no rate limiting) + Rate limits apply on Onyx Cloud only: + - Subscribed users: CRAFT_PAID_USER_RATE_LIMIT messages per week + (configurable, default 25) + - Non-subscribed users: 5 messages (lifetime total) + - Per-tenant/user overrides via the `craft-has-usage-limits` PostHog flag + + All non-cloud deployments (self-hosted, managed) are unlimited. Args: user: The authenticated user @@ -79,8 +83,9 @@ def get_user_rate_limit_status( Returns: RateLimitResponse with current limit status """ - # Self-hosted deployments have no rate limits - if not MULTI_TENANT: + # Rate limits apply only on Onyx Cloud (multi-tenant); all other deployments + # are unlimited. + if not MULTI_TENANT or AUTH_TYPE != AuthType.CLOUD: return RateLimitResponse( is_limited=False, limit_type="weekly", diff --git a/backend/onyx/server/features/build/utils.py b/backend/onyx/server/features/build/utils.py index be365793d5..d124157ccc 100644 --- a/backend/onyx/server/features/build/utils.py +++ b/backend/onyx/server/features/build/utils.py @@ -13,6 +13,7 @@ from onyx.feature_flags.interface import NoOpFeatureFlagProvider from onyx.server.features.build.configs import ENABLE_CRAFT from onyx.server.features.build.configs import MAX_UPLOAD_FILE_SIZE_BYTES from onyx.utils.logger import setup_logger +from shared_configs.contextvars import get_current_tenant_id logger = setup_logger() @@ -126,10 +127,10 @@ def is_onyx_craft_enabled(user: User) -> bool: if isinstance(feature_flag_provider, NoOpFeatureFlagProvider): return ENABLE_CRAFT - # Use the feature flag provider - is_enabled = feature_flag_provider.feature_enabled( + is_enabled = feature_flag_provider.feature_enabled_for_user_tenant( ONYX_CRAFT_ENABLED_FLAG, - user.id, + user, + get_current_tenant_id(), ) if is_enabled: diff --git a/backend/tests/unit/onyx/server/features/build/test_feature_gate.py b/backend/tests/unit/onyx/server/features/build/test_feature_gate.py index bd0c5c96ee..201cc0a090 100644 --- a/backend/tests/unit/onyx/server/features/build/test_feature_gate.py +++ b/backend/tests/unit/onyx/server/features/build/test_feature_gate.py @@ -23,26 +23,32 @@ from onyx.server.features.build.utils import is_onyx_craft_enabled class _StubPostHogProvider(FeatureFlagProvider): - """A non-NoOp provider that returns a fixed answer for the craft flag.""" + """A non-NoOp provider that returns a fixed answer for the craft flag. + + Inherits the base `feature_enabled_for_user_tenant`, which is the method the + gate calls — so `calls` captures the `user_properties` (including + `tenant_id`) that get forwarded to PostHog. + """ def __init__(self, enabled: bool) -> None: self._enabled = enabled - self.calls: list[tuple[str, UUID]] = [] + self.calls: list[tuple[str, UUID, dict[str, Any] | None]] = [] def feature_enabled( self, flag_key: str, user_id: UUID, - user_properties: dict[str, Any] | None = None, # noqa: ARG002 + user_properties: dict[str, Any] | None = None, ) -> bool: - self.calls.append((flag_key, user_id)) + self.calls.append((flag_key, user_id, user_properties)) return self._enabled def _make_user() -> MagicMock: - """Build a minimal stand-in for `User` - only `.id` is read.""" + """Build a minimal stand-in for `User` - only `.id` and `.email` are read.""" user = MagicMock() user.id = uuid4() + user.email = "user@tenant-dev.example" return user @@ -70,6 +76,15 @@ def test_enabled_via_env_when_no_flag_provider( "get_default_feature_flag_provider", lambda: NoOpFeatureFlagProvider(), ) + # The env/NoOp path must not touch the tenant contextvar, which can raise + # in multi-tenant mode when no tenant is set. + monkeypatch.setattr( + build_utils, + "get_current_tenant_id", + lambda: pytest.fail( + "get_current_tenant_id should not be called on the env path" + ), + ) assert is_onyx_craft_enabled(_make_user()) is True @@ -77,6 +92,7 @@ def test_enabled_via_env_when_no_flag_provider( def test_posthog_flag_overrides_env(monkeypatch: pytest.MonkeyPatch) -> None: """A real provider's verdict wins regardless of ENABLE_CRAFT.""" monkeypatch.setattr(build_utils, "ENABLE_CRAFT", False) + monkeypatch.setattr(build_utils, "get_current_tenant_id", lambda: "tenant_dev") provider = _StubPostHogProvider(enabled=True) monkeypatch.setattr( build_utils, "get_default_feature_flag_provider", lambda: provider @@ -85,4 +101,30 @@ def test_posthog_flag_overrides_env(monkeypatch: pytest.MonkeyPatch) -> None: user = _make_user() assert is_onyx_craft_enabled(user) is True # The provider was consulted with the craft-enabled flag key for this user. - assert provider.calls == [("onyx-craft-enabled", user.id)] + assert provider.calls == [ + ( + "onyx-craft-enabled", + user.id, + {"tenant_id": "tenant_dev", "email": "user@tenant-dev.example"}, + ) + ] + + +def test_posthog_flag_disables_when_false(monkeypatch: pytest.MonkeyPatch) -> None: + """A real provider returning False disables Craft even if ENABLE_CRAFT=True.""" + monkeypatch.setattr(build_utils, "ENABLE_CRAFT", True) + monkeypatch.setattr(build_utils, "get_current_tenant_id", lambda: "tenant_other") + provider = _StubPostHogProvider(enabled=False) + monkeypatch.setattr( + build_utils, "get_default_feature_flag_provider", lambda: provider + ) + + user = _make_user() + assert is_onyx_craft_enabled(user) is False + assert provider.calls == [ + ( + "onyx-craft-enabled", + user.id, + {"tenant_id": "tenant_other", "email": "user@tenant-dev.example"}, + ) + ]