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 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 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_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.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 @@ async def on_after_forgot_password(
             )
             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 @@
 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 = "[email protected]"
 

backend/onyx/redis/redis_pool.py

@@ -25,6 +25,7 @@
 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 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 @@
 
 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 @@ def do_request(self, query: str) -> dict:
         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 @@
 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 @@ def login_as_user(test_user: DATestUser) -> DATestUser:
         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")

web/src/app/auth/forgot-password/utils.ts

@@ -28,6 +28,12 @@ export const resetPassword = async (
   });
 
   if (!response.ok) {
-    throw new Error("Failed to reset password");
+    const error = await response.json();
+    if (error?.detail?.code === "RESET_PASSWORD_INVALID_PASSWORD") {
+      throw new Error(error.detail.reason || "Invalid password");
+    }
+    const errorMessage =
+      error?.detail || "An error occurred during password reset.";
+    throw new Error(errorMessage);
   }
 };

web/src/app/auth/reset-password/page.tsx

@@ -1,5 +1,5 @@
 "use client";
-import React, { useState } from "react";
+import React, { useState, useEffect } from "react";
 import { resetPassword } from "../forgot-password/utils";
 import AuthFlowContainer from "@/components/auth/AuthFlowContainer";
 import CardSection from "@/components/admin/CardSection";
@@ -13,13 +13,28 @@ import { TextFormField } from "@/components/admin/connectors/Field";
 import { usePopup } from "@/components/admin/connectors/Popup";
 import { Spinner } from "@/components/Spinner";
 import { redirect, useSearchParams } from "next/navigation";
-import { NEXT_PUBLIC_FORGOT_PASSWORD_ENABLED } from "@/lib/constants";
+import {
+  NEXT_PUBLIC_FORGOT_PASSWORD_ENABLED,
+  TENANT_ID_COOKIE_NAME,
+} from "@/lib/constants";
+import Cookies from "js-cookie";
 
 const ResetPasswordPage: React.FC = () => {
   const { popup, setPopup } = usePopup();
   const [isWorking, setIsWorking] = useState(false);
   const searchParams = useSearchParams();
   const token = searchParams.get("token");
+  const tenantId = searchParams.get(TENANT_ID_COOKIE_NAME);
+  // Keep search param same name as cookie for simplicity
+
+  useEffect(() => {
+    if (tenantId) {
+      Cookies.set(TENANT_ID_COOKIE_NAME, tenantId, {
+        path: "/",
+        expires: 1 / 24,
+      }); // Expires in 1 hour
+    }
+  }, [tenantId]);
 
   if (!NEXT_PUBLIC_FORGOT_PASSWORD_ENABLED) {
     redirect("/auth/login");
@@ -63,10 +78,18 @@ const ResetPasswordPage: React.FC = () => {
                 redirect("/auth/login");
               }, 1000);
             } catch (error) {
-              setPopup({
-                type: "error",
-                message: "An error occurred. Please try again.",
-              });
+              if (error instanceof Error) {
+                setPopup({
+                  type: "error",
+                  message:
+                    error.message || "An error occurred during password reset.",
+                });
+              } else {
+                setPopup({
+                  type: "error",
+                  message: "An unexpected error occurred. Please try again.",
+                });
+              }
             } finally {
               setIsWorking(false);
             }

web/src/lib/constants.ts

@@ -18,6 +18,8 @@ export const NEXT_PUBLIC_DO_NOT_USE_TOGGLE_OFF_DANSWER_POWERED =
   process.env.NEXT_PUBLIC_DO_NOT_USE_TOGGLE_OFF_DANSWER_POWERED?.toLowerCase() ===
   "true";
 
+export const TENANT_ID_COOKIE_NAME = "onyx_tid";
+
 export const GMAIL_AUTH_IS_ADMIN_COOKIE_NAME = "gmail_auth_is_admin";
 
 export const GOOGLE_DRIVE_AUTH_IS_ADMIN_COOKIE_NAME =