refactor: enhance billing info response handling (#34340)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py

@@ -607,19 +607,15 @@ class PublishedRagPipelineApi(Resource):
         # The role of the current user in the ta table must be admin, owner, or editor
         current_user, _ = current_account_with_tenant()
         rag_pipeline_service = RagPipelineService()
-        with Session(db.engine) as session:
-            pipeline = session.merge(pipeline)
-            workflow = rag_pipeline_service.publish_workflow(
-                session=session,
-                pipeline=pipeline,
-                account=current_user,
-            )
-            pipeline.is_published = True
-            pipeline.workflow_id = workflow.id
-            session.add(pipeline)
-            workflow_created_at = TimestampField().format(workflow.created_at)
-
-            session.commit()
+        workflow = rag_pipeline_service.publish_workflow(
+            session=db.session,  # type: ignore[reportArgumentType,arg-type]
+            pipeline=pipeline,
+            account=current_user,
+        )
+        pipeline.is_published = True
+        pipeline.workflow_id = workflow.id
+        db.session.commit()
+        workflow_created_at = TimestampField().format(workflow.created_at)
 
         return {
             "result": "success",

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py

@@ -288,26 +288,27 @@ class TidbOnQdrantVector(BaseVector):
         if not ids:
             return
 
-        try:
-            filter = models.Filter(
-                must=[
-                    models.FieldCondition(
-                        key="metadata.doc_id",
-                        match=models.MatchAny(any=ids),
-                    ),
-                ],
-            )
-            self._client.delete(
-                collection_name=self._collection_name,
-                points_selector=FilterSelector(filter=filter),
-            )
-        except UnexpectedResponse as e:
-            # Collection does not exist, so return
-            if e.status_code == 404:
-                return
-            # Some other error occurred, so re-raise the exception
-            else:
-                raise e
+        batch_size = 1000
+        for i in range(0, len(ids), batch_size):
+            batch = ids[i : i + batch_size]
+
+            try:
+                filter = models.Filter(
+                    must=[
+                        models.FieldCondition(
+                            key="metadata.doc_id",
+                            match=models.MatchAny(any=batch),
+                        ),
+                    ],
+                )
+                self._client.delete(
+                    collection_name=self._collection_name,
+                    points_selector=FilterSelector(filter=filter),
+                )
+            except UnexpectedResponse as e:
+                # Collection does not exist, so return
+                if e.status_code != 404:
+                    raise e
 
     def text_exists(self, id: str) -> bool:
         all_collection_name = []

api/models/enums.py

@@ -113,6 +113,7 @@ class DataSourceType(StrEnum):
     WEBSITE_CRAWL = "website_crawl"
     LOCAL_FILE = "local_file"
     ONLINE_DOCUMENT = "online_document"
+    ONLINE_DRIVE = "online_drive"
 
 
 class ProcessRuleMode(StrEnum):

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):

api/services/credit_pool_service.py

@@ -1,7 +1,7 @@
 import logging
 
-from sqlalchemy import update
-from sqlalchemy.orm import Session
+from sqlalchemy import select, update
+from sqlalchemy.orm import Session, sessionmaker
 
 from configs import dify_config
 from core.errors.error import QuotaExceededError
@@ -29,14 +29,15 @@ class CreditPoolService:
     @classmethod
     def get_pool(cls, tenant_id: str, pool_type: str = "trial") -> TenantCreditPool | None:
         """get tenant credit pool"""
-        return (
-            db.session.query(TenantCreditPool)
-            .filter_by(
-                tenant_id=tenant_id,
-                pool_type=pool_type,
+        with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
+            return session.scalar(
+                select(TenantCreditPool)
+                .where(
+                    TenantCreditPool.tenant_id == tenant_id,
+                    TenantCreditPool.pool_type == pool_type,
+                )
+                .limit(1)
             )
-            .first()
-        )
 
     @classmethod
     def check_credits_available(

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"]

api/tasks/async_workflow_tasks.py

@@ -156,7 +156,12 @@ def _execute_workflow_common(
                 state_owner_user_id=workflow.created_by,
             )
 
-            # Execute the workflow with the trigger type
+            # NOTE (hj24)
+            # Release the transaction before the blocking generate() call,
+            # otherwise the connection stays "idle in transaction" for hours.
+            session.commit()
+            # NOTE END
+
             generator.generate(
                 app_model=app_model,
                 workflow=workflow,

api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_generator.py

@@ -8,6 +8,7 @@ import core.app.apps.pipeline.pipeline_generator as module
 from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.datasource.entities.datasource_entities import DatasourceProviderType
+from models.enums import DataSourceType
 
 
 class FakeRagPipelineGenerateEntity(SimpleNamespace):
@@ -558,6 +559,24 @@ def test_build_document_sets_metadata_for_builtin_fields(generator, mocker):
     assert document.doc_metadata
 
 
+def test_build_document_supports_online_drive_datasource_type(generator):
+    document = generator._build_document(
+        tenant_id="tenant",
+        dataset_id="ds",
+        built_in_field_enabled=True,
+        datasource_type=DatasourceProviderType.ONLINE_DRIVE,
+        datasource_info={"id": "file-1", "bucket": "bucket-1", "name": "drive.pdf", "type": "file"},
+        created_from="rag-pipeline",
+        position=1,
+        account=_build_user(),
+        batch="batch",
+        document_form="text",
+    )
+
+    assert DataSourceType(document.data_source_type) == DataSourceType.ONLINE_DRIVE
+    assert document.name == "drive.pdf"
+
+
 def test_build_document_invalid_datasource_type(generator):
     with pytest.raises(ValueError):
         generator._build_document(

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")

web/app/components/plugins/plugin-auth/authorize/api-key-modal.tsx

@@ -142,7 +142,7 @@ const ApiKeyModal = ({
       onExtraButtonClick={onRemove}
       disabled={disabled || isLoading || doingAction}
       clickOutsideNotClose={true}
-      wrapperClassName="!z-[101]"
+      wrapperClassName="!z-[1002]"
     >
       {pluginPayload.detail && (
         <ReadmeEntrance pluginDetail={pluginPayload.detail} showType={ReadmeShowType.modal} />

web/app/components/plugins/plugin-auth/authorize/oauth-client-settings.tsx

@@ -157,7 +157,7 @@ const OAuthClientSettings = ({
         )
       }
       containerClassName="pt-0"
-      wrapperClassName="!z-[101]"
+      wrapperClassName="!z-[1002]"
       clickOutsideNotClose={true}
     >
       {pluginPayload.detail && (

web/package.json

@@ -125,15 +125,15 @@
     "mime": "4.1.0",
     "mitt": "3.0.1",
     "negotiator": "1.0.0",
-    "next": "16.2.1",
+    "next": "16.2.3",
     "next-themes": "0.4.6",
     "nuqs": "2.8.9",
     "pinyin-pro": "3.28.0",
     "qrcode.react": "4.2.0",
     "qs": "6.15.0",
-    "react": "19.2.4",
+    "react": "19.2.5",
     "react-18-input-autosize": "3.0.0",
-    "react-dom": "19.2.4",
+    "react-dom": "19.2.5",
     "react-easy-crop": "5.5.6",
     "react-hotkeys-hook": "5.2.4",
     "react-i18next": "16.6.1",
@@ -173,8 +173,8 @@
     "@mdx-js/loader": "3.1.1",
     "@mdx-js/react": "3.1.1",
     "@mdx-js/rollup": "3.1.1",
-    "@next/eslint-plugin-next": "16.2.1",
-    "@next/mdx": "16.2.1",
+    "@next/eslint-plugin-next": "16.2.3",
+    "@next/mdx": "16.2.3",
     "@rgrove/parse-xml": "4.2.0",
     "@storybook/addon-docs": "10.3.1",
     "@storybook/addon-links": "10.3.1",
@@ -231,7 +231,7 @@
     "nock": "14.0.11",
     "postcss": "8.5.8",
     "postcss-js": "5.1.0",
-    "react-server-dom-webpack": "19.2.4",
+    "react-server-dom-webpack": "19.2.5",
     "sass": "1.98.0",
     "storybook": "10.3.1",
     "tailwindcss": "3.4.19",