fix: oidc self deadlock (#14181)

backend/onyx/auth/users.py

@@ -541,24 +541,29 @@ async def resolve_tenant_for_user(email: str, request: Request | None = None) ->
     )(email=email, request=request)
 
 
-@asynccontextmanager
-async def _tenant_session_with_context(
-    tenant_id: str,
-) -> AsyncGenerator[AsyncSession, None]:
-    token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
-    try:
-        async with get_async_session_context_manager(tenant_id) as db_session:
-            yield db_session
-    finally:
-        CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
-
-
 class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
     reset_password_token_secret = USER_AUTH_SECRET
     verification_token_secret = USER_AUTH_SECRET
     verification_token_lifetime_seconds = AUTH_COOKIE_EXPIRE_TIME_SECONDS
     user_db: SQLAlchemyUserDatabase[User, uuid.UUID]
 
+    @asynccontextmanager
+    async def _tenant_session_with_bound_user_db(
+        self, tenant_id: str
+    ) -> AsyncGenerator[AsyncSession, None]:
+        """Run tenant work on one AsyncSession; user_db writes use that connection."""
+        token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
+        previous_user_db = self.user_db
+        try:
+            async with get_async_session_context_manager(tenant_id) as db_session:
+                self.user_db = SQLAlchemyUserDatabase[User, uuid.UUID](
+                    db_session, User, OAuthAccount
+                )
+                yield db_session
+        finally:
+            self.user_db = previous_user_db
+            CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
+
     async def get_by_email(self, user_email: str) -> User:
         tenant_id = fetch_ee_implementation_or_noop(
             "onyx.db.user_tenant_mapping", "get_tenant_id_for_email", None
@@ -963,8 +968,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
         if not tenant_id:
             raise HTTPException(status_code=401, detail="User not found")
 
-        # Proceed with the tenant context
-        async with _tenant_session_with_context(tenant_id) as db_session:
+        async with self._tenant_session_with_bound_user_db(tenant_id) as db_session:
             verify_email_in_whitelist(account_email, tenant_id, oauth_name, account_id)
             oauth_security_settings = get_security_settings()
             effective_valid_email_domains = (
@@ -991,14 +995,6 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
                         "single sign-on.",
                     )
 
-            # NOTE(rkuo): If this UserManager is instantiated per connection
-            # should we even be doing this here?
-            if MULTI_TENANT:
-                tenant_user_db = SQLAlchemyUserDatabase[User, uuid.UUID](
-                    db_session, User, OAuthAccount
-                )
-                self.user_db = tenant_user_db
-
             oauth_account_dict = {
                 "oauth_name": oauth_name,
                 "access_token": access_token,
@@ -1140,14 +1136,10 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
 
             # Handle case where user has used product outside of web and is now creating an account through web
             if not user.account_type.is_web_login():
-                # We must use the existing user in the session if it matches
-                # the user we just got by email/oauth. Note that this only applies
-                # to multi-tenant, due to the overwriting of the user_db
-                if MULTI_TENANT:
-                    if user.id:
-                        user_by_session = await db_session.get(User, user.id)
-                        if user_by_session:
-                            user = user_by_session
+                if user.id:
+                    user_by_session = await db_session.get(User, user.id)
+                    if user_by_session:
+                        user = user_by_session
 
                 # Lock + check + upgrade in one transaction.
                 was_inactive = not user.is_active

backend/onyx/utils/timing.py

@@ -1,3 +1,4 @@
+import inspect
 import time
 from collections.abc import Callable, Generator, Iterator
 from functools import wraps
@@ -40,13 +41,8 @@ def log_function_time(
     """
 
     def decorator(func: F) -> F:
-        @wraps(func)
-        def wrapped_func(*args: Any, **kwargs: Any) -> Any:
-            # Elapsed time should use monotonic.
-            start_time = time.monotonic()
-            result = func(*args, **kwargs)
-            elapsed_time = time.monotonic() - start_time
-            elapsed_time_str = f"{elapsed_time:.3f}"
+        def _log_elapsed(start_time: float, *args: Any, **kwargs: Any) -> None:
+            elapsed_time_str = f"{time.monotonic() - start_time:.3f}"
             log_name = func_name or func.__name__  # ty: ignore[unresolved-attribute]
             args_str = ""
             if include_args:
@@ -63,8 +59,6 @@ def log_function_time(
             if debug_only:
                 logger.debug(final_log)
             else:
-                # These are generally more important logs so the level is a bit
-                # higher.
                 logger.notice(final_log)
 
             if not print_only:
@@ -75,6 +69,22 @@ def log_function_time(
                     user_id=str(user.id) if user else "Unknown",
                 )
 
+        if inspect.iscoroutinefunction(func):
+
+            @wraps(func)
+            async def wrapped_async(*args: Any, **kwargs: Any) -> Any:
+                start_time = time.monotonic()
+                result = await func(*args, **kwargs)
+                _log_elapsed(start_time, *args, **kwargs)
+                return result
+
+            return cast(F, wrapped_async)
+
+        @wraps(func)
+        def wrapped_func(*args: Any, **kwargs: Any) -> Any:
+            start_time = time.monotonic()
+            result = func(*args, **kwargs)
+            _log_elapsed(start_time, *args, **kwargs)
             return result
 
         return cast(F, wrapped_func)

backend/tests/external_dependency_unit/auth/test_oauth_callback_oidc_expiry_session.py

@@ -0,0 +1,205 @@
+"""Single-tenant OIDC callback must persist `oidc_expiry` on the existing
+OAuth user without blocking on another connection's row lock, and must not
+leave a second transaction idle on `"user"`.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from datetime import datetime, timezone
+from uuid import uuid4
+
+import pytest
+from fastapi_users_db_sqlalchemy import SQLAlchemyUserDatabase
+from sqlalchemy import delete, text
+from sqlalchemy.orm import Session
+
+import onyx.auth.users as users_module
+from onyx.auth.users import UserManager
+from onyx.db.engine.async_sql_engine import get_async_session_context_manager
+from onyx.db.engine.sql_engine import get_session_with_current_tenant
+from onyx.db.models import OAuthAccount, User
+from onyx.server.security.store import _build_env_defaults
+from tests.external_dependency_unit.conftest import create_test_user, delete_test_user
+
+_CALLBACK_TIMEOUT_SECONDS = 5
+_POLL_INTERVAL_SECONDS = 0.2
+_OAUTH_NAME = "oidc"
+
+
+def _idle_user_lock_pids(exclude: set[int]) -> list[int]:
+    """Backends sitting idle while still holding a lock on `"user"`."""
+    with get_session_with_current_tenant() as session:
+        checker_pid = session.execute(text("SELECT pg_backend_pid()")).scalar()
+        exclude = exclude | {int(checker_pid)}
+        pids = session.execute(
+            text(
+                "SELECT DISTINCT a.pid "
+                "FROM pg_stat_activity a "
+                "JOIN pg_locks l ON l.pid = a.pid "
+                "JOIN pg_class c ON c.oid = l.relation "
+                "WHERE a.datname = current_database() "
+                "AND a.state = 'idle in transaction' "
+                "AND c.relname = 'user'"
+            )
+        ).scalars()
+        return [int(pid) for pid in pids if int(pid) not in exclude]
+
+
+def _assert_no_idle_user_lock(exclude: set[int]) -> None:
+    deadline = time.monotonic() + 1.0
+    leftover: list[int] = []
+    while time.monotonic() < deadline:
+        leftover = _idle_user_lock_pids(exclude)
+        if not leftover:
+            return
+        time.sleep(0.05)
+    pytest.fail(f'leftover idle-in-transaction lock on "user"; pids={leftover}')
+
+
+def _attach_oauth_account(db_session: Session, user: User) -> str:
+    account_id = f"entra-{uuid4().hex}"
+    db_session.add(
+        OAuthAccount(
+            user_id=user.id,
+            oauth_name=_OAUTH_NAME,
+            account_id=account_id,
+            account_email=user.email,
+            access_token="test-access-token",
+            refresh_token="",
+        )
+    )
+    db_session.commit()
+    return account_id
+
+
+def _delete_oauth_accounts(db_session: Session, user: User) -> None:
+    db_session.execute(
+        delete(OAuthAccount).where(OAuthAccount.__table__.c.user_id == user.id)
+    )
+
+
+async def _run_oauth_callback(
+    *,
+    user_email: str,
+    account_id: str,
+    expires_at: int | None,
+    track_external_idp_expiry: bool,
+    monkeypatch: pytest.MonkeyPatch,
+    exclude_pids: set[int],
+) -> User:
+    monkeypatch.setattr(users_module, "MULTI_TENANT", False)
+    monkeypatch.setattr(
+        users_module,
+        "get_security_settings",
+        lambda: _build_env_defaults().model_copy(
+            update={"track_external_idp_expiry": track_external_idp_expiry}
+        ),
+    )
+
+    async with get_async_session_context_manager() as injected_session:
+        injected_pid = (
+            await injected_session.execute(text("SELECT pg_backend_pid()"))
+        ).scalar()
+        assert injected_pid is not None
+        watched_pids = exclude_pids | {int(injected_pid)}
+
+        manager = UserManager(
+            SQLAlchemyUserDatabase(injected_session, User, OAuthAccount)
+        )
+        task = asyncio.create_task(
+            manager.oauth_callback(
+                oauth_name=_OAUTH_NAME,
+                access_token="rotated-access-token",
+                account_id=account_id,
+                account_email=user_email,
+                expires_at=expires_at,
+                is_verified_by_default=True,
+            )
+        )
+        deadline = time.monotonic() + _CALLBACK_TIMEOUT_SECONDS
+        idle_pids: list[int] = []
+        while not task.done():
+            if time.monotonic() >= deadline:
+                idle_pids = _idle_user_lock_pids(watched_pids)
+                task.cancel()
+                await asyncio.gather(task, return_exceptions=True)
+                pytest.fail(
+                    "oauth_callback blocked instead of writing oidc_expiry; "
+                    f"idle-in-transaction pids={idle_pids}"
+                )
+            await asyncio.sleep(_POLL_INTERVAL_SECONDS)
+            idle_pids = _idle_user_lock_pids(watched_pids)
+        return task.result()
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("tenant_context")
+async def test_existing_oauth_login_writes_oidc_expiry(
+    db_session: Session, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    user = create_test_user(db_session, "oidc_expiry_write")
+    account_id = _attach_oauth_account(db_session, user)
+    expires_at = int(datetime.now(timezone.utc).timestamp()) + 3600
+
+    sync_pid = db_session.execute(text("SELECT pg_backend_pid()")).scalar()
+    try:
+        result = await _run_oauth_callback(
+            user_email=user.email,
+            account_id=account_id,
+            expires_at=expires_at,
+            track_external_idp_expiry=True,
+            monkeypatch=monkeypatch,
+            exclude_pids={int(sync_pid)},
+        )
+        assert result.id == user.id
+
+        db_session.expire_all()
+        persisted = db_session.get(User, user.id)
+        assert persisted is not None
+        assert persisted.oidc_expiry is not None
+        assert persisted.oidc_expiry == datetime.fromtimestamp(
+            expires_at, tz=timezone.utc
+        )
+        sync_pid = db_session.execute(text("SELECT pg_backend_pid()")).scalar()
+        _assert_no_idle_user_lock({int(sync_pid)})
+    finally:
+        _delete_oauth_accounts(db_session, user)
+        delete_test_user(db_session, user)
+        db_session.commit()
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("tenant_context")
+async def test_existing_oauth_login_clears_oidc_expiry_when_tracking_disabled(
+    db_session: Session, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    user = create_test_user(db_session, "oidc_expiry_clear")
+    user.oidc_expiry = datetime.now(timezone.utc)
+    db_session.commit()
+    account_id = _attach_oauth_account(db_session, user)
+
+    sync_pid = db_session.execute(text("SELECT pg_backend_pid()")).scalar()
+    try:
+        result = await _run_oauth_callback(
+            user_email=user.email,
+            account_id=account_id,
+            expires_at=int(datetime.now(timezone.utc).timestamp()) + 3600,
+            track_external_idp_expiry=False,
+            monkeypatch=monkeypatch,
+            exclude_pids={int(sync_pid)},
+        )
+        assert result.id == user.id
+        assert result.oidc_expiry is None
+
+        db_session.expire_all()
+        persisted = db_session.get(User, user.id)
+        assert persisted is not None
+        assert persisted.oidc_expiry is None
+        sync_pid = db_session.execute(text("SELECT pg_backend_pid()")).scalar()
+        _assert_no_idle_user_lock({int(sync_pid)})
+    finally:
+        _delete_oauth_accounts(db_session, user)
+        delete_test_user(db_session, user)
+        db_session.commit()

backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_per_user_key.py

@@ -18,6 +18,7 @@ from starlette.responses import JSONResponse, Response
 
 DEFAULT_PORT = 8003
 MCP_PATH_PREFIX = "/mcp"
+USER_EMAIL_HEADER = "X-User-Email"
 
 
 # ---- pretend database --------------------------------------------------------
@@ -122,7 +123,10 @@ class RequireHeadersMiddleware(BaseHTTPMiddleware):
             return await call_next(request)
 
         for header in self.required_headers:
-            if not request.headers.get(header):
+            value = request.headers.get(header)
+            if header.lower() == USER_EMAIL_HEADER.lower():
+                print(f"{USER_EMAIL_HEADER}: {value!r}", flush=True)
+            if not value:
                 return JSONResponse(
                     {"error": f"Missing required header '{header}'"},
                     status_code=401,

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

@@ -47,12 +47,16 @@ def mock_async_session() -> MagicMock:
     session.commit = AsyncMock()
     session.rollback = AsyncMock()
     session.run_sync = AsyncMock(return_value=None)
+    session.get = AsyncMock(return_value=None)
     return session
 
 
 class _AsyncSessionContextManager:
     def __init__(self, session: MagicMock) -> None:
         self._session = session
+        # oauth_callback awaits session.get on the placeholder upgrade path.
+        if not isinstance(session.get, AsyncMock):
+            session.get = AsyncMock(return_value=None)
 
     async def __aenter__(self) -> MagicMock:
         return self._session
@@ -79,6 +83,14 @@ def _mock_user_manager_methods(user_manager: UserManager) -> None:
     user_manager.validate_password = AsyncMock()
 
 
+def _bind_sqlalchemy_user_db_cls(
+    mock_user_db_cls: MagicMock, mock_user_db: MagicMock
+) -> None:
+    # SQLAlchemyUserDatabase[User, uuid.UUID](...) hits __getitem__, then call.
+    mock_user_db_cls.return_value = mock_user_db
+    mock_user_db_cls.__getitem__.return_value = mock_user_db_cls
+
+
 class TestDisposableEmailValidation:
     """Test disposable email validation before tenant provisioning."""
 
@@ -580,8 +592,10 @@ class TestOAuthDottedGmail:
     @patch("onyx.auth.users.fetch_ee_implementation_or_noop")
     @patch("onyx.auth.users.get_async_session_context_manager")
     @patch("onyx.auth.users.remove_user_from_invited_users")
+    @patch("onyx.auth.users.SQLAlchemyUserDatabase")
     async def test_oauth_create_does_not_block_dotted_gmail(
         self,
+        mock_user_db_cls: MagicMock,
         mock_remove_invited: MagicMock,  # noqa: ARG002
         mock_session_manager: MagicMock,
         mock_fetch_ee: MagicMock,
@@ -617,6 +631,7 @@ class TestOAuthDottedGmail:
         mock_user_db.session = MagicMock()
         mock_user_db.session.run_sync = AsyncMock()
         user_manager.user_db = mock_user_db
+        _bind_sqlalchemy_user_db_cls(mock_user_db_cls, mock_user_db)
 
         await user_manager.oauth_callback(
             oauth_name="google",
@@ -661,7 +676,9 @@ class TestOAuthNoAutoLinkExemptions:
         )
 
     @staticmethod
-    def _manager_with_existing(existing_user: MagicMock) -> UserManager:
+    def _manager_with_existing(
+        existing_user: MagicMock, mock_user_db_cls: MagicMock | None = None
+    ) -> UserManager:
         user_manager = UserManager(MagicMock())
         _mock_user_manager_methods(user_manager)
         user_manager.on_after_register = AsyncMock()
@@ -675,6 +692,8 @@ class TestOAuthNoAutoLinkExemptions:
         mock_user_db.update = AsyncMock(return_value=existing_user)
         mock_user_db.session = MagicMock()
         user_manager.user_db = mock_user_db
+        if mock_user_db_cls is not None:
+            _bind_sqlalchemy_user_db_cls(mock_user_db_cls, mock_user_db)
         return user_manager
 
     @pytest.mark.asyncio
@@ -690,8 +709,10 @@ class TestOAuthNoAutoLinkExemptions:
     @patch("onyx.auth.users.assign_user_to_default_groups__no_commit")
     @patch("onyx.auth.users._upgrade_will_add_seat", return_value=False)
     @patch("onyx.auth.users.get_session_with_current_tenant")
+    @patch("onyx.auth.users.SQLAlchemyUserDatabase")
     async def test_placeholder_promoted_without_auto_link(
         self,
+        mock_user_db_cls: MagicMock,
         mock_sync_session_factory: MagicMock,
         mock_will_add_seat: MagicMock,  # noqa: ARG002
         mock_assign_groups: MagicMock,
@@ -717,7 +738,8 @@ class TestOAuthNoAutoLinkExemptions:
             account_type=AccountType.EXT_PERM_USER,
             is_active=is_active,
         )
-        user_manager = self._manager_with_existing(placeholder)
+        user_manager = self._manager_with_existing(placeholder, mock_user_db_cls)
+        mock_async_session.get = AsyncMock(return_value=placeholder)
 
         sync_user = MagicMock(is_active=is_active)
         mock_sync_db = MagicMock()
@@ -757,8 +779,10 @@ class TestOAuthNoAutoLinkExemptions:
     @patch("onyx.auth.users.fetch_ee_implementation_or_noop")
     @patch("onyx.auth.users.get_async_session_context_manager")
     @patch("onyx.auth.users.remove_user_from_invited_users")
+    @patch("onyx.auth.users.SQLAlchemyUserDatabase")
     async def test_unclaimed_row_is_claimed_by_first_login(
         self,
+        mock_user_db_cls: MagicMock,
         mock_remove_invited: MagicMock,  # noqa: ARG002
         mock_session_manager: MagicMock,
         mock_fetch_ee: MagicMock,
@@ -774,7 +798,7 @@ class TestOAuthNoAutoLinkExemptions:
         mock_fetch_ee.return_value = AsyncMock(return_value="test_tenant")
 
         provisioned = self._unclaimed()
-        user_manager = self._manager_with_existing(provisioned)
+        user_manager = self._manager_with_existing(provisioned, mock_user_db_cls)
 
         result = await user_manager.oauth_callback(
             oauth_name="okta",
@@ -804,8 +828,10 @@ class TestOAuthNoAutoLinkExemptions:
     @patch("onyx.auth.users.fetch_ee_implementation_or_noop")
     @patch("onyx.auth.users.get_async_session_context_manager")
     @patch("onyx.auth.users.remove_user_from_invited_users")
+    @patch("onyx.auth.users.SQLAlchemyUserDatabase")
     async def test_spoken_for_row_is_rejected(
         self,
+        mock_user_db_cls: MagicMock,
         mock_remove_invited: MagicMock,  # noqa: ARG002
         mock_session_manager: MagicMock,
         mock_fetch_ee: MagicMock,
@@ -819,7 +845,9 @@ class TestOAuthNoAutoLinkExemptions:
         )
         mock_fetch_ee.return_value = AsyncMock(return_value="test_tenant")
 
-        user_manager = self._manager_with_existing(self._unclaimed(**overrides))
+        user_manager = self._manager_with_existing(
+            self._unclaimed(**overrides), mock_user_db_cls
+        )
 
         with pytest.raises(exceptions.UserAlreadyExists):
             await user_manager.oauth_callback(
@@ -839,8 +867,10 @@ class TestOAuthNoAutoLinkExemptions:
     @patch("onyx.auth.users.fetch_ee_implementation_or_noop")
     @patch("onyx.auth.users.get_async_session_context_manager")
     @patch("onyx.auth.users.remove_user_from_invited_users")
+    @patch("onyx.auth.users.SQLAlchemyUserDatabase")
     async def test_deactivated_row_is_not_linked(
         self,
+        mock_user_db_cls: MagicMock,
         mock_remove_invited: MagicMock,  # noqa: ARG002
         mock_session_manager: MagicMock,
         mock_fetch_ee: MagicMock,
@@ -856,7 +886,7 @@ class TestOAuthNoAutoLinkExemptions:
         mock_fetch_ee.return_value = AsyncMock(return_value="test_tenant")
 
         deactivated = self._unclaimed(is_active=False)
-        user_manager = self._manager_with_existing(deactivated)
+        user_manager = self._manager_with_existing(deactivated, mock_user_db_cls)
 
         result = await user_manager.oauth_callback(
             oauth_name="okta",