From 7db3a521e19dbb9cefa1efb3123f02e27eb3147b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E7=8E=AE=20=28Jade=20Lin=29?= Date: Mon, 25 May 2026 16:58:42 +0800 Subject: [PATCH] fix: member invite limits with dedup, locking, and accurate new-member counting (#36512) --- api/controllers/console/workspace/members.py | 115 ++++++++++++------ .../console/test_workspace_members.py | 20 ++- .../console/workspace/test_members.py | 55 +++++++++ web/app/components/billing/type.ts | 1 + web/context/provider-context-provider.tsx | 34 +++++- 5 files changed, 184 insertions(+), 41 deletions(-) diff --git a/api/controllers/console/workspace/members.py b/api/controllers/console/workspace/members.py index d068c15667..c7769f4e6a 100644 --- a/api/controllers/console/workspace/members.py +++ b/api/controllers/console/workspace/members.py @@ -4,6 +4,7 @@ from uuid import UUID from flask import abort, request from flask_restx import Resource from pydantic import BaseModel, Field, TypeAdapter +from sqlalchemy import func, select import services from configs import dify_config @@ -22,15 +23,15 @@ from controllers.console.auth.error import ( from controllers.console.error import EmailSendIpLimitError, WorkspaceMembersLimitExceeded from controllers.console.wraps import ( account_initialization_required, - cloud_edition_billing_resource_check, is_allow_transfer_owner, setup_required, ) from extensions.ext_database import db +from extensions.ext_redis import redis_client from fields.member_fields import AccountWithRole, AccountWithRoleList from libs.helper import extract_remote_ip from libs.login import current_account_with_tenant, login_required -from models.account import Account, TenantAccountRole +from models.account import Account, TenantAccountJoin, TenantAccountRole from services.account_service import AccountService, RegisterService, TenantService from services.enterprise import rbac_service as enterprise_rbac_service from services.errors.account import AccountAlreadyInTenantError @@ -93,6 +94,54 @@ def _normalize_enum_value(value: object) -> str: return str(normalized) if normalized is not None else "" +def _normalize_invitee_emails(emails: list[str]) -> list[str]: + return list(dict.fromkeys(email.lower() for email in emails)) + + +def _count_new_member_invites(tenant_id: str, emails: list[str]) -> int: + new_member_count = 0 + for email in emails: + account = AccountService.get_account_by_email_with_case_fallback(email) + if not account: + new_member_count += 1 + continue + + exists = db.session.scalar( + select(TenantAccountJoin.id) + .where(TenantAccountJoin.tenant_id == tenant_id, TenantAccountJoin.account_id == account.id) + .limit(1) + ) + if not exists: + new_member_count += 1 + + return new_member_count + + +def _count_current_members(tenant_id: str) -> int: + return ( + db.session.scalar(select(func.count(TenantAccountJoin.id)).where(TenantAccountJoin.tenant_id == tenant_id)) or 0 + ) + + +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) + + if dify_config.ENTERPRISE_ENABLED: + workspace_members = features.workspace_members + if workspace_members.enabled is True and not workspace_members.is_available(new_member_count): + raise WorkspaceMembersLimitExceeded() + return + + if dify_config.BILLING_ENABLED and features.billing.enabled is True: + members = features.members + current_member_count = _count_current_members(tenant_id) + if 0 < members.limit < current_member_count + new_member_count: + raise WorkspaceMembersLimitExceeded() + + @console_ns.route("/workspaces/current/members") class MemberListApi(Resource): """List all members of current tenant.""" @@ -148,12 +197,11 @@ class MemberInviteEmailApi(Resource): @setup_required @login_required @account_initialization_required - @cloud_edition_billing_resource_check("members") def post(self): payload = console_ns.payload or {} args = MemberInvitePayload.model_validate(payload) - invitee_emails = args.emails + invitee_emails = _normalize_invitee_emails(args.emails) invitee_role = args.role interface_language = args.language if not dify_config.RBAC_ENABLED: @@ -174,37 +222,36 @@ class MemberInviteEmailApi(Resource): invitation_results = [] console_web_url = dify_config.CONSOLE_WEB_URL - workspace_members = FeatureService.get_features(tenant_id=inviter.current_tenant.id).workspace_members + tenant_id = inviter.current_tenant.id + with redis_client.lock(f"workspace_member_invite:{tenant_id}", timeout=60): + new_member_count = _count_new_member_invites(tenant_id, invitee_emails) + _check_member_invite_limits(tenant_id, new_member_count) - if not workspace_members.is_available(len(invitee_emails)): - raise WorkspaceMembersLimitExceeded() - - for invitee_email in invitee_emails: - normalized_invitee_email = invitee_email.lower() - try: - if not inviter.current_tenant: - raise ValueError("No current tenant") - token = RegisterService.invite_new_member( - tenant=inviter.current_tenant, - email=invitee_email, - language=interface_language, - role=invitee_role, - inviter=inviter, - ) - encoded_invitee_email = parse.quote(normalized_invitee_email) - invitation_results.append( - { - "status": "success", - "email": normalized_invitee_email, - "url": f"{console_web_url}/activate?email={encoded_invitee_email}&token={token}", - } - ) - except AccountAlreadyInTenantError: - invitation_results.append( - {"status": "success", "email": normalized_invitee_email, "url": f"{console_web_url}/signin"} - ) - except Exception as e: - invitation_results.append({"status": "failed", "email": normalized_invitee_email, "message": str(e)}) + for invitee_email in invitee_emails: + try: + if not inviter.current_tenant: + raise ValueError("No current tenant") + token = RegisterService.invite_new_member( + tenant=inviter.current_tenant, + email=invitee_email, + language=interface_language, + role=invitee_role, + inviter=inviter, + ) + encoded_invitee_email = parse.quote(invitee_email) + invitation_results.append( + { + "status": "success", + "email": invitee_email, + "url": f"{console_web_url}/activate?email={encoded_invitee_email}&token={token}", + } + ) + except AccountAlreadyInTenantError: + invitation_results.append( + {"status": "success", "email": invitee_email, "url": f"{console_web_url}/signin"} + ) + except Exception as e: + invitation_results.append({"status": "failed", "email": invitee_email, "message": str(e)}) return { "result": "success", diff --git a/api/tests/unit_tests/controllers/console/test_workspace_members.py b/api/tests/unit_tests/controllers/console/test_workspace_members.py index 4f73a5b8c5..152a41d730 100644 --- a/api/tests/unit_tests/controllers/console/test_workspace_members.py +++ b/api/tests/unit_tests/controllers/console/test_workspace_members.py @@ -1,3 +1,4 @@ +from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import patch @@ -18,7 +19,7 @@ def app(): def _build_feature_flags(): placeholder_quota = SimpleNamespace(limit=0, size=0) - workspace_members = SimpleNamespace(is_available=lambda count: True) + workspace_members = SimpleNamespace(enabled=False, is_available=lambda count: True) return SimpleNamespace( billing=SimpleNamespace(enabled=False), workspace_members=workspace_members, @@ -31,6 +32,11 @@ def _build_feature_flags(): class TestMemberInviteEmailApi: + @pytest.fixture(autouse=True) + def _mock_member_invite_lock(self): + with patch("controllers.console.workspace.members.redis_client.lock", return_value=nullcontext()): + yield + @patch("controllers.console.workspace.members.FeatureService.get_features") @patch("controllers.console.workspace.members.RegisterService.invite_new_member") @patch("controllers.console.workspace.members.current_account_with_tenant") @@ -52,9 +58,13 @@ class TestMemberInviteEmailApi: inviter = SimpleNamespace(email="Owner@Example.com", current_tenant=tenant, status="active") mock_current_account.return_value = (inviter, tenant.id) - with patch("controllers.console.workspace.members.dify_config") as mock_config: - mock_config.RBAC_ENABLED = False - mock_config.CONSOLE_WEB_URL = "https://console.example.com" + with ( + patch("controllers.console.workspace.members.dify_config.RBAC_ENABLED", False), + patch("controllers.console.workspace.members.dify_config.CONSOLE_WEB_URL", "https://console.example.com"), + patch("controllers.console.workspace.members._count_new_member_invites", return_value=1), + patch("controllers.console.workspace.members.dify_config.ENTERPRISE_ENABLED", False), + patch("controllers.console.workspace.members.dify_config.BILLING_ENABLED", False), + ): with app.test_request_context( "/workspaces/current/members/invite-email", method="POST", @@ -72,7 +82,7 @@ class TestMemberInviteEmailApi: assert mock_invite_member.call_count == 1 call_args = mock_invite_member.call_args assert call_args.kwargs["tenant"] == tenant - assert call_args.kwargs["email"] == "User@Example.com" + assert call_args.kwargs["email"] == "user@example.com" assert call_args.kwargs["language"] == "en-US" assert call_args.kwargs["role"] == TenantAccountRole.EDITOR assert call_args.kwargs["inviter"] == inviter diff --git a/api/tests/unit_tests/controllers/console/workspace/test_members.py b/api/tests/unit_tests/controllers/console/workspace/test_members.py index c207aedd23..2cb4c54eb4 100644 --- a/api/tests/unit_tests/controllers/console/workspace/test_members.py +++ b/api/tests/unit_tests/controllers/console/workspace/test_members.py @@ -1,3 +1,4 @@ +from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -123,6 +124,11 @@ class TestMemberListApi: class TestMemberInviteEmailApi: + @pytest.fixture(autouse=True) + def _mock_member_invite_lock(self): + with patch("controllers.console.workspace.members.redis_client.lock", return_value=nullcontext()): + yield + def test_invite_success(self, app: Flask): api = MemberInviteEmailApi() method = unwrap(api.post) @@ -130,6 +136,8 @@ class TestMemberInviteEmailApi: tenant = MagicMock(id="t1") user = MagicMock(current_tenant=tenant) features = MagicMock() + features.billing.enabled = False + features.workspace_members.enabled = False features.workspace_members.is_available.return_value = True payload = { @@ -142,8 +150,11 @@ class TestMemberInviteEmailApi: app.test_request_context("/", json=payload), patch("controllers.console.workspace.members.current_account_with_tenant", return_value=(user, "t1")), patch("controllers.console.workspace.members.FeatureService.get_features", return_value=features), + patch("controllers.console.workspace.members._count_new_member_invites", return_value=1), patch("controllers.console.workspace.members.RegisterService.invite_new_member", return_value="token"), patch("controllers.console.workspace.members.dify_config.CONSOLE_WEB_URL", "http://x"), + patch("controllers.console.workspace.members.dify_config.ENTERPRISE_ENABLED", False), + patch("controllers.console.workspace.members.dify_config.BILLING_ENABLED", False), ): result, status = method(api) @@ -157,6 +168,8 @@ class TestMemberInviteEmailApi: tenant = MagicMock(id="t1") user = MagicMock(current_tenant=tenant) features = MagicMock() + features.billing.enabled = False + features.workspace_members.enabled = True features.workspace_members.is_available.return_value = False payload = { @@ -168,6 +181,38 @@ class TestMemberInviteEmailApi: app.test_request_context("/", json=payload), patch("controllers.console.workspace.members.current_account_with_tenant", return_value=(user, "t1")), patch("controllers.console.workspace.members.FeatureService.get_features", return_value=features), + patch("controllers.console.workspace.members._count_new_member_invites", return_value=1), + patch("controllers.console.workspace.members.dify_config.ENTERPRISE_ENABLED", True), + patch("controllers.console.workspace.members.dify_config.BILLING_ENABLED", False), + ): + with pytest.raises(WorkspaceMembersLimitExceeded): + method(api) + + def test_invite_billing_limit_exceeded(self, app: Flask): + api = MemberInviteEmailApi() + method = unwrap(api.post) + + tenant = MagicMock(id="t1") + user = MagicMock(current_tenant=tenant) + features = MagicMock() + features.billing.enabled = True + features.members.size = 9 + features.members.limit = 10 + features.workspace_members.enabled = False + + payload = { + "emails": ["a@test.com", "b@test.com"], + "role": "normal", + } + + with ( + app.test_request_context("/", json=payload), + patch("controllers.console.workspace.members.current_account_with_tenant", return_value=(user, "t1")), + patch("controllers.console.workspace.members.FeatureService.get_features", return_value=features), + patch("controllers.console.workspace.members._count_new_member_invites", return_value=2), + patch("controllers.console.workspace.members._count_current_members", return_value=9), + patch("controllers.console.workspace.members.dify_config.ENTERPRISE_ENABLED", False), + patch("controllers.console.workspace.members.dify_config.BILLING_ENABLED", True), ): with pytest.raises(WorkspaceMembersLimitExceeded): method(api) @@ -179,6 +224,8 @@ class TestMemberInviteEmailApi: tenant = MagicMock(id="t1") user = MagicMock(current_tenant=tenant) features = MagicMock() + features.billing.enabled = False + features.workspace_members.enabled = False features.workspace_members.is_available.return_value = True payload = { @@ -190,11 +237,14 @@ class TestMemberInviteEmailApi: app.test_request_context("/", json=payload), patch("controllers.console.workspace.members.current_account_with_tenant", return_value=(user, "t1")), patch("controllers.console.workspace.members.FeatureService.get_features", return_value=features), + patch("controllers.console.workspace.members._count_new_member_invites", return_value=0), patch( "controllers.console.workspace.members.RegisterService.invite_new_member", side_effect=AccountAlreadyInTenantError(), ), patch("controllers.console.workspace.members.dify_config.CONSOLE_WEB_URL", "http://x"), + patch("controllers.console.workspace.members.dify_config.ENTERPRISE_ENABLED", False), + patch("controllers.console.workspace.members.dify_config.BILLING_ENABLED", False), ): result, status = method(api) @@ -222,6 +272,8 @@ class TestMemberInviteEmailApi: tenant = MagicMock(id="t1") user = MagicMock(current_tenant=tenant) features = MagicMock() + features.billing.enabled = False + features.workspace_members.enabled = False features.workspace_members.is_available.return_value = True payload = { @@ -233,11 +285,14 @@ class TestMemberInviteEmailApi: app.test_request_context("/", json=payload), patch("controllers.console.workspace.members.current_account_with_tenant", return_value=(user, "t1")), patch("controllers.console.workspace.members.FeatureService.get_features", return_value=features), + patch("controllers.console.workspace.members._count_new_member_invites", return_value=1), patch( "controllers.console.workspace.members.RegisterService.invite_new_member", side_effect=Exception("boom"), ), patch("controllers.console.workspace.members.dify_config.CONSOLE_WEB_URL", "http://x"), + patch("controllers.console.workspace.members.dify_config.ENTERPRISE_ENABLED", False), + patch("controllers.console.workspace.members.dify_config.BILLING_ENABLED", False), ): result, _ = method(api) diff --git a/web/app/components/billing/type.ts b/web/app/components/billing/type.ts index e40c89f1a7..168bc53893 100644 --- a/web/app/components/billing/type.ts +++ b/web/app/components/billing/type.ts @@ -91,6 +91,7 @@ export type CurrentPlanInfoBackend = { } webapp_copyright_enabled: boolean workspace_members: { + enabled?: boolean size: number limit: number } diff --git a/web/context/provider-context-provider.tsx b/web/context/provider-context-provider.tsx index 160a559a77..b69a395a76 100644 --- a/web/context/provider-context-provider.tsx +++ b/web/context/provider-context-provider.tsx @@ -28,6 +28,37 @@ type ProviderContextProviderProps = { children: ReactNode } +type MemberInviteLimit = { + size: number + limit: number +} + +const unlimitedMemberInviteLimit: MemberInviteLimit = { + size: 0, + limit: 0, +} + +const resolveMemberInviteLimit = (data: Awaited>): MemberInviteLimit => { + if (!data) + return unlimitedMemberInviteLimit + + if (data.workspace_members?.enabled) { + return { + size: data.workspace_members.size, + limit: data.workspace_members.limit, + } + } + + if (data.billing?.enabled && data.members?.limit > 0) { + return { + size: data.members.size, + limit: data.members.limit, + } + } + + return unlimitedMemberInviteLimit +} + export const ProviderContextProvider = ({ children, }: ProviderContextProviderProps) => { @@ -87,8 +118,7 @@ export const ProviderContextProvider = ({ setDatasetOperatorEnabled(true) if (data.webapp_copyright_enabled) setWebappCopyrightEnabled(true) - if (data.workspace_members) - setLicenseLimit({ workspace_members: data.workspace_members }) + setLicenseLimit({ workspace_members: resolveMemberInviteLimit(data) }) if (data.is_allow_transfer_workspace) setIsAllowTransferWorkspace(data.is_allow_transfer_workspace) if (data.knowledge_pipeline?.publish_enabled)