refactor: enhance billing info response handling (#34340)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-05-04 17:38:04 +08:00 · 2026-04-08 11:49:35 +08:00
parent c3e05045bb
commit 1a9a1e821f
3 changed files with 276 additions and 20 deletions
--- a/api/services/billing_service.py
+++ b/api/services/billing_service.py
@ -2,12 +2,11 @@ import json
 import logging
 import os
 from collections.abc import Sequence
-from typing import Literal
+from typing import Literal, NotRequired, TypedDict

 import httpx
 from pydantic import TypeAdapter
 from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed
-from typing_extensions import TypedDict
 from werkzeug.exceptions import InternalServerError

 from enums.cloud_plan import CloudPlan
@ -26,6 +25,103 @@ class SubscriptionPlan(TypedDict):
    expiration_date: int


+class _BillingQuota(TypedDict):
+    size: int
+    limit: int
+
+
+class _VectorSpaceQuota(TypedDict):
+    size: float
+    limit: int
+
+
+class _KnowledgeRateLimit(TypedDict):
+    # NOTE (hj24):
+    # 1. Return for sandbox users but is null for other plans, it's defined but never used.
+    # 2. Keep it for compatibility for now, can be deprecated in future versions.
+    size: NotRequired[int]
+    # NOTE END
+    limit: int
+
+
+class _BillingSubscription(TypedDict):
+    plan: str
+    interval: str
+    education: bool
+
+
+class BillingInfo(TypedDict):
+    """Response of /subscription/info.
+
+    NOTE (hj24):
+    - Fields not listed here (e.g. trigger_event, api_rate_limit) are stripped by TypeAdapter.validate_python()
+    - To ensure the precision, billing may convert fields like int as str, be careful when use TypeAdapter:
+        1. validate_python in non-strict mode will coerce it to the expected type
+        2. In strict mode, it will raise ValidationError
+        3. To preserve compatibility, always keep non-strict mode here and avoid strict mode
+    """
+
+    enabled: bool
+    subscription: _BillingSubscription
+    members: _BillingQuota
+    apps: _BillingQuota
+    vector_space: _VectorSpaceQuota
+    knowledge_rate_limit: _KnowledgeRateLimit
+    documents_upload_quota: _BillingQuota
+    annotation_quota_limit: _BillingQuota
+    docs_processing: str
+    can_replace_logo: bool
+    model_load_balancing_enabled: bool
+    knowledge_pipeline_publish_enabled: bool
+    next_credit_reset_date: NotRequired[int]
+
+
+_billing_info_adapter = TypeAdapter(BillingInfo)
+
+
+class KnowledgeRateLimitDict(TypedDict):
+    limit: int
+    subscription_plan: str
+
+
+class TenantFeaturePlanUsageDict(TypedDict):
+    result: str
+    history_id: str
+
+
+class LangContentDict(TypedDict):
+    lang: str
+    title: str
+    subtitle: str
+    body: str
+    title_pic_url: str
+
+
+class NotificationDict(TypedDict):
+    notification_id: str
+    contents: dict[str, LangContentDict]
+    frequency: Literal["once", "every_page_load"]
+
+
+class AccountNotificationDict(TypedDict, total=False):
+    should_show: bool
+    notification: NotificationDict
+    shouldShow: bool
+    notifications: list[dict]
+
+
+class UpsertNotificationDict(TypedDict):
+    notification_id: str
+
+
+class BatchAddNotificationAccountsDict(TypedDict):
+    count: int
+
+
+class DismissNotificationDict(TypedDict):
+    success: bool
+
+
 class BillingService:
    base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL")
    secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY")
@ -38,11 +134,11 @@ class BillingService:
    _PLAN_CACHE_TTL = 600

    @classmethod
-    def get_info(cls, tenant_id: str):
+    def get_info(cls, tenant_id: str) -> BillingInfo:
        params = {"tenant_id": tenant_id}

        billing_info = cls._send_request("GET", "/subscription/info", params=params)
-        return billing_info
+        return _billing_info_adapter.validate_python(billing_info)

    @classmethod
    def get_tenant_feature_plan_usage_info(cls, tenant_id: str):
--- a/api/services/feature_service.py
+++ b/api/services/feature_service.py
@ -312,7 +312,10 @@ class FeatureService:
            features.apps.limit = billing_info["apps"]["limit"]

        if "vector_space" in billing_info:
-            features.vector_space.size = billing_info["vector_space"]["size"]
+            # NOTE (hj24): billing API returns vector_space.size as float (e.g. 0.0)
+            # but LimitationModel.size is int; truncate here for compatibility
+            features.vector_space.size = int(billing_info["vector_space"]["size"])
+            # NOTE END
            features.vector_space.limit = billing_info["vector_space"]["limit"]

        if "documents_upload_quota" in billing_info:
@ -333,7 +336,11 @@ class FeatureService:
            features.model_load_balancing_enabled = billing_info["model_load_balancing_enabled"]

        if "knowledge_rate_limit" in billing_info:
+            # NOTE (hj24):
+            # 1. knowledge_rate_limit size is nullable, currently it's defined but never used, only limit is used.
+            # 2. So be careful if later we decide to use [size], we cannot assume it is always present.
            features.knowledge_rate_limit = billing_info["knowledge_rate_limit"]["limit"]
+            # NOTE END

        if "knowledge_pipeline_publish_enabled" in billing_info:
            features.knowledge_pipeline.publish_enabled = billing_info["knowledge_pipeline_publish_enabled"]
--- a/api/tests/unit_tests/services/test_billing_service.py
+++ b/api/tests/unit_tests/services/test_billing_service.py
@ -290,9 +290,19 @@ class TestBillingServiceSubscriptionInfo:
        # Arrange
        tenant_id = "tenant-123"
        expected_response = {
-            "subscription_plan": "professional",
-            "billing_cycle": "monthly",
-            "status": "active",
+            "enabled": True,
+            "subscription": {"plan": "professional", "interval": "month", "education": False},
+            "members": {"size": 1, "limit": 50},
+            "apps": {"size": 1, "limit": 200},
+            "vector_space": {"size": 0.0, "limit": 20480},
+            "knowledge_rate_limit": {"limit": 1000},
+            "documents_upload_quota": {"size": 0, "limit": 1000},
+            "annotation_quota_limit": {"size": 0, "limit": 5000},
+            "docs_processing": "top-priority",
+            "can_replace_logo": True,
+            "model_load_balancing_enabled": True,
+            "knowledge_pipeline_publish_enabled": True,
+            "next_credit_reset_date": 1775952000,
        }
        mock_send_request.return_value = expected_response

@ -1009,17 +1019,14 @@ class TestBillingServiceEdgeCases:
            yield mock

    def test_get_info_empty_response(self, mock_send_request):
-        """Test handling of empty billing info response."""
-        # Arrange
+        """Empty response from billing API should raise ValidationError due to missing required fields."""
+        from pydantic import ValidationError
+
        tenant_id = "tenant-empty"
        mock_send_request.return_value = {}

-        # Act
-        result = BillingService.get_info(tenant_id)
-
-        # Assert
-        assert result == {}
-        mock_send_request.assert_called_once()
+        with pytest.raises(ValidationError):
+            BillingService.get_info(tenant_id)

    def test_update_tenant_feature_plan_usage_zero_delta(self, mock_send_request):
        """Test updating tenant feature usage with zero delta (no change)."""
@ -1434,12 +1441,21 @@ class TestBillingServiceIntegrationScenarios:

        # Step 1: Get current billing info
        mock_send_request.return_value = {
-            "subscription_plan": "sandbox",
-            "billing_cycle": "monthly",
-            "status": "active",
+            "enabled": True,
+            "subscription": {"plan": "sandbox", "interval": "", "education": False},
+            "members": {"size": 0, "limit": 1},
+            "apps": {"size": 0, "limit": 5},
+            "vector_space": {"size": 0.0, "limit": 50},
+            "knowledge_rate_limit": {"limit": 10},
+            "documents_upload_quota": {"size": 0, "limit": 50},
+            "annotation_quota_limit": {"size": 0, "limit": 10},
+            "docs_processing": "standard",
+            "can_replace_logo": False,
+            "model_load_balancing_enabled": False,
+            "knowledge_pipeline_publish_enabled": False,
        }
        current_info = BillingService.get_info(tenant_id)
-        assert current_info["subscription_plan"] == "sandbox"
+        assert current_info["subscription"]["plan"] == "sandbox"

        # Step 2: Get payment link for upgrade
        mock_send_request.return_value = {"payment_link": "https://payment.example.com/upgrade"}
@ -1553,3 +1569,140 @@ class TestBillingServiceIntegrationScenarios:
            mock_send_request.return_value = {"result": "success", "activated": True}
            activate_result = BillingService.EducationIdentity.activate(account, "token-123", "MIT", "student")
            assert activate_result["activated"] is True
+
+
+class TestBillingServiceSubscriptionInfoDataType:
+    """Unit tests for data type coercion in BillingService.get_info
+
+    1. Verifies the get_info returns correct Python types for numeric fields
+    2. Ensure the compatibility regardless of what results the upstream billing API returns
+    """
+
+    @pytest.fixture
+    def mock_send_request(self):
+        with patch.object(BillingService, "_send_request") as mock:
+            yield mock
+
+    @pytest.fixture
+    def normal_billing_response(self) -> dict:
+        return {
+            "enabled": True,
+            "subscription": {
+                "plan": "team",
+                "interval": "year",
+                "education": False,
+            },
+            "members": {"size": 10, "limit": 50},
+            "apps": {"size": 80, "limit": 200},
+            "vector_space": {"size": 5120.75, "limit": 20480},
+            "knowledge_rate_limit": {"limit": 1000},
+            "documents_upload_quota": {"size": 450, "limit": 1000},
+            "annotation_quota_limit": {"size": 1200, "limit": 5000},
+            "docs_processing": "top-priority",
+            "can_replace_logo": True,
+            "model_load_balancing_enabled": True,
+            "knowledge_pipeline_publish_enabled": True,
+            "next_credit_reset_date": 1745971200,
+        }
+
+    @pytest.fixture
+    def string_billing_response(self) -> dict:
+        return {
+            "enabled": True,
+            "subscription": {
+                "plan": "team",
+                "interval": "year",
+                "education": False,
+            },
+            "members": {"size": "10", "limit": "50"},
+            "apps": {"size": "80", "limit": "200"},
+            "vector_space": {"size": 5120.75, "limit": "20480"},
+            "knowledge_rate_limit": {"limit": "1000"},
+            "documents_upload_quota": {"size": "450", "limit": "1000"},
+            "annotation_quota_limit": {"size": "1200", "limit": "5000"},
+            "docs_processing": "top-priority",
+            "can_replace_logo": True,
+            "model_load_balancing_enabled": True,
+            "knowledge_pipeline_publish_enabled": True,
+            "next_credit_reset_date": "1745971200",
+        }
+
+    @staticmethod
+    def _assert_billing_info_types(result: dict):
+        assert isinstance(result["enabled"], bool)
+        assert isinstance(result["subscription"]["plan"], str)
+        assert isinstance(result["subscription"]["interval"], str)
+        assert isinstance(result["subscription"]["education"], bool)
+
+        assert isinstance(result["members"]["size"], int)
+        assert isinstance(result["members"]["limit"], int)
+
+        assert isinstance(result["apps"]["size"], int)
+        assert isinstance(result["apps"]["limit"], int)
+
+        assert isinstance(result["vector_space"]["size"], float)
+        assert isinstance(result["vector_space"]["limit"], int)
+
+        assert isinstance(result["knowledge_rate_limit"]["limit"], int)
+
+        assert isinstance(result["documents_upload_quota"]["size"], int)
+        assert isinstance(result["documents_upload_quota"]["limit"], int)
+
+        assert isinstance(result["annotation_quota_limit"]["size"], int)
+        assert isinstance(result["annotation_quota_limit"]["limit"], int)
+
+        assert isinstance(result["docs_processing"], str)
+        assert isinstance(result["can_replace_logo"], bool)
+        assert isinstance(result["model_load_balancing_enabled"], bool)
+        assert isinstance(result["knowledge_pipeline_publish_enabled"], bool)
+        if "next_credit_reset_date" in result:
+            assert isinstance(result["next_credit_reset_date"], int)
+
+    def test_get_info_with_normal_types(self, mock_send_request, normal_billing_response):
+        """When the billing API returns native numeric types, get_info should preserve them."""
+        mock_send_request.return_value = normal_billing_response
+
+        result = BillingService.get_info("tenant-type-test")
+
+        self._assert_billing_info_types(result)
+        mock_send_request.assert_called_once_with("GET", "/subscription/info", params={"tenant_id": "tenant-type-test"})
+
+    def test_get_info_with_string_types(self, mock_send_request, string_billing_response):
+        """When the billing API returns numeric values as strings, get_info should coerce them."""
+        mock_send_request.return_value = string_billing_response
+
+        result = BillingService.get_info("tenant-type-test")
+
+        self._assert_billing_info_types(result)
+        mock_send_request.assert_called_once_with("GET", "/subscription/info", params={"tenant_id": "tenant-type-test"})
+
+    def test_get_info_without_optional_fields(self, mock_send_request, string_billing_response):
+        """NotRequired fields can be absent without raising."""
+        del string_billing_response["next_credit_reset_date"]
+        mock_send_request.return_value = string_billing_response
+
+        result = BillingService.get_info("tenant-type-test")
+
+        assert "next_credit_reset_date" not in result
+        self._assert_billing_info_types(result)
+
+    def test_get_info_with_extra_fields(self, mock_send_request, string_billing_response):
+        """Undefined fields are silently stripped by validate_python."""
+        string_billing_response["new_feature"] = "something"
+        mock_send_request.return_value = string_billing_response
+
+        result = BillingService.get_info("tenant-type-test")
+
+        # extra fields are dropped by TypeAdapter on TypedDict
+        assert "new_feature" not in result
+        self._assert_billing_info_types(result)
+
+    def test_get_info_missing_required_field_raises(self, mock_send_request, string_billing_response):
+        """Missing a required field should raise ValidationError."""
+        from pydantic import ValidationError
+
+        del string_billing_response["members"]
+        mock_send_request.return_value = string_billing_response
+
+        with pytest.raises(ValidationError):
+            BillingService.get_info("tenant-type-test")