Redirect user back to KO Review page if that is where they came from

atst/routes/task_orders/new.py

@@ -207,6 +207,33 @@ def new(screen, task_order_id=None, portfolio_id=None):
         if not TaskOrders.all_sections_complete(task_order):
             flash("task_order_draft")
 
+    if workflow.task_order:
+        ko_review_url = url_for(
+            "portfolios.ko_review",
+            portfolio_id=workflow.task_order.portfolio.id,
+            task_order_id=task_order_id,
+            _external=True,
+        )
+        redirect_url = url_for(
+            "portfolios.ko_review",
+            portfolio_id=workflow.task_order.portfolio.id,
+            task_order_id=task_order_id,
+        )
+
+        if http_request.referrer == ko_review_url:
+            return render_template(
+                workflow.template,
+                current=screen,
+                task_order_id=task_order_id,
+                task_order=workflow.task_order,
+                portfolio_id=portfolio_id,
+                screens=workflow.display_screens,
+                form=workflow.form,
+                complete=workflow.is_complete,
+                from_ko_review=True,
+                next=redirect_url,
+            )
+
     return render_template(
         workflow.template,
         current=screen,
@@ -229,8 +256,12 @@ def update(screen, task_order_id=None, portfolio_id=None):
     workflow = UpdateTaskOrderWorkflow(
         g.current_user, form_data, screen, task_order_id, portfolio_id
     )
+
     if workflow.validate():
         workflow.update()
+        if http_request.args.get("next"):
+            return redirect(http_request.args.get("next"))
+        else:
             return redirect(
                 url_for(
                     "task_orders.new",

templates/task_orders/_new.html

@@ -9,6 +9,10 @@
   {% include "fragments/flash.html" %}
 
   {% block form_action %}
+    {% if from_ko_review %}
+      <form method='POST' action="{{ url_for('task_orders.update', screen=current, task_order_id=task_order_id, next=next) }}" autocomplete="off" enctype="multipart/form-data">
+    {% endif %}
+
     {% if task_order_id %}
       <form method='POST' action="{{ url_for('task_orders.new', screen=current, task_order_id=task_order_id) }}" autocomplete="off" enctype="multipart/form-data">
     {% else %}
@@ -39,7 +43,11 @@
     {% block next %}
 
       <div class='action-group'>
+        {% if from_ko_review %}
+          <input type='submit' class='usa-button usa-button-primary' value='Save' />
+        {% else %}
           <input type='submit' class='usa-button usa-button-primary' value='Save & Continue' />
+        {% endif %}
       </div>
 
     {% endblock %}