diff --git a/api/controllers/console/auth/forgot_password.py b/api/controllers/console/auth/forgot_password.py index a9c4300b9a..a9c4cc77fd 100644 --- a/api/controllers/console/auth/forgot_password.py +++ b/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"} diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index 41362e9fa2..c475466a79 100644 --- a/api/controllers/console/auth/login.py +++ b/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()} diff --git a/api/controllers/console/error.py b/api/controllers/console/error.py index ee87138a44..47317dfed6 100644 --- a/api/controllers/console/error.py +++ b/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." diff --git a/api/controllers/console/workspace/members.py b/api/controllers/console/workspace/members.py index e0c2386963..4f93406de2 100644 --- a/api/controllers/console/workspace/members.py +++ b/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 { - "code": "limit-exceeded", - "message": "Limit exceeded", - }, 400 + raise WorkspaceMembersLimitExceeded() for invitee_email in invitee_emails: try: diff --git a/api/services/account_service.py b/api/services/account_service.py index 4e2ec13375..6ba1bc2dfb 100644 --- a/api/services/account_service.py +++ b/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( - language=language, - to=account_email, - code=code, - ) + # send_reset_password_mail_task.delay( + # language=language, + # to=account_email, + # 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( - language=language, - to=account.email if account else email, - code=code, - ) + # send_email_code_login_mail_task.delay( + # language=language, + # to=account.email if account else email, + # 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) diff --git a/api/services/errors/workspace.py b/api/services/errors/workspace.py index 714064ffdf..f2973674e9 100644 --- a/api/services/errors/workspace.py +++ b/api/services/errors/workspace.py @@ -7,3 +7,7 @@ class WorkSpaceNotAllowedCreateError(BaseServiceError): class WorkSpaceNotFoundError(BaseServiceError): pass + + +class WorkspacesLimitExceededError(BaseServiceError): + pass \ No newline at end of file