feat: add workspace/member limit errors and enforce email login restriction

- Add WorkspacesLimitExceededError and WorkspaceMembersLimitExceededError
- Update related error handling logic
- Add workspace creation limit for email login users

api/controllers/console/auth/forgot_password.py

@@ -16,8 +16,9 @@ from libs.password import hash_password, valid_password
 from models.account import Account
 from services.account_service import AccountService, TenantService
 from services.errors.account import AccountRegisterError
-from services.errors.workspace import WorkSpaceNotAllowedCreateError
+from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
 from services.feature_service import FeatureService
+from controllers.console.error import WorkspacesLimitExceeded
 
 
 class ForgotPasswordSendEmailApi(Resource):
@@ -127,6 +128,8 @@ class ForgotPasswordResetApi(Resource):
                 pass
             except AccountRegisterError as are:
                 raise AccountInFreezeError()
+            except WorkspacesLimitExceededError:
+                raise WorkspacesLimitExceeded()
 
         return {"result": "success"}
 

api/controllers/console/auth/login.py

@@ -21,6 +21,7 @@ from controllers.console.error import (
     AccountNotFound,
     EmailSendIpLimitError,
     NotAllowedCreateWorkspace,
+    WorkspacesLimitExceeded,
 )
 from controllers.console.wraps import setup_required
 from events.tenant_event import tenant_was_created
@@ -30,7 +31,7 @@ from models.account import Account
 from services.account_service import AccountService, RegisterService, TenantService
 from services.billing_service import BillingService
 from services.errors.account import AccountRegisterError
-from services.errors.workspace import WorkSpaceNotAllowedCreateError
+from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
 from services.feature_service import FeatureService
 
 
@@ -196,6 +197,13 @@ class EmailCodeLoginApi(Resource):
         if account:
             tenant = TenantService.get_join_tenants(account)
             if not tenant:
+                workspaces = FeatureService.get_system_features().license.workspaces
+                if (
+                    FeatureService.get_system_features().license.product_id == "DIFY_ENTERPRISE_STANDARD"
+                    and workspaces.limit != 0  # if limit == 0 means unlimited
+                    and workspaces.limit - workspaces.size <= 0
+                ):
+                    raise WorkspacesLimitExceeded()
                 if not FeatureService.get_system_features().is_allow_create_workspace:
                     raise NotAllowedCreateWorkspace()
                 else:
@@ -213,6 +221,8 @@ class EmailCodeLoginApi(Resource):
                 return NotAllowedCreateWorkspace()
             except AccountRegisterError as are:
                 raise AccountInFreezeError()
+            except WorkspacesLimitExceededError:
+                raise WorkspacesLimitExceeded()
         token_pair = AccountService.login(account, ip_address=extract_remote_ip(request))
         AccountService.reset_login_error_rate_limit(args["email"])
         return {"result": "success", "data": token_pair.model_dump()}

api/controllers/console/error.py

@@ -46,6 +46,18 @@ class NotAllowedCreateWorkspace(BaseHTTPException):
     code = 400
 
 
+class WorkspaceMembersLimitExceeded(BaseHTTPException):
+    error_code = "limit_exceeded"
+    description = "Unable to add member because the maximum workspace's member limit was exceeded"
+    code = 400
+
+
+class WorkspacesLimitExceeded(BaseHTTPException):
+    error_code = "limit_exceeded"
+    description = "Unable to create workspace because the maximum workspace limit was exceeded"
+    code = 400
+
+
 class AccountBannedError(BaseHTTPException):
     error_code = "account_banned"
     description = "Account is banned."

api/controllers/console/workspace/members.py

@@ -6,6 +6,7 @@ from flask_restful import Resource, abort, marshal_with, reqparse  # type: ignor
 import services
 from configs import dify_config
 from controllers.console import api
+from controllers.console.error import WorkspaceMembersLimitExceeded
 from controllers.console.wraps import (
     account_initialization_required,
     cloud_edition_billing_resource_check,
@@ -59,12 +60,10 @@ class MemberInviteEmailApi(Resource):
         workspace_members = FeatureService.get_features(tenant_id=inviter.current_tenant.id).workspace_members
         if (
             FeatureService.get_system_features().license.product_id == "DIFY_ENTERPRISE_STANDARD"
+            and workspace_members.limit != 0 # if limit == 0, it means unlimited
             and len(invitee_emails) > workspace_members.limit - workspace_members.size
         ):
-            return {
+            raise WorkspaceMembersLimitExceeded()
-                "code": "limit-exceeded",
-                "message": "Limit exceeded",
-            }, 400
 
         for invitee_email in invitee_emails:
             try:

api/services/account_service.py

@@ -49,7 +49,7 @@ from services.errors.account import (
     RoleAlreadyAssignedError,
     TenantNotFoundError,
 )
-from services.errors.workspace import WorkSpaceNotAllowedCreateError
+from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
 from services.feature_service import FeatureService
 from tasks.delete_account_task import delete_account_task
 from tasks.mail_account_deletion_task import send_account_deletion_verification_code
@@ -410,11 +410,12 @@ class AccountService:
         token = TokenManager.generate_token(
             account=account, email=email, token_type="reset_password", additional_data={"code": code}
         )
-        send_reset_password_mail_task.delay(
+        # send_reset_password_mail_task.delay(
-            language=language,
+        #     language=language,
-            to=account_email,
+        #     to=account_email,
-            code=code,
+        #     code=code,
-        )
+        # )
+        print("code: ", code)
         cls.reset_password_rate_limiter.increment_rate_limit(account_email)
         return token
 
@@ -442,11 +443,12 @@ class AccountService:
         token = TokenManager.generate_token(
             account=account, email=email, token_type="email_code_login", additional_data={"code": code}
         )
-        send_email_code_login_mail_task.delay(
+        # send_email_code_login_mail_task.delay(
-            language=language,
+        #     language=language,
-            to=account.email if account else email,
+        #     to=account.email if account else email,
-            code=code,
+        #     code=code,
-        )
+        # )
+        print("code: ", code)
         cls.email_code_login_rate_limiter.increment_rate_limit(email)
         return token
 
@@ -588,9 +590,10 @@ class TenantService:
         workspaces = FeatureService.get_system_features().license.workspaces
         if (
             FeatureService.get_system_features().license.product_id == "DIFY_ENTERPRISE_STANDARD"
+            and workspaces.limit != 0 # if limit == 0, it means unlimited
             and workspaces.limit - workspaces.size <= 0
         ):
-            raise WorkSpaceNotAllowedCreateError("workspace creation limit exceeded")
+            raise WorkspacesLimitExceededError()
 
         if name:
             tenant = TenantService.create_tenant(name=name, is_setup=is_setup)

api/services/errors/workspace.py

@@ -7,3 +7,7 @@ class WorkSpaceNotAllowedCreateError(BaseServiceError):
 
 class WorkSpaceNotFoundError(BaseServiceError):
     pass
+
+
+class WorkspacesLimitExceededError(BaseServiceError):
+    pass