diff --git a/core/modules/workflows/src/Form/DialogFormTrait.php b/core/modules/workflows/src/Form/DialogFormTrait.php
index 30a848fa6a..ee64932b9c 100644
--- a/core/modules/workflows/src/Form/DialogFormTrait.php
+++ b/core/modules/workflows/src/Form/DialogFormTrait.php
@@ -1,14 +1,13 @@
 <?php
 
-namespace Drupal\Core\Form;
+namespace Drupal\workflows\Form;
 
-use Drupal\Component\Utility\UrlHelper;
 use Drupal\Core\Ajax\RedirectCommand;
 use Drupal\Core\Ajax\ReplaceCommand;
 use Drupal\Core\EventSubscriber\MainContentViewSubscriber;
 use Drupal\Core\Ajax\AjaxResponse;
 use Drupal\Core\Ajax\CloseDialogCommand;
-use Drupal\Core\Url;
+use Drupal\Core\Form\FormStateInterface;
 
 /**
  * Provides utilities for forms that want to be rendered in dialogs.
@@ -120,7 +119,9 @@ public function submitFormDialog(array &$form, FormStateInterface $form_state) {
       $command = new ReplaceCommand('#dialog-form', $form);
     }
     else {
-      $command = new RedirectCommand($this->getEntity()->toUrl('edit-form')->toString());
+      /** @var \Drupal\workflows\WorkflowInterface $workflow */
+      $workflow = $this->workflow ?: $this->getEntity();
+      $command = new RedirectCommand($workflow->toUrl('edit-form')->toString());
     }
     return $response->addCommand($command);
   }
diff --git a/core/modules/workflows/src/Form/WorkflowStateAddForm.php b/core/modules/workflows/src/Form/WorkflowStateAddForm.php
index 21afc28db4..c1cc580088 100644
--- a/core/modules/workflows/src/Form/WorkflowStateAddForm.php
+++ b/core/modules/workflows/src/Form/WorkflowStateAddForm.php
@@ -4,14 +4,11 @@
 
 use Drupal\Core\Entity\EntityInterface;
 use Drupal\Core\Form\FormStateInterface;
-use Drupal\Core\Form\DialogFormTrait;
 
 /**
  * Class WorkflowStateAddForm.
  */
-class WorkflowStateAddForm extends WorkflowEntityForm {
-
-  use DialogFormTrait;
+class WorkflowStateAddForm extends WorkflowEntityFormBase {
 
   /**
    * {@inheritdoc}
diff --git a/core/modules/workflows/src/Form/WorkflowStateDeleteForm.php b/core/modules/workflows/src/Form/WorkflowStateDeleteForm.php
index 4461a35939..71d2785b63 100644
--- a/core/modules/workflows/src/Form/WorkflowStateDeleteForm.php
+++ b/core/modules/workflows/src/Form/WorkflowStateDeleteForm.php
@@ -2,7 +2,6 @@
 
 namespace Drupal\workflows\Form;
 
-use Drupal\Core\Form\DialogFormTrait;
 use Drupal\workflows\WorkflowInterface;
 use Drupal\Core\Form\ConfirmFormBase;
 use Drupal\Core\Form\FormStateInterface;
diff --git a/core/modules/workflows/src/Form/WorkflowStateEditForm.php b/core/modules/workflows/src/Form/WorkflowStateEditForm.php
index aaf0c11393..bd640f7538 100644
--- a/core/modules/workflows/src/Form/WorkflowStateEditForm.php
+++ b/core/modules/workflows/src/Form/WorkflowStateEditForm.php
@@ -4,13 +4,11 @@
 
 use Drupal\Core\Entity\EntityInterface;
 use Drupal\Core\Form\FormStateInterface;
-use Drupal\Core\Form\DialogFormTrait;
+
 /**
  * Class WorkflowStateEditForm.
  */
-class WorkflowStateEditForm extends WorkflowEntityForm {
-
-  use DialogFormTrait;
+class WorkflowStateEditForm extends WorkflowEntityFormBase {
 
   /**
    * The ID of the state that is being edited.
diff --git a/core/modules/workflows/src/Form/WorkflowTransitionAddForm.php b/core/modules/workflows/src/Form/WorkflowTransitionAddForm.php
index 4363388c21..f2173f2873 100644
--- a/core/modules/workflows/src/Form/WorkflowTransitionAddForm.php
+++ b/core/modules/workflows/src/Form/WorkflowTransitionAddForm.php
@@ -5,14 +5,11 @@
 use Drupal\Core\Entity\EntityInterface;
 use Drupal\Core\Form\FormStateInterface;
 use Drupal\workflows\State;
-use Drupal\Core\Form\DialogFormTrait;
 
 /**
  * Class WorkflowTransitionAddForm.
  */
-class WorkflowTransitionAddForm extends WorkflowEntityForm {
-
-  use DialogFormTrait;
+class WorkflowTransitionAddForm extends WorkflowEntityFormBase {
 
   /**
    * {@inheritdoc}
diff --git a/core/modules/workflows/src/Form/WorkflowTransitionDeleteForm.php b/core/modules/workflows/src/Form/WorkflowTransitionDeleteForm.php
index 199bf4c8c8..7b3763cb5b 100644
--- a/core/modules/workflows/src/Form/WorkflowTransitionDeleteForm.php
+++ b/core/modules/workflows/src/Form/WorkflowTransitionDeleteForm.php
@@ -2,7 +2,6 @@
 
 namespace Drupal\workflows\Form;
 
-use Drupal\Core\Form\DialogFormTrait;
 use Drupal\workflows\WorkflowInterface;
 use Drupal\Core\Form\ConfirmFormBase;
 use Drupal\Core\Form\FormStateInterface;
diff --git a/core/modules/workflows/src/Form/WorkflowTransitionEditForm.php b/core/modules/workflows/src/Form/WorkflowTransitionEditForm.php
index a2704a33a9..ec59533b9f 100644
--- a/core/modules/workflows/src/Form/WorkflowTransitionEditForm.php
+++ b/core/modules/workflows/src/Form/WorkflowTransitionEditForm.php
@@ -4,15 +4,12 @@
 
 use Drupal\Core\Entity\EntityInterface;
 use Drupal\Core\Form\FormStateInterface;
-use Drupal\Core\Form\DialogFormTrait;
 use Drupal\workflows\State;
 
 /**
  * Class WorkflowTransitionEditForm.
  */
-class WorkflowTransitionEditForm extends WorkflowEntityForm {
-
-  use DialogFormTrait;
+class WorkflowTransitionEditForm extends WorkflowEntityFormBase {
 
   /**
    * The ID of the transition that is being edited.