Billing fixes (#3976)

* misct billing_fixes
g

* various billing updates

* improvements

* quick nit

* k

* update

* k

* k

* k

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

@@ -18,11 +18,15 @@
 from ee.onyx.server.tenants.anonymous_user_path import modify_anonymous_user_path
 from ee.onyx.server.tenants.anonymous_user_path import validate_anonymous_user_path
 from ee.onyx.server.tenants.billing import fetch_billing_information
+from ee.onyx.server.tenants.billing import fetch_stripe_checkout_session
 from ee.onyx.server.tenants.billing import fetch_tenant_stripe_information
 from ee.onyx.server.tenants.models import AnonymousUserPath
 from ee.onyx.server.tenants.models import BillingInformation
 from ee.onyx.server.tenants.models import ImpersonateRequest
 from ee.onyx.server.tenants.models import ProductGatingRequest
+from ee.onyx.server.tenants.models import ProductGatingResponse
+from ee.onyx.server.tenants.models import SubscriptionSessionResponse
+from ee.onyx.server.tenants.models import SubscriptionStatusResponse
 from ee.onyx.server.tenants.provisioning import delete_user_from_control_plane
 from ee.onyx.server.tenants.user_mapping import get_tenant_id_for_email
 from ee.onyx.server.tenants.user_mapping import remove_all_users_from_tenant
@@ -39,7 +43,6 @@
 from onyx.db.engine import get_current_tenant_id
 from onyx.db.engine import get_session
 from onyx.db.engine import get_session_with_tenant
-from onyx.db.notification import create_notification
 from onyx.db.users import delete_user_from_db
 from onyx.db.users import get_user_by_email
 from onyx.server.manage.models import UserByEmail
@@ -126,37 +129,38 @@ async def login_as_anonymous_user(
 @router.post("/product-gating")
 def gate_product(
     product_gating_request: ProductGatingRequest, _: None = Depends(control_plane_dep)
-) -> None:
+) -> ProductGatingResponse:
     """
     Gating the product means that the product is not available to the tenant.
     They will be directed to the billing page.
     We gate the product when
     1) User has ended free trial without adding payment method
     2) User's card has declined
     """
-    tenant_id = product_gating_request.tenant_id
-    token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
+    try:
+        tenant_id = product_gating_request.tenant_id
+        token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
 
-    settings = load_settings()
-    settings.product_gating = product_gating_request.product_gating
-    store_settings(settings)
+        settings = load_settings()
+        settings.application_status = product_gating_request.application_status
+        store_settings(settings)
 
-    if product_gating_request.notification:
-        with get_session_with_tenant(tenant_id) as db_session:
-            create_notification(None, product_gating_request.notification, db_session)
+        if token is not None:
+            CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
 
-    if token is not None:
-        CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
+        return ProductGatingResponse(updated=True, error=None)
+
+    except Exception as e:
+        logger.exception("Failed to gate product")
+        return ProductGatingResponse(updated=False, error=str(e))
 
 
-@router.get("/billing-information", response_model=BillingInformation)
+@router.get("/billing-information")
 async def billing_information(
     _: User = Depends(current_admin_user),
-) -> BillingInformation:
+) -> BillingInformation | SubscriptionStatusResponse:
     logger.info("Fetching billing information")
-    return BillingInformation(
-        **fetch_billing_information(CURRENT_TENANT_ID_CONTEXTVAR.get())
-    )
+    return fetch_billing_information(CURRENT_TENANT_ID_CONTEXTVAR.get())
 
 
 @router.post("/create-customer-portal-session")
@@ -169,9 +173,10 @@ async def create_customer_portal_session(_: User = Depends(current_admin_user))
         if not stripe_customer_id:
             raise HTTPException(status_code=400, detail="Stripe customer ID not found")
         logger.info(stripe_customer_id)
+
         portal_session = stripe.billing_portal.Session.create(
             customer=stripe_customer_id,
-            return_url=f"{WEB_DOMAIN}/admin/cloud-settings",
+            return_url=f"{WEB_DOMAIN}/admin/billing",
         )
         logger.info(portal_session)
         return {"url": portal_session.url}
@@ -180,6 +185,20 @@ async def create_customer_portal_session(_: User = Depends(current_admin_user))
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@router.post("/create-subscription-session")
+async def create_subscription_session(
+    _: User = Depends(current_admin_user),
+) -> SubscriptionSessionResponse:
+    try:
+        tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()
+        session_id = fetch_stripe_checkout_session(tenant_id)
+        return SubscriptionSessionResponse(sessionId=session_id)
+
+    except Exception as e:
+        logger.exception("Failed to create resubscription session")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @router.post("/impersonate")
 async def impersonate_user(
     impersonate_request: ImpersonateRequest,

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

@@ -6,6 +6,7 @@
 from ee.onyx.configs.app_configs import STRIPE_PRICE_ID
 from ee.onyx.configs.app_configs import STRIPE_SECRET_KEY
 from ee.onyx.server.tenants.access import generate_data_plane_token
+from ee.onyx.server.tenants.models import BillingInformation
 from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL
 from onyx.utils.logger import setup_logger
 
@@ -14,6 +15,19 @@
 logger = setup_logger()
 
 
+def fetch_stripe_checkout_session(tenant_id: str) -> str:
+    token = generate_data_plane_token()
+    headers = {
+        "Authorization": f"Bearer {token}",
+        "Content-Type": "application/json",
+    }
+    url = f"{CONTROL_PLANE_API_BASE_URL}/create-checkout-session"
+    params = {"tenant_id": tenant_id}
+    response = requests.post(url, headers=headers, params=params)
+    response.raise_for_status()
+    return response.json()["sessionId"]
+
+
 def fetch_tenant_stripe_information(tenant_id: str) -> dict:
     token = generate_data_plane_token()
     headers = {
@@ -27,7 +41,7 @@ def fetch_tenant_stripe_information(tenant_id: str) -> dict:
     return response.json()
 
 
-def fetch_billing_information(tenant_id: str) -> dict:
+def fetch_billing_information(tenant_id: str) -> BillingInformation:
     logger.info("Fetching billing information")
     token = generate_data_plane_token()
     headers = {
@@ -38,7 +52,7 @@ def fetch_billing_information(tenant_id: str) -> dict:
     params = {"tenant_id": tenant_id}
     response = requests.get(url, headers=headers, params=params)
     response.raise_for_status()
-    billing_info = response.json()
+    billing_info = BillingInformation(**response.json())
     return billing_info
 
 

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

@@ -1,7 +1,8 @@
+from datetime import datetime
+
 from pydantic import BaseModel
 
-from onyx.configs.constants import NotificationType
-from onyx.server.settings.models import GatingType
+from onyx.server.settings.models import ApplicationStatus
 
 
 class CheckoutSessionCreationRequest(BaseModel):
@@ -15,15 +16,24 @@ class CreateTenantRequest(BaseModel):
 
 class ProductGatingRequest(BaseModel):
     tenant_id: str
-    product_gating: GatingType
-    notification: NotificationType | None = None
+    application_status: ApplicationStatus
+
+
+class SubscriptionStatusResponse(BaseModel):
+    subscribed: bool
 
 
 class BillingInformation(BaseModel):
+    stripe_subscription_id: str
+    status: str
+    current_period_start: datetime
+    current_period_end: datetime
+    number_of_seats: int
+    cancel_at_period_end: bool
+    canceled_at: datetime | None
+    trial_start: datetime | None
+    trial_end: datetime | None
     seats: int
-    subscription_status: str
-    billing_start: str
-    billing_end: str
     payment_method_enabled: bool
 
 
@@ -48,3 +58,12 @@ class TenantDeletionPayload(BaseModel):
 
 class AnonymousUserPath(BaseModel):
     anonymous_user_path: str | None
+
+
+class ProductGatingResponse(BaseModel):
+    updated: bool
+    error: str | None
+
+
+class SubscriptionSessionResponse(BaseModel):
+    sessionId: str

backend/onyx/server/settings/models.py

@@ -12,10 +12,10 @@ class PageType(str, Enum):
     SEARCH = "search"
 
 
-class GatingType(str, Enum):
-    FULL = "full"  # Complete restriction of access to the product or service
-    PARTIAL = "partial"  # Full access but warning (no credit card on file)
-    NONE = "none"  # No restrictions, full access to all features
+class ApplicationStatus(str, Enum):
+    PAYMENT_REMINDER = "payment_reminder"
+    GATED_ACCESS = "gated_access"
+    ACTIVE = "active"
 
 
 class Notification(BaseModel):
@@ -43,7 +43,7 @@ class Settings(BaseModel):
 
     maximum_chat_retention_days: int | None = None
     gpu_enabled: bool | None = None
-    product_gating: GatingType = GatingType.NONE
+    application_status: ApplicationStatus = ApplicationStatus.ACTIVE
     anonymous_user_enabled: bool | None = None
     pro_search_disabled: bool | None = None
     auto_scroll: bool | None = None

web/src/app/admin/settings/interfaces.ts

@@ -1,7 +1,7 @@
-export enum GatingType {
-  FULL = "full",
-  PARTIAL = "partial",
-  NONE = "none",
+export enum ApplicationStatus {
+  PAYMENT_REMINDER = "payment_reminder",
+  GATED_ACCESS = "gated_access",
+  ACTIVE = "active",
 }
 
 export interface Settings {
@@ -11,7 +11,7 @@ export interface Settings {
   needs_reindexing: boolean;
   gpu_enabled: boolean;
   pro_search_disabled: boolean | null;
-  product_gating: GatingType;
+  application_status: ApplicationStatus;
   auto_scroll: boolean;
 }
 

web/src/app/chat/ChatPage.tsx

@@ -2291,8 +2291,6 @@ export function ChatPage({
                 bg-opacity-80
                 duration-300
                 ease-in-out
-
-                
                 ${
                   !untoggled && (showHistorySidebar || sidebarVisible)
                     ? "opacity-100 w-[250px] translate-x-0"

web/src/app/ee/admin/billing/BillingAlerts.tsx

@@ -0,0 +1,73 @@
+import React from "react";
+import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert";
+import { CircleAlert, Info } from "lucide-react";
+import { BillingInformation, BillingStatus } from "./interfaces";
+
+export function BillingAlerts({
+  billingInformation,
+}: {
+  billingInformation: BillingInformation;
+}) {
+  const isTrialing = billingInformation.status === BillingStatus.TRIALING;
+  const isCancelled = billingInformation.cancel_at_period_end;
+  const isExpired =
+    new Date(billingInformation.current_period_end) < new Date();
+  const noPaymentMethod = !billingInformation.payment_method_enabled;
+
+  const messages: string[] = [];
+
+  if (isExpired) {
+    messages.push(
+      "Your subscription has expired. Please resubscribe to continue using the service."
+    );
+  }
+  if (isCancelled && !isExpired) {
+    messages.push(
+      `Your subscription will cancel on ${new Date(
+        billingInformation.current_period_end
+      ).toLocaleDateString()}. You can resubscribe before this date to remain uninterrupted.`
+    );
+  }
+  if (isTrialing) {
+    messages.push(
+      `You're currently on a trial. Your trial ends on ${
+        billingInformation.trial_end
+          ? new Date(billingInformation.trial_end).toLocaleDateString()
+          : "N/A"
+      }.`
+    );
+  }
+  if (noPaymentMethod) {
+    messages.push(
+      "You currently have no payment method on file. Please add one to avoid service interruption."
+    );
+  }
+
+  const variant = isExpired || noPaymentMethod ? "destructive" : "default";
+
+  if (messages.length === 0) return null;
+
+  return (
+    <Alert variant={variant}>
+      <AlertTitle className="flex items-center space-x-2">
+        {variant === "destructive" ? (
+          <CircleAlert className="h-4 w-4" />
+        ) : (
+          <Info className="h-4 w-4" />
+        )}
+        <span>
+          {variant === "destructive"
+            ? "Important Subscription Notice"
+            : "Subscription Notice"}
+        </span>
+      </AlertTitle>
+      <AlertDescription>
+        <ul className="list-disc list-inside space-y-1 mt-2">
+          {messages.map((msg, idx) => (
+            <li key={idx}>{msg}</li>
+          ))}
+        </ul>
+      </AlertDescription>
+    </Alert>
+  );
+}