How to retrieve 'login theme' from a client for each email sent?

Hi everyone,

I’m currently working on customizing the email templates in Keycloak to use client-specific themes. I need to retrieve the ‘login theme’ attribute from a client every time that client sends an email. Here’s my current implementation for Keycloak 24.0.2 Version:

package com.custom.keycloak;

import org.jboss.logging.Logger;
import org.keycloak.email.freemarker.FreeMarkerEmailTemplateProvider;
import org.keycloak.models.ClientModel;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.KeycloakContext;
import org.keycloak.sessions.AuthenticationSessionModel;
import org.keycloak.theme.Theme;

import java.io.IOException;

public class ClientThemeEmailTemplateProvider extends FreeMarkerEmailTemplateProvider {

    private static final Logger LOG = Logger.getLogger(ClientThemeEmailTemplateProvider.class);

    public ClientThemeEmailTemplateProvider(KeycloakSession session) {
        super(session);
    }

    @Override
    protected Theme getTheme() throws IOException {
        ClientModel client = getClientFromContextOrSession();
        if (client == null) {
            LOG.warn("No client context available. Using default theme.");
            return super.getTheme();
        }

        String loginTheme = client.getAttribute("loginTheme");
        if (loginTheme != null) {
            LOG.infof("Using theme for client: %s", loginTheme);
            return session.theme().getTheme(loginTheme, Theme.Type.EMAIL);
        } else {
            LOG.info("No login theme attribute found for client. Using default theme.");
            return super.getTheme();
        }
    }

    private ClientModel getClientFromContextOrSession() {
        KeycloakContext context = session.getContext();
        AuthenticationSessionModel authSession = context.getAuthenticationSession();

        if (authSession != null && authSession.getClient() != null) {
            return authSession.getClient();
        }
        return context.getClient();
    }
}

Despite this implementation, I am not able to retrieve any client information from KeycloakSession or AuthenticationSessionModel, and I am currently stuck.

Is there a more efficient or recommended way to achieve this in Keycloak, or am I missing something in my approach?

Thanks in advance for your help!