diff --git a/lib/galaxy/workflow/modules.py b/lib/galaxy/workflow/modules.py
index d0428ae0ee7..3ea170baeee 100644
--- a/lib/galaxy/workflow/modules.py
+++ b/lib/galaxy/workflow/modules.py
@@ -1029,13 +1029,6 @@ class InputModule(WorkflowModule):
                     step_outputs["input_ds_copy"] = new_hdca
                 else:
                     raise Exception("Unknown history content encountered")
-        # If coming from UI - we haven't registered invocation inputs yet,
-        # so do that now so dependent steps can be recalculated. In the future
-        # everything should come in from the API and this can be eliminated.
-        if not invocation.has_input_for_step(step.id):
-            content = next(iter(step_outputs.values()))
-            if content and content is not NO_REPLACEMENT:
-                invocation.add_input(content, step.id)
         progress.set_outputs_for_input(invocation_step, step_outputs)
         return None
 
diff --git a/lib/galaxy/workflow/run.py b/lib/galaxy/workflow/run.py
index 6d4bd6d1eb0..84594022bf6 100644
--- a/lib/galaxy/workflow/run.py
+++ b/lib/galaxy/workflow/run.py
@@ -623,6 +623,12 @@ class WorkflowProgress:
 
         if step.label and step.type == "parameter_input" and "output" in outputs:
             self.runtime_replacements[step.label] = str(outputs["output"])
+        invocation = invocation_step.workflow_invocation
+        if not invocation.has_input_for_step(step.id):
+            content = outputs.get("output", NO_REPLACEMENT)
+            if content is not NO_REPLACEMENT:
+                log.info("ADDING INPUT FOR STEP %s: %s", step.id, content, exc_info=True)
+                invocation.add_input(content, step.id)
         self.set_step_outputs(invocation_step, outputs, already_persisted=already_persisted)
 
     def effective_replacement_dict(self):
diff --git a/test/unit/workflows/test_workflow_progress.py b/test/unit/workflows/test_workflow_progress.py
index 1a6147c95ca..08cc8c057f4 100644
--- a/test/unit/workflows/test_workflow_progress.py
+++ b/test/unit/workflows/test_workflow_progress.py
@@ -94,8 +94,7 @@ class TestWorkflowProgress(TestCase):
                 workflow_invocation_step.state = "scheduled"
                 workflow_invocation_step.workflow_step = self._step(i)
                 assert step_id == self._step(i).id
-                # workflow_invocation_step.workflow_invocation = self.invocation
-                self.invocation.steps.append(workflow_invocation_step)
+                workflow_invocation_step.workflow_invocation = self.invocation
 
             workflow_invocation_step_state = model.WorkflowRequestStepState()
             workflow_invocation_step_state.workflow_step_id = step_id
@@ -111,6 +110,7 @@ class TestWorkflowProgress(TestCase):
         else:
             workflow_invocation_step = model.WorkflowInvocationStep()
             workflow_invocation_step.workflow_step = self._step(index)
+            workflow_invocation_step.workflow_invocation = self.invocation
             return workflow_invocation_step
 
     def test_connect_data_input(self):
@@ -211,6 +211,7 @@ class TestWorkflowProgress(TestCase):
         subworkflow_invocation_step.workflow_step_id = subworkflow_input_step.id
         subworkflow_invocation_step.state = "new"
         subworkflow_invocation_step.workflow_step = subworkflow_input_step
+        subworkflow_invocation_step.workflow_invocation = subworkflow_invocation
 
         subworkflow_progress.set_outputs_for_input(subworkflow_invocation_step)