diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 8dcc49c2fd8..3817935bf71 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -198,7 +198,9 @@ async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | No return loaded if isinstance(loaded, str) else None -async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResolutionFailure": +async def load_active_user_by_id( + user_id: str, *, sso_user_id: str | None = None, user_email: str | None = None +) -> "LiteLLM_UserTable | _KeyResolutionFailure": """Load a live litellm user by id, returning the record when the user is active or a precise failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a @@ -232,6 +234,8 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, + sso_user_id=sso_user_id, + user_email=user_email, ) except (ProxyException, HTTPException): return "no_active_key" @@ -352,7 +356,7 @@ async def _extract_jwt_user_id(token: str) -> str | None: return None if await _key_owner_scim_deactivated(mapped) else _active_key_user_id(mapped) if mapped is not None: return None - user_id, _, valid_email = await JWTAuthManager.get_user_info(jwt_handler, claims) + user_id, user_email, valid_email = await JWTAuthManager.get_user_info(jwt_handler, claims) object_id: Final = jwt_handler.get_object_id(token=claims, default_value=None) owner_id: Final = ( object_id @@ -361,7 +365,7 @@ async def _extract_jwt_user_id(token: str) -> str | None: ) if not owner_id or valid_email is False: return None - owner: Final = await load_active_user_by_id(owner_id) + owner: Final = await load_active_user_by_id(owner_id, sso_user_id=owner_id, user_email=user_email) return None if isinstance(owner, str) else owner.user_id except Exception as exc: # noqa: BLE001 # public OAuth exchange stays available; unvalidated identities never write credentials verbose_logger.debug("OAuth JWT identity could not be validated (%s)", type(exc).__name__) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index af4c9eea770..ce985064eb0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11428,6 +11428,7 @@ def _oauth_identity_jwt( { "sub": "not-the-configured-user-id", "identity": {"user_id": owner}, + "email": "owner@example.test", "iss": issuer, "aud": audience, "exp": int(time.time()) + expires_in, @@ -11649,3 +11650,38 @@ async def test_oauth_jwt_uses_rbac_user_object_id(jwt_oauth_identity: tuple["JWT ] request: Final = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}"}) assert await _extract_user_id_from_request(request) == "jwt-owner" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", ["sso", "email"]) +@pytest.mark.parametrize("inactive", [False, True]) +async def test_oauth_jwt_resolves_canonical_owner_without_cached_identity( + jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"], + monkeypatch: pytest.MonkeyPatch, + identity: str, + inactive: bool, +) -> None: + from litellm.models.user import LiteLLM_UserTable + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _extract_user_id_from_request + + handler, signing_key = jwt_oauth_identity + external_id: Final = f"external-{identity}-{inactive}" + handler.litellm_jwtauth.user_email_jwt_field = "email" + owner: Final = LiteLLM_UserTable( + user_id="canonical-oauth-owner", + user_email="owner@example.test", + metadata={"scim_active": not inactive}, + organization_memberships=[], + ) + database: Final = MagicMock() + table: Final = database.db.litellm_usertable + table.find_unique = AsyncMock(side_effect=[None, owner if identity == "sso" else None]) + table.find_first = AsyncMock(return_value=owner) + table.update = AsyncMock(return_value=owner) + monkeypatch.setattr(proxy_server, "prisma_client", database) + request: Final = _token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key, owner=external_id)}"}) + assert await _extract_user_id_from_request(request) == (None if inactive else "canonical-oauth-owner") + assert table.find_unique.await_count == 2 + if identity == "email": + table.find_first.assert_awaited_once()