fix: ensure workflow node executions are tracked after addition to se…

…ssion Signed-off-by: -LAN- <laipz8200@outlook.com>
langgenius · Jan 3, 2025 · d3161ff · d3161ff
1 parent db1b05a
commit d3161ff
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/api/core/app/task_pipeline/workflow_cycle_manage.py b/api/core/app/task_pipeline/workflow_cycle_manage.py
@@ -307,6 +307,8 @@ def _handle_node_execution_start(
         workflow_node_execution.created_at = datetime.now(UTC).replace(tzinfo=None)
 
         session.add(workflow_node_execution)
+
+        self._workflow_node_executions[workflow_node_execution.id] = workflow_node_execution
         return workflow_node_execution
 
     def _handle_workflow_node_execution_success(
@@ -424,6 +426,8 @@ def _handle_workflow_node_execution_retried(
         workflow_node_execution.index = event.node_run_index
 
         session.add(workflow_node_execution)
+
+        self._workflow_node_executions[workflow_node_execution.id] = workflow_node_execution
         return workflow_node_execution
 
     #################################################