From ba33f4c2a4bdc4d7530b46d313c9fb06ac2307a9 Mon Sep 17 00:00:00 2001
From: zhangx1n <sjhpzx@gmail.com>
Date: Wed, 23 Apr 2025 17:57:36 +0800
Subject: [PATCH] wip

---
 api/services/enterprise/base.py               |  3 +--
 api/services/enterprise/enterprise_service.py |  8 ++++++--
 api/services/feature_service.py               | 20 +++++++++----------
 3 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/api/services/enterprise/base.py b/api/services/enterprise/base.py
index ad526d5f34..6120de5351 100644
--- a/api/services/enterprise/base.py
+++ b/api/services/enterprise/base.py
@@ -15,8 +15,7 @@ class EnterpriseRequest:
     @classmethod
     def send_request(cls, method, endpoint, json=None, params=None):
         headers = {"Content-Type": "application/json", "Enterprise-Api-Secret-Key": cls.secret_key}
-        if params["tenant_id"]:
-            headers["X-Tenant-ID"] = params["tenant_id"]
         url = f"{cls.base_url}{endpoint}"
+        print("url: ", url)
         response = requests.request(method, url, json=json, params=params, headers=headers, proxies=cls.proxies)
         return response.json()
diff --git a/api/services/enterprise/enterprise_service.py b/api/services/enterprise/enterprise_service.py
index 0e1e6c2500..5e3db10ae9 100644
--- a/api/services/enterprise/enterprise_service.py
+++ b/api/services/enterprise/enterprise_service.py
@@ -15,8 +15,12 @@ class WebAppSettings(BaseModel):
 
 class EnterpriseService:
     @classmethod
-    def get_info(cls, tenant_id=None):
-        return EnterpriseRequest.send_request("GET", "/info", params={"tenant_id": tenant_id})
+    def get_info(cls):
+        return EnterpriseRequest.send_request("GET", "/info")
+
+    @classmethod
+    def get_workspace_info(cls, tenant_id:str):
+        return EnterpriseRequest.send_request("GET", f"/workspace/{tenant_id}/info")
 
     class WebAppAuth:
         @classmethod
diff --git a/api/services/feature_service.py b/api/services/feature_service.py
index 199b63e679..c87b28e9dd 100644
--- a/api/services/feature_service.py
+++ b/api/services/feature_service.py
@@ -35,8 +35,6 @@ class LicenseModel(BaseModel):
     status: LicenseStatus = LicenseStatus.NONE
     expired_at: str = ""
     product_id: str = ""
-    workspaces: LimitationModel = LimitationModel(size=0, limit=0)
-
 
 class BrandingModel(BaseModel):
     enabled: bool = False
@@ -86,6 +84,7 @@ class SystemFeatureModel(BaseModel):
     is_allow_create_workspace: bool = False
     is_email_setup: bool = False
     license: LicenseModel = LicenseModel()
+    workspaces: LimitationModel = LimitationModel(size=0, limit=0)
     branding: BrandingModel = BrandingModel()
     webapp_auth: WebAppAuthModel = WebAppAuthModel()
 
@@ -102,7 +101,7 @@ class FeatureService:
 
         if dify_config.ENTERPRISE_ENABLED:
             features.webapp_copyright_enabled = True
-            cls._fulfill_parms_from_license_info(features, tenant_id)
+            cls._fulfill_params_from_workspaces_quota(features, tenant_id)
 
         return features
 
@@ -135,10 +134,10 @@ class FeatureService:
         features.dataset_operator_enabled = dify_config.DATASET_OPERATOR_ENABLED
 
     @classmethod
-    def _fulfill_parms_from_license_info(cls, features: FeatureModel, tenant_id: str):
-        license_info = EnterpriseService.get_info(tenant_id)["License"]
-        features.workspace_members.limit = license_info["workspaceMembers"]["limit"]
-        features.workspace_members.size = license_info["workspaceMembers"]["used"]
+    def _fulfill_params_from_workspaces_quota(cls, features: FeatureModel, tenant_id: str):
+        workspace_members_quota = EnterpriseService.get_workspace_info(tenant_id)["WorkspaceMembersQuota"]
+        features.workspace_members.limit = workspace_members_quota["limit"]
+        features.workspace_members.size = workspace_members_quota["used"]
 
     @classmethod
     def _fulfill_params_from_billing_api(cls, features: FeatureModel, tenant_id: str):
@@ -230,6 +229,7 @@ class FeatureService:
             if "productId" in license_info:
                 features.license.product_id = license_info["productId"]
 
-            if "workspaces" in license_info:
-                features.license.workspaces.limit = license_info["workspaces"]["limit"]
-                features.license.workspaces.size = license_info["workspaces"]["used"]
+        if "WorkspacesQuota" in enterprise_info:
+            features.workspaces.limit =enterprise_info["WorkspacesQuota"]["limit"]
+            features.workspaces.size = enterprise_info["WorkspacesQuota"]["used"]
+