Password reset tenant (#3895)

* nots

* functional

* minor naming cleanup

* nit

* update constant

* k

backend/ee/onyx/server/middleware/tenant_tracking.py

@@ -10,6 +10,7 @@ from fastapi import Response
 from ee.onyx.auth.users import decode_anonymous_user_jwt_token
 from ee.onyx.configs.app_configs import ANONYMOUS_USER_COOKIE_NAME
 from onyx.auth.api_key import extract_tenant_from_api_key_header
+from onyx.configs.constants import TENANT_ID_COOKIE_NAME
 from onyx.db.engine import is_valid_schema_name
 from onyx.redis.redis_pool import retrieve_auth_token_data_from_redis
 from shared_configs.configs import MULTI_TENANT
@@ -43,6 +44,7 @@ async def _get_tenant_id_from_request(
     Attempt to extract tenant_id from:
     1) The API key header
     2) The Redis-based token (stored in Cookie: fastapiusersauth)
+    3)  Reset token cookie
     Fallback: POSTGRES_DEFAULT_SCHEMA
     """
     # Check for API key
@@ -90,3 +92,12 @@ async def _get_tenant_id_from_request(
     except Exception as e:
         logger.error(f"Unexpected error in _get_tenant_id_from_request: {str(e)}")
         raise HTTPException(status_code=500, detail="Internal server error")
+
+    finally:
+        # As a final step, check for explicit tenant_id cookie
+        tenant_id_cookie = request.cookies.get(TENANT_ID_COOKIE_NAME)
+        if tenant_id_cookie and is_valid_schema_name(tenant_id_cookie):
+            return tenant_id_cookie
+
+        # If we've reached this point, return the default schema
+        return POSTGRES_DEFAULT_SCHEMA

backend/ee/onyx/server/tenants/api.py

@@ -34,6 +34,7 @@ from onyx.auth.users import get_redis_strategy
 from onyx.auth.users import optional_user
 from onyx.auth.users import User
 from onyx.configs.app_configs import WEB_DOMAIN
+from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
 from onyx.db.auth import get_user_count
 from onyx.db.engine import get_current_tenant_id
 from onyx.db.engine import get_session
@@ -111,7 +112,7 @@ async def login_as_anonymous_user(
     token = generate_anonymous_user_jwt_token(tenant_id)
 
     response = Response()
-    response.delete_cookie("fastapiusersauth")
+    response.delete_cookie(FASTAPI_USERS_AUTH_COOKIE_NAME)
     response.set_cookie(
         key=ANONYMOUS_USER_COOKIE_NAME,
         value=token,

backend/onyx/auth/email_utils.py

@@ -10,6 +10,7 @@ from onyx.configs.app_configs import SMTP_PORT
 from onyx.configs.app_configs import SMTP_SERVER
 from onyx.configs.app_configs import SMTP_USER
 from onyx.configs.app_configs import WEB_DOMAIN
+from onyx.configs.constants import TENANT_ID_COOKIE_NAME
 from onyx.db.models import User
 
 
@@ -65,9 +66,13 @@ def send_forgot_password_email(
     user_email: str,
     token: str,
     mail_from: str = EMAIL_FROM,
+    tenant_id: str | None = None,
 ) -> None:
     subject = "Onyx Forgot Password"
     link = f"{WEB_DOMAIN}/auth/reset-password?token={token}"
+    if tenant_id:
+        link += f"&{TENANT_ID_COOKIE_NAME}={tenant_id}"
+        # Keep search param same name as cookie for simplicity
     body = f"Click the following link to reset your password: {link}"
     send_email(user_email, subject, body, mail_from)
 

backend/onyx/auth/users.py

@@ -73,6 +73,7 @@ from onyx.configs.app_configs import WEB_DOMAIN
 from onyx.configs.constants import AuthType
 from onyx.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN
 from onyx.configs.constants import DANSWER_API_KEY_PREFIX
+from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
 from onyx.configs.constants import MilestoneRecordType
 from onyx.configs.constants import OnyxRedisLocks
 from onyx.configs.constants import PASSWORD_SPECIAL_CHARS
@@ -218,6 +219,24 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
     verification_token_lifetime_seconds = AUTH_COOKIE_EXPIRE_TIME_SECONDS
     user_db: SQLAlchemyUserDatabase[User, uuid.UUID]
 
+    async def get_by_email(self, user_email: str) -> User:
+        tenant_id = fetch_ee_implementation_or_noop(
+            "onyx.server.tenants.user_mapping", "get_tenant_id_for_email", None
+        )(user_email)
+        async with get_async_session_with_tenant(tenant_id) as db_session:
+            if MULTI_TENANT:
+                tenant_user_db = SQLAlchemyUserAdminDB[User, uuid.UUID](
+                    db_session, User, OAuthAccount
+                )
+                user = await tenant_user_db.get_by_email(user_email)
+            else:
+                user = await self.user_db.get_by_email(user_email)
+
+        if not user:
+            raise exceptions.UserNotExists()
+
+        return user
+
     async def create(
         self,
         user_create: schemas.UC | UserCreate,
@@ -504,9 +523,15 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             )
             raise HTTPException(
                 status.HTTP_500_INTERNAL_SERVER_ERROR,
-                "Your admin has not enbaled this feature.",
+                "Your admin has not enabled this feature.",
             )
-        send_forgot_password_email(user.email, token)
+        tenant_id = await fetch_ee_implementation_or_noop(
+            "onyx.server.tenants.provisioning",
+            "get_or_provision_tenant",
+            async_return_default_schema,
+        )(email=user.email)
+
+        send_forgot_password_email(user.email, token, tenant_id=tenant_id)
 
     async def on_after_request_verify(
         self, user: User, token: str, request: Optional[Request] = None
@@ -580,6 +605,7 @@ async def get_user_manager(
 cookie_transport = CookieTransport(
     cookie_max_age=SESSION_EXPIRE_TIME_SECONDS,
     cookie_secure=WEB_DOMAIN.startswith("https"),
+    cookie_name=FASTAPI_USERS_AUTH_COOKIE_NAME,
 )
 
 

backend/onyx/configs/constants.py

@@ -15,6 +15,12 @@ ID_SEPARATOR = ":;:"
 DEFAULT_BOOST = 0
 SESSION_KEY = "session"
 
+# Cookies
+FASTAPI_USERS_AUTH_COOKIE_NAME = (
+    "fastapiusersauth"  # Currently a constant, but logic allows for configuration
+)
+TENANT_ID_COOKIE_NAME = "onyx_tid"  # tenant id - for workaround cases
+
 NO_AUTH_USER_ID = "__no_auth_user__"
 NO_AUTH_USER_EMAIL = "anonymous@onyx.app"
 

backend/onyx/redis/redis_pool.py

@@ -25,6 +25,7 @@ from onyx.configs.app_configs import REDIS_REPLICA_HOST
 from onyx.configs.app_configs import REDIS_SSL
 from onyx.configs.app_configs import REDIS_SSL_CA_CERTS
 from onyx.configs.app_configs import REDIS_SSL_CERT_REQS
+from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
 from onyx.configs.constants import REDIS_SOCKET_KEEPALIVE_OPTIONS
 from onyx.utils.logger import setup_logger
 
@@ -287,7 +288,7 @@ async def get_async_redis_connection() -> aioredis.Redis:
 
 
 async def retrieve_auth_token_data_from_redis(request: Request) -> dict | None:
-    token = request.cookies.get("fastapiusersauth")
+    token = request.cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME)
     if not token:
         logger.debug("No auth token cookie found")
         return None

backend/onyx/server/manage/users.py

@@ -38,6 +38,7 @@ from onyx.configs.app_configs import ENABLE_EMAIL_INVITES
 from onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS
 from onyx.configs.app_configs import VALID_EMAIL_DOMAINS
 from onyx.configs.constants import AuthType
+from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
 from onyx.db.api_key import is_api_key_email_address
 from onyx.db.auth import get_total_users_count
 from onyx.db.engine import CURRENT_TENANT_ID_CONTEXTVAR
@@ -479,7 +480,7 @@ def get_current_token_expiration_jwt(
 
     try:
         # Get the JWT from the cookie
-        jwt_token = request.cookies.get("fastapiusersauth")
+        jwt_token = request.cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME)
         if not jwt_token:
             logger.error("No JWT token found in cookies")
             return None

backend/scripts/sources_selection_analysis.py

@@ -11,6 +11,8 @@ from typing import Optional
 
 import requests
 
+from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
+
 parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.append(parent_dir)
 
@@ -374,7 +376,11 @@ class SelectionAnalysis:
         Returns:
             dict: The Onyx API response content
         """
-        cookies = {"fastapiusersauth": self._auth_cookie} if self._auth_cookie else {}
+        cookies = (
+            {FASTAPI_USERS_AUTH_COOKIE_NAME: self._auth_cookie}
+            if self._auth_cookie
+            else {}
+        )
 
         endpoint = f"http://127.0.0.1:{self._web_port}/api/direct-qa"
         query_json = {

backend/tests/integration/common_utils/managers/user.py

@@ -7,6 +7,7 @@ import requests
 from requests import HTTPError
 
 from onyx.auth.schemas import UserRole
+from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
 from onyx.server.documents.models import PaginatedReturn
 from onyx.server.models import FullUserSnapshot
 from tests.integration.common_utils.constants import API_SERVER_URL
@@ -82,7 +83,7 @@ class UserManager:
         response.raise_for_status()
 
         cookies = response.cookies.get_dict()
-        session_cookie = cookies.get("fastapiusersauth")
+        session_cookie = cookies.get(FASTAPI_USERS_AUTH_COOKIE_NAME)
 
         if not session_cookie:
             raise Exception("Failed to login")