From f5d664887b34bae60025a01bf35f36c4736eeb7a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9D=9E=E6=B3=95=E6=93=8D=E4=BD=9C?= Date: Tue, 26 May 2026 16:50:54 +0800 Subject: [PATCH] chore: backend feature api exclude_vector_space (#36642) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- api/controllers/console/workspace/members.py | 4 +- .../console/workspace/workspace.py | 4 +- api/controllers/console/wraps.py | 23 +++++--- api/controllers/service_api/wraps.py | 17 ++++-- api/controllers/web/site.py | 4 +- api/core/provider_manager.py | 4 +- api/libs/workspace_permission.py | 2 +- api/schedule/clean_unused_datasets_task.py | 2 +- .../mail_clean_document_notify_task.py | 2 +- api/services/annotation_service.py | 2 +- api/services/billing_service.py | 3 + api/services/dataset_service.py | 6 +- api/services/document_indexing_proxy/base.py | 2 +- api/services/feature_service.py | 5 +- .../human_input_delivery_test_service.py | 2 +- .../rag_pipeline/rag_pipeline_task_proxy.py | 2 +- api/services/workspace_service.py | 2 +- api/tasks/document_indexing_task.py | 1 + api/tasks/duplicate_document_indexing_task.py | 1 + api/tasks/mail_human_input_delivery_task.py | 2 +- api/tasks/retry_document_indexing_task.py | 1 + .../sync_website_document_indexing_task.py | 1 + .../test_human_input_delivery_test_service.py | 10 ++-- .../controllers/console/test_wraps.py | 34 ++++++++++- .../console/workspace/test_workspace.py | 4 +- .../dataset/test_dataset_segment.py | 13 ++++ .../service_api/dataset/test_document.py | 12 +++- .../controllers/service_api/test_wraps.py | 59 +++++++++++++++++++ .../controllers/web/test_human_input_form.py | 4 +- .../libs/test_workspace_permission.py | 2 +- .../services/test_batch_indexing_base.py | 2 +- .../services/test_billing_service.py | 31 ++++++++++ .../test_dataset_service_lock_not_owned.py | 2 +- .../test_document_indexing_task_proxy.py | 2 +- ..._duplicate_document_indexing_task_proxy.py | 2 +- .../test_feature_service_vector_space.py | 37 ++++++++++++ .../services/test_rag_pipeline_task_proxy.py | 2 +- .../test_mail_human_input_delivery_task.py | 8 +-- 38 files changed, 262 insertions(+), 54 deletions(-) create mode 100644 api/tests/unit_tests/services/test_feature_service_vector_space.py diff --git a/api/controllers/console/workspace/members.py b/api/controllers/console/workspace/members.py index 1f770faa11..93a830e8ee 100644 --- a/api/controllers/console/workspace/members.py +++ b/api/controllers/console/workspace/members.py @@ -77,7 +77,7 @@ register_response_schema_models(console_ns, SimpleResultDataResponse, Verificati def _is_role_enabled(role: TenantAccountRole | str, tenant_id: str) -> bool: if role != TenantAccountRole.DATASET_OPERATOR: return True - return FeatureService.get_features(tenant_id=tenant_id).dataset_operator_enabled + return FeatureService.get_features(tenant_id=tenant_id, exclude_vector_space=True).dataset_operator_enabled def _normalize_invitee_emails(emails: list[str]) -> list[str]: @@ -113,7 +113,7 @@ def _check_member_invite_limits(tenant_id: str, new_member_count: int) -> None: if new_member_count <= 0: return - features = FeatureService.get_features(tenant_id=tenant_id) + features = FeatureService.get_features(tenant_id=tenant_id, exclude_vector_space=True) if dify_config.ENTERPRISE_ENABLED: workspace_members = features.workspace_members diff --git a/api/controllers/console/workspace/workspace.py b/api/controllers/console/workspace/workspace.py index 77501eed72..57a37f20f1 100644 --- a/api/controllers/console/workspace/workspace.py +++ b/api/controllers/console/workspace/workspace.py @@ -166,10 +166,10 @@ class TenantListApi(Resource): if tenant_plan: plan = tenant_plan["plan"] or CloudPlan.SANDBOX else: - features = FeatureService.get_features(tenant.id) + features = FeatureService.get_features(tenant.id, exclude_vector_space=True) plan = features.billing.subscription.plan or CloudPlan.SANDBOX elif not is_enterprise_only: - features = FeatureService.get_features(tenant.id) + features = FeatureService.get_features(tenant.id, exclude_vector_space=True) plan = features.billing.subscription.plan or CloudPlan.SANDBOX # Create a dictionary with tenant attributes diff --git a/api/controllers/console/wraps.py b/api/controllers/console/wraps.py index ad406f2a9e..603645278f 100644 --- a/api/controllers/console/wraps.py +++ b/api/controllers/console/wraps.py @@ -96,21 +96,28 @@ def cloud_edition_billing_resource_check[**P, R](resource: str) -> Callable[[Cal @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): _, current_tenant_id = current_account_with_tenant() - features = FeatureService.get_features(current_tenant_id) + if resource == "vector_space": + if not dify_config.BILLING_ENABLED: + return view(*args, **kwargs) + + vector_space = FeatureService.get_vector_space(current_tenant_id) + if 0 < vector_space.limit <= vector_space.size: + abort( + 403, + "The capacity of the knowledge storage space has reached the limit of your subscription.", + ) + return view(*args, **kwargs) + + features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True) if features.billing.enabled: members = features.members apps = features.apps - vector_space = features.vector_space documents_upload_quota = features.documents_upload_quota annotation_quota_limit = features.annotation_quota_limit if resource == "members" and 0 < members.limit <= members.size: abort(403, "The number of members has reached the limit of your subscription.") elif resource == "apps" and 0 < apps.limit <= apps.size: abort(403, "The number of apps has reached the limit of your subscription.") - elif resource == "vector_space" and 0 < vector_space.limit <= vector_space.size: - abort( - 403, "The capacity of the knowledge storage space has reached the limit of your subscription." - ) elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size: # The api of file upload is used in the multiple places, # so we need to check the source of the request from datasets @@ -140,7 +147,7 @@ def cloud_edition_billing_knowledge_limit_check[**P, R]( @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): _, current_tenant_id = current_account_with_tenant() - features = FeatureService.get_features(current_tenant_id) + features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True) if features.billing.enabled: if resource == "add_segment": if features.billing.subscription.plan == CloudPlan.SANDBOX: @@ -291,7 +298,7 @@ def knowledge_pipeline_publish_enabled[**P, R](view: Callable[P, R]) -> Callable @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): _, current_tenant_id = current_account_with_tenant() - features = FeatureService.get_features(current_tenant_id) + features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True) if features.knowledge_pipeline.publish_enabled: return view(*args, **kwargs) abort(403) diff --git a/api/controllers/service_api/wraps.py b/api/controllers/service_api/wraps.py index b9389ccc47..013ea34a6a 100644 --- a/api/controllers/service_api/wraps.py +++ b/api/controllers/service_api/wraps.py @@ -13,6 +13,7 @@ from pydantic import BaseModel from sqlalchemy import select from werkzeug.exceptions import Forbidden, NotFound, Unauthorized +from configs import dify_config from enums.cloud_plan import CloudPlan from extensions.ext_database import db from extensions.ext_redis import redis_client @@ -140,20 +141,26 @@ def cloud_edition_billing_resource_check[**P, R]( def interceptor(view: Callable[P, R]): def decorated(*args: P.args, **kwargs: P.kwargs): api_token = validate_and_get_api_token(api_token_type) - features = FeatureService.get_features(api_token.tenant_id) + if resource == "vector_space": + if not dify_config.BILLING_ENABLED: + return view(*args, **kwargs) + + vector_space = FeatureService.get_vector_space(api_token.tenant_id) + if 0 < vector_space.limit <= vector_space.size: + raise Forbidden("The capacity of the vector space has reached the limit of your subscription.") + return view(*args, **kwargs) + + features = FeatureService.get_features(api_token.tenant_id, exclude_vector_space=True) if features.billing.enabled: members = features.members apps = features.apps - vector_space = features.vector_space documents_upload_quota = features.documents_upload_quota if resource == "members" and 0 < members.limit <= members.size: raise Forbidden("The number of members has reached the limit of your subscription.") elif resource == "apps" and 0 < apps.limit <= apps.size: raise Forbidden("The number of apps has reached the limit of your subscription.") - elif resource == "vector_space" and 0 < vector_space.limit <= vector_space.size: - raise Forbidden("The capacity of the vector space has reached the limit of your subscription.") elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size: raise Forbidden("The number of documents has reached the limit of your subscription.") else: @@ -174,7 +181,7 @@ def cloud_edition_billing_knowledge_limit_check[**P, R]( @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): api_token = validate_and_get_api_token(api_token_type) - features = FeatureService.get_features(api_token.tenant_id) + features = FeatureService.get_features(api_token.tenant_id, exclude_vector_space=True) if features.billing.enabled: if resource == "add_segment": if features.billing.subscription.plan == CloudPlan.SANDBOX: diff --git a/api/controllers/web/site.py b/api/controllers/web/site.py index 7d2080dd91..bd21632b05 100644 --- a/api/controllers/web/site.py +++ b/api/controllers/web/site.py @@ -81,7 +81,7 @@ class AppSiteApi(WebApiResource): if app_model.tenant.status == TenantStatus.ARCHIVE: raise Forbidden() - can_replace_logo = FeatureService.get_features(app_model.tenant_id).can_replace_logo + can_replace_logo = FeatureService.get_features(app_model.tenant_id, exclude_vector_space=True).can_replace_logo return AppSiteInfo(app_model.tenant, app_model, site, end_user.id, can_replace_logo) @@ -119,6 +119,6 @@ def serialize_site(site: Site) -> dict[str, Any]: def serialize_app_site_payload(app_model: App, site: Site, end_user_id: str | None) -> dict[str, Any]: - can_replace_logo = FeatureService.get_features(app_model.tenant_id).can_replace_logo + can_replace_logo = FeatureService.get_features(app_model.tenant_id, exclude_vector_space=True).can_replace_logo app_site_info = AppSiteInfo(app_model.tenant, app_model, site, end_user_id, can_replace_logo) return cast(dict[str, Any], marshal(app_site_info, AppSiteApi.app_fields)) diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index 9faa70a0b8..0ba668a5e8 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -534,7 +534,9 @@ class ProviderManager: cache_key = f"tenant:{tenant_id}:model_load_balancing_enabled" cache_result = redis_client.get(cache_key) if cache_result is None: - model_load_balancing_enabled = FeatureService.get_features(tenant_id).model_load_balancing_enabled + model_load_balancing_enabled = FeatureService.get_features( + tenant_id, exclude_vector_space=True + ).model_load_balancing_enabled redis_client.setex(cache_key, 120, str(model_load_balancing_enabled)) else: cache_result = cache_result.decode("utf-8") diff --git a/api/libs/workspace_permission.py b/api/libs/workspace_permission.py index dd42a7facf..435b07dd6e 100644 --- a/api/libs/workspace_permission.py +++ b/api/libs/workspace_permission.py @@ -58,7 +58,7 @@ def check_workspace_owner_transfer_permission(workspace_id: str) -> None: Raises: Forbidden: If either billing plan or workspace policy prohibits ownership transfer """ - features = FeatureService.get_features(workspace_id) + features = FeatureService.get_features(workspace_id, exclude_vector_space=True) if not features.is_allow_transfer_workspace: raise Forbidden("Your current plan does not allow workspace ownership transfer") diff --git a/api/schedule/clean_unused_datasets_task.py b/api/schedule/clean_unused_datasets_task.py index 0b0fc1b229..849274311a 100644 --- a/api/schedule/clean_unused_datasets_task.py +++ b/api/schedule/clean_unused_datasets_task.py @@ -112,7 +112,7 @@ def clean_unused_datasets_task(): features_cache_key = f"features:{dataset.tenant_id}" plan_cache = redis_client.get(features_cache_key) if plan_cache is None: - features = FeatureService.get_features(dataset.tenant_id) + features = FeatureService.get_features(dataset.tenant_id, exclude_vector_space=True) redis_client.setex(features_cache_key, 600, features.billing.subscription.plan) plan = features.billing.subscription.plan else: diff --git a/api/schedule/mail_clean_document_notify_task.py b/api/schedule/mail_clean_document_notify_task.py index 2cc0192a4a..1a76a4aa30 100644 --- a/api/schedule/mail_clean_document_notify_task.py +++ b/api/schedule/mail_clean_document_notify_task.py @@ -45,7 +45,7 @@ def mail_clean_document_notify_task(): dataset_auto_disable_logs_map[dataset_auto_disable_log.tenant_id].append(dataset_auto_disable_log) url = f"{dify_config.CONSOLE_WEB_URL}/datasets" for tenant_id, tenant_dataset_auto_disable_logs in dataset_auto_disable_logs_map.items(): - features = FeatureService.get_features(tenant_id) + features = FeatureService.get_features(tenant_id, exclude_vector_space=True) plan = features.billing.subscription.plan if plan != CloudPlan.SANDBOX: knowledge_details = [] diff --git a/api/services/annotation_service.py b/api/services/annotation_service.py index aa6b8ffc6e..e1762c686f 100644 --- a/api/services/annotation_service.py +++ b/api/services/annotation_service.py @@ -521,7 +521,7 @@ class AppAnnotationService: ) # Check annotation quota limit - features = FeatureService.get_features(current_tenant_id) + features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True) if features.billing.enabled: annotation_quota_limit = features.annotation_quota_limit if annotation_quota_limit.limit < len(result) + annotation_quota_limit.size: diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 6021d46c72..5c59a5f8dd 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -192,6 +192,9 @@ class BillingService: params["exclude_vector_space"] = "true" billing_info = cls._send_request("GET", "/subscription/info", params=params) + if exclude_vector_space and billing_info.get("vector_space") is None: + # Unset proto message fields can be serialized as null; the light billing contract treats it as absent. + billing_info.pop("vector_space", None) return _billing_info_adapter.validate_python(billing_info) @classmethod diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py index 1332390d46..2e593ea71f 100644 --- a/api/services/dataset_service.py +++ b/api/services/dataset_service.py @@ -1325,7 +1325,7 @@ class DatasetService: def get_dataset_auto_disable_logs(dataset_id: str) -> AutoDisableLogsDict: assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None - features = FeatureService.get_features(current_user.current_tenant_id) + features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True) if not features.billing.enabled or features.billing.subscription.plan == CloudPlan.SANDBOX: return { "document_ids": [], @@ -2007,7 +2007,7 @@ class DocumentService: assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None - features = FeatureService.get_features(current_user.current_tenant_id) + features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True) if features.billing.enabled: if not knowledge_config.original_document_id: @@ -2798,7 +2798,7 @@ class DocumentService: assert current_user.current_tenant_id is not None assert knowledge_config.data_source - features = FeatureService.get_features(current_user.current_tenant_id) + features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True) if features.billing.enabled: count = 0 diff --git a/api/services/document_indexing_proxy/base.py b/api/services/document_indexing_proxy/base.py index 56e47857c9..02df6752f3 100644 --- a/api/services/document_indexing_proxy/base.py +++ b/api/services/document_indexing_proxy/base.py @@ -41,7 +41,7 @@ class DocumentTaskProxyBase(ABC): @cached_property def features(self): - return FeatureService.get_features(self._tenant_id) + return FeatureService.get_features(self._tenant_id, exclude_vector_space=True) @abstractmethod def _send_to_direct_queue(self, task_func: Callable[..., Any]): diff --git a/api/services/feature_service.py b/api/services/feature_service.py index c4f723fde8..7137a9c8c9 100644 --- a/api/services/feature_service.py +++ b/api/services/feature_service.py @@ -130,7 +130,7 @@ class FeatureModel(FeatureResponseModel): education: EducationModel = EducationModel() members: LimitationModel = LimitationModel(size=0, limit=1) apps: LimitationModel = LimitationModel(size=0, limit=10) - vector_space: LimitationModel = LimitationModel(size=0, limit=5) + vector_space: LimitationModel | None = LimitationModel(size=0, limit=5) knowledge_rate_limit: int = 10 annotation_quota_limit: LimitationModel = LimitationModel(size=0, limit=10) documents_upload_quota: LimitationModel = LimitationModel(size=0, limit=50) @@ -188,6 +188,8 @@ class FeatureService: @classmethod def get_features(cls, tenant_id: str, exclude_vector_space: bool = False) -> FeatureModel: features = FeatureModel() + if exclude_vector_space: + features.vector_space = None cls._fulfill_params_from_env(features) @@ -347,6 +349,7 @@ class FeatureService: features.apps.limit = billing_info["apps"]["limit"] if not exclude_vector_space: + assert features.vector_space is not None cls._fulfill_vector_space_from_billing_info(features.vector_space, billing_info) if "documents_upload_quota" in billing_info: diff --git a/api/services/human_input_delivery_test_service.py b/api/services/human_input_delivery_test_service.py index 8b4983e5f7..c266d4f958 100644 --- a/api/services/human_input_delivery_test_service.py +++ b/api/services/human_input_delivery_test_service.py @@ -136,7 +136,7 @@ class EmailDeliveryTestHandler: ) -> DeliveryTestResult: if not isinstance(method, EmailDeliveryMethod): raise DeliveryTestUnsupportedError("Delivery method does not support test send.") - features = FeatureService.get_features(context.tenant_id) + features = FeatureService.get_features(context.tenant_id, exclude_vector_space=True) if not features.human_input_email_delivery_enabled: raise DeliveryTestError("Email delivery is not available for current plan.") if not mail.is_inited(): diff --git a/api/services/rag_pipeline/rag_pipeline_task_proxy.py b/api/services/rag_pipeline/rag_pipeline_task_proxy.py index 1a7b104a70..52ebbce65a 100644 --- a/api/services/rag_pipeline/rag_pipeline_task_proxy.py +++ b/api/services/rag_pipeline/rag_pipeline_task_proxy.py @@ -29,7 +29,7 @@ class RagPipelineTaskProxy: @cached_property def features(self): - return FeatureService.get_features(self._dataset_tenant_id) + return FeatureService.get_features(self._dataset_tenant_id, exclude_vector_space=True) def _upload_invoke_entities(self) -> str: text = [item.model_dump() for item in self._rag_pipeline_invoke_entities] diff --git a/api/services/workspace_service.py b/api/services/workspace_service.py index eb4671cfaa..70114a83f0 100644 --- a/api/services/workspace_service.py +++ b/api/services/workspace_service.py @@ -33,7 +33,7 @@ class WorkspaceService: assert tenant_account_join is not None, "TenantAccountJoin not found" tenant_info["role"] = tenant_account_join.role - feature = FeatureService.get_features(tenant.id) + feature = FeatureService.get_features(tenant.id, exclude_vector_space=True) can_replace_logo = feature.can_replace_logo if can_replace_logo and TenantService.has_roles(tenant, [TenantAccountRole.OWNER, TenantAccountRole.ADMIN]): diff --git a/api/tasks/document_indexing_task.py b/api/tasks/document_indexing_task.py index 31dad7937c..c173606d74 100644 --- a/api/tasks/document_indexing_task.py +++ b/api/tasks/document_indexing_task.py @@ -66,6 +66,7 @@ def _document_indexing(dataset_id: str, document_ids: Sequence[str]): try: if features.billing.enabled: vector_space = features.vector_space + assert vector_space is not None count = len(document_ids) batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT) if features.billing.subscription.plan == CloudPlan.SANDBOX and count > 1: diff --git a/api/tasks/duplicate_document_indexing_task.py b/api/tasks/duplicate_document_indexing_task.py index 71f367c5e7..69747ba497 100644 --- a/api/tasks/duplicate_document_indexing_task.py +++ b/api/tasks/duplicate_document_indexing_task.py @@ -92,6 +92,7 @@ def _duplicate_document_indexing_task(dataset_id: str, document_ids: Sequence[st try: if features.billing.enabled: vector_space = features.vector_space + assert vector_space is not None count = len(document_ids) if features.billing.subscription.plan == CloudPlan.SANDBOX and count > 1: raise ValueError("Your current plan does not support batch upload, please upgrade your plan.") diff --git a/api/tasks/mail_human_input_delivery_task.py b/api/tasks/mail_human_input_delivery_task.py index 2a60be7762..8ed50071df 100644 --- a/api/tasks/mail_human_input_delivery_task.py +++ b/api/tasks/mail_human_input_delivery_task.py @@ -157,7 +157,7 @@ def dispatch_human_input_email_task(form_id: str, node_title: str | None = None, if form is None: logger.warning("Human input form not found, form_id=%s", form_id) return - features = FeatureService.get_features(form.tenant_id) + features = FeatureService.get_features(form.tenant_id, exclude_vector_space=True) if not features.human_input_email_delivery_enabled: logger.info( "Human input email delivery is not available for tenant=%s, form_id=%s", diff --git a/api/tasks/retry_document_indexing_task.py b/api/tasks/retry_document_indexing_task.py index 0df5896ce3..fa02afda15 100644 --- a/api/tasks/retry_document_indexing_task.py +++ b/api/tasks/retry_document_indexing_task.py @@ -52,6 +52,7 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str], user_ try: if features.billing.enabled: vector_space = features.vector_space + assert vector_space is not None if 0 < vector_space.limit <= vector_space.size: raise ValueError( "Your total number of documents plus the number of uploads have over the limit of " diff --git a/api/tasks/sync_website_document_indexing_task.py b/api/tasks/sync_website_document_indexing_task.py index 06eb460311..5bdeac7e9d 100644 --- a/api/tasks/sync_website_document_indexing_task.py +++ b/api/tasks/sync_website_document_indexing_task.py @@ -39,6 +39,7 @@ def sync_website_document_indexing_task(dataset_id: str, document_id: str): try: if features.billing.enabled: vector_space = features.vector_space + assert vector_space is not None if 0 < vector_space.limit <= vector_space.size: raise ValueError( "Your total number of documents plus the number of uploads have over the limit of " diff --git a/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test_service.py b/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test_service.py index 454b8096d1..718ff05d22 100644 --- a/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test_service.py +++ b/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test_service.py @@ -127,7 +127,7 @@ class TestEmailDeliveryTestHandler: monkeypatch.setattr( service_module.FeatureService, "get_features", - lambda _tenant_id: SimpleNamespace(human_input_email_delivery_enabled=False), + lambda _tenant_id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=False), ) handler = EmailDeliveryTestHandler(session_factory=MagicMock()) context = DeliveryTestContext( @@ -142,7 +142,7 @@ class TestEmailDeliveryTestHandler: monkeypatch.setattr( service_module.FeatureService, "get_features", - lambda _id: SimpleNamespace(human_input_email_delivery_enabled=True), + lambda _id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True), ) monkeypatch.setattr(service_module.mail, "is_inited", lambda: False) @@ -159,7 +159,7 @@ class TestEmailDeliveryTestHandler: monkeypatch.setattr( service_module.FeatureService, "get_features", - lambda _id: SimpleNamespace(human_input_email_delivery_enabled=True), + lambda _id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True), ) monkeypatch.setattr(service_module.mail, "is_inited", lambda: True) @@ -178,7 +178,7 @@ class TestEmailDeliveryTestHandler: monkeypatch.setattr( service_module.FeatureService, "get_features", - lambda _id: SimpleNamespace(human_input_email_delivery_enabled=True), + lambda _id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True), ) monkeypatch.setattr(service_module.mail, "is_inited", lambda: True) mock_mail_send = MagicMock() @@ -214,7 +214,7 @@ class TestEmailDeliveryTestHandler: monkeypatch.setattr( service_module.FeatureService, "get_features", - lambda _id: SimpleNamespace(human_input_email_delivery_enabled=True), + lambda _id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True), ) monkeypatch.setattr(service_module.mail, "is_inited", lambda: True) mock_mail_send = MagicMock() diff --git a/api/tests/unit_tests/controllers/console/test_wraps.py b/api/tests/unit_tests/controllers/console/test_wraps.py index 6ddb1748d6..c392ffc69d 100644 --- a/api/tests/unit_tests/controllers/console/test_wraps.py +++ b/api/tests/unit_tests/controllers/console/test_wraps.py @@ -255,11 +255,43 @@ class TestBillingResourceLimits: with patch( "controllers.console.wraps.current_account_with_tenant", return_value=(MockUser("test_user"), "tenant123") ): - with patch("controllers.console.wraps.FeatureService.get_features", return_value=mock_features): + with patch( + "controllers.console.wraps.FeatureService.get_features", return_value=mock_features + ) as get_features: result = add_member() # Assert assert result == "member_added" + get_features.assert_called_once_with("tenant123", exclude_vector_space=True) + + def test_should_load_vector_space_from_dedicated_quota_api(self): + """Test vector-space limit checks avoid loading the full feature payload.""" + # Arrange + mock_vector_space = MagicMock() + mock_vector_space.limit = 10 + mock_vector_space.size = 5 + + @cloud_edition_billing_resource_check("vector_space") + def add_segment(): + return "segment_added" + + # Act + with patch( + "controllers.console.wraps.current_account_with_tenant", return_value=(MockUser("test_user"), "tenant123") + ): + with ( + patch("controllers.console.wraps.dify_config.BILLING_ENABLED", True), + patch( + "controllers.console.wraps.FeatureService.get_vector_space", return_value=mock_vector_space + ) as get_vector_space, + patch("controllers.console.wraps.FeatureService.get_features") as get_features, + ): + result = add_segment() + + # Assert + assert result == "segment_added" + get_vector_space.assert_called_once_with("tenant123") + get_features.assert_not_called() def test_should_reject_when_over_resource_limit(self): """Test that requests are rejected when over resource limits""" diff --git a/api/tests/unit_tests/controllers/console/workspace/test_workspace.py b/api/tests/unit_tests/controllers/console/workspace/test_workspace.py index 0ee1877630..95c69d30c2 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_workspace.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_workspace.py @@ -139,7 +139,7 @@ class TestTenantListApi: assert result["workspaces"][0]["plan"] == CloudPlan.TEAM assert result["workspaces"][1]["plan"] == CloudPlan.PROFESSIONAL get_plan_bulk_mock.assert_called_once_with(["t1", "t2"]) - get_features_mock.assert_called_once_with("t2") + get_features_mock.assert_called_once_with("t2", exclude_vector_space=True) def test_get_saas_path_falls_back_to_legacy_feature_path_on_bulk_error(self, app: Flask): """Test fallback to FeatureService when bulk billing returns empty result. @@ -235,7 +235,7 @@ class TestTenantListApi: assert status == 200 assert result["workspaces"][0]["plan"] == CloudPlan.SANDBOX - get_features_mock.assert_called_once_with("t1") + get_features_mock.assert_called_once_with("t1", exclude_vector_space=True) def test_get_enterprise_only_skips_feature_service(self, app: Flask): api = TenantListApi() diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py index fe8fc02548..2e1051ab6b 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py @@ -872,6 +872,11 @@ class TestSegmentApiPost: mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features + mock_vector_space = Mock() + mock_vector_space.limit = 10 + mock_vector_space.size = 0 + mock_feature_svc.get_vector_space.return_value = mock_vector_space + mock_rate_limit = Mock() mock_rate_limit.enabled = False mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit @@ -1209,6 +1214,10 @@ class TestDatasetSegmentApiUpdate: mock_features = Mock() mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features + mock_vector_space = Mock() + mock_vector_space.limit = 10 + mock_vector_space.size = 0 + mock_feature_svc.get_vector_space.return_value = mock_vector_space mock_rate_limit = Mock() mock_rate_limit.enabled = False mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit @@ -1710,6 +1719,10 @@ class TestChildChunkApiPost: mock_features = Mock() mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features + mock_vector_space = Mock() + mock_vector_space.limit = 10 + mock_vector_space.size = 0 + mock_feature_svc.get_vector_space.return_value = mock_vector_space mock_rate_limit = Mock() mock_rate_limit.enabled = False mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_document.py b/api/tests/unit_tests/controllers/service_api/dataset/test_document.py index 61ec397193..2185e65326 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_document.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_document.py @@ -950,7 +950,8 @@ class TestDocumentAddByTextApi: """Configure mocks to neutralise billing/auth decorators. ``cloud_edition_billing_resource_check`` calls - ``FeatureService.get_features`` and + ``FeatureService.get_vector_space`` for vector-space checks and + ``FeatureService.get_features`` for other resource checks. ``cloud_edition_billing_rate_limit_check`` calls ``FeatureService.get_knowledge_rate_limit``. Both call ``validate_and_get_api_token`` first. @@ -963,6 +964,11 @@ class TestDocumentAddByTextApi: mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features + mock_vector_space = Mock() + mock_vector_space.limit = 10 + mock_vector_space.size = 0 + mock_feature_svc.get_vector_space.return_value = mock_vector_space + mock_rate_limit = Mock() mock_rate_limit.enabled = False mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit @@ -1140,6 +1146,10 @@ def _setup_billing_mocks(mock_validate_token, mock_feature_svc, tenant_id: str): mock_features = Mock() mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features + mock_vector_space = Mock() + mock_vector_space.limit = 10 + mock_vector_space.size = 0 + mock_feature_svc.get_vector_space.return_value = mock_vector_space mock_rate_limit = Mock() mock_rate_limit.enabled = False mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit diff --git a/api/tests/unit_tests/controllers/service_api/test_wraps.py b/api/tests/unit_tests/controllers/service_api/test_wraps.py index 6e8d971c0d..0b5c5d95b6 100644 --- a/api/tests/unit_tests/controllers/service_api/test_wraps.py +++ b/api/tests/unit_tests/controllers/service_api/test_wraps.py @@ -265,6 +265,65 @@ class TestCloudEditionBillingResourceCheck: # Assert assert result == "member_added" + mock_get_features.assert_called_once_with("tenant123", exclude_vector_space=True) + + @patch("controllers.service_api.wraps.validate_and_get_api_token") + @patch("controllers.service_api.wraps.FeatureService.get_features") + @patch("controllers.service_api.wraps.FeatureService.get_vector_space") + def test_loads_vector_space_from_dedicated_quota_api( + self, mock_get_vector_space, mock_get_features, mock_validate_token, app: Flask + ): + """Test vector-space resource checks avoid loading the full feature payload.""" + # Arrange + mock_validate_token.return_value = Mock(tenant_id="tenant123") + + mock_vector_space = Mock() + mock_vector_space.limit = 10 + mock_vector_space.size = 5 + mock_get_vector_space.return_value = mock_vector_space + + @cloud_edition_billing_resource_check("vector_space", "dataset") + def add_segment(): + return "segment_added" + + # Act + with ( + app.test_request_context("/", method="GET"), + patch("controllers.service_api.wraps.dify_config.BILLING_ENABLED", True), + ): + result = add_segment() + + # Assert + assert result == "segment_added" + mock_get_vector_space.assert_called_once_with("tenant123") + mock_get_features.assert_not_called() + + @patch("controllers.service_api.wraps.validate_and_get_api_token") + @patch("controllers.service_api.wraps.FeatureService.get_features") + def test_loads_features_when_checking_non_vector_space_limit( + self, mock_get_features, mock_validate_token, app: Flask + ): + """Test non-vector-space resource checks keep using the light feature payload.""" + # Arrange + mock_validate_token.return_value = Mock(tenant_id="tenant123") + + mock_features = Mock() + mock_features.billing.enabled = True + mock_features.documents_upload_quota.limit = 10 + mock_features.documents_upload_quota.size = 5 + mock_get_features.return_value = mock_features + + @cloud_edition_billing_resource_check("documents", "dataset") + def upload_document(): + return "document_uploaded" + + # Act + with app.test_request_context("/", method="GET"): + result = upload_document() + + # Assert + assert result == "document_uploaded" + mock_get_features.assert_called_once_with("tenant123", exclude_vector_space=True) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") diff --git a/api/tests/unit_tests/controllers/web/test_human_input_form.py b/api/tests/unit_tests/controllers/web/test_human_input_form.py index 5f2dc19aab..f9d49237b7 100644 --- a/api/tests/unit_tests/controllers/web/test_human_input_form.py +++ b/api/tests/unit_tests/controllers/web/test_human_input_form.py @@ -126,7 +126,7 @@ def test_get_form_includes_site(monkeypatch: pytest.MonkeyPatch, app: Flask): monkeypatch.setattr( site_module.FeatureService, "get_features", - lambda tenant_id: SimpleNamespace(can_replace_logo=True), + lambda tenant_id, **_kwargs: SimpleNamespace(can_replace_logo=True), ) with app.test_request_context("/api/form/human_input/token-1", method="GET"): @@ -245,7 +245,7 @@ def test_get_form_allows_backstage_token(monkeypatch: pytest.MonkeyPatch, app: F monkeypatch.setattr( site_module.FeatureService, "get_features", - lambda tenant_id: SimpleNamespace(can_replace_logo=True), + lambda tenant_id, **_kwargs: SimpleNamespace(can_replace_logo=True), ) with app.test_request_context("/api/form/human_input/token-1", method="GET"): diff --git a/api/tests/unit_tests/libs/test_workspace_permission.py b/api/tests/unit_tests/libs/test_workspace_permission.py index 89586ccf26..48d9b351a4 100644 --- a/api/tests/unit_tests/libs/test_workspace_permission.py +++ b/api/tests/unit_tests/libs/test_workspace_permission.py @@ -36,7 +36,7 @@ class TestWorkspacePermissionHelper: # Should not raise check_workspace_owner_transfer_permission("test-workspace-id") - mock_feature_service.get_features.assert_called_once_with("test-workspace-id") + mock_feature_service.get_features.assert_called_once_with("test-workspace-id", exclude_vector_space=True) @patch("libs.workspace_permission.EnterpriseService") @patch("libs.workspace_permission.dify_config") diff --git a/api/tests/unit_tests/services/test_batch_indexing_base.py b/api/tests/unit_tests/services/test_batch_indexing_base.py index bd68b67d89..8d07fc3050 100644 --- a/api/tests/unit_tests/services/test_batch_indexing_base.py +++ b/api/tests/unit_tests/services/test_batch_indexing_base.py @@ -344,7 +344,7 @@ class TestDispatchRouting: proxy._dispatch() # Assert - mock_features.assert_called_once_with(TENANT_ID) + mock_features.assert_called_once_with(TENANT_ID, exclude_vector_space=True) class TestBaseRouterHelpers: diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index e7a195a472..67d1cc0291 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/api/tests/unit_tests/services/test_billing_service.py @@ -343,6 +343,37 @@ class TestBillingServiceSubscriptionInfo: params={"tenant_id": tenant_id, "exclude_vector_space": "true"}, ) + def test_get_info_exclude_vector_space_normalizes_null_field(self, mock_send_request): + """When billing serializes skipped vector_space as null, get_info treats it as absent.""" + # Arrange + tenant_id = "tenant-123" + expected_response = { + "enabled": True, + "subscription": {"plan": "professional", "interval": "month", "education": False}, + "members": {"size": 1, "limit": 50}, + "apps": {"size": 1, "limit": 200}, + "vector_space": None, + "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, + } + mock_send_request.return_value = expected_response + + # Act + result = BillingService.get_info(tenant_id, exclude_vector_space=True) + + # Assert + assert "vector_space" not in result + mock_send_request.assert_called_once_with( + "GET", + "/subscription/info", + params={"tenant_id": tenant_id, "exclude_vector_space": "true"}, + ) + def test_get_vector_space_success(self, mock_send_request): """Test successful retrieval of vector-space usage and limit.""" # Arrange diff --git a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py index f5879d973d..352a765de2 100644 --- a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py +++ b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py @@ -39,7 +39,7 @@ def fake_features(monkeypatch: pytest.MonkeyPatch): ) monkeypatch.setattr( "services.dataset_service.FeatureService.get_features", - lambda tenant_id: features, + lambda tenant_id, **_kwargs: features, ) return features diff --git a/api/tests/unit_tests/services/test_document_indexing_task_proxy.py b/api/tests/unit_tests/services/test_document_indexing_task_proxy.py index 98c30c3722..28de9efa57 100644 --- a/api/tests/unit_tests/services/test_document_indexing_task_proxy.py +++ b/api/tests/unit_tests/services/test_document_indexing_task_proxy.py @@ -75,7 +75,7 @@ class TestDocumentIndexingTaskProxy: assert features1 == mock_features assert features2 == mock_features assert features1 is features2 # Should be the same instance due to caching - mock_feature_service.get_features.assert_called_once_with("tenant-123") + mock_feature_service.get_features.assert_called_once_with("tenant-123", exclude_vector_space=True) @patch("services.document_indexing_proxy.document_indexing_task_proxy.normal_document_indexing_task") def test_send_to_direct_queue(self, mock_task): diff --git a/api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py b/api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py index 68bafe3d5e..20358d6a0c 100644 --- a/api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py +++ b/api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py @@ -94,7 +94,7 @@ class TestDuplicateDocumentIndexingTaskProxy: assert features1 == mock_features assert features2 == mock_features assert features1 is features2 # Should be the same instance due to caching - mock_feature_service.get_features.assert_called_once_with("tenant-123") + mock_feature_service.get_features.assert_called_once_with("tenant-123", exclude_vector_space=True) @patch( "services.document_indexing_proxy.duplicate_document_indexing_task_proxy.normal_duplicate_document_indexing_task" diff --git a/api/tests/unit_tests/services/test_feature_service_vector_space.py b/api/tests/unit_tests/services/test_feature_service_vector_space.py new file mode 100644 index 0000000000..dc0e7ddb54 --- /dev/null +++ b/api/tests/unit_tests/services/test_feature_service_vector_space.py @@ -0,0 +1,37 @@ +from unittest.mock import patch + +from services.feature_service import FeatureService + + +def test_get_features_exclude_vector_space_sets_vector_space_to_none(): + tenant_id = "tenant-id" + billing_info = { + "enabled": True, + "subscription": {"plan": "pro", "interval": "monthly", "education": False}, + "members": {"size": 1, "limit": 10}, + "apps": {"size": 2, "limit": 20}, + "documents_upload_quota": {"size": 3, "limit": 100}, + "annotation_quota_limit": {"size": 4, "limit": 50}, + "docs_processing": "standard", + "can_replace_logo": True, + "model_load_balancing_enabled": True, + "knowledge_rate_limit": {"limit": 100}, + "knowledge_pipeline_publish_enabled": True, + } + + with ( + patch("services.feature_service.dify_config") as mock_config, + patch("services.feature_service.BillingService.get_info", return_value=billing_info) as get_info, + patch("services.feature_service.BillingService.get_quota_info", return_value={}), + ): + mock_config.BILLING_ENABLED = True + mock_config.ENTERPRISE_ENABLED = False + mock_config.CAN_REPLACE_LOGO = False + mock_config.MODEL_LB_ENABLED = False + mock_config.DATASET_OPERATOR_ENABLED = False + mock_config.EDUCATION_ENABLED = False + + features = FeatureService.get_features(tenant_id, exclude_vector_space=True) + + assert features.vector_space is None + get_info.assert_called_once_with(tenant_id, exclude_vector_space=True) diff --git a/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py b/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py index f5a48b1416..cfc685e4cb 100644 --- a/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py +++ b/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py @@ -144,7 +144,7 @@ class TestRagPipelineTaskProxy: assert features1 == mock_features assert features2 == mock_features assert features1 is features2 # Should be the same instance due to caching - mock_feature_service.get_features.assert_called_once_with("tenant-123") + mock_feature_service.get_features.assert_called_once_with("tenant-123", exclude_vector_space=True) @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") diff --git a/api/tests/unit_tests/tasks/test_mail_human_input_delivery_task.py b/api/tests/unit_tests/tasks/test_mail_human_input_delivery_task.py index 37b7a85451..f8f9ec9971 100644 --- a/api/tests/unit_tests/tasks/test_mail_human_input_delivery_task.py +++ b/api/tests/unit_tests/tasks/test_mail_human_input_delivery_task.py @@ -54,7 +54,7 @@ def test_dispatch_human_input_email_task_sends_to_each_recipient(monkeypatch: py monkeypatch.setattr( task_module.FeatureService, "get_features", - lambda _tenant_id: SimpleNamespace(human_input_email_delivery_enabled=True), + lambda _tenant_id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True), ) jobs: Sequence[task_module._EmailDeliveryJob] = [_build_job(recipient_count=2)] monkeypatch.setattr(task_module, "_load_email_jobs", lambda _session, _form: jobs) @@ -78,7 +78,7 @@ def test_dispatch_human_input_email_task_skips_when_feature_disabled(monkeypatch monkeypatch.setattr( task_module.FeatureService, "get_features", - lambda _tenant_id: SimpleNamespace(human_input_email_delivery_enabled=False), + lambda _tenant_id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=False), ) monkeypatch.setattr(task_module, "_load_email_jobs", lambda _session, _form: []) @@ -109,7 +109,7 @@ def test_dispatch_human_input_email_task_replaces_body_variables(monkeypatch: py monkeypatch.setattr( task_module.FeatureService, "get_features", - lambda _tenant_id: SimpleNamespace(human_input_email_delivery_enabled=True), + lambda _tenant_id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True), ) monkeypatch.setattr(task_module, "_load_email_jobs", lambda _session, _form: [job]) monkeypatch.setattr(task_module, "_load_variable_pool", lambda _workflow_run_id: variable_pool) @@ -142,7 +142,7 @@ def test_dispatch_human_input_email_task_sanitizes_subject( monkeypatch.setattr( task_module.FeatureService, "get_features", - lambda _tenant_id: SimpleNamespace(human_input_email_delivery_enabled=True), + lambda _tenant_id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True), ) monkeypatch.setattr(task_module, "_load_email_jobs", lambda _session, _form: [job]) monkeypatch.setattr(task_module, "_load_variable_pool", lambda _workflow_run_id: None)