feat(auth): cloud SSO login core (workspace pin + session override) (#13936)

backend/ee/onyx/db/user_tenant_mapping.py

@@ -97,6 +97,47 @@ def get_tenant_id_for_email(email: str) -> str:
     return tenant_id
 
 
+def lookup_tenant_id_for_login(email: str) -> str | None:
+    """Workspace this address may sign in to, resolved without side effects.
+
+    SSO discovery runs before the user has authenticated, so unlike
+    ``get_tenant_id_for_email`` it must not accept a pending invitation on the
+    caller's behalf. An address that names no single workspace resolves to
+    None, which the caller surfaces as "no SSO here" rather than an error.
+    """
+    if not MULTI_TENANT:
+        return POSTGRES_DEFAULT_SCHEMA
+
+    email = email.lower()
+    with get_catalog_session() as db_session:
+        active_tenant_ids = (
+            db_session.scalars(
+                select(UserTenantMapping.tenant_id).where(
+                    UserTenantMapping.email == email,
+                    UserTenantMapping.active.is_(True),
+                )
+            )
+            .unique()
+            .all()
+        )
+        if len(active_tenant_ids) == 1:
+            return active_tenant_ids[0]
+        if active_tenant_ids:
+            return None
+
+        pending_tenant_ids = (
+            db_session.scalars(
+                select(UserTenantMapping.tenant_id).where(
+                    UserTenantMapping.email == email,
+                    UserTenantMapping.active.is_(False),
+                )
+            )
+            .unique()
+            .all()
+        )
+        return pending_tenant_ids[0] if len(pending_tenant_ids) == 1 else None
+
+
 def resolve_tenant_id(
     email: str, oauth_name: str | None = None, account_id: str | None = None
 ) -> str | None:

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

@@ -10,11 +10,13 @@ from onyx.db.engine.sql_engine import get_session_with_tenant
 from onyx.db.users import get_user_by_email
 from onyx.utils.audit import (
     AuditAction,
+    AuditActor,
     AuditOutcome,
     actor_from_user,
     emit_audit_event,
 )
 from onyx.utils.logger import setup_logger
+from shared_configs.contextvars import SESSION_TENANT_OVERRIDE_CONTEXTVAR
 
 logger = setup_logger()
 
@@ -42,6 +44,25 @@ async def impersonate_user(
         )
         raise HTTPException(status_code=422, detail=detail)
 
+    # The token must name the impersonated user's workspace. Left to the ambient
+    # request tenant it would name the superuser's own.
+    override_token = SESSION_TENANT_OVERRIDE_CONTEXTVAR.set(tenant_id)
+    try:
+        return await _issue_impersonation_token(
+            impersonate_request=impersonate_request,
+            tenant_id=tenant_id,
+            actor=actor,
+        )
+    finally:
+        SESSION_TENANT_OVERRIDE_CONTEXTVAR.reset(override_token)
+
+
+async def _issue_impersonation_token(
+    *,
+    impersonate_request: ImpersonateRequest,
+    tenant_id: str,
+    actor: AuditActor | None,
+) -> Response:
     with get_session_with_tenant(tenant_id=tenant_id) as tenant_session:
         user_to_impersonate = get_user_by_email(
             impersonate_request.email, tenant_session

backend/onyx/auth/sso_tenant_token.py

@@ -0,0 +1,53 @@
+"""Signed workspace pin for an in-flight cloud SSO login.
+
+Cloud serves every workspace from one domain, so the SSO login routes have no
+tenant until the user identifies themselves. Discovery resolves the workspace
+from the catalog and returns this token. The authorize call presents it to
+select the schema holding the provider row. Signed because the tenant it names
+decides which workspace's IdP configuration is read, and short-lived because it
+only has to survive the click between discovery and authorize.
+"""
+
+from typing import Any
+
+import jwt
+from fastapi_users.jwt import decode_jwt, generate_jwt
+
+from onyx.configs.app_configs import USER_AUTH_SECRET
+from onyx.db.engine.sql_engine import is_valid_schema_name
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError
+
+SSO_TENANT_TOKEN_AUDIENCE = "onyx:sso-login-tenant"
+SSO_TENANT_TOKEN_LIFETIME_SECONDS = 600
+SSO_TENANT_TOKEN_PARAM = "workspace_token"
+
+_TENANT_CLAIM = "tenant_id"
+
+
+def generate_sso_tenant_token(tenant_id: str) -> str:
+    data: dict[str, Any] = {
+        _TENANT_CLAIM: tenant_id,
+        "aud": SSO_TENANT_TOKEN_AUDIENCE,
+    }
+    return generate_jwt(data, USER_AUTH_SECRET, SSO_TENANT_TOKEN_LIFETIME_SECONDS)
+
+
+def decode_sso_tenant_token(token: str) -> str:
+    """The tenant this login step is pinned to. Raises when the token is
+    unsigned by us, expired, or names something that is not a schema."""
+    try:
+        payload = decode_jwt(token, USER_AUTH_SECRET, [SSO_TENANT_TOKEN_AUDIENCE])
+    except jwt.PyJWTError as e:
+        raise OnyxError(
+            OnyxErrorCode.UNAUTHORIZED,
+            "Sign-in link has expired. Return to the login page and try again.",
+        ) from e
+
+    tenant_id = payload.get(_TENANT_CLAIM)
+    if not isinstance(tenant_id, str) or not is_valid_schema_name(tenant_id):
+        raise OnyxError(
+            OnyxErrorCode.UNAUTHORIZED,
+            "Sign-in link is not valid for any workspace.",
+        )
+    return tenant_id

backend/onyx/auth/users.py

@@ -173,6 +173,7 @@ from shared_configs.contextvars import (
     CURRENT_TENANT_ID_CONTEXTVAR,
     CURRENT_USAGE_CREDENTIAL_CONTEXTVAR,
     CURRENT_USER_ID_CONTEXTVAR,
+    SESSION_TENANT_OVERRIDE_CONTEXTVAR,
     UsageCredentialIdentity,
     get_current_tenant_id,
 )
@@ -524,6 +525,20 @@ def _invalidate_license_cache_after_seat_change() -> None:
     )()
 
 
+async def resolve_tenant_for_user(email: str, request: Request | None = None) -> str:
+    """Workspace to bind a user to. An explicit override wins, since an SSO login
+    knows its workspace before the user row exists."""
+    override = SESSION_TENANT_OVERRIDE_CONTEXTVAR.get()
+    if override is not None:
+        return override
+
+    return await fetch_ee_implementation_or_noop(
+        "onyx.server.tenants.provisioning",
+        "get_or_provision_tenant",
+        async_return_default_schema,
+    )(email=email, request=request)
+
+
 @asynccontextmanager
 async def _tenant_session_with_context(
     tenant_id: str,
@@ -959,7 +974,11 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             getattr(request.state, "referral_source", None) if request else None
         )
 
-        tenant_id = await fetch_ee_implementation_or_noop(
+        # A workspace-configured provider vouches for who someone is, never for
+        # where they belong, so a pinned login (override set) skips provisioning
+        # entirely and lets the invite gate below decide whether it admits them.
+        override = SESSION_TENANT_OVERRIDE_CONTEXTVAR.get()
+        tenant_id = override or await fetch_ee_implementation_or_noop(
             "onyx.server.tenants.provisioning",
             "get_or_provision_tenant",
             async_return_default_schema,
@@ -1219,14 +1238,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
     async def on_after_register(
         self, user: User, request: Optional[Request] = None
     ) -> None:
-        tenant_id = await fetch_ee_implementation_or_noop(
-            "onyx.server.tenants.provisioning",
-            "get_or_provision_tenant",
-            async_return_default_schema,
-        )(
-            email=user.email,
-            request=request,
-        )
+        tenant_id = await resolve_tenant_for_user(user.email, request)
 
         user_count = None
         token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
@@ -1598,11 +1610,9 @@ class TenantAwareRedisStrategy(RedisStrategy[User, uuid.UUID]):
     async def write_token(self, user: User) -> str:
         redis = await get_async_redis_connection()
 
-        tenant_id = await fetch_ee_implementation_or_noop(
-            "onyx.server.tenants.provisioning",
-            "get_or_provision_tenant",
-            async_return_default_schema,
-        )(email=user.email)
+        # The token names the workspace every later request runs against, so it
+        # has to agree with the one the login actually entered.
+        tenant_id = await resolve_tenant_for_user(user.email)
 
         now = datetime.now(timezone.utc)
         token = secrets.token_urlsafe()
@@ -2543,7 +2553,14 @@ async def complete_login_flow(
     allowed_email_domains_override: Sequence[str] | None = None,
 ) -> RedirectResponse:
     """Shared post-token OAuth/OIDC login: read the verified identity, create or
-    authenticate the user, and return a web or mobile redirect."""
+    authenticate the user, and return a web or mobile redirect.
+
+    Runs inside whatever tenant context the caller set. An SSO login pins its
+    workspace via SESSION_TENANT_OVERRIDE_CONTEXTVAR, which keeps a pinned login
+    inside that one workspace: session issuance and the post-register hook read
+    the same override rather than re-deriving a workspace from the address, which
+    a first-time member does not yet answer to.
+    """
     # Convert a failed or unverified userinfo fetch into a controlled login
     # rejection. OnyxError has a global handler, GetIdEmailError would 500.
     try:
@@ -2566,9 +2583,14 @@ async def complete_login_flow(
     referral_source = state_data.get("referral_source", None)
     # Drives the new_team redirect below. Resolving differently from the login
     # itself would greet a returning user as a brand new signup.
-    tenant_id = fetch_ee_implementation_or_noop(
-        "onyx.db.user_tenant_mapping", "resolve_tenant_id", None
-    )(account_email, oauth_client.name, account_id)
+    tenant_id = (
+        SESSION_TENANT_OVERRIDE_CONTEXTVAR.get()
+        or fetch_ee_implementation_or_noop(
+            "onyx.db.user_tenant_mapping", "resolve_tenant_id", None
+        )(account_email, oauth_client.name, account_id)
+    )
+
+    request.state.referral_source = referral_source
 
     # Snapshot the raw IdP claims for directory-profile enrichment and the admin
     # "OAuth Test" page. The subject-resolved tenant keeps capture working after
@@ -2577,8 +2599,6 @@ async def complete_login_flow(
         oauth_client, account_email, token, tenant_id=tenant_id
     )
 
-    request.state.referral_source = referral_source
-
     try:
         user = await user_manager.oauth_callback(  # ty: ignore[invalid-argument-type]
             oauth_client.name,

backend/onyx/main.py

@@ -159,6 +159,7 @@ from onyx.server.saml_multi import router as saml_multi_router
 from onyx.server.security.api import admin_router as security_admin_router
 from onyx.server.settings.api import admin_router as settings_admin_router
 from onyx.server.settings.api import basic_router as settings_router
+from onyx.server.sso_discovery import router as sso_discovery_router
 from onyx.server.utils import BasicAuthenticationError
 from onyx.setup import setup_multitenant_onyx, setup_onyx
 from onyx.tracing.setup import setup_tracing
@@ -696,6 +697,13 @@ def get_application(lifespan_override: Lifespan | None = None) -> FastAPI:
         oidc_multi_router,
     )
 
+    # Resolves which workspace's providers to offer before any of the above can
+    # run. Mounted unconditionally: single-tenant answers from its one schema.
+    include_auth_router_with_prefix(
+        application,
+        sso_discovery_router,
+    )
+
     include_auth_router_with_prefix(
         application,
         fastapi_users.get_refresh_router(auth_backend),

backend/onyx/server/auth_check.py

@@ -78,6 +78,10 @@ PUBLIC_ENDPOINT_SPECS = [
     # db-backed multi-provider oidc/google (oidc_multi router)
     ("/auth/oidc/{provider_name}/authorize", {"GET"}),
     ("/auth/oidc/{provider_name}/callback", {"GET"}),
+    # Resolves a workspace's SSO buttons before the user has any session. Public
+    # only because it answers uniformly for unknown, ambiguous, and SSO-less
+    # addresses, and is rate limited.
+    ("/auth/sso/discover", {"POST"}),
     # saml (single router: legacy-compatible + parametric authorize, one
     # issuer-resolved callback)
     ("/auth/saml/authorize", {"GET"}),

backend/onyx/server/oidc_multi.py

@@ -3,11 +3,18 @@
 Resolves the enabled provider row from the database on each request so one
 deployment can serve multiple Google and generic OIDC IdPs. Ships dark when no
 matching provider rows exist.
+
+Provider rows are per-workspace, and on cloud a login request carries no session
+to resolve the workspace from. Both halves of the flow therefore name it
+explicitly: authorize takes the signed pin discovery issued, and the callback
+reads it back out of the OAuth state Onyx signed on the way out.
 """
 
 import hashlib
 import json
 import uuid
+from collections.abc import AsyncGenerator, Generator
+from contextlib import contextmanager
 from typing import Any
 
 from cachetools import TTLCache
@@ -21,6 +28,10 @@ from httpx_oauth.oauth2 import BaseOAuth2, GetAccessTokenError
 from sqlalchemy.orm import Session
 
 from onyx.auth.oidc_client import VerifiedEmailOpenID
+from onyx.auth.sso_tenant_token import (
+    SSO_TENANT_TOKEN_PARAM,
+    decode_sso_tenant_token,
+)
 from onyx.auth.sso_url_guard import (
     UnsafeSSOUrl,
     validate_discovered_endpoints,
@@ -50,7 +61,10 @@ from onyx.configs.app_configs import (
     USER_AUTH_SECRET,
     WEB_DOMAIN,
 )
-from onyx.db.engine.sql_engine import get_session
+from onyx.db.engine.sql_engine import (
+    get_session_with_current_tenant,
+    get_session_with_tenant,
+)
 from onyx.db.enums import SSOProviderType
 from onyx.db.models import SSOProvider, User
 from onyx.db.sso_provider import (
@@ -61,10 +75,24 @@ from onyx.db.sso_provider import (
 from onyx.error_handling.error_codes import OnyxErrorCode
 from onyx.error_handling.exceptions import OnyxError
 from onyx.utils.url import sanitize_next_url
-from shared_configs.contextvars import get_current_tenant_id
+from shared_configs.configs import MULTI_TENANT
+from shared_configs.contextvars import (
+    CURRENT_TENANT_ID_CONTEXTVAR,
+    SESSION_TENANT_OVERRIDE_CONTEXTVAR,
+    get_current_tenant_id,
+)
 
 router = APIRouter(prefix="/auth/oidc")
 
+_NO_WORKSPACE_DETAIL = (
+    "Sign-in did not identify a workspace. Return to the login page and enter "
+    "your email to continue."
+)
+
+# Named in the OAuth state so the callback can reach the workspace whose
+# provider row started the flow.
+_STATE_TENANT_KEY = "tenant_id"
+
 _CLIENT_CACHE_TTL_SECONDS = 600
 # Hardcoded off: linking a second IdP to an existing account by verified email is
 # an account-takeover vector when two IdPs can assert one domain.
@@ -76,6 +104,60 @@ _CLIENT_CACHE: TTLCache[tuple[str, str, SSOProviderType, str], BaseOAuth2[Any]]
 _COOKIE_SECURE = WEB_DOMAIN.startswith("https")
 
 
+@contextmanager
+def _workspace_session(tenant_id: str | None) -> Generator[Session, None, None]:
+    # Context var travels with the session so downstream login code agrees on
+    # the workspace this request belongs to.
+    if not MULTI_TENANT:
+        with get_session_with_current_tenant() as db_session:
+            yield db_session
+        return
+
+    if not tenant_id:
+        raise OnyxError(OnyxErrorCode.UNAUTHORIZED, _NO_WORKSPACE_DETAIL)
+    context_token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
+    try:
+        with get_session_with_tenant(tenant_id=tenant_id) as db_session:
+            yield db_session
+    finally:
+        CURRENT_TENANT_ID_CONTEXTVAR.reset(context_token)
+
+
+@contextmanager
+def _pinned_workspace(tenant_id: str | None) -> Generator[None, None, None]:
+    """Bind the callback's whole tail to one workspace. The override is what
+    keeps a pinned login there: session issuance and post-register read it
+    instead of re-deriving from the address."""
+    if not MULTI_TENANT:
+        yield
+        return
+    if not tenant_id:
+        raise OnyxError(OnyxErrorCode.UNAUTHORIZED, _NO_WORKSPACE_DETAIL)
+    tenant_token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
+    override_token = SESSION_TENANT_OVERRIDE_CONTEXTVAR.set(tenant_id)
+    try:
+        yield
+    finally:
+        SESSION_TENANT_OVERRIDE_CONTEXTVAR.reset(override_token)
+        CURRENT_TENANT_ID_CONTEXTVAR.reset(tenant_token)
+
+
+async def get_authorize_session(
+    request: Request,
+) -> AsyncGenerator[Session, None]:
+    """Session for the authorize call, which has no session cookie to resolve a
+    workspace from and instead presents the pin discovery issued."""
+    workspace_token = request.query_params.get(SSO_TENANT_TOKEN_PARAM)
+    tenant_id = decode_sso_tenant_token(workspace_token) if workspace_token else None
+    with _workspace_session(tenant_id) as db_session:
+        yield db_session
+
+
+def _state_workspace(state_data: dict[str, Any]) -> str | None:
+    tenant_id = state_data.get(_STATE_TENANT_KEY)
+    return tenant_id if isinstance(tenant_id, str) and tenant_id else None
+
+
 def _resolve_oidc_provider(
     db_session: Session, provider_name: str
 ) -> tuple[SSOProvider, dict[str, Any]]:
@@ -229,7 +311,7 @@ def _callback_uri(provider: SSOProvider, config: dict[str, Any]) -> str:
 async def oidc_login_for_provider(
     provider_name: str,
     request: Request,
-    db_session: Session = Depends(get_session),
+    db_session: Session = Depends(get_authorize_session),
 ) -> Response:
     provider, config = _resolve_oidc_provider(db_session, provider_name)
     client = await _get_oauth_client(provider, config)
@@ -244,6 +326,9 @@ async def oidc_login_for_provider(
             # Pins this flow's PKCE mode, so a provider edit mid-login cannot
             # make the callback disagree with the authorization request.
             "pkce": use_pkce,
+            # The IdP redirect arrives with no session, so the state is the only
+            # thing carrying the workspace across the round trip.
+            _STATE_TENANT_KEY: get_current_tenant_id(),
             CSRF_TOKEN_KEY: csrf_token,
         },
         USER_AUTH_SECRET,
@@ -297,7 +382,6 @@ async def oidc_login_callback(
     code: str | None = None,
     state: str | None = None,
     error: str | None = None,
-    db_session: Session = Depends(get_session),
     strategy: Strategy[User, uuid.UUID] = Depends(auth_backend.get_strategy),
     user_manager: UserManager = Depends(get_user_manager),
 ) -> Response:
@@ -326,7 +410,6 @@ async def oidc_login_callback(
         code=code,
         state=state,
         error=error,
-        db_session=db_session,
         strategy=strategy,
         user_manager=user_manager,
     )
@@ -340,14 +423,9 @@ async def oidc_login_callback_for_provider(
     code: str | None = None,
     state: str | None = None,
     error: str | None = None,
-    db_session: Session = Depends(get_session),
     strategy: Strategy[User, uuid.UUID] = Depends(auth_backend.get_strategy),
     user_manager: UserManager = Depends(get_user_manager),
 ) -> Response:
-    provider, config = _resolve_oidc_provider(db_session, provider_name)
-    client = await _get_oauth_client(provider, config)
-    redirect_uri = _callback_uri(provider, config)
-
     if error is not None:
         raise OnyxError(
             OnyxErrorCode.VALIDATION_ERROR,
@@ -364,47 +442,59 @@ async def oidc_login_callback_for_provider(
             "Missing state parameter in OAuth callback",
         )
 
+    # Validated before the workspace is read out of it, so a tampered or expired
+    # state is rejected here before any workspace is selected.
     state_data = decode_and_validate_oauth_state(
         request=request,
         state_value=state,
         state_secret=USER_AUTH_SECRET,
         expected_provider_name=provider_name,
     )
+    pinned_tenant_id = _state_workspace(state_data)
 
-    # The state pins the flow's PKCE mode. States without the claim fall back
-    # to the row's current setting so logins in flight across a deploy complete.
-    use_pkce = (
-        bool(state_data["pkce"]) if "pkce" in state_data else _pkce_enabled(config)
-    )
-    code_verifier: str | None = None
-    if use_pkce:
-        code_verifier = request.cookies.get(get_pkce_cookie_name(state))
-        if not code_verifier:
+    # Everything below runs inside the pinned workspace: the client cache is
+    # keyed on the current tenant, and session issuance reads the override.
+    with _pinned_workspace(pinned_tenant_id):
+        with get_session_with_current_tenant() as db_session:
+            provider, config = _resolve_oidc_provider(db_session, provider_name)
+            client = await _get_oauth_client(provider, config)
+            redirect_uri = _callback_uri(provider, config)
+            allowed_email_domains = list(provider.allowed_email_domains)
+
+        # The state pins the flow's PKCE mode. States without the claim fall
+        # back to the row's setting so logins in flight across a deploy complete.
+        use_pkce = (
+            bool(state_data["pkce"]) if "pkce" in state_data else _pkce_enabled(config)
+        )
+        code_verifier: str | None = None
+        if use_pkce:
+            code_verifier = request.cookies.get(get_pkce_cookie_name(state))
+            if not code_verifier:
+                raise OnyxError(
+                    OnyxErrorCode.VALIDATION_ERROR,
+                    "Missing PKCE verifier cookie in OAuth callback",
+                )
+
+        try:
+            token = await client.get_access_token(code, redirect_uri, code_verifier)
+        except GetAccessTokenError as e:
             raise OnyxError(
                 OnyxErrorCode.VALIDATION_ERROR,
-                "Missing PKCE verifier cookie in OAuth callback",
-            )
+                "Authorization code exchange failed",
+            ) from e
 
-    try:
-        token = await client.get_access_token(code, redirect_uri, code_verifier)
-    except GetAccessTokenError as e:
-        raise OnyxError(
-            OnyxErrorCode.VALIDATION_ERROR,
-            "Authorization code exchange failed",
-        ) from e
-
-    redirect_response = await complete_login_flow(
-        oauth_client=client,
-        token=token,
-        state_data=state_data,
-        request=request,
-        user_manager=user_manager,
-        backend=auth_backend,
-        strategy=strategy,
-        associate_by_email=_ALLOW_AUTO_LINK,
-        is_verified_by_default=True,
-        allowed_email_domains_override=provider.allowed_email_domains,
-    )
+        redirect_response = await complete_login_flow(
+            oauth_client=client,
+            token=token,
+            state_data=state_data,
+            request=request,
+            user_manager=user_manager,
+            backend=auth_backend,
+            strategy=strategy,
+            associate_by_email=_ALLOW_AUTO_LINK,
+            is_verified_by_default=True,
+            allowed_email_domains_override=allowed_email_domains,
+        )
 
     if use_pkce:
         delete_pkce_cookie(redirect_response, state)

backend/onyx/server/sso_discovery.py

@@ -0,0 +1,139 @@
+"""Workspace discovery for the cloud login page.
+
+Cloud serves every workspace from one domain, so the login page cannot know
+which IdP to offer until the visitor names themselves. This endpoint takes an
+address, resolves the workspace from the shared catalog, and returns that
+workspace's sign-in buttons with a signed pin the authorize call uses to reach
+the right schema.
+
+The response is deliberately uniform: an address that maps nowhere, maps to
+several workspaces, or maps to a workspace with no SSO all return an empty
+list, so the endpoint cannot be used to tell those cases apart.
+"""
+
+import time
+
+from fastapi import APIRouter, Request
+from fastapi.concurrency import run_in_threadpool
+from pydantic import BaseModel, EmailStr
+
+from onyx.auth.sso_tenant_token import (
+    SSO_TENANT_TOKEN_PARAM,
+    generate_sso_tenant_token,
+)
+from onyx.db.engine.sql_engine import get_session_with_tenant
+from onyx.db.models import SSOProvider
+from onyx.db.sso_provider import (
+    fetch_sso_providers,
+    sso_authorize_path,
+    sso_provider_type_supported,
+)
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError
+from onyx.redis.redis_pool import get_async_redis_connection
+from onyx.server.manage.models import SSOProviderOption
+from onyx.utils.client_ip import get_client_ip
+from onyx.utils.logger import setup_logger
+from onyx.utils.url import add_url_params
+from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop
+from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA
+
+logger = setup_logger()
+
+router = APIRouter(prefix="/auth/sso")
+
+# A whole office shares one public address, so this has to clear a company
+# signing in on a Monday morning while still bounding a script walking an
+# address list. It bounds the rate of enumeration, not the possibility.
+_PER_IP_PER_HOUR = 300
+_BUCKET_SECONDS = 3600
+_REDIS_KEY_PREFIX = "sso_discovery_rate:"
+
+
+class SSODiscoveryRequest(BaseModel):
+    email: EmailStr
+
+
+class SSODiscoveryResponse(BaseModel):
+    providers: list[SSOProviderOption]
+
+
+async def _enforce_discovery_rate_limit(request: Request) -> None:
+    # Single-tenant has one workspace to enumerate and runs without Redis on
+    # Onyx-lite, so the bound applies only on cloud, where it is the sole limit
+    # on address enumeration.
+    if not MULTI_TENANT:
+        return
+
+    ip = get_client_ip(request) or "unknown"
+    key = f"{_REDIS_KEY_PREFIX}{ip}:{int(time.time() // _BUCKET_SECONDS)}"
+    try:
+        redis = await get_async_redis_connection()
+        pipe = redis.pipeline()
+        pipe.incr(key)
+        pipe.expire(key, _BUCKET_SECONDS)
+        incr_result, _ = await pipe.execute()
+        count = int(incr_result)
+    except Exception as e:
+        # Fails closed. This endpoint answers about addresses it has not
+        # authenticated, so losing Redis must cost the lookup, not the probing
+        # bound on it.
+        logger.exception("SSO discovery rate-limit unavailable, refusing lookup")
+        raise OnyxError(
+            OnyxErrorCode.RATE_LIMITED,
+            "Sign-in lookup is briefly unavailable. Use Google or your password.",
+        ) from e
+
+    if count > _PER_IP_PER_HOUR:
+        logger.warning(
+            "SSO discovery rate limit exceeded for ip=%s count=%s", ip, count
+        )
+        raise OnyxError(
+            OnyxErrorCode.RATE_LIMITED,
+            "Too many sign-in lookups from this network. Please wait before trying again.",
+        )
+
+
+def _authorize_url(provider: SSOProvider, workspace_token: str | None) -> str:
+    path = sso_authorize_path(provider)
+    if workspace_token is None:
+        return path
+    return add_url_params(path, {SSO_TENANT_TOKEN_PARAM: workspace_token})
+
+
+def _login_options(
+    tenant_id: str, workspace_token: str | None
+) -> list[SSOProviderOption]:
+    with get_session_with_tenant(tenant_id=tenant_id) as db_session:
+        return [
+            SSOProviderOption(
+                name=provider.name,
+                display_name=provider.display_name,
+                provider_type=provider.provider_type,
+                authorize_url=_authorize_url(provider, workspace_token),
+            )
+            for provider in fetch_sso_providers(db_session, enabled_only=True)
+            if sso_provider_type_supported(provider.provider_type)
+        ]
+
+
+def _discover(email: str) -> list[SSOProviderOption]:
+    if not MULTI_TENANT:
+        return _login_options(POSTGRES_DEFAULT_SCHEMA, workspace_token=None)
+
+    tenant_id = fetch_ee_implementation_or_noop(
+        "onyx.db.user_tenant_mapping", "lookup_tenant_id_for_login", None
+    )(email)
+    if not tenant_id or tenant_id == POSTGRES_DEFAULT_SCHEMA:
+        return []
+
+    return _login_options(tenant_id, generate_sso_tenant_token(tenant_id))
+
+
+@router.post("/discover")
+async def discover_sso_providers(
+    payload: SSODiscoveryRequest, request: Request
+) -> SSODiscoveryResponse:
+    await _enforce_discovery_rate_limit(request)
+    providers = await run_in_threadpool(_discover, payload.email)
+    return SSODiscoveryResponse(providers=providers)

backend/shared_configs/contextvars.py

@@ -11,6 +11,13 @@ CURRENT_TENANT_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = (
     )
 )
 
+# Workspace a session must be issued against, set only by a caller that already
+# decided it. Not CURRENT_TENANT_ID_CONTEXTVAR: that is whatever cookie the
+# request carried, which can name a workspace this user does not belong to.
+SESSION_TENANT_OVERRIDE_CONTEXTVAR: contextvars.ContextVar[str | None] = (
+    contextvars.ContextVar("session_tenant_override", default=None)
+)
+
 # set by every route in the API server
 INDEXING_REQUEST_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = (
     contextvars.ContextVar("indexing_request_id", default=None)

backend/tests/external_dependency_unit/auth/test_sso_login_discovery.py

@@ -0,0 +1,200 @@
+"""Guard the cloud SSO discovery contract.
+
+Discovery runs before the visitor has authenticated, so it must answer only
+about workspaces the address already belongs to, must not accept an invitation
+on their behalf, and must not distinguish "no such workspace" from "workspace
+without SSO" in what it returns.
+"""
+
+from collections.abc import Generator
+from typing import Any, cast
+from unittest.mock import AsyncMock, patch
+from uuid import uuid4
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from sqlalchemy import Table, delete, inspect, select
+from sqlalchemy.orm import Session
+
+from ee.onyx.db.user_tenant_mapping import lookup_tenant_id_for_login
+from onyx.auth.sso_tenant_token import (
+    decode_sso_tenant_token,
+    generate_sso_tenant_token,
+)
+from onyx.db.models import UserTenantMapping
+from onyx.error_handling.error_codes import OnyxErrorCode
+from onyx.error_handling.exceptions import OnyxError, register_onyx_exception_handlers
+from onyx.server.sso_discovery import router as sso_discovery_router
+
+_TEST_SECRET = "test-user-auth-secret-at-least-32-bytes-long"
+
+
+@pytest.fixture()
+def client() -> Generator[TestClient, None, None]:
+    app = FastAPI()
+    register_onyx_exception_handlers(app)
+    app.include_router(sso_discovery_router)
+    with TestClient(app) as test_client:
+        yield test_client
+
+
+@pytest.fixture()
+def catalog_session(db_session: Session) -> Generator[Session, None, None]:
+    """`user_tenant_mapping` lives in the multi-tenant catalog migrations, which
+    a single-tenant dev database never runs. Create it for the test and leave a
+    real catalog untouched."""
+    bind = db_session.get_bind()
+    table = cast(Table, UserTenantMapping.__table__)
+    created_here = not inspect(bind).has_table(table.name, schema="public")
+    if created_here:
+        table.create(bind=bind)
+        db_session.commit()
+    try:
+        yield db_session
+    finally:
+        if created_here:
+            db_session.rollback()
+            table.drop(bind=bind)
+            db_session.commit()
+
+
+def _new_email() -> str:
+    return f"discovery-{uuid4().hex[:10]}@example.com"
+
+
+def _add_mapping(session: Session, email: str, tenant_id: str, active: bool) -> None:
+    session.add(UserTenantMapping(email=email, tenant_id=tenant_id, active=active))
+    session.commit()
+
+
+def _cleanup(session: Session, email: str) -> None:
+    session.execute(delete(UserTenantMapping).where(UserTenantMapping.email == email))
+    session.commit()
+
+
+def _mapping_active(session: Session, email: str) -> bool:
+    session.expire_all()
+    return bool(
+        session.scalar(
+            select(UserTenantMapping.active).where(UserTenantMapping.email == email)
+        )
+    )
+
+
+@patch("ee.onyx.db.user_tenant_mapping.MULTI_TENANT", True)
+def test_lookup_does_not_accept_a_pending_invitation(
+    catalog_session: Session,
+) -> None:
+    """An unauthenticated lookup that flipped `active` would enroll someone in a
+    workspace by typing their address into a login form."""
+    email = _new_email()
+    tenant_id = f"tenant_{uuid4().hex[:12]}"
+    _add_mapping(catalog_session, email, tenant_id, active=False)
+    try:
+        assert lookup_tenant_id_for_login(email) == tenant_id
+        assert _mapping_active(catalog_session, email) is False
+    finally:
+        _cleanup(catalog_session, email)
+
+
+@patch("ee.onyx.db.user_tenant_mapping.MULTI_TENANT", True)
+def test_lookup_refuses_to_choose_between_invitations(
+    catalog_session: Session,
+) -> None:
+    email = _new_email()
+    _add_mapping(catalog_session, email, f"tenant_{uuid4().hex[:12]}", active=False)
+    _add_mapping(catalog_session, email, f"tenant_{uuid4().hex[:12]}", active=False)
+    try:
+        assert lookup_tenant_id_for_login(email) is None
+    finally:
+        _cleanup(catalog_session, email)
+
+
+@pytest.mark.usefixtures("catalog_session")
+@patch("ee.onyx.db.user_tenant_mapping.MULTI_TENANT", True)
+def test_lookup_returns_none_for_unknown_address() -> None:
+    assert lookup_tenant_id_for_login(_new_email()) is None
+
+
+@patch("onyx.server.sso_discovery.MULTI_TENANT", True)
+@patch("onyx.server.sso_discovery._enforce_discovery_rate_limit", new=AsyncMock())
+def test_unknown_address_returns_an_empty_list(client: TestClient) -> None:
+    """Same shape an SSO-less workspace returns, so the response cannot be used
+    to tell whether an address belongs to a customer."""
+    response = client.post("/auth/sso/discover", json={"email": _new_email()})
+    assert response.status_code == 200
+    assert response.json() == {"providers": []}
+
+
+def test_malformed_address_is_rejected(client: TestClient) -> None:
+    response = client.post("/auth/sso/discover", json={"email": "not-an-email"})
+    assert response.status_code == 422
+
+
+@patch("onyx.server.sso_discovery.MULTI_TENANT", True)
+@patch(
+    "onyx.server.sso_discovery.get_async_redis_connection",
+    side_effect=ConnectionError("redis is down"),
+)
+def test_lookup_refuses_when_the_limiter_is_unavailable(
+    _redis: object, client: TestClient
+) -> None:
+    """The limiter is the only bound on probing this endpoint, so losing it has
+    to cost the lookup rather than the bound."""
+    response = client.post("/auth/sso/discover", json={"email": _new_email()})
+    assert response.status_code == OnyxErrorCode.RATE_LIMITED.status_code
+
+
+def test_workspace_pin_round_trips_and_rejects_another_signer() -> None:
+    """The pin decides which workspace's IdP configuration a login reads, so a
+    token this deployment did not sign must not select one."""
+    tenant_id = f"tenant_{uuid4().hex[:12]}"
+
+    with patch("onyx.auth.sso_tenant_token.USER_AUTH_SECRET", _TEST_SECRET):
+        token = generate_sso_tenant_token(tenant_id)
+        assert decode_sso_tenant_token(token) == tenant_id
+
+    # Same token, a different signing secret. Flipping a character instead would
+    # be flaky: base64url characters can differ while the decoded bytes match.
+    with patch("onyx.auth.sso_tenant_token.USER_AUTH_SECRET", _TEST_SECRET + "x"):
+        with pytest.raises(OnyxError):
+            decode_sso_tenant_token(token)
+
+
+@patch("onyx.auth.sso_tenant_token.USER_AUTH_SECRET", _TEST_SECRET)
+@patch("onyx.server.sso_discovery.MULTI_TENANT", True)
+@patch("onyx.server.sso_discovery._enforce_discovery_rate_limit", new=AsyncMock())
+def test_resolved_workspace_authorize_urls_carry_a_workspace_pin(
+    client: TestClient,
+) -> None:
+    """Authorize has no session to read the workspace from, so discovery has to
+    hand it over in the URL it returns."""
+    tenant_id = f"tenant_{uuid4().hex[:12]}"
+    fake_provider: Any = type(
+        "FakeProvider",
+        (),
+        {
+            "name": "okta",
+            "display_name": "Okta",
+            "provider_type": "OIDC",
+        },
+    )()
+
+    with (
+        patch(
+            "onyx.server.sso_discovery.fetch_ee_implementation_or_noop",
+            return_value=lambda _email: tenant_id,
+        ),
+        patch("onyx.server.sso_discovery.get_session_with_tenant"),
+        patch(
+            "onyx.server.sso_discovery.fetch_sso_providers",
+            return_value=[fake_provider],
+        ),
+    ):
+        response = client.post("/auth/sso/discover", json={"email": _new_email()})
+
+    assert response.status_code == 200
+    [provider] = response.json()["providers"]
+    assert provider["authorize_url"].startswith("/api/auth/oidc/okta/authorize?")
+    assert "workspace_token=" in provider["authorize_url"]

backend/tests/unit/ee/onyx/server/tenants/test_impersonation_tenant_binding.py

@@ -0,0 +1,74 @@
+"""The impersonation token must name the impersonated user's workspace.
+
+Session issuance takes an explicit override when the caller has already decided
+the workspace. Without one it would fall back to resolving the superuser's own
+address, minting a session for the wrong workspace.
+"""
+
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import AsyncMock, patch
+from uuid import uuid4
+
+import pytest
+
+from ee.onyx.server.tenants.admin_api import impersonate_user
+from shared_configs.contextvars import SESSION_TENANT_OVERRIDE_CONTEXTVAR
+
+_SUPERUSER_TENANT = "tenant_superuser"
+_TARGET_TENANT = "tenant_target"
+
+
+@pytest.mark.asyncio
+async def test_token_is_minted_against_the_target_workspace() -> None:
+    target_user = SimpleNamespace(id=uuid4(), email="member@target.example")
+    superuser = SimpleNamespace(id=uuid4(), email="ops@onyx.app")
+    tenant_at_issuance: dict[str, str | None] = {}
+
+    async def capture_write_token(_user: Any) -> str:
+        tenant_at_issuance["value"] = SESSION_TENANT_OVERRIDE_CONTEXTVAR.get()
+        return "token"
+
+    # A stale override from earlier work on this task must not leak through.
+    context_token = SESSION_TENANT_OVERRIDE_CONTEXTVAR.set(_SUPERUSER_TENANT)
+    try:
+        with (
+            patch(
+                "ee.onyx.server.tenants.admin_api.get_tenant_id_for_email",
+                return_value=_TARGET_TENANT,
+            ),
+            patch("ee.onyx.server.tenants.admin_api.get_session_with_tenant"),
+            patch(
+                "ee.onyx.server.tenants.admin_api.get_user_by_email",
+                return_value=target_user,
+            ),
+            patch(
+                "ee.onyx.server.tenants.admin_api.get_redis_strategy",
+                return_value=SimpleNamespace(write_token=capture_write_token),
+            ),
+            patch(
+                "ee.onyx.server.tenants.admin_api.auth_backend",
+                SimpleNamespace(
+                    transport=SimpleNamespace(
+                        get_login_response=AsyncMock(
+                            return_value=SimpleNamespace(set_cookie=lambda **_kw: None)
+                        )
+                    )
+                ),
+            ),
+            patch("ee.onyx.server.tenants.admin_api.emit_audit_event"),
+            patch(
+                "ee.onyx.server.tenants.admin_api.actor_from_user", return_value=None
+            ),
+        ):
+            await impersonate_user(
+                impersonate_request=SimpleNamespace(email=target_user.email),  # ty: ignore[invalid-argument-type]
+                superuser=superuser,  # ty: ignore[invalid-argument-type]
+            )
+            # Checked before the test's own reset, so this is the handler's
+            # restoration and not this fixture's.
+            assert SESSION_TENANT_OVERRIDE_CONTEXTVAR.get() == _SUPERUSER_TENANT
+    finally:
+        SESSION_TENANT_OVERRIDE_CONTEXTVAR.reset(context_token)
+
+    assert tenant_at_issuance["value"] == _TARGET_TENANT

backend/tests/unit/onyx/auth/test_session_tenant_resolution.py

@@ -0,0 +1,61 @@
+"""Which workspace a freshly issued session names.
+
+Only a caller that has already decided the workspace may set it. The request's
+own tenant context is whatever cookie the browser sent, which on a re-login can
+name a workspace the user being authenticated does not belong to, so reading it
+here would hand them a session for someone else's.
+"""
+
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+
+from onyx.auth.users import resolve_tenant_for_user
+from shared_configs.contextvars import (
+    CURRENT_TENANT_ID_CONTEXTVAR,
+    SESSION_TENANT_OVERRIDE_CONTEXTVAR,
+)
+
+_COOKIE_TENANT = "tenant_from_stale_cookie"
+_MEMBER_TENANT = "tenant_the_user_belongs_to"
+_PINNED_TENANT = "tenant_pinned_by_sso"
+
+
+def _catalog_returns(tenant_id: str) -> Any:
+    async def _resolve(**_kwargs: Any) -> str:
+        return tenant_id
+
+    return lambda *_args, **_kwargs: _resolve
+
+
+@pytest.mark.asyncio
+async def test_ambient_request_tenant_does_not_decide_the_session() -> None:
+    token = CURRENT_TENANT_ID_CONTEXTVAR.set(_COOKIE_TENANT)
+    try:
+        with patch(
+            "onyx.auth.users.fetch_ee_implementation_or_noop",
+            _catalog_returns(_MEMBER_TENANT),
+        ):
+            resolved = await resolve_tenant_for_user("member@example.com")
+    finally:
+        CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
+
+    assert resolved == _MEMBER_TENANT
+
+
+@pytest.mark.asyncio
+async def test_explicit_override_wins_over_the_catalog() -> None:
+    """An SSO login pins its workspace before the user row exists, so the
+    catalog cannot answer for a first-time member."""
+    token = SESSION_TENANT_OVERRIDE_CONTEXTVAR.set(_PINNED_TENANT)
+    try:
+        with patch(
+            "onyx.auth.users.fetch_ee_implementation_or_noop",
+            _catalog_returns(_MEMBER_TENANT),
+        ):
+            resolved = await resolve_tenant_for_user("newcomer@example.com")
+    finally:
+        SESSION_TENANT_OVERRIDE_CONTEXTVAR.reset(token)
+
+    assert resolved == _PINNED_TENANT

backend/tests/unit/onyx/server/test_oidc_multi.py

@@ -444,7 +444,6 @@ def test_fixed_callback_rejects_missing_state() -> None:
                 code="code",
                 state=None,
                 error=None,
-                db_session=_DB,
                 strategy=cast(Any, None),
                 user_manager=cast(Any, None),
             )
@@ -472,7 +471,6 @@ def test_fixed_callback_rejects_state_without_provider(
                 code="code",
                 state=state,
                 error=None,
-                db_session=_DB,
                 strategy=cast(Any, None),
                 user_manager=cast(Any, None),
             )