refactor: enhance billing info response handling

api/services/billing_service.py

@@ -2,7 +2,7 @@ import json
 import logging
 import os
 from collections.abc import Sequence
-from typing import Literal
+from typing import Literal, NotRequired
 
 import httpx
 from pydantic import TypeAdapter
@@ -26,6 +26,56 @@ class SubscriptionPlan(TypedDict):
     expiration_date: int
 
 
+class _BillingQuota(TypedDict):
+    size: int
+    limit: int
+
+
+class _VectorSpaceQuota(TypedDict):
+    size: float
+    limit: int
+
+
+class _KnowledgeRateLimit(TypedDict):
+    size: NotRequired[int]
+    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 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 +88,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):

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:

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",
+            "enabled": True,
-            "billing_cycle": "monthly",
+            "subscription": {"plan": "professional", "interval": "month", "education": False},
-            "status": "active",
+            "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."""
+        """Empty response from billing API should raise ValidationError due to missing required fields."""
-        # Arrange
+        from pydantic import ValidationError
+
         tenant_id = "tenant-empty"
         mock_send_request.return_value = {}
 
-        # Act
+        with pytest.raises(ValidationError):
-        result = BillingService.get_info(tenant_id)
+            BillingService.get_info(tenant_id)
-
-        # Assert
-        assert result == {}
-        mock_send_request.assert_called_once()
 
     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",
+            "enabled": True,
-            "billing_cycle": "monthly",
+            "subscription": {"plan": "sandbox", "interval": "", "education": False},
-            "status": "active",
+            "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")