Cannot remove media when editing Layout Builder form

I'm using it with Layout Builder but I cannot remove media I can only change it.

It seems that function valueCallback not working anymore because $form_state->getTriggeringElement() is empty.

// Remove the default value propery in case of AJAX removal.
if ($form_state->isSubmitted() && end($form_state->getTriggeringElement()['#parents']) === 'remove_button') {
  $element['#default_value'] = NULL;
}