Variable tax rate on wc checkout depending on a legal person selection field only works for the first purchase

Question

Found this LoicTheAztec code and works great but only for first time customers. If a returning customer tries to make a purchase (logged in), the tax rate applied is the same of the first purchase, regardless the legal person option chosen in the subsequent purchases (the checkout review appears to be updating but gets the same result of the first purchase).

add_filter('woocommerce_billing_fields', 'add_legal_person_billing_field');
function add_legal_person_billing_field( $fields ) {
    $fields['billing_company']['class'][] = 'hidden'; 

    $fields['billing_legal_person'] = array(
        'label'       => __( 'Legal person', 'woocommerce' ),
        'type'        => 'select',
        'options'     => [
            ''        => __("Choose an option",  "woocommerce"),
            'Person'  => __("Person",  "woocommerce"),
            'Company' => __("Company", "woocommerce"),
        ],
        'required'    => true,
        'priority'    => 25,
    );
    return $fields;
}

// Validate Billing company field
add_action( 'woocommerce_checkout_process', 'validate_custom_checkout_fields' );
function validate_custom_checkout_fields() {
    if ( isset($_POST['billing_legal_person']) && $_POST['billing_legal_person'] === 'Company' 
    && isset($_POST['billing_company']) && empty($_POST['billing_company']) ) {
        wc_add_notice('Billing company is a required field.', 'error');
    }
}

// PHP: Replace "(optional)" with required for Billing company
add_filter( 'woocommerce_form_field' , 'remove_checkout_optional_fields_label', 10, 4 );
function remove_checkout_optional_fields_label( $field, $key, $args, $value ) {
    // Only on checkout page
    if( ( ( is_checkout() && ! is_wc_endpoint_url() ) 
    || is_wc_endpoint_url('edit-address') ) && $key === 'billing_company' ) {
        $optional = '&nbsp;<span class="optional">(' . esc_html__( 'optional', 'woocommerce' ) . ')</span>';
        $required = '&nbsp;<abbr class="required" title="required">*</abbr>';
        $field = str_replace( $optional, $required, $field );
    }
    return $field;
}

// Hide Billing company field on checkout page load 
add_action('wp_head', 'checkout_legal_person_css');
function checkout_legal_person_css() {
    if ( ( is_checkout() && ! is_wc_endpoint_url() ) || is_wc_endpoint_url('edit-address') ) : ?>
    <style>
        #billing_company_field.hidden {display:none !important;}
    </style>
    <?php endif;
}

// JavaScript / Ajax: Show/hide billing company and set the tax rate via Ajax
add_action( 'template_redirect' , 'checkout_legal_person_js' );
function checkout_legal_person_js() {
    if ( ( is_checkout() && ! is_wc_endpoint_url() ) || is_wc_endpoint_url('edit-address') ) :

    $legal_person = WC()->customer->get_meta('billing_legal_person');
    if ( is_checkout() && $legal_person ) {
        WC()->session->set('legal_person', $legal_person);
    }
    
    wc_enqueue_js("
        const company         = '#billing_company',       
              companyField    = company+'_field', 
              legalPers       = '#billing_legal_person',
              isCheckout      = ".( is_checkout() ? 'true' : 'false' ).";

        function showHideCompanyField( legalPers, companyField ) {
            if ( $(legalPers).val() === 'Company' ) {
                $(companyField).show();
            } else {
                $(companyField).hide();
            }
        }

        // On start:
        $(legalPers).selectWoo();
        $(companyField).addClass('validate-required woocommerce-invalid woocommerce-invalid-required-field');
        showHideCompanyField( legalPers, companyField );
        $(companyField).removeClass('hidden');

        // On Legal Person change: Show/Hide company field
        $('form.woocommerce-checkout,.woocommerce-edit-address form').on('change', legalPers, function() {
            showHideCompanyField( legalPers, companyField );
            
            if ( isCheckout ) {
                $.ajax({
                    type: 'POST',
                    url: '".admin_url('admin-ajax.php')."',
                    data: {
                        'action': 'legal_person_tax',
                        'legal_person' : $(legalPers).val(), 
                    },
                    success: function (response) {
                        $(document.body).trigger('update_checkout');
                        console.log(response);
                    }
                });
            }
        });
    ");
    endif;
}

// PHP AJAX receiver: Set "legal_person" session variable
add_action('wp_ajax_legal_person_tax', 'set_legal_person_tax' );
add_action('wp_ajax_nopriv_legal_person_tax', 'set_legal_person_tax' );
function set_legal_person_tax() {
    if ( isset($_POST['legal_person']) && ! empty($_POST['legal_person']) ) {
        WC()->session->set('legal_person', esc_attr($_POST['legal_person']));
    }
    wp_die(esc_attr($_POST['legal_person']));
}

// Dynamically change tax rate based on chosen "legal_person"
add_action('woocommerce_before_calculate_totals', 'define_cart_item_tax_class_for_companies');
function define_cart_item_tax_class_for_companies( $cart )
{
    if ((is_admin() && !defined('DOING_AJAX')))
        return;

    if (did_action('woocommerce_before_calculate_totals') >= 2)
        return;

    $legal_person = WC()->session->get('legal_person');
    $legal_person = $legal_person ? $legal_person : WC()->customer->get_meta('billing_legal_person');

    foreach ($cart->get_cart() as $item) {
        if ( $legal_person === 'Company' ) {
            $item['data']->set_tax_class('companies'); // Set tax class
        } else {
            $item['data']->set_tax_class('');
        }
    }
}

// Remove session variable when order is placed
add_action('woocommerce_checkout_order_created', 'unset_legal_person_session_variable');
function unset_legal_person_session_variable() {
    WC()->session->__unset('legal_person');
}

// Change tax rate if customer is a a company
add_filter( 'woocommerce_product_get_tax_class', 'tax_rate_based_on_chosen_legal_person', 10, 2 );
add_filter( 'woocommerce_product_variation_get_tax_class', 'tax_rate_based_on_chosen_legal_person', 10, 2 );
function tax_rate_based_on_chosen_legal_person( $tax_class, $product ) {
    if ( is_admin() ) return $tax_class;
    
    $legal_person = WC()->customer->get_meta('billing_legal_person');

    if ( $legal_person === 'Company' ) {
        $tax_class = 'companies';
    }
    return $tax_class;
}

I assume the code adds a legal_person label to the user in the first purchase and overrides the subsequent different legal_person selections of the same user, but I can't find a way to manage the label directly from the user data or a way to avoid to link the legal_person seleccion to the user and only use that legal_person selection on the checkout for the tax rate selection.