diff --git a/src/Form/IndexForm.php b/src/Form/IndexForm.php
index 732f63b..c4509ee 100644
--- a/src/Form/IndexForm.php
+++ b/src/Form/IndexForm.php
@@ -552,9 +552,9 @@ public function actions(array $form, FormStateInterface $form_state) {
         '#value' => $this->t('Save and add fields'),
         '#submit' => $actions['submit']['#submit'],
         '#button_type' => 'primary',
-        // Work around for submit callbacks after save()
-        // not being called due to batch operations.
-        '#redirect_to_fields' => TRUE,
+        // Work around for submit callbacks after save() not being called due to
+        // batch operations.
+        '#redirect_to_url' => $this->getEntity()->toUrl('add-fields'),
       );
     }
 
@@ -610,8 +610,8 @@ public function save(array $form, FormStateInterface $form_state) {
         $index->save();
         drupal_set_message($this->t('The index was successfully saved.'));
         $button = $form_state->getTriggeringElement();
-        if (isset($button['#redirect_to_fields'])) {
-          $form_state->setRedirectUrl($index->toUrl('add-fields'));
+        if (!empty($button['#redirect_to_url'])) {
+          $form_state->setRedirectUrl($button['#redirect_to_url']);
         }
         else {
           $form_state->setRedirect('entity.search_api_index.canonical', array('search_api_index' => $index->id()));
diff --git a/src/Tests/IntegrationTest.php b/src/Tests/IntegrationTest.php
index bd872a0..cef9daa 100644
--- a/src/Tests/IntegrationTest.php
+++ b/src/Tests/IntegrationTest.php
@@ -302,7 +302,7 @@ protected function createIndex() {
       throw new SearchApiException();
     }
 
-    // Test the "Save and edit" button.
+    // Test the "Save and add fields" button.
     $index2_id = 'test_index2';
     $edit['id'] = $index2_id;
     unset($edit['server']);