feat: cloud usage limits (#7192)

Author: evan-onyx

backend/alembic/versions/2b90f3af54b8_usage_limits.py

@@ -0,0 +1,46 @@
+"""usage_limits
+
+Revision ID: 2b90f3af54b8
+Revises: 9a0296d7421e
+Create Date: 2026-01-03 16:55:30.449692
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "2b90f3af54b8"
+down_revision = "9a0296d7421e"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "tenant_usage",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column(
+            "window_start", sa.DateTime(timezone=True), nullable=False, index=True
+        ),
+        sa.Column("llm_cost_cents", sa.Float(), nullable=False, server_default="0.0"),
+        sa.Column("chunks_indexed", sa.Integer(), nullable=False, server_default="0"),
+        sa.Column("api_calls", sa.Integer(), nullable=False, server_default="0"),
+        sa.Column(
+            "non_streaming_api_calls", sa.Integer(), nullable=False, server_default="0"
+        ),
+        sa.Column(
+            "updated_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.func.now(),
+            nullable=True,
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("window_start", name="uq_tenant_usage_window"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_index("ix_tenant_usage_window_start", table_name="tenant_usage")
+    op.drop_table("tenant_usage")

backend/ee/onyx/configs/app_configs.py

@@ -111,17 +111,6 @@
 STRIPE_SECRET_KEY = os.environ.get("STRIPE_SECRET_KEY")
 STRIPE_PRICE_ID = os.environ.get("STRIPE_PRICE")
 
-OPENAI_DEFAULT_API_KEY = os.environ.get("OPENAI_DEFAULT_API_KEY")
-ANTHROPIC_DEFAULT_API_KEY = os.environ.get("ANTHROPIC_DEFAULT_API_KEY")
-COHERE_DEFAULT_API_KEY = os.environ.get("COHERE_DEFAULT_API_KEY")
-
-# Vertex AI default credentials (JSON content or path to JSON file)
-VERTEXAI_DEFAULT_CREDENTIALS = os.environ.get("VERTEXAI_DEFAULT_CREDENTIALS")
-VERTEXAI_DEFAULT_LOCATION = os.environ.get("VERTEXAI_DEFAULT_LOCATION", "global")
-
-# OpenRouter default API key
-OPENROUTER_DEFAULT_API_KEY = os.environ.get("OPENROUTER_DEFAULT_API_KEY")
-
 # JWT Public Key URL
 JWT_PUBLIC_KEY_URL: str | None = os.getenv("JWT_PUBLIC_KEY_URL", None)
 

backend/ee/onyx/server/tenants/provisioning.py

@@ -9,13 +9,7 @@
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
-from ee.onyx.configs.app_configs import ANTHROPIC_DEFAULT_API_KEY
-from ee.onyx.configs.app_configs import COHERE_DEFAULT_API_KEY
 from ee.onyx.configs.app_configs import HUBSPOT_TRACKING_URL
-from ee.onyx.configs.app_configs import OPENAI_DEFAULT_API_KEY
-from ee.onyx.configs.app_configs import OPENROUTER_DEFAULT_API_KEY
-from ee.onyx.configs.app_configs import VERTEXAI_DEFAULT_CREDENTIALS
-from ee.onyx.configs.app_configs import VERTEXAI_DEFAULT_LOCATION
 from ee.onyx.server.tenants.access import generate_data_plane_token
 from ee.onyx.server.tenants.models import TenantByDomainResponse
 from ee.onyx.server.tenants.models import TenantCreationPayload
@@ -27,8 +21,14 @@
 from ee.onyx.server.tenants.user_mapping import get_tenant_id_for_email
 from ee.onyx.server.tenants.user_mapping import user_owns_a_tenant
 from onyx.auth.users import exceptions
+from onyx.configs.app_configs import ANTHROPIC_DEFAULT_API_KEY
+from onyx.configs.app_configs import COHERE_DEFAULT_API_KEY
 from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL
 from onyx.configs.app_configs import DEV_MODE
+from onyx.configs.app_configs import OPENAI_DEFAULT_API_KEY
+from onyx.configs.app_configs import OPENROUTER_DEFAULT_API_KEY
+from onyx.configs.app_configs import VERTEXAI_DEFAULT_CREDENTIALS
+from onyx.configs.app_configs import VERTEXAI_DEFAULT_LOCATION
 from onyx.configs.constants import MilestoneRecordType
 from onyx.db.engine.sql_engine import get_session_with_shared_schema
 from onyx.db.engine.sql_engine import get_session_with_tenant

backend/ee/onyx/server/usage_limits.py

@@ -0,0 +1,47 @@
+"""EE Usage limits - trial detection via billing information."""
+
+from datetime import datetime
+from datetime import timezone
+
+from ee.onyx.server.tenants.billing import fetch_billing_information
+from ee.onyx.server.tenants.models import BillingInformation
+from ee.onyx.server.tenants.models import SubscriptionStatusResponse
+from onyx.utils.logger import setup_logger
+from shared_configs.configs import MULTI_TENANT
+
+logger = setup_logger()
+
+
+def is_tenant_on_trial(tenant_id: str) -> bool:
+    """
+    Determine if a tenant is currently on a trial subscription.
+
+    In multi-tenant mode, we fetch billing information from the control plane
+    to determine if the tenant has an active trial.
+    """
+    if not MULTI_TENANT:
+        return False
+
+    try:
+        billing_info = fetch_billing_information(tenant_id)
+
+        # If not subscribed at all, check if we have trial information
+        if isinstance(billing_info, SubscriptionStatusResponse):
+            # No subscription means they're likely on trial (new tenant)
+            return True
+
+        if isinstance(billing_info, BillingInformation):
+            # Check if trial is active
+            if billing_info.trial_end is not None:
+                now = datetime.now(timezone.utc)
+                # Trial active if trial_end is in the future
+                # and subscription status indicates trialing
+                if billing_info.trial_end > now and billing_info.status == "trialing":
+                    return True
+
+        return False
+
+    except Exception as e:
+        logger.warning(f"Failed to fetch billing info for trial check: {e}")
+        # Default to trial limits on error (more restrictive = safer)
+        return True

backend/onyx/auth/captcha.py

@@ -0,0 +1,107 @@
+"""Captcha verification for user registration."""
+
+import httpx
+from pydantic import BaseModel
+from pydantic import Field
+
+from onyx.configs.app_configs import CAPTCHA_ENABLED
+from onyx.configs.app_configs import RECAPTCHA_SCORE_THRESHOLD
+from onyx.configs.app_configs import RECAPTCHA_SECRET_KEY
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+RECAPTCHA_VERIFY_URL = "https://www.google.com/recaptcha/api/siteverify"
+
+
+class CaptchaVerificationError(Exception):
+    """Raised when captcha verification fails."""
+
+
+class RecaptchaResponse(BaseModel):
+    """Response from Google reCAPTCHA verification API."""
+
+    success: bool
+    score: float | None = None  # Only present for reCAPTCHA v3
+    action: str | None = None
+    challenge_ts: str | None = None
+    hostname: str | None = None
+    error_codes: list[str] | None = Field(default=None, alias="error-codes")
+
+
+def is_captcha_enabled() -> bool:
+    """Check if captcha verification is enabled."""
+    return CAPTCHA_ENABLED and bool(RECAPTCHA_SECRET_KEY)
+
+
+async def verify_captcha_token(
+    token: str,
+    expected_action: str = "signup",
+) -> None:
+    """
+    Verify a reCAPTCHA token with Google's API.
+
+    Args:
+        token: The reCAPTCHA response token from the client
+        expected_action: Expected action name for v3 verification
+
+    Raises:
+        CaptchaVerificationError: If verification fails
+    """
+    if not is_captcha_enabled():
+        return
+
+    if not token:
+        raise CaptchaVerificationError("Captcha token is required")
+
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                RECAPTCHA_VERIFY_URL,
+                data={
+                    "secret": RECAPTCHA_SECRET_KEY,
+                    "response": token,
+                },
+                timeout=10.0,
+            )
+            response.raise_for_status()
+
+            data = response.json()
+            result = RecaptchaResponse(**data)
+
+            if not result.success:
+                error_codes = result.error_codes or ["unknown-error"]
+                logger.warning(f"Captcha verification failed: {error_codes}")
+                raise CaptchaVerificationError(
+                    f"Captcha verification failed: {', '.join(error_codes)}"
+                )
+
+            # For reCAPTCHA v3, also check the score
+            if result.score is not None:
+                if result.score < RECAPTCHA_SCORE_THRESHOLD:
+                    logger.warning(
+                        f"Captcha score too low: {result.score} < {RECAPTCHA_SCORE_THRESHOLD}"
+                    )
+                    raise CaptchaVerificationError(
+                        "Captcha verification failed: suspicious activity detected"
+                    )
+
+                # Optionally verify the action matches
+                if result.action and result.action != expected_action:
+                    logger.warning(
+                        f"Captcha action mismatch: {result.action} != {expected_action}"
+                    )
+                    raise CaptchaVerificationError(
+                        "Captcha verification failed: action mismatch"
+                    )
+
+            logger.debug(
+                f"Captcha verification passed: score={result.score}, "
+                f"action={result.action}"
+            )
+
+    except httpx.HTTPError as e:
+        logger.error(f"Captcha API request failed: {e}")
+        # In case of API errors, we might want to allow registration
+        # to prevent blocking legitimate users. This is a policy decision.
+        raise CaptchaVerificationError("Captcha verification service unavailable")

backend/onyx/auth/schemas.py

@@ -40,6 +40,8 @@ class UserRead(schemas.BaseUser[uuid.UUID]):
 class UserCreate(schemas.BaseUserCreate):
     role: UserRole = UserRole.BASIC
     tenant_id: str | None = None
+    # Captcha token for cloud signup protection (optional, only used when captcha is enabled)
+    captcha_token: str | None = None
 
 
 class UserUpdateWithRole(schemas.BaseUserUpdate):

backend/onyx/auth/users.py

@@ -292,6 +292,31 @@ async def create(
         safe: bool = False,
         request: Optional[Request] = None,
     ) -> User:
+        # Verify captcha if enabled (for cloud signup protection)
+        from onyx.auth.captcha import CaptchaVerificationError
+        from onyx.auth.captcha import is_captcha_enabled
+        from onyx.auth.captcha import verify_captcha_token
+
+        if is_captcha_enabled() and request is not None:
+            # Get captcha token from request body or headers
+            captcha_token = None
+            if hasattr(user_create, "captcha_token"):
+                captcha_token = getattr(user_create, "captcha_token", None)
+
+            # Also check headers as a fallback
+            if not captcha_token:
+                captcha_token = request.headers.get("X-Captcha-Token")
+
+            try:
+                await verify_captcha_token(
+                    captcha_token or "", expected_action="signup"
+                )
+            except CaptchaVerificationError as e:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail={"reason": str(e)},
+                )
+
         # We verify the password here to make sure it's valid before we proceed
         await self.validate_password(
             user_create.password, cast(schemas.UC, user_create)

backend/onyx/background/celery/tasks/docprocessing/tasks.py

@@ -86,6 +86,7 @@
 from onyx.db.search_settings import get_current_search_settings
 from onyx.db.search_settings import get_secondary_search_settings
 from onyx.db.swap_index import check_and_perform_index_swap
+from onyx.db.usage import UsageLimitExceededError
 from onyx.document_index.factory import get_default_document_index
 from onyx.file_store.document_batch_storage import DocumentBatchStorage
 from onyx.file_store.document_batch_storage import get_document_batch_storage
@@ -112,6 +113,7 @@
 from shared_configs.configs import INDEXING_MODEL_SERVER_HOST
 from shared_configs.configs import INDEXING_MODEL_SERVER_PORT
 from shared_configs.configs import MULTI_TENANT
+from shared_configs.configs import USAGE_LIMITS_ENABLED
 from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
 from shared_configs.contextvars import INDEX_ATTEMPT_INFO_CONTEXTVAR
 
@@ -1279,6 +1281,31 @@ def docprocessing_task(
         INDEX_ATTEMPT_INFO_CONTEXTVAR.reset(token)
 
 
+def _check_chunk_usage_limit(tenant_id: str) -> None:
+    """Check if chunk indexing usage limit has been exceeded.
+
+    Raises UsageLimitExceededError if the limit is exceeded.
+    """
+    if not USAGE_LIMITS_ENABLED:
+        return
+
+    from onyx.db.usage import check_usage_limit
+    from onyx.db.usage import UsageType
+    from onyx.server.usage_limits import get_limit_for_usage_type
+    from onyx.server.usage_limits import is_tenant_on_trial
+
+    is_trial = is_tenant_on_trial(tenant_id)
+    limit = get_limit_for_usage_type(UsageType.CHUNKS_INDEXED, is_trial)
+
+    with get_session_with_current_tenant() as db_session:
+        check_usage_limit(
+            db_session=db_session,
+            usage_type=UsageType.CHUNKS_INDEXED,
+            limit=limit,
+            pending_amount=0,  # Just check current usage
+        )
+
+
 def _docprocessing_task(
     index_attempt_id: int,
     cc_pair_id: int,
@@ -1290,6 +1317,25 @@ def _docprocessing_task(
     if tenant_id:
         CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
 
+    # Check if chunk indexing usage limit has been exceeded before processing
+    if USAGE_LIMITS_ENABLED:
+        try:
+            _check_chunk_usage_limit(tenant_id)
+        except UsageLimitExceededError as e:
+            # Log the error and fail the indexing attempt
+            task_logger.error(
+                f"Chunk indexing usage limit exceeded for tenant {tenant_id}: {e}"
+            )
+            with get_session_with_current_tenant() as db_session:
+                from onyx.db.index_attempt import mark_attempt_failed
+
+                mark_attempt_failed(
+                    index_attempt_id=index_attempt_id,
+                    db_session=db_session,
+                    failure_reason=str(e),
+                )
+            raise
+
     task_logger.info(
         f"Processing document batch: "
         f"attempt={index_attempt_id} "
@@ -1434,6 +1480,23 @@ def _docprocessing_task(
                 adapter=adapter,
             )
 
+        # Track chunk indexing usage for cloud usage limits
+        if USAGE_LIMITS_ENABLED and index_pipeline_result.total_chunks > 0:
+            try:
+                from onyx.db.usage import increment_usage
+                from onyx.db.usage import UsageType
+
+                with get_session_with_current_tenant() as usage_db_session:
+                    increment_usage(
+                        db_session=usage_db_session,
+                        usage_type=UsageType.CHUNKS_INDEXED,
+                        amount=index_pipeline_result.total_chunks,
+                    )
+                    usage_db_session.commit()
+            except Exception as e:
+                # Log but don't fail indexing if usage tracking fails
+                task_logger.warning(f"Failed to track chunk indexing usage: {e}")
+
         # Update batch completion and document counts atomically using database coordination
 
         with get_session_with_current_tenant() as db_session, cross_batch_db_lock: