fix(api): let a removed active member sign in with a re-invite (#43667)

Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Co-authored-by: Stream <Stream_2@qq.com>

api/extensions/application_services/account.py

@@ -196,22 +196,23 @@ def build_account_services(
         sessions=RedisAccountSessionGateway(redis=redis),
         now=naive_utc_now,
     )
+    activation = AccountActivationService(
+        tokens=invitation_tokens,
+        accounts=activation_accounts,
+        workspace_policy=DeploymentWorkspaceInvitePolicy(),
+        eligibility=BillingAccountActivationEligibility(
+            enabled=deployment_edition == DeploymentEdition.CLOUD,
+        ),
+        membership_cache=BillingWorkspaceMembershipCache(
+            enabled=deployment_edition == DeploymentEdition.CLOUD,
+        ),
+        member_access_sync=RBACWorkspaceMemberAccessSync(
+            enabled=dify_config.RBAC_ENABLED,
+        ),
+    )
     return AccountServices(
         lifecycle=lifecycle,
-        activation=AccountActivationService(
-            tokens=invitation_tokens,
-            accounts=activation_accounts,
-            workspace_policy=DeploymentWorkspaceInvitePolicy(),
-            eligibility=BillingAccountActivationEligibility(
-                enabled=deployment_edition == DeploymentEdition.CLOUD,
-            ),
-            membership_cache=BillingWorkspaceMembershipCache(
-                enabled=deployment_edition == DeploymentEdition.CLOUD,
-            ),
-            member_access_sync=RBACWorkspaceMemberAccessSync(
-                enabled=dify_config.RBAC_ENABLED,
-            ),
-        ),
+        activation=activation,
         identity=AccountIdentityGateway(accounts=accounts, redis=redis),
         access=AccountAccessService(
             accounts=accounts,
@@ -226,6 +227,7 @@ def build_account_services(
             invitations=AccountActivationConsoleAuthInvitationGateway(
                 tokens=invitation_tokens,
                 accounts=activation_accounts,
+                activation=activation,
             ),
             policies=DeploymentConsoleAuthPolicyGateway(
                 billing_enabled=deployment_edition == DeploymentEdition.CLOUD,

api/extensions/ext_redis.py

@@ -149,6 +149,9 @@ class RedisClientWrapper:
     def ttl(self, name: str | bytes) -> Any:
         return self._require_client().ttl(_serialize_redis_name_arg(name, self._get_prefix()))
 
+    def pttl(self, name: str | bytes) -> Any:
+        return self._require_client().pttl(_serialize_redis_name_arg(name, self._get_prefix()))
+
     def getdel(self, name: str | bytes) -> Any:
         return self._require_client().getdel(_serialize_redis_name_arg(name, self._get_prefix()))
 

api/repositories/account_activation_repository.py

@@ -58,6 +58,7 @@ class SQLAlchemyAccountActivationRepository(AccountActivationRepository):
         *,
         role: str,
         setup: AccountSetup | None,
+        create_membership: bool = True,
     ) -> ActivationPersistenceResult | None:
         with self._session_factory.begin() as session:
             tenant_id = session.scalar(
@@ -83,6 +84,8 @@ class SQLAlchemyAccountActivationRepository(AccountActivationRepository):
                     TenantAccountJoin.account_id == account.id,
                 )
             )
+            if membership is None and not create_membership:
+                return None
             membership_created = membership is None
             if membership is None:
                 membership = TenantAccountJoin(

api/services/account/adapters.py

@@ -185,6 +185,18 @@ class RedisInvitationTokenStore(InvitationTokenStore):
             return None
         return _invitation_token_adapter.validate_json(data)
 
+    @override
+    def consume_membership(self, invitation: InvitationLookup) -> bool:
+        key = self._invitation_token_key(invitation.token)
+        ttl = self._redis.pttl(key)
+        if ttl <= 0:
+            return False
+        return bool(self._redis.set(self._membership_used_key(invitation.token), "1", nx=True, px=ttl))
+
+    @override
+    def is_membership_consumed(self, invitation: InvitationLookup) -> bool:
+        return self._redis.get(self._membership_used_key(invitation.token)) == b"1"
+
     @override
     def revoke(self, invitation: InvitationLookup) -> None:
         if invitation.workspace_id is not None and invitation.email is not None:
@@ -196,6 +208,10 @@ class RedisInvitationTokenStore(InvitationTokenStore):
     def _invitation_token_key(token: str) -> str:
         return f"member_invite:token:{token}"
 
+    @classmethod
+    def _membership_used_key(cls, token: str) -> str:
+        return f"{cls._invitation_token_key(token)}:membership_used"
+
     @staticmethod
     def _workspace_invitation_key(invitation: InvitationLookup) -> str:
         assert invitation.workspace_id is not None

api/services/account/login_adapters.py

@@ -27,7 +27,11 @@ from services.account.login_service import (
     HumanVerificationGateway,
     ResetPasswordEmailGateway,
 )
-from services.account_activation_service import AccountActivationRepository, InvitationTokenStore
+from services.account_activation_service import (
+    AccountActivationRepository,
+    AccountActivationService,
+    InvitationTokenStore,
+)
 from services.account_security_gateway import RedisAccountEmailSecurityGateway
 from services.billing_service import BillingService
 from services.email_code_login_challenge import EmailCodeLoginChallengeStore, EmailCodeLoginChallengeUnavailableError
@@ -62,11 +66,18 @@ _RESET_PASSWORD_RATE_LIMIT_SECONDS = 60
 
 
 class AccountActivationConsoleAuthInvitationGateway(ConsoleAuthInvitationGateway):
-    """Resolve login invitations through the shared activation contracts."""
+    """Resolve login invitations and restore invited membership through activation."""
 
-    def __init__(self, *, tokens: InvitationTokenStore, accounts: AccountActivationRepository) -> None:
+    def __init__(
+        self,
+        *,
+        tokens: InvitationTokenStore,
+        accounts: AccountActivationRepository,
+        activation: AccountActivationService,
+    ) -> None:
         self._tokens = tokens
         self._accounts = accounts
+        self._activation = activation
 
     @override
     def resolve(self, *, email: str, token: str) -> LoginInvitation | None:
@@ -82,6 +93,10 @@ class AccountActivationConsoleAuthInvitationGateway(ConsoleAuthInvitationGateway
                 return LoginInvitation(email=invitation.account_email)
         return None
 
+    @override
+    def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool:
+        return self._activation.ensure_membership(account_id=account_id, email=email, token=token)
+
 
 class DeploymentConsoleAuthPolicyGateway(ConsoleAuthPolicyGateway):
     def __init__(self, *, billing_enabled: bool) -> None:

api/services/account/login_service.py

@@ -41,6 +41,8 @@ class ConsoleAuthWorkspaceQuery(Protocol):
 class ConsoleAuthInvitationGateway(Protocol):
     def resolve(self, *, email: str, token: str) -> LoginInvitation | None: ...
 
+    def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool: ...
+
 
 class ConsoleAuthPolicyGateway(Protocol):
     def get_email_freeze_type(self, email: str) -> str | None: ...
@@ -200,9 +202,12 @@ class ConsoleAuthenticationService:
             self._invalid_credentials(normalized_email, command.ip_address)
 
         if not self._workspaces.has_active_for_account(account.id):
-            if self._policies.is_workspace_creation_allowed() and not self._policies.has_workspace_capacity():
-                raise account_errors.LoginWorkspaceLimitError
-            return PasswordLoginResult(tokens=None, workspace_found=False)
+            if invitation_token is not None:
+                self._invitations.ensure_membership(account_id=account.id, email=command.email, token=invitation_token)
+            if not self._workspaces.has_active_for_account(account.id):
+                if self._policies.is_workspace_creation_allowed() and not self._policies.has_workspace_capacity():
+                    raise account_errors.LoginWorkspaceLimitError
+                return PasswordLoginResult(tokens=None, workspace_found=False)
 
         tokens = self._issue_session(
             account_id=account.id,

api/services/account_activation_service.py

@@ -25,6 +25,10 @@ _PENDING_ACCOUNT_STATUS = "pending"
 
 
 class InvitationTokenStore(Protocol):
+    def consume_membership(self, invitation: InvitationLookup) -> bool: ...
+
+    def is_membership_consumed(self, invitation: InvitationLookup) -> bool: ...
+
     def find(self, invitation: InvitationLookup) -> InvitationToken | None: ...
 
     def revoke(self, invitation: InvitationLookup) -> None: ...
@@ -39,6 +43,7 @@ class AccountActivationRepository(Protocol):
         *,
         role: str,
         setup: AccountSetup | None,
+        create_membership: bool = True,
     ) -> ActivationPersistenceResult | None: ...
 
 
@@ -108,8 +113,9 @@ class AccountActivationService:
             raise FrozenAccountError
 
         setup = self._resolve_setup(invitation, command)
-        raw_role = invitation.role
-        role = raw_role if raw_role is not None and raw_role in _NON_OWNER_ROLES else _DEFAULT_ROLE
+        role = self._membership_role(invitation.role)
+        # An invitation that already restored membership at sign-in must not grant it again after removal.
+        membership_consumed = self._tokens.is_membership_consumed(command.invitation)
 
         normalized_email = command.invitation.email.lower() if command.invitation.email else None
         self._tokens.revoke(
@@ -119,13 +125,29 @@ class AccountActivationService:
                 token=command.invitation.token,
             )
         )
-        result = self._accounts.activate(invitation, role=role, setup=setup)
+        result = self._accounts.activate(invitation, role=role, setup=setup, create_membership=not membership_consumed)
         if result is None:
             raise InvalidInvitationError
         if result.membership_created:
             self._membership_cache.invalidate(invitation.workspace_id)
         self._member_access_sync.sync(invitation.workspace_id, invitation.account_id)
 
+    def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool:
+        lookup = InvitationLookup(workspace_id=None, email=email, token=token)
+        invitation = self._resolve(lookup)
+        if invitation is None or invitation.account_id != account_id:
+            return False
+        if not self._tokens.consume_membership(lookup):
+            return False
+
+        result = self._accounts.activate(invitation, role=self._membership_role(invitation.role), setup=None)
+        if result is None:
+            return False
+        if result.membership_created:
+            self._membership_cache.invalidate(invitation.workspace_id)
+            self._member_access_sync.sync(invitation.workspace_id, invitation.account_id)
+        return True
+
     def _resolve(self, invitation: InvitationLookup) -> AccountInvitation | None:
         token = self._tokens.find(invitation)
         resolved = self._accounts.resolve(token) if token is not None else None
@@ -146,6 +168,12 @@ class AccountActivationService:
             return None
         return self._accounts.resolve(token)
 
+    @staticmethod
+    def _membership_role(role: str | None) -> str:
+        if role is not None and role in _NON_OWNER_ROLES:
+            return role
+        return _DEFAULT_ROLE
+
     @staticmethod
     def _requires_setup(invitation: AccountInvitation) -> bool:
         if invitation.requires_setup is not None:

api/tests/unit_tests/repositories/test_account_activation_repository.py

@@ -176,3 +176,52 @@ class TestPersistActivation:
         assert membership is not None
         assert membership.role == TenantAccountRole.EDITOR
         assert membership.current is True
+
+    def test_does_not_create_membership_when_creation_is_disallowed(
+        self,
+        sqlite_session: Session,
+        sqlite_session_factory: sessionmaker[Session],
+    ) -> None:
+        account, _ = _persist_invitation_state(sqlite_session)
+        repository = SQLAlchemyAccountActivationRepository(sqlite_session_factory)
+
+        result = repository.activate(
+            _invitation(),
+            role="admin",
+            setup=AccountSetup(name="Name", interface_language="en-US", timezone="UTC"),
+            create_membership=False,
+        )
+
+        assert result is None
+        sqlite_session.expire_all()
+        assert sqlite_session.scalar(select(TenantAccountJoin)) is None
+        reloaded = sqlite_session.get(Account, account.id)
+        assert reloaded is not None
+        assert reloaded.status == AccountStatus.PENDING
+
+    def test_keeps_existing_membership_when_creation_is_disallowed(
+        self,
+        sqlite_session: Session,
+        sqlite_session_factory: sessionmaker[Session],
+    ) -> None:
+        account, tenant = _persist_invitation_state(sqlite_session)
+        sqlite_session.add(
+            TenantAccountJoin(
+                tenant_id=tenant.id,
+                account_id=account.id,
+                role=TenantAccountRole.EDITOR,
+            )
+        )
+        account.status = AccountStatus.ACTIVE
+        sqlite_session.commit()
+        repository = SQLAlchemyAccountActivationRepository(sqlite_session_factory)
+
+        result = repository.activate(_invitation(), role="admin", setup=None, create_membership=False)
+
+        assert result is not None
+        assert result.membership_created is False
+        sqlite_session.expire_all()
+        membership = sqlite_session.scalar(select(TenantAccountJoin))
+        assert membership is not None
+        assert membership.role == TenantAccountRole.EDITOR
+        assert membership.current is True

api/tests/unit_tests/services/account/test_adapters.py

@@ -84,6 +84,53 @@ def test_invitation_token_store_revokes_its_redis_key(redis_transport: tuple[Red
     assert commands.call_count == 1
 
 
+def test_invitation_token_store_marks_membership_use_with_token_ttl(
+    redis_transport: tuple[RedisClientWrapper, MagicMock],
+) -> None:
+    redis, commands = redis_transport
+    commands.side_effect = [1200, b"OK"]
+    lookup = InvitationLookup(workspace_id=None, email="invitee@example.com", token="token-1")
+
+    assert RedisInvitationTokenStore(redis=redis).consume_membership(lookup) is True
+
+    assert commands.call_args_list[0].args == ("PTTL", "member_invite:token:token-1")
+    set_args = commands.call_args_list[1].args
+    assert set_args[:3] == ("SET", "member_invite:token:token-1:membership_used", "1")
+    assert "NX" in set_args
+    assert set_args[set_args.index("PX") + 1] == 1200
+
+
+@pytest.mark.parametrize("ttl", [-2, -1, 0])
+def test_invitation_token_store_refuses_membership_use_without_live_token(
+    redis_transport: tuple[RedisClientWrapper, MagicMock], ttl: int
+) -> None:
+    redis, commands = redis_transport
+    commands.return_value = ttl
+    lookup = InvitationLookup(workspace_id=None, email="invitee@example.com", token="token-1")
+
+    assert RedisInvitationTokenStore(redis=redis).consume_membership(lookup) is False
+
+    assert commands.call_count == 1
+    assert commands.call_args.args == ("PTTL", "member_invite:token:token-1")
+
+
+@pytest.mark.parametrize(("stored", "expected"), [(b"1", True), (None, False)])
+def test_invitation_token_store_reports_membership_use(
+    redis_transport: tuple[RedisClientWrapper, MagicMock], stored: bytes | None, expected: bool
+) -> None:
+    redis, commands = redis_transport
+    commands.return_value = stored
+    lookup = InvitationLookup(workspace_id="workspace-1", email="invitee@example.com", token="token-1")
+
+    assert RedisInvitationTokenStore(redis=redis).is_membership_consumed(lookup) is expected
+
+    commands.assert_called_once_with(
+        "GET",
+        "member_invite:token:token-1:membership_used",
+        keys=["member_invite:token:token-1:membership_used"],
+    )
+
+
 def test_billing_eligibility_skips_gateway_when_disabled() -> None:
     with patch("services.account.adapters.BillingService.get_email_freeze_type") as get_freeze_type:
         result = BillingAccountActivationEligibility(enabled=False).get_freeze_type("invitee@example.com")

api/tests/unit_tests/services/account/test_login_adapters.py

@@ -16,17 +16,33 @@ from extensions.ext_redis import RedisClientWrapper
 from libs.helper import RateLimiter
 from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole
 from repositories.account.repository import SQLAlchemyAccountRepository
+from repositories.account_activation_repository import SQLAlchemyAccountActivationRepository
+from repositories.workspace.workspace_repository import WorkspaceRepository
 from services import account_errors
 from services.account import login_adapters as adapters
+from services.account.adapters import (
+    BillingAccountActivationEligibility,
+    BillingWorkspaceMembershipCache,
+    DeploymentWorkspaceInvitePolicy,
+    RBACWorkspaceMemberAccessSync,
+    RedisInvitationTokenStore,
+)
+from services.account_activation_service import AccountActivationService
 from services.email_code_login_challenge import (
     EmailCodeLoginChallengeResult,
     EmailCodeLoginChallengeStatus,
     EmailCodeLoginChallengeUnavailableError,
 )
+from services.entities.account_activation_entities import InvitationToken
 from services.entities.account_entities import AccountSessionTokens
-from services.entities.account_login_entities import EmailCodeChallengeStatus, RefreshAccountStatus
+from services.entities.account_login_entities import (
+    EmailCodeChallengeStatus,
+    PasswordLoginCommand,
+    RefreshAccountStatus,
+)
 from services.turnstile_service import TurnstileChallengeRejectedError, TurnstileUpstreamError
 from tests.unit_tests.account_domain import AccountDomain
+from tests.unit_tests.services.account.test_login_service import Dependencies
 
 
 class FakeRateLimiter(RateLimiter):
@@ -67,6 +83,134 @@ def _persist_account(session: Session) -> Account:
     return account
 
 
+def test_invitation_gateway_restores_editor_membership_without_revoking_token(
+    sqlite_session_factory: sessionmaker[Session],
+    redis_transport: tuple[RedisClientWrapper, MagicMock],
+) -> None:
+    redis, commands = redis_transport
+    commands.return_value = (
+        b'{"account_id":"account-1","email":"user@example.com","workspace_id":"workspace-1",'
+        b'"role":"editor","requires_setup":false}'
+    )
+    with sqlite_session_factory.begin() as session:
+        account = Account(name="User", email="user@example.com", status=AccountStatus.ACTIVE)
+        account.id = "account-1"
+        tenant = Tenant(name="Workspace")
+        tenant.id = "workspace-1"
+        session.add_all([account, tenant])
+    tokens = RedisInvitationTokenStore(redis=redis)
+    accounts = SQLAlchemyAccountActivationRepository(sqlite_session_factory)
+    activation = AccountActivationService(
+        tokens=tokens,
+        accounts=accounts,
+        workspace_policy=DeploymentWorkspaceInvitePolicy(),
+        eligibility=BillingAccountActivationEligibility(enabled=False),
+        membership_cache=BillingWorkspaceMembershipCache(enabled=False),
+        member_access_sync=RBACWorkspaceMemberAccessSync(enabled=False),
+    )
+    gateway = adapters.AccountActivationConsoleAuthInvitationGateway(
+        tokens=tokens,
+        accounts=accounts,
+        activation=activation,
+    )
+
+    commands.side_effect = lambda command, *_args, **_kwargs: (
+        3600 if command == "PTTL" else b"OK" if command == "SET" else commands.return_value
+    )
+    assert gateway.ensure_membership(account_id="different-account", email="user@example.com", token="token-1") is False
+    with sqlite_session_factory() as session:
+        assert session.scalar(select(TenantAccountJoin)) is None
+    assert gateway.ensure_membership(account_id="account-1", email="user@example.com", token="token-1") is True
+
+    with sqlite_session_factory() as session:
+        persisted = session.get(Account, "account-1")
+        membership = session.scalar(select(TenantAccountJoin))
+        assert persisted is not None
+        assert persisted.id == "account-1"
+        assert persisted.status == AccountStatus.ACTIVE
+        assert persisted.name == "User"
+        assert membership is not None
+        assert membership.role == TenantAccountRole.EDITOR
+        assert membership.current is True
+    assert not any(command.args and command.args[0] == "DEL" for command in commands.call_args_list)
+
+
+def test_login_restores_membership_when_invitation_is_stored_under_prefixed_key(
+    sqlite_session_factory: sessionmaker[Session],
+    redis_transport: tuple[RedisClientWrapper, MagicMock],
+    config_overrides: Callable[..., None],
+) -> None:
+    redis, commands = redis_transport
+    config_overrides(REDIS_KEY_PREFIX="enterprise-a")
+    stored: dict[str, object] = {}
+
+    def execute(command: object, *args: object, **_kwargs: object) -> object:
+        operation, key = str(command).upper(), str(args[0])
+        if operation == "SETEX":
+            stored[key] = args[2]
+            return True
+        if operation == "GET":
+            return stored.get(key)
+        if operation == "PTTL":
+            return 3600 if key in stored else -2
+        if operation == "SET" and "NX" in args and key not in stored:
+            stored[key] = b"1"
+            return True
+        return None
+
+    commands.side_effect = execute
+    with sqlite_session_factory.begin() as session:
+        account = Account(
+            name="User",
+            email="user@example.com",
+            status=AccountStatus.ACTIVE,
+            password="hash",
+            password_salt="salt",
+        )
+        account.id = "account-1"
+        workspace = Tenant(name="Workspace")
+        workspace.id = "workspace-1"
+        session.add_all([account, workspace])
+    tokens = RedisInvitationTokenStore(redis=redis)
+    token = tokens.create(
+        InvitationToken(
+            account_id="account-1",
+            email="user@example.com",
+            workspace_id="workspace-1",
+            role="editor",
+            requires_setup=False,
+        )
+    )
+    activation_accounts = SQLAlchemyAccountActivationRepository(sqlite_session_factory)
+    gateway = adapters.AccountActivationConsoleAuthInvitationGateway(
+        tokens=tokens,
+        accounts=activation_accounts,
+        activation=AccountActivationService(
+            tokens=tokens,
+            accounts=activation_accounts,
+            workspace_policy=DeploymentWorkspaceInvitePolicy(),
+            eligibility=BillingAccountActivationEligibility(enabled=False),
+            membership_cache=BillingWorkspaceMembershipCache(enabled=False),
+            member_access_sync=RBACWorkspaceMemberAccessSync(enabled=False),
+        ),
+    )
+    dependencies = Dependencies()
+    dependencies.accounts = SQLAlchemyAccountRepository(sqlite_session_factory)  # type: ignore[assignment]
+    dependencies.workspaces = WorkspaceRepository(sqlite_session_factory)  # type: ignore[assignment]
+    dependencies.invitations = gateway  # type: ignore[assignment]
+    result = dependencies.service().login_with_password(
+        PasswordLoginCommand(email="user@example.com", password="password", invite_token=token, ip_address="127.0.0.1")
+    )
+
+    assert f"enterprise-a:member_invite:token:{token}" in stored
+    assert result.workspace_found is True
+    assert result.tokens == AccountSessionTokens(access_token="access", refresh_token="refresh", csrf_token="csrf")
+    with sqlite_session_factory() as session:
+        membership = session.scalar(select(TenantAccountJoin))
+        assert membership is not None
+        assert membership.role == TenantAccountRole.EDITOR
+
+
 def test_security_gateway_owns_login_failure_state(redis_transport: tuple[RedisClientWrapper, MagicMock]) -> None:
     redis, commands = redis_transport
     commands.side_effect = [b"6", b"2", True, 1]

api/tests/unit_tests/services/account/test_login_service.py

@@ -49,20 +49,30 @@ class FakeAccounts:
 @dataclass
 class FakeWorkspaces:
     active: bool = False
+    checks: list[str] = field(default_factory=list)
 
     def has_active_for_account(self, account_id: str) -> bool:
-        _ = account_id
+        self.checks.append(account_id)
         return self.active
 
 
 @dataclass
 class FakeInvitations:
     invitation: LoginInvitation | None = None
+    ensured: list[tuple[str, str, str]] = field(default_factory=list)
+    restored_workspace: FakeWorkspaces | None = None
 
     def resolve(self, *, email: str, token: str) -> LoginInvitation | None:
         _ = email, token
         return self.invitation
 
+    def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool:
+        self.ensured.append((account_id, email, token))
+        if self.restored_workspace is None:
+            return False
+        self.restored_workspace.active = True
+        return True
+
 
 @dataclass
 class FakePolicies:
@@ -410,6 +420,7 @@ def test_invitation_login_rejects_empty_password_without_completing_account() ->
     assert dependencies.passwords.hashed_passwords == []
     assert dependencies.accounts.completions == []
     assert dependencies.security.failures == ["user@example.com"]
+    assert dependencies.invitations.ensured == []
 
 
 def test_invitation_login_does_not_initialize_password_when_token_is_unknown() -> None:
@@ -421,6 +432,7 @@ def test_invitation_login_does_not_initialize_password_when_token_is_unknown() -
 
     assert dependencies.passwords.hashed_passwords == []
     assert dependencies.accounts.completions == []
+    assert dependencies.invitations.ensured == []
 
 
 def test_invitation_login_applies_new_password_policy_before_completion() -> None:
@@ -436,6 +448,7 @@ def test_invitation_login_applies_new_password_policy_before_completion() -> Non
 
     assert dependencies.passwords.hashed_passwords == ["letters-only"]
     assert dependencies.accounts.completions == []
+    assert dependencies.invitations.ensured == []
 
 
 def test_invalid_password_records_rate_limit_and_audit() -> None:
@@ -448,6 +461,20 @@ def test_invalid_password_records_rate_limit_and_audit() -> None:
 
     assert dependencies.security.failures == ["user@example.com"]
     assert dependencies.audit.failures == [("user@example.com", LoginFailureReason.INVALID_CREDENTIALS, "127.0.0.1")]
+    assert dependencies.invitations.ensured == []
+
+
+def test_wrong_password_does_not_restore_invited_membership() -> None:
+    dependencies = Dependencies()
+    dependencies.accounts.candidates = [_account()]
+    dependencies.invitations.invitation = LoginInvitation(email="user@example.com")
+    dependencies.passwords.valid = False
+
+    with pytest.raises(account_errors.InvalidLoginCredentialsError):
+        dependencies.service().login_with_password(_password_command(invite_token="invite"))
+
+    assert dependencies.invitations.ensured == []
+    assert dependencies.security.failures == ["user@example.com"]
 
 
 def test_password_login_rejects_locked_account_before_loading_candidates() -> None:
@@ -504,6 +531,7 @@ def test_password_login_rejects_invitation_for_another_email() -> None:
     assert dependencies.audit.failures == [
         ("user@example.com", LoginFailureReason.INVALID_INVITATION_EMAIL, "127.0.0.1")
     ]
+    assert dependencies.invitations.ensured == []
 
 
 def test_password_login_tries_lowercase_candidate_after_exact_candidate_password_fails() -> None:
@@ -532,6 +560,40 @@ def test_password_login_returns_no_workspace_without_creating_one() -> None:
     assert not result.workspace_found
     assert result.tokens is None
     assert dependencies.workspaces_provisioning.account_ids == []
+    assert dependencies.invitations.ensured == []
+
+
+def test_password_login_restores_invited_membership_when_account_has_no_workspace() -> None:
+    dependencies = Dependencies()
+    dependencies.accounts.candidates = [_account()]
+    dependencies.invitations.invitation = LoginInvitation(email="user@example.com")
+    dependencies.invitations.restored_workspace = dependencies.workspaces
+    dependencies.policies.workspace_capacity = False
+
+    result = dependencies.service().login_with_password(_password_command(invite_token="invite"))
+
+    assert dependencies.invitations.ensured == [("account-1", "User@Example.com", "invite")]
+    assert dependencies.workspaces.checks == ["account-1", "account-1"]
+    assert result.workspace_found
+    assert result.tokens == TOKENS
+    assert dependencies.sessions.issued == ["account-1"]
+    assert dependencies.security.resets == ["user@example.com"]
+
+
+@pytest.mark.parametrize("invite_token", [None, "invite"])
+def test_password_login_does_not_restore_membership_for_an_existing_workspace(invite_token: str | None) -> None:
+    dependencies = Dependencies()
+    dependencies.accounts.candidates = [_account()]
+    dependencies.workspaces.active = True
+    if invite_token is not None:
+        dependencies.invitations.invitation = LoginInvitation(email="user@example.com")
+
+    result = dependencies.service().login_with_password(_password_command(invite_token=invite_token))
+
+    assert result.workspace_found
+    assert result.tokens == TOKENS
+    assert dependencies.invitations.ensured == []
+    assert dependencies.workspaces.checks == ["account-1"]
 
 
 def test_password_login_enforces_workspace_capacity() -> None:
@@ -542,6 +604,22 @@ def test_password_login_enforces_workspace_capacity() -> None:
     with pytest.raises(account_errors.LoginWorkspaceLimitError):
         dependencies.service().login_with_password(_password_command())
 
+    assert dependencies.invitations.ensured == []
+
+
+def test_password_login_enforces_workspace_capacity_when_invitation_does_not_restore_membership() -> None:
+    dependencies = Dependencies()
+    dependencies.accounts.candidates = [_account()]
+    dependencies.invitations.invitation = LoginInvitation(email="user@example.com")
+    dependencies.policies.workspace_capacity = False
+
+    with pytest.raises(account_errors.LoginWorkspaceLimitError):
+        dependencies.service().login_with_password(_password_command(invite_token="invite"))
+
+    assert dependencies.invitations.ensured == [("account-1", "User@Example.com", "invite")]
+    assert dependencies.workspaces.checks == ["account-1", "account-1"]
+    assert dependencies.sessions.issued == []
+
 
 def test_send_email_code_normalizes_identity_and_preserves_account_recipient_address() -> None:
     dependencies = Dependencies()

api/tests/unit_tests/services/test_account_activation_service.py

@@ -220,6 +220,7 @@ class TestActivateInvitation:
             _invitation(role=role, requires_setup=True if role == "owner" else None),
             role="normal",
             setup=AccountSetup(name="John Doe", interface_language="en-US", timezone="UTC"),
+            create_membership=True,
         )
         boundaries["cache"].assert_called_once_with("workspace-1")
         boundaries["sync"].assert_called_once_with("workspace-1", "account-1", operator_account_id=None)
@@ -264,7 +265,10 @@ class TestActivateInvitation:
         )
 
         boundaries["activate"].assert_called_once_with(
-            _invitation(account_status="active", role="editor", requires_setup=False), role="editor", setup=None
+            _invitation(account_status="active", role="editor", requires_setup=False),
+            role="editor",
+            setup=None,
+            create_membership=True,
         )
         boundaries["cache"].assert_not_called()
         boundaries["sync"].assert_called_once_with("workspace-1", "account-1", operator_account_id=None)
@@ -276,3 +280,227 @@ class TestActivateInvitation:
             assert len(memberships) == 1
             assert memberships[0].role == TenantAccountRole.EDITOR
             assert memberships[0].current is True
+
+
+class TestEnsureMembership:
+    def test_missing_token_returns_false_without_revoking(
+        self,
+        service: AccountActivationService,
+        boundaries: dict[str, MagicMock],
+        redis_transport: tuple[RedisClientWrapper, MagicMock],
+    ) -> None:
+        redis_transport[1].return_value = None
+
+        assert service.ensure_membership(account_id="account-1", email="invitee@example.com", token="missing") is False
+
+        boundaries["revoke"].assert_not_called()
+        boundaries["activate"].assert_not_called()
+        boundaries["cache"].assert_not_called()
+        boundaries["sync"].assert_not_called()
+
+    @pytest.mark.parametrize(
+        ("role", "stored_role"),
+        [
+            ("editor", TenantAccountRole.EDITOR),
+            (None, TenantAccountRole.NORMAL),
+            ("owner", TenantAccountRole.NORMAL),
+        ],
+    )
+    def test_restores_membership_once_and_preserves_activation(
+        self,
+        service: AccountActivationService,
+        boundaries: dict[str, MagicMock],
+        sqlite_session_factory: sessionmaker[Session],
+        redis_transport: tuple[RedisClientWrapper, MagicMock],
+        role: str | None,
+        stored_role: TenantAccountRole,
+    ) -> None:
+        with sqlite_session_factory.begin() as session:
+            account = session.get(Account, "account-1")
+            assert account is not None
+            account.status = AccountStatus.ACTIVE
+        role_json = f',"role":"{role}"' if role is not None else ""
+        redis_transport[1].return_value = (
+            '{"account_id":"account-1","email":"invitee@example.com","workspace_id":"workspace-1"'
+            f'{role_json},"requires_setup":false}}'
+        ).encode()
+
+        used: set[object] = set()
+
+        def execute(command: object, key: object, *args: object, **_kwargs: object) -> object:
+            if command == "PTTL":
+                return 3600
+            if command == "SET":
+                assert "NX" in args
+                assert "PX" in args
+                if key in used:
+                    return None
+                used.add(key)
+                return b"OK"
+            if command == "GET" and str(key).endswith(":membership_used"):
+                return b"1" if key in used else None
+            return redis_transport[1].return_value
+
+        redis_transport[1].side_effect = execute
+        restored = service.ensure_membership(account_id="account-1", email="Invitee@Example.com", token="token-1")
+
+        assert restored is True
+        boundaries["revoke"].assert_not_called()
+        boundaries["activate"].assert_called_once_with(
+            _invitation(account_status="active", role=role, requires_setup=False),
+            role=stored_role,
+            setup=None,
+        )
+        boundaries["cache"].assert_called_once_with("workspace-1")
+        boundaries["sync"].assert_called_once_with("workspace-1", "account-1", operator_account_id=None)
+        with sqlite_session_factory() as session:
+            account = session.get(Account, "account-1")
+            membership = session.scalar(select(TenantAccountJoin))
+            assert account is not None
+            assert account.id == "account-1"
+            assert account.status == AccountStatus.ACTIVE
+            assert account.name == "Test User"
+            assert membership is not None
+            assert membership.role == stored_role
+            assert membership.current is True
+        check = service.check(InvitationLookup(workspace_id=None, email=None, token="token-1"))
+        assert check.is_valid is True
+        boundaries["revoke"].assert_not_called()
+
+        with sqlite_session_factory.begin() as session:
+            session.execute(delete(TenantAccountJoin))
+        boundaries["activate"].reset_mock()
+        assert service.ensure_membership(account_id="account-1", email="invitee@example.com", token="token-1") is False
+        boundaries["activate"].assert_not_called()
+        with sqlite_session_factory() as session:
+            assert session.scalar(select(TenantAccountJoin)) is None
+
+        with pytest.raises(InvalidInvitationError):
+            service.activate(
+                ActivationCommand(invitation=InvitationLookup(workspace_id=None, email=None, token="token-1")),
+                authenticated_account_id="account-1",
+            )
+        boundaries["revoke"].assert_called_once()
+        boundaries["activate"].assert_called_once_with(
+            _invitation(account_status="active", role=role, requires_setup=False),
+            role=stored_role,
+            setup=None,
+            create_membership=False,
+        )
+        with sqlite_session_factory() as session:
+            assert session.scalar(select(TenantAccountJoin)) is None
+
+    def test_activation_after_restore_keeps_existing_membership(
+        self,
+        service: AccountActivationService,
+        boundaries: dict[str, MagicMock],
+        sqlite_session_factory: sessionmaker[Session],
+        redis_transport: tuple[RedisClientWrapper, MagicMock],
+    ) -> None:
+        with sqlite_session_factory.begin() as session:
+            account = session.get(Account, "account-1")
+            assert account is not None
+            account.status = AccountStatus.ACTIVE
+        redis_transport[1].return_value = (
+            b'{"account_id":"account-1","email":"invitee@example.com","workspace_id":"workspace-1",'
+            b'"role":"editor","requires_setup":false}'
+        )
+        used: set[object] = set()
+
+        def execute(command: object, key: object, *_args: object, **_kwargs: object) -> object:
+            if command == "PTTL":
+                return 3600
+            if command == "SET":
+                used.add(key)
+                return b"OK"
+            if command == "GET" and str(key).endswith(":membership_used"):
+                return b"1" if key in used else None
+            return redis_transport[1].return_value
+
+        redis_transport[1].side_effect = execute
+        assert service.ensure_membership(account_id="account-1", email="invitee@example.com", token="token-1") is True
+        boundaries["cache"].reset_mock()
+        boundaries["sync"].reset_mock()
+
+        service.activate(
+            ActivationCommand(invitation=InvitationLookup(workspace_id=None, email=None, token="token-1")),
+            authenticated_account_id="account-1",
+        )
+
+        boundaries["revoke"].assert_called_once()
+        boundaries["cache"].assert_not_called()
+        boundaries["sync"].assert_called_once_with("workspace-1", "account-1", operator_account_id=None)
+        with sqlite_session_factory() as session:
+            memberships = session.scalars(select(TenantAccountJoin)).all()
+            assert len(memberships) == 1
+            assert memberships[0].role == TenantAccountRole.EDITOR
+            assert memberships[0].current is True
+
+    def test_existing_membership_is_not_rewritten_or_resynced(
+        self,
+        service: AccountActivationService,
+        boundaries: dict[str, MagicMock],
+        sqlite_session_factory: sessionmaker[Session],
+        redis_transport: tuple[RedisClientWrapper, MagicMock],
+    ) -> None:
+        with sqlite_session_factory.begin() as session:
+            session.add(
+                TenantAccountJoin(tenant_id="workspace-1", account_id="account-1", role=TenantAccountRole.ADMIN)
+            )
+        redis_transport[1].side_effect = lambda command, *_args, **_kwargs: (
+            3600
+            if command == "PTTL"
+            else b"OK"
+            if command == "SET"
+            else b'{"account_id":"account-1","email":"invitee@example.com","workspace_id":"workspace-1"}'
+        )
+
+        assert service.ensure_membership(account_id="account-1", email="invitee@example.com", token="token-1") is True
+
+        boundaries["activate"].assert_called_once()
+        boundaries["cache"].assert_not_called()
+        boundaries["sync"].assert_not_called()
+        with sqlite_session_factory() as session:
+            memberships = session.scalars(select(TenantAccountJoin)).all()
+            assert len(memberships) == 1
+            assert memberships[0].role == TenantAccountRole.ADMIN
+
+    def test_returns_false_when_activation_finds_no_workspace(
+        self,
+        service: AccountActivationService,
+        boundaries: dict[str, MagicMock],
+        mocker: MockerFixture,
+        redis_transport: tuple[RedisClientWrapper, MagicMock],
+    ) -> None:
+        redis_transport[1].side_effect = lambda command, *_args, **_kwargs: (
+            3600
+            if command == "PTTL"
+            else b"OK"
+            if command == "SET"
+            else b'{"account_id":"account-1","email":"invitee@example.com","workspace_id":"workspace-1"}'
+        )
+        mocker.patch.object(service._accounts, "activate", return_value=None)
+
+        assert service.ensure_membership(account_id="account-1", email="invitee@example.com", token="token-1") is False
+
+        boundaries["cache"].assert_not_called()
+        boundaries["sync"].assert_not_called()
+
+    def test_account_mismatch_does_not_write_membership_or_consume_token(
+        self,
+        service: AccountActivationService,
+        boundaries: dict[str, MagicMock],
+        redis_transport: tuple[RedisClientWrapper, MagicMock],
+    ) -> None:
+        redis_transport[
+            1
+        ].return_value = b'{"account_id":"account-1","email":"invitee@example.com","workspace_id":"workspace-1"}'
+        assert (
+            service.ensure_membership(account_id="different-account", email="invitee@example.com", token="token-1")
+            is False
+        )
+        boundaries["activate"].assert_not_called()
+        boundaries["revoke"].assert_not_called()
+        boundaries["cache"].assert_not_called()
+        boundaries["sync"].assert_not_called()
+        assert all(args[0] == "GET" for args, _ in redis_transport[1].call_args_list)