work with menu link edit as well as add

menu_save_add_another.info.yml

@@ -4,3 +4,5 @@ description: Make menu item stubbing easier. Fix the menu add item redirect and
 package: Menu
 core: 8.x
 core_version_requirement: ^8 || ^9
+dependencies:
+  - drupal:menu_ui

menu_save_add_another.module

@@ -7,6 +7,7 @@
 
 use Drupal\Core\Form\FormStateInterface;
 use Drupal\Core\Url;
+use Drupal\Core\Render\Element;
 
 /**
  * Implements hook_menu_local_actions_alter().
@@ -21,16 +22,26 @@ function menu_save_add_another_menu_local_actions_alter(array &$local_actions) {
 /**
  * Implements hook_form_alter().
  */
-function menu_save_add_another_form_menu_link_content_form_alter(array &$form, FormStateInterface $form_state, $form_id) {
-  // Set redirect on the original save button (since we've removed the ?destination).
-  $form['actions']['submit']['#submit'][] = '_menu_save_add_another_redirect_menu';
+function menu_save_add_another_form_menu_edit_form_alter(array &$form, FormStateInterface $form_state, $form_id) {
+  // Unset the edit link redirects.
+  foreach (Element::children($form['links']['links']) as $link) {
+    $form['links']['links'][$link]['operations']['#links']['edit']['query'] = [];
+  }
+}
 
+/**
+ * Implements hook_form_alter().
+ */
+function menu_save_add_another_form_menu_link_content_form_alter(array &$form, FormStateInterface $form_state, $form_id) {
   // Add another save button.
   $form['actions']['submit_and_add'] = $form['actions']['submit'];
   $form['actions']['submit_and_add']['#value'] = t('Save and Add Another');
   $form['actions']['submit_and_add']['#weight'] = 6;
   $form['actions']['submit_and_add']['#button_type'] = '';
   $form['actions']['submit_and_add']['#submit'][] = '_menu_save_add_another_redirect_add';
+
+  // Set redirect on the original save button (since we've removed the ?destination).
+  $form['actions']['submit']['#submit'][] = '_menu_save_add_another_redirect_menu';
 }
 
 /**
@@ -49,7 +60,16 @@ function _menu_save_add_another_redirect_add(array $form, FormStateInterface $fo
 
 function _menu_save_add_another_set_redirect(FormStateInterface $form_state, $route_name) {
   $route_match = \Drupal::routeMatch();
-  // Route params are the same for the menu and for the menu item.
-  $options = $route_match->getRawParameters()->all();
-  $form_state->setRedirect($route_name, $options);
+  // Route params might either include a menu (for new items) or the menu_link_content (for existing).
+  $menu = $route_match->getRawParameter('menu');
+  if (!$menu) {
+    $menu_link = $route_match->getParameter('menu_link_content');
+    if ($menu_link) {
+      // Look up the menu for the link.
+      $menu = $menu_link->getMenuName();
+    }
+  }
+  if ($menu) {
+    $form_state->setRedirect($route_name, ['menu' => $menu]);
+  }
 }
\ No newline at end of file