interaction between form element plugins and config entity plugin collections

I'm making a form element plugin which has compound elements, but which should return a single string value. This is the same idea as core's password confirm form element (https://api.drupal.org/api/drupal/core%21lib%21Drupal%21Core%21Render%21Element%21PasswordConfirm.php/10).

I've imitated what password confirm does: an #element_validate callback gets the value from the inner elements, and sets it as the value for the element as a whole:

    // Password field must be converted from a two-element array into a single
    // string regardless of validation results.
    $form_state->setValueForElement($element['pass1'], NULL);
    $form_state->setValueForElement($element['pass2'], NULL);
    $form_state->setValueForElement($element, $pass1);

This works fine in a plain form.

However, when I then try to use this form element in a config entity's form to select a plugin, it crashes on submit.

This is because when the form is rebuilt on submission, EntityForm::copyFormValuesToEntity() is called and calls $entity->getPluginCollections(), and at that point, the value of the main element is an array, which causes a crash in the plugin collection system because it tries to use that as an array key.

The value is an array because after main element is processed and the #element_validate callback sets sets its value to a string, the nested element is processed, and sets its value, and that causes NestedArray to create an array in the form values:

Step 1, form values are:

[
  main_element => value,
]

Step 2, inner element's value sets set on form state values and we get:

[
  main_element => [
    inner_element => value,
  ],
]

I don't understand why during rebuild we get this effect, but in the final submission, the value of the element is a string as expected.

Am I doing something wrong, or is this just not going to work?

Snippet of code for the form element:

  /**
   * Process callback.
   */
  public static function processPlugin(&$element, FormStateInterface $form_state, &$complete_form) {
    $element['#tree'] = TRUE;

    // This needs to be a nested element so the radio or select element
    // processing and theming takes place.
    $element['plugin_id'] = [
      '#type' => $element['#options_element_type'],
      '#title' => $element['#title'],
      '#options' => $options, // SNIP defined elsewhere
      '#empty_value' => '',
      '#required' => $element['#required'],
      '#default_value' => $element['#default_value'] ?? '',
    ];

    $element['#element_validate'] = [[static::class, 'validatePlugin']];

    return $element;
  }

  /**
   * {@inheritdoc}
   */
  public static function valueCallback(&$element, $input, FormStateInterface $form_state) {
    if ($input === FALSE) {
      return $element['#default_value'] ?? '';
    }
    elseif (is_null($input)) {
      return $element['#default_value'] ?? '';
    }
    else {
      return $input['plugin_id'];
    }
  }

  public static function validatePlugin(&$element, FormStateInterface $form_state, &$complete_form) {
    $plugin_id = $element['plugin_id']['#value'];

    // Set the value at the top level of the element.
    $form_state->setValueForElement($element['plugin_id'], NULL);
    $form_state->setValueForElement($element, $plugin_id);

    return $element;
  }

(Full thing is in this MR https://git.drupalcode.org/project/plugin/-/merge_requests/3#note_144104 for this issue https://www.drupal.org/project/plugin/issues/3197304.)

Snippet of code for the form that uses it:

    $form['link_style'] = [
      '#type' => 'action_link_plugin',
      '#title' => $this->t('Link style'),
      '#required' => TRUE,
      '#default_value' => $action_link->get('link_style'),
      '#plugin_type' => 'action_link.link_style',
      '#options_element_type' => 'radios',
    ];

And in the config entity class:

public function getPluginCollections() {
  $collections = [];
  if ($this->getActionLinkPluginCollection()) {
    $collections['plugin_config'] = $this->getActionLinkPluginCollection();
  }
  if ($this->getLinkStylePluginCollection()) {
    $collections['link_style_collection'] = $this->getLinkStylePluginCollection();
  }
  return $collections;
}

protected function getActionLinkPluginCollection() {
  if (!$this->actionLinkPluginCollection && $this->plugin_id) {
    $this->actionLinkPluginCollection = new DefaultSingleLazyPluginCollection(
      \Drupal::service('plugin.manager.action_link_state_action'),
      $this->plugin_id, $this->plugin_config
    );
  }
  return $this->actionLinkPluginCollection;
}

protected function getLinkStylePluginCollection() {
  // Horrible workaround for the form element's inner element's value getting
  // set and then the resulting value *array* for the outer element being used
  // by copyFormValuesToEntity().
  if (is_array($this->link_style)) {
    return NULL;
  }
  if (!$this->linkStylePluginCollection && $this->link_style) {
    $this->linkStylePluginCollection = new DefaultSingleLazyPluginCollection(
      \Drupal::service('plugin.manager.action_link_style'),
      $this->link_style,
      []
    );
  }
  return $this->linkStylePluginCollection;
}