feat: no login experience and prem tokens

surfsense_backend/app/routes/stripe_routes.py

@@ -13,13 +13,24 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from stripe import SignatureVerificationError, StripeClient, StripeError
 
 from app.config import config
-from app.db import PagePurchase, PagePurchaseStatus, User, get_async_session
+from app.db import (
+    PagePurchase,
+    PagePurchaseStatus,
+    PremiumTokenPurchase,
+    PremiumTokenPurchaseStatus,
+    User,
+    get_async_session,
+)
 from app.schemas.stripe import (
     CreateCheckoutSessionRequest,
     CreateCheckoutSessionResponse,
+    CreateTokenCheckoutSessionRequest,
+    CreateTokenCheckoutSessionResponse,
     PagePurchaseHistoryResponse,
     StripeStatusResponse,
     StripeWebhookResponse,
+    TokenPurchaseHistoryResponse,
+    TokenStripeStatusResponse,
 )
 from app.users import current_active_user
 
@@ -151,6 +162,26 @@ async def _mark_purchase_failed(
     return StripeWebhookResponse()
 
 
+async def _mark_token_purchase_failed(
+    db_session: AsyncSession, checkout_session_id: str
+) -> StripeWebhookResponse:
+    purchase = (
+        await db_session.execute(
+            select(PremiumTokenPurchase)
+            .where(
+                PremiumTokenPurchase.stripe_checkout_session_id == checkout_session_id
+            )
+            .with_for_update()
+        )
+    ).scalar_one_or_none()
+
+    if purchase is not None and purchase.status == PremiumTokenPurchaseStatus.PENDING:
+        purchase.status = PremiumTokenPurchaseStatus.FAILED
+        await db_session.commit()
+
+    return StripeWebhookResponse()
+
+
 async def _fulfill_completed_purchase(
     db_session: AsyncSession, checkout_session: Any
 ) -> StripeWebhookResponse:
@@ -201,6 +232,86 @@ async def _fulfill_completed_purchase(
     return StripeWebhookResponse()
 
 
+async def _fulfill_completed_token_purchase(
+    db_session: AsyncSession, checkout_session: Any
+) -> StripeWebhookResponse:
+    """Grant premium tokens to the user after a confirmed Stripe payment."""
+    checkout_session_id = str(checkout_session.id)
+    purchase = (
+        await db_session.execute(
+            select(PremiumTokenPurchase)
+            .where(
+                PremiumTokenPurchase.stripe_checkout_session_id == checkout_session_id
+            )
+            .with_for_update()
+        )
+    ).scalar_one_or_none()
+
+    if purchase is None:
+        metadata = _get_metadata(checkout_session)
+        user_id = metadata.get("user_id")
+        quantity = int(metadata.get("quantity", "0"))
+        tokens_per_unit = int(metadata.get("tokens_per_unit", "0"))
+
+        if not user_id or quantity <= 0 or tokens_per_unit <= 0:
+            logger.error(
+                "Skipping token fulfillment for session %s: incomplete metadata %s",
+                checkout_session_id,
+                metadata,
+            )
+            return StripeWebhookResponse()
+
+        purchase = PremiumTokenPurchase(
+            user_id=uuid.UUID(user_id),
+            stripe_checkout_session_id=checkout_session_id,
+            stripe_payment_intent_id=_normalize_optional_string(
+                getattr(checkout_session, "payment_intent", None)
+            ),
+            quantity=quantity,
+            tokens_granted=quantity * tokens_per_unit,
+            amount_total=getattr(checkout_session, "amount_total", None),
+            currency=getattr(checkout_session, "currency", None),
+            status=PremiumTokenPurchaseStatus.PENDING,
+        )
+        db_session.add(purchase)
+        await db_session.flush()
+
+    if purchase.status == PremiumTokenPurchaseStatus.COMPLETED:
+        return StripeWebhookResponse()
+
+    user = (
+        (
+            await db_session.execute(
+                select(User).where(User.id == purchase.user_id).with_for_update(of=User)
+            )
+        )
+        .unique()
+        .scalar_one_or_none()
+    )
+    if user is None:
+        logger.error(
+            "Skipping token fulfillment for session %s: user %s not found",
+            purchase.stripe_checkout_session_id,
+            purchase.user_id,
+        )
+        return StripeWebhookResponse()
+
+    purchase.status = PremiumTokenPurchaseStatus.COMPLETED
+    purchase.completed_at = datetime.now(UTC)
+    purchase.amount_total = getattr(checkout_session, "amount_total", None)
+    purchase.currency = getattr(checkout_session, "currency", None)
+    purchase.stripe_payment_intent_id = _normalize_optional_string(
+        getattr(checkout_session, "payment_intent", None)
+    )
+    user.premium_tokens_limit = (
+        max(user.premium_tokens_used, user.premium_tokens_limit)
+        + purchase.tokens_granted
+    )
+
+    await db_session.commit()
+    return StripeWebhookResponse()
+
+
 @router.post("/create-checkout-session", response_model=CreateCheckoutSessionResponse)
 async def create_checkout_session(
     body: CreateCheckoutSessionRequest,
@@ -333,6 +444,10 @@ async def stripe_webhook(
             )
             return StripeWebhookResponse()
 
+        metadata = _get_metadata(checkout_session)
+        purchase_type = metadata.get("purchase_type", "page_packs")
+        if purchase_type == "premium_tokens":
+            return await _fulfill_completed_token_purchase(db_session, checkout_session)
         return await _fulfill_completed_purchase(db_session, checkout_session)
 
     if event.type in {
@@ -340,6 +455,12 @@ async def stripe_webhook(
         "checkout.session.expired",
     }:
         checkout_session = event.data.object
+        metadata = _get_metadata(checkout_session)
+        purchase_type = metadata.get("purchase_type", "page_packs")
+        if purchase_type == "premium_tokens":
+            return await _mark_token_purchase_failed(
+                db_session, str(checkout_session.id)
+            )
         return await _mark_purchase_failed(db_session, str(checkout_session.id))
 
     return StripeWebhookResponse()
@@ -369,3 +490,146 @@ async def get_page_purchases(
     )
 
     return PagePurchaseHistoryResponse(purchases=purchases)
+
+
+# =============================================================================
+# Premium Token Purchase Routes
+# =============================================================================
+
+
+def _ensure_token_buying_enabled() -> None:
+    if not config.STRIPE_TOKEN_BUYING_ENABLED:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="Premium token purchases are temporarily unavailable.",
+        )
+
+
+def _get_token_checkout_urls(search_space_id: int) -> tuple[str, str]:
+    if not config.NEXT_FRONTEND_URL:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="NEXT_FRONTEND_URL is not configured.",
+        )
+    base_url = config.NEXT_FRONTEND_URL.rstrip("/")
+    success_url = f"{base_url}/dashboard/{search_space_id}/purchase-success"
+    cancel_url = f"{base_url}/dashboard/{search_space_id}/purchase-cancel"
+    return success_url, cancel_url
+
+
+def _get_required_token_price_id() -> str:
+    if not config.STRIPE_PREMIUM_TOKEN_PRICE_ID:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="STRIPE_PREMIUM_TOKEN_PRICE_ID is not configured.",
+        )
+    return config.STRIPE_PREMIUM_TOKEN_PRICE_ID
+
+
+@router.post("/create-token-checkout-session")
+async def create_token_checkout_session(
+    body: CreateTokenCheckoutSessionRequest,
+    user: User = Depends(current_active_user),
+    db_session: AsyncSession = Depends(get_async_session),
+):
+    """Create a Stripe Checkout Session for buying premium token packs."""
+    _ensure_token_buying_enabled()
+    stripe_client = get_stripe_client()
+    price_id = _get_required_token_price_id()
+    success_url, cancel_url = _get_token_checkout_urls(body.search_space_id)
+    tokens_granted = body.quantity * config.STRIPE_TOKENS_PER_UNIT
+
+    try:
+        checkout_session = stripe_client.v1.checkout.sessions.create(
+            params={
+                "mode": "payment",
+                "success_url": success_url,
+                "cancel_url": cancel_url,
+                "line_items": [
+                    {
+                        "price": price_id,
+                        "quantity": body.quantity,
+                    }
+                ],
+                "client_reference_id": str(user.id),
+                "customer_email": user.email,
+                "metadata": {
+                    "user_id": str(user.id),
+                    "quantity": str(body.quantity),
+                    "tokens_per_unit": str(config.STRIPE_TOKENS_PER_UNIT),
+                    "purchase_type": "premium_tokens",
+                },
+            }
+        )
+    except StripeError as exc:
+        logger.exception("Failed to create token checkout session for user %s", user.id)
+        raise HTTPException(
+            status_code=status.HTTP_502_BAD_GATEWAY,
+            detail="Unable to create Stripe checkout session.",
+        ) from exc
+
+    checkout_url = getattr(checkout_session, "url", None)
+    if not checkout_url:
+        raise HTTPException(
+            status_code=status.HTTP_502_BAD_GATEWAY,
+            detail="Stripe checkout session did not return a URL.",
+        )
+
+    db_session.add(
+        PremiumTokenPurchase(
+            user_id=user.id,
+            stripe_checkout_session_id=str(checkout_session.id),
+            stripe_payment_intent_id=_normalize_optional_string(
+                getattr(checkout_session, "payment_intent", None)
+            ),
+            quantity=body.quantity,
+            tokens_granted=tokens_granted,
+            amount_total=getattr(checkout_session, "amount_total", None),
+            currency=getattr(checkout_session, "currency", None),
+            status=PremiumTokenPurchaseStatus.PENDING,
+        )
+    )
+    await db_session.commit()
+
+    return CreateTokenCheckoutSessionResponse(checkout_url=checkout_url)
+
+
+@router.get("/token-status")
+async def get_token_status(
+    user: User = Depends(current_active_user),
+):
+    """Return token-buying availability and current premium quota for frontend."""
+    used = user.premium_tokens_used
+    limit = user.premium_tokens_limit
+    return TokenStripeStatusResponse(
+        token_buying_enabled=config.STRIPE_TOKEN_BUYING_ENABLED,
+        premium_tokens_used=used,
+        premium_tokens_limit=limit,
+        premium_tokens_remaining=max(0, limit - used),
+    )
+
+
+@router.get("/token-purchases")
+async def get_token_purchases(
+    user: User = Depends(current_active_user),
+    db_session: AsyncSession = Depends(get_async_session),
+    offset: int = 0,
+    limit: int = 50,
+):
+    """Return the authenticated user's premium token purchase history."""
+    limit = min(limit, 100)
+    purchases = (
+        (
+            await db_session.execute(
+                select(PremiumTokenPurchase)
+                .where(PremiumTokenPurchase.user_id == user.id)
+                .order_by(PremiumTokenPurchase.created_at.desc())
+                .offset(offset)
+                .limit(limit)
+            )
+        )
+        .scalars()
+        .all()
+    )
+
+    return TokenPurchaseHistoryResponse(purchases=purchases)