Multi-plan subscriptions, quotas, and token usage APIs (#16 #17 #36) (#37)

## Summary
Implements [#16](#16), [#17](#17), and [#36](#36) in one backend PR.

- **#36 Multi-plan catalog**: Founders ($10, public), Standard ($15), Pro ($40), Business ($99), Backer ($0). Future tiers seeded but hidden/`is_selectable=false`. Backer email whitelist auto-assigns Founders-level access with no checkout.
- **#36 Feature + prompt gating**: plan feature flags (text vs image); rolling **6h** prompt windows (100 / 200 / 300 / 300 / 300). Enforced in both chat consumers when `ENFORCE_SUBSCRIPTION_GATES=true`.
- **#17 Token-period quotas**: optional `monthly_token_quota` on plans + per-user override; calendar-month aggregation from `PromptMetric`; warn/block when reported token totals exceed cap. Null provider usage never fabricated as 0; tracked via `turns_missing_token_usage`.
- **#16 Token API exposure**: `tokens_in` / `tokens_out` on conversation + prompt serializers (null when unknown). `GET /api/finance/subscription/` returns plan + usage snapshot for the FE.
- Checkout defaults to **Founders**; Stripe paid webhooks assign Founders. Registration/OAuth redeem Backer whitelist and return `needs_checkout`.

Companion FE PR: `chat_web_app` branch `feature/plans-quotas-token-usage`.

## Test plan
- [ ] `manage.py migrate` seeds five plans; admin can add Backer emails
- [ ] Public `GET /api/finance/plans/` returns only Founders
- [ ] Register with Backer email → active Backer, `needs_checkout=false`, checkout rejected
- [ ] Founders checkout + paid webhook → active Founders subscription
- [ ] Chat turn blocked without subscription / when prompt window exceeded / when token period exceeded
- [ ] Standard plan denies image feature; Pro/Founders/Backer allow
- [ ] Conversation/prompt API returns `null` tokens when unreported, sums when present
- [ ] `finance.tests.test_plans_quotas` + existing finance/checkout tests passReviewed-on: #37

llm_be/chat_backend/consumers.py

@@ -27,7 +27,17 @@ from .services.title_generator import title_generator
 from .services.moderation_classifier import moderation_classifier, ModerationLabel
 from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
 from .services.data_analysis_service import AsyncDataAnalysisService
-from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message
+from .utils import (
+    extract_token_usage,
+    has_usable_user_prompt,
+    is_heartbeat_payload,
+    normalize_user_message,
+)
+from finance.services.quotas import (
+    FeatureNotAllowed,
+    QuotaExceeded,
+    check_generation_allowed,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -47,6 +57,35 @@ def create_conversation(prompt, email, title):
     return conversation.id
 
 
+@database_sync_to_async
+def resolve_chat_user(email=None, conversation_id=None):
+    if email:
+        user = CustomUser.objects.filter(email__iexact=email).first()
+        if user:
+            return user
+    if conversation_id:
+        conversation = (
+            Conversation.objects.select_related("user")
+            .filter(id=conversation_id)
+            .first()
+        )
+        if conversation and conversation.user_id:
+            return conversation.user
+    return None
+
+
+@database_sync_to_async
+def enforce_generation_gates(user, feature="text_generation"):
+    return check_generation_allowed(user, feature=feature)
+
+
+@database_sync_to_async
+def enforce_feature_gate(user, feature):
+    from finance.services.quotas import assert_feature_allowed
+
+    assert_feature_allowed(user, feature)
+
+
 @database_sync_to_async
 def get_workspace(conversation_id):
     conversation = Conversation.objects.get(id=conversation_id)
@@ -262,6 +301,36 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
                 )
                 return
 
+            chat_user = await resolve_chat_user(
+                email=email, conversation_id=conversation_id
+            )
+            if chat_user is None:
+                await self.send_json_message(
+                    json.dumps(
+                        {
+                            "type": "error",
+                            "code": "user_not_found",
+                            "content": "Unable to resolve user for this chat session.",
+                        }
+                    )
+                )
+                return
+
+            try:
+                await enforce_generation_gates(chat_user, feature="text_generation")
+            except (QuotaExceeded, FeatureNotAllowed) as exc:
+                await self.send_json_message(
+                    json.dumps(
+                        {
+                            "type": "error",
+                            "code": exc.code,
+                            "content": exc.message,
+                            "details": getattr(exc, "details", {}),
+                        }
+                    )
+                )
+                return
+
             if not conversation_id:
                 # we need to create a new conversation
                 # we will generate a name for it too
@@ -334,11 +403,20 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
                     decoded_file = input_dict.get("decoded_file")
                     file_type = input_dict.get("file_type")
                     
-                    # Feature Flag: Image Generation
+                    # Feature Flag + plan gate: Image Generation
                     if prompt_type == PromptType.IMAGE_GENERATION:
                         if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
                              return {"type": "text", "content": "Image Generation is disabled."}
-                        # If enabled, proceed (assuming implementation exists, but user said "have it set to false for now")
+                        try:
+                            await enforce_feature_gate(
+                                chat_user, "image_generation"
+                            )
+                        except FeatureNotAllowed as exc:
+                            return {
+                                "type": "error",
+                                "code": exc.code,
+                                "content": exc.message,
+                            }
                         return {"type": "text", "content": "Image Generation is not supported at this time, but it will be soon."}
 
                     if prompt_type == PromptType.SEARCH:
@@ -446,7 +524,17 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
                 await self.send("END_OF_THE_STREAM_ENDER_GAME_42")
                 
                 await save_generated_message(conversation_id, full_response)
-                await finish_prompt_metric(prompt_metric, len(full_response))
+                tokens_in, tokens_out = extract_token_usage(
+                    response_generator_or_dict
+                    if isinstance(response_generator_or_dict, dict)
+                    else None
+                )
+                await finish_prompt_metric(
+                    prompt_metric,
+                    len(full_response),
+                    tokens_in=tokens_in,
+                    tokens_out=tokens_out,
+                )
 
         if bytes_data:
             logger.info("we have byte data")

llm_be/chat_backend/consumers_graph.py

@@ -22,7 +22,13 @@ from .services.title_generator import title_generator
 from .services.moderation_classifier import moderation_classifier, ModerationLabel
 from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
 from .services.data_analysis_service import AsyncDataAnalysisService
-from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message
+from .utils import (
+    extract_token_usage,
+    has_usable_user_prompt,
+    is_heartbeat_payload,
+    normalize_user_message,
+)
+from finance.services.quotas import FeatureNotAllowed, QuotaExceeded, check_generation_allowed
 
 logger = logging.getLogger(__name__)
 
@@ -40,6 +46,35 @@ def create_conversation(prompt, email, title):
     conversation.save()
     return conversation.id
 
+
+@database_sync_to_async
+def resolve_chat_user(email=None, conversation_id=None):
+    if email:
+        user = CustomUser.objects.filter(email__iexact=email).first()
+        if user:
+            return user
+    if conversation_id:
+        conversation = (
+            Conversation.objects.select_related("user")
+            .filter(id=conversation_id)
+            .first()
+        )
+        if conversation and conversation.user_id:
+            return conversation.user
+    return None
+
+
+@database_sync_to_async
+def enforce_generation_gates(user, feature="text_generation"):
+    return check_generation_allowed(user, feature=feature)
+
+
+@database_sync_to_async
+def enforce_feature_gate(user, feature):
+    from finance.services.quotas import assert_feature_allowed
+
+    assert_feature_allowed(user, feature)
+
 @database_sync_to_async
 def get_workspace(conversation_id):
     conversation = Conversation.objects.get(id=conversation_id)
@@ -172,6 +207,7 @@ class ChatState(TypedDict):
     response_generator: Any # AsyncGenerator or dict
     error: Union[str, None]
     model_name: str
+    chat_user: Any
 
 
 # --- LangGraph Nodes ---
@@ -210,10 +246,22 @@ async def generation_node(state: ChatState) -> ChatState:
     decoded_file = state.get("decoded_file")
     file_type = state.get("file_type")
     
-    # Feature Flag: Image Generation
+    # Feature Flag + plan gate: Image Generation
     if prompt_type == PromptType.IMAGE_GENERATION:
         if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
              return {"response_generator": {"type": "text", "content": "Image Generation is disabled."}}
+        chat_user = state.get("chat_user")
+        if chat_user is not None:
+            try:
+                await enforce_feature_gate(chat_user, "image_generation")
+            except FeatureNotAllowed as exc:
+                return {
+                    "response_generator": {
+                        "type": "error",
+                        "code": exc.code,
+                        "content": exc.message,
+                    }
+                }
         return {"response_generator": {"type": "text", "content": "Image Generation is not supported at this time, but it will be soon."}}
 
     # Feature Flag: Internet Access
@@ -312,6 +360,36 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
                 )
                 return
 
+            chat_user = await resolve_chat_user(
+                email=email, conversation_id=conversation_id
+            )
+            if chat_user is None:
+                await self.send_json_message(
+                    json.dumps(
+                        {
+                            "type": "error",
+                            "code": "user_not_found",
+                            "content": "Unable to resolve user for this chat session.",
+                        }
+                    )
+                )
+                return
+
+            try:
+                await enforce_generation_gates(chat_user, feature="text_generation")
+            except (QuotaExceeded, FeatureNotAllowed) as exc:
+                await self.send_json_message(
+                    json.dumps(
+                        {
+                            "type": "error",
+                            "code": exc.code,
+                            "content": exc.message,
+                            "details": getattr(exc, "details", {}),
+                        }
+                    )
+                )
+                return
+
             if not conversation_id:
                 title = await title_generator.generate_async(message)
                 conversation_id = await create_conversation(message, email, title)
@@ -357,7 +435,8 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
                     "prompt_type": None,
                     "response_generator": None,
                     "error": None,
-                    "model_name": model
+                    "model_name": model,
+                    "chat_user": chat_user,
                 }
                 print("Initial State: ", initial_state)
 
@@ -387,4 +466,14 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
                 await self.send("END_OF_THE_STREAM_ENDER_GAME_42")
                 
                 await save_generated_message(conversation_id, full_response)
-                await finish_prompt_metric(prompt_metric, len(full_response))
+                tokens_in, tokens_out = extract_token_usage(
+                    response_generator_or_dict
+                    if isinstance(response_generator_or_dict, dict)
+                    else None
+                )
+                await finish_prompt_metric(
+                    prompt_metric,
+                    len(full_response),
+                    tokens_in=tokens_in,
+                    tokens_out=tokens_out,
+                )

llm_be/chat_backend/oauth.py

@@ -307,6 +307,8 @@ def upsert_identity(user: CustomUser, profile: ProviderProfile) -> OAuthIdentity
 
 
 def _create_sso_user(profile: ProviderProfile) -> CustomUser:
+    from finance.services.plans import try_redeem_backer_email
+
     company = Company.objects.create(
         name=f"{profile.email}'s workspace",
         state="NA",
@@ -323,6 +325,7 @@ def _create_sso_user(profile: ProviderProfile) -> CustomUser:
     )
     user.set_unusable_password()
     user.save()
+    try_redeem_backer_email(user)
     return user
 
 

llm_be/chat_backend/serializers.py

@@ -1,11 +1,14 @@
 from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
 from rest_framework import serializers
+from django.db.models import Count, Q, Sum
+
 from .models import (
     CustomUser,
     Announcement,
     Company,
     Conversation,
     Prompt,
+    PromptMetric,
     Feedback,
     FEEDBACK_CATEGORIES,
     DocumentWorkspace,
@@ -48,12 +51,33 @@ class CustomUserSerializer(serializers.ModelSerializer):
     password = serializers.CharField(min_length=8, write_only=True)
     company = CompanySerializer()
     has_usable_password = serializers.BooleanField()
+    subscription = serializers.SerializerMethodField()
 
     class Meta:
         model = CustomUser
         fields = "__all__"
         extra_kwargs = {"password": {"write_only": True}}
 
+    def get_subscription(self, obj):
+        from finance.services.plans import needs_checkout, plan_to_dict
+        from finance.models import UserSubscription
+
+        try:
+            sub = obj.subscription
+        except UserSubscription.DoesNotExist:
+            return {
+                "plan": None,
+                "status": UserSubscription.Status.NONE,
+                "source": UserSubscription.Source.NONE,
+                "needs_checkout": True,
+            }
+        return {
+            "plan": plan_to_dict(sub.plan) if sub.plan_id else None,
+            "status": sub.status,
+            "source": sub.source,
+            "needs_checkout": needs_checkout(obj),
+        }
+
 
 class SelfServeRegistrationSerializer(serializers.Serializer):
     """Minimal payload for public self-serve sign-up (gated by settings)."""
@@ -79,6 +103,8 @@ class SelfServeRegistrationSerializer(serializers.Serializer):
         return email
 
     def create(self, validated_data):
+        from finance.services.plans import try_redeem_backer_email
+
         email = validated_data["email"]
         password = validated_data["password"]
         first_name = (validated_data.get("first_name") or "").strip()
@@ -102,16 +128,73 @@ class SelfServeRegistrationSerializer(serializers.Serializer):
             company=company,
             is_company_manager=True,
         )
+        try_redeem_backer_email(user)
         return user
 
 
+def _conversation_token_totals(conversation_id: int):
+    """
+    Sum PromptMetric tokens for a conversation.
+
+    Returns nulls when the provider never reported usage (never fabricate 0).
+    """
+    agg = PromptMetric.objects.filter(conversation_id=conversation_id).aggregate(
+        tin=Sum("tokens_in"),
+        tout=Sum("tokens_out"),
+        with_in=Count("id", filter=Q(tokens_in__isnull=False)),
+        with_out=Count("id", filter=Q(tokens_out__isnull=False)),
+    )
+    return (
+        agg["tin"] if agg["with_in"] else None,
+        agg["tout"] if agg["with_out"] else None,
+    )
+
+
+def _prompt_token_pair(prompt_id: int):
+    metric = (
+        PromptMetric.objects.filter(prompt_id=prompt_id)
+        .order_by("-created")
+        .only("tokens_in", "tokens_out")
+        .first()
+    )
+    if metric is None:
+        return None, None
+    return metric.tokens_in, metric.tokens_out
+
+
 class ConversationSerializer(serializers.ModelSerializer):
+    tokens_in = serializers.SerializerMethodField()
+    tokens_out = serializers.SerializerMethodField()
+
     class Meta:
         model = Conversation
-        fields = ("title", "created", "last_modified", "id")
+        fields = (
+            "title",
+            "created",
+            "last_modified",
+            "id",
+            "tokens_in",
+            "tokens_out",
+        )
+
+    def _token_pair(self, obj):
+        cache = self.context.setdefault("_conversation_token_cache", {})
+        if obj.id not in cache:
+            cache[obj.id] = _conversation_token_totals(obj.id)
+        return cache[obj.id]
+
+    def get_tokens_in(self, obj):
+        tin, _ = self._token_pair(obj)
+        return tin
+
+    def get_tokens_out(self, obj):
+        _, tout = self._token_pair(obj)
+        return tout
 
 
 class PromptSerializer(serializers.ModelSerializer):
+    tokens_in = serializers.SerializerMethodField()
+    tokens_out = serializers.SerializerMethodField()
 
     class Meta:
         model = Prompt
@@ -120,8 +203,24 @@ class PromptSerializer(serializers.ModelSerializer):
             "user_created",
             "created",
             "id",
+            "tokens_in",
+            "tokens_out",
         )
 
+    def _token_pair(self, obj):
+        cache = self.context.setdefault("_prompt_token_cache", {})
+        if obj.id not in cache:
+            cache[obj.id] = _prompt_token_pair(obj.id)
+        return cache[obj.id]
+
+    def get_tokens_in(self, obj):
+        tin, _ = self._token_pair(obj)
+        return tin
+
+    def get_tokens_out(self, obj):
+        _, tout = self._token_pair(obj)
+        return tout
+
     def validate_message(self, value: str) -> str:
         if value is None or not str(value).strip():
             raise serializers.ValidationError("Message text cannot be empty.")

llm_be/chat_backend/tests/test_serializers.py

@@ -38,7 +38,10 @@ class ConversationSerializerTestCase(TestCase):
 
         data = ConversationSerializer(conversation).data
 
-        self.assertEqual(set(data.keys()), {"title", "created", "last_modified", "id"})
+        self.assertEqual(
+            set(data.keys()),
+            {"title", "created", "last_modified", "id", "tokens_in", "tokens_out"},
+        )
         self.assertEqual(data["title"], "Weather Inquiry")
 
 

llm_be/chat_backend/views.py

@@ -136,6 +136,8 @@ class CustomUserCreate(APIView):
 
         user = serializer.save()
         refresh = RefreshToken.for_user(user)
+        from finance.services.plans import needs_checkout
+
         return Response(
             {
                 "email": user.email,
@@ -143,6 +145,7 @@ class CustomUserCreate(APIView):
                 "last_name": user.last_name,
                 "access": str(refresh.access_token),
                 "refresh": str(refresh),
+                "needs_checkout": needs_checkout(user),
             },
             status=status.HTTP_201_CREATED,
         )

llm_be/chat_backend/views_oauth.py

@@ -139,7 +139,9 @@ class OAuthCallbackView(APIView):
             return _redirect_error("server_error", "Unexpected OAuth error.")
 
         refresh = RefreshToken.for_user(user)
-        needs_checkout = "1" if created else "0"
+        from finance.services.plans import needs_checkout as user_needs_checkout
+
+        needs_checkout = "1" if (created and user_needs_checkout(user)) else "0"
         return HttpResponseRedirect(
             _frontend_callback_url(
                 access=str(refresh.access_token),