diff --git a/core/modules/media_library/src/Form/AddFormBase.php b/core/modules/media_library/src/Form/AddFormBase.php
index 482f1ec684..a9356ffbc4 100644
--- a/core/modules/media_library/src/Form/AddFormBase.php
+++ b/core/modules/media_library/src/Form/AddFormBase.php
@@ -422,7 +422,7 @@ public function updateFormCallback(array &$form, FormStateInterface $form_state)
       // shift focus back to the first tabbable element (which should be the
       // source field).
       if (empty($added_media)) {
-        $response->addCommand(new ReplaceCommand('#media-library-add-form-wrapper', $this->buildMediaLibraryUi()));
+        $response->addCommand(new ReplaceCommand('#media-library-add-form-wrapper', $this->buildMediaLibraryUi($form_state)));
         $response->addCommand(new InvokeCommand('#media-library-add-form-wrapper :tabbable', 'focus'));
       }
       // When there are still more items, the AJAX system automatically sets
@@ -505,17 +505,20 @@ public function updateWidget(array &$form, FormStateInterface $form_state) {
 
     $response = new AjaxResponse();
     $response->addCommand(new UpdateSelectionCommand($media_ids));
-    $response->addCommand(new ReplaceCommand('#media-library-add-form-wrapper', $this->buildMediaLibraryUi()));
+    $response->addCommand(new ReplaceCommand('#media-library-add-form-wrapper', $this->buildMediaLibraryUi($form_state)));
     return $response;
   }
 
   /**
    * Build the render array of the media library UI.
    *
+   * @param \Drupal\Core\Form\FormStateInterface $form_state
+   *   The current form state.
+   *
    * @return array
    *   The render array for the media library.
    */
-  protected function buildMediaLibraryUi() {
+  protected function buildMediaLibraryUi(FormStateInterface $form_state) {
     // Get the render array for the media library. The media library state might
     // contain the 'media_library_content' when it has been opened from a
     // vertical tab. We need to remove that to make sure the render array
diff --git a/core/modules/media_library/tests/src/FunctionalJavascript/MediaLibraryTest.php b/core/modules/media_library/tests/src/FunctionalJavascript/MediaLibraryTest.php
index 5c1326a439..f3f62dd06e 100644
--- a/core/modules/media_library/tests/src/FunctionalJavascript/MediaLibraryTest.php
+++ b/core/modules/media_library/tests/src/FunctionalJavascript/MediaLibraryTest.php
@@ -820,7 +820,7 @@ public function testWidgetUpload() {
     $assert_session->pageTextContains($jpg_image->filename);
 
     // Assert removing an uploaded media item before save works as expected.
-    $assert_session->elementExists('css', '.media-library-open-button[href*="field_unlimited_media"]')->click();
+    $assert_session->elementExists('css', '.media-library-open-button[name^="field_unlimited_media"]')->click();
     $assert_session->assertWaitOnAjaxRequest();
     $assert_session->pageTextContains('Add or select media');
     $page->clickLink('Type Three');
@@ -946,7 +946,7 @@ public function testWidgetOEmbed() {
     $assert_session->pageTextContains('Custom video title');
 
     // Assert removing an uploaded media item before save works as expected.
-    $assert_session->elementExists('css', '.media-library-open-button[href*="field_unlimited_media"]')->click();
+    $assert_session->elementExists('css', '.media-library-open-button[name^="field_unlimited_media"]')->click();
     $assert_session->assertWaitOnAjaxRequest();
     $assert_session->pageTextContains('Add or select media');
     $page->clickLink('Type Five');