I finally found a Rule script that worked:
function (user, context, callback) {
const namespace = 'http://schemas.microsoft.com/ws/2008/06/identity/claims';
const assignedRoles = (context.authorization || {}).roles;
let idTokenClaims = context.idToken || {};
let accessTokenClaims = context.accessToken || {};
idTokenClaims[`${namespace}/role`] = assignedRoles;
accessTokenClaims[`${namespace}/role`] = assignedRoles;
context.idToken = idTokenClaims;
context.accessToken = accessTokenClaims;
callback(null, user, context);
}