ensure we remove users in time

backend/danswer/auth/users.py

@@ -77,6 +77,7 @@ from danswer.configs.constants import DANSWER_API_KEY_PREFIX
 from danswer.configs.constants import UNNAMED_KEY_PLACEHOLDER
 from danswer.db.auth import get_access_token_db
 from danswer.db.auth import get_default_admin_user_emails
+from danswer.db.auth import get_total_users_count
 from danswer.db.auth import get_user_count
 from danswer.db.auth import get_user_db
 from danswer.db.auth import SQLAlchemyUserAdminDB
@@ -93,6 +94,7 @@ from danswer.utils.logger import setup_logger
 from danswer.utils.telemetry import optional_telemetry
 from danswer.utils.telemetry import RecordType
 from danswer.utils.variable_functionality import fetch_versioned_implementation
+from ee.danswer.server.tenants.billing import register_tenant_users
 from shared_configs.configs import MULTI_TENANT
 from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
@@ -422,9 +424,20 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             return user
 
     async def on_after_register(
-        self, user: User, request: Optional[Request] = None
+        self,
+        user: User,
+        request: Optional[Request] = None,
     ) -> None:
         logger.notice(f"User {user.id} has registered.")
+
+        if MULTI_TENANT:
+            # We will always have set the context var (in `create` and `oauth_callback`)
+            with get_session_with_tenant() as db_session:
+                register_tenant_users(
+                    CURRENT_TENANT_ID_CONTEXTVAR.get(),
+                    get_total_users_count(db_session),
+                )
+
         optional_telemetry(
             record_type=RecordType.SIGN_UP,
             data={"action": "create"},

backend/danswer/db/auth.py

@@ -12,7 +12,6 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.future import select
 from sqlalchemy.orm import Session
 
-from danswer.auth.invited_users import get_invited_users
 from danswer.auth.schemas import UserRole
 from danswer.db.engine import get_async_session
 from danswer.db.engine import get_async_session_with_tenant
@@ -44,10 +43,11 @@ def get_total_users_count(db_session: Session) -> int:
     user_count = (
         db_session.query(User)
         .filter(~User.email.endswith(get_api_key_email_pattern()))  # type: ignore
+        .filter(User.is_active)  # type: ignore
         .count()
     )
-    invited_users = len(get_invited_users())
-    return user_count + invited_users
+
+    return user_count
 
 
 async def get_user_count() -> int:

backend/danswer/server/manage/users.py

@@ -41,6 +41,7 @@ from danswer.configs.constants import AuthType
 from danswer.db.auth import get_total_users_count
 from danswer.db.engine import CURRENT_TENANT_ID_CONTEXTVAR
 from danswer.db.engine import get_session
+from danswer.db.engine import get_session_with_tenant
 from danswer.db.models import AccessToken
 from danswer.db.models import DocumentSet__User
 from danswer.db.models import Persona__User
@@ -225,10 +226,6 @@ def bulk_invite_users(
     if not MULTI_TENANT:
         return number_of_invited_users
     try:
-        logger.info("Registering tenant users")
-        register_tenant_users(
-            CURRENT_TENANT_ID_CONTEXTVAR.get(), get_total_users_count(db_session)
-        )
         if ENABLE_EMAIL_INVITES:
             try:
                 for email in all_emails:
@@ -260,18 +257,6 @@ def remove_invited_user(
     remove_users_from_tenant([user_email.user_email], tenant_id)
     number_of_invited_users = write_invited_users(remaining_users)
 
-    try:
-        if MULTI_TENANT:
-            register_tenant_users(
-                CURRENT_TENANT_ID_CONTEXTVAR.get(), get_total_users_count(db_session)
-            )
-    except Exception:
-        logger.error(
-            "Request to update number of seats taken in control plane failed. "
-            "This may cause synchronization issues/out of date enforcement of seat limits."
-        )
-        raise
-
     return number_of_invited_users
 
 
@@ -303,6 +288,14 @@ def deactivate_user(
     db_session.add(user_to_deactivate)
     db_session.commit()
 
+    if MULTI_TENANT:
+        # We will always have set the context var (in `create` and `oauth_callback`)
+        with get_session_with_tenant() as db_session:
+            register_tenant_users(
+                CURRENT_TENANT_ID_CONTEXTVAR.get(),
+                get_total_users_count(db_session),
+            )
+
 
 @router.delete("/manage/admin/delete-user")
 async def delete_user(