Add user input into the pstat in thread start and resume

Author: tjeerddie

orchestrator/services/executors/threadpool.py

@@ -18,8 +18,10 @@
 
 from oauth2_lib.fastapi import OIDCUserModel
 from orchestrator.db import ProcessTable, db
+from orchestrator.services.input_state import retrieve_input_state
 from orchestrator.services.processes import (
     SYSTEM_USER,
+    StateMerger,
     _get_process,
     _run_process_async,
     create_process,
@@ -52,6 +54,10 @@ def thread_start_process(
     db.session.add(process)
     db.session.commit()
 
+    pstat = load_process(process)
+    input_data = retrieve_input_state(process.process_id, "initial_state")
+    pstat.update(state=pstat.state.map(lambda state: StateMerger.merge(state, input_data.input_state)))
+
     _safe_logstep_with_func = partial(safe_logstep, broadcast_func=broadcast_func)
     return _run_process_async(pstat.process_id, lambda: runwf(pstat, _safe_logstep_with_func))
 
@@ -71,6 +77,9 @@ def thread_resume_process(
     if user:
         pstat.update(current_user=user)
 
+    input_data = retrieve_input_state(process.process_id, "user_input")
+    pstat.update(state=pstat.state.map(lambda state: StateMerger.merge(state, input_data.input_state)))
+
     # enforce an update to the process status to properly show the process
     process.last_status = ProcessStatus.RUNNING
     db.session.add(process)

orchestrator/services/processes.py

@@ -486,6 +486,29 @@ def start_process(
     return start_func(pstat, user=user, user_model=user_model, broadcast_func=broadcast_func)
 
 
+def restart_process(
+    process: ProcessTable,
+    *,
+    user: str | None = None,
+    broadcast_func: BroadcastFunc | None = None,
+) -> UUID:
+    """Start a process for workflow.
+
+    Args:
+        process: Process from database
+        user: user who resumed this process
+        broadcast_func: Optional function to broadcast process data
+
+    Returns:
+        process id
+
+    """
+    pstat = load_process(process)
+
+    start_func = get_execution_context()["start"]
+    return start_func(pstat, user=user, broadcast_func=broadcast_func)
+
+
 def resume_process(
     process: ProcessTable,
     *,

orchestrator/workflows/tasks/resume_workflows.py

@@ -50,9 +50,10 @@ def resume_found_workflows(
     created_state_process_ids: list[UUIDstr],
     resumed_state_process_ids: list[UUIDstr],
 ) -> State:
-    all_processes = waiting_process_ids + created_state_process_ids + resumed_state_process_ids
+    resume_processes = waiting_process_ids + resumed_state_process_ids
+
     resumed_process_ids = []
-    for process_id in all_processes:
+    for process_id in resume_processes:
         try:
             process = db.session.get(ProcessTable, process_id)
             if not process:
@@ -67,7 +68,28 @@ def resume_found_workflows(
             # Make sure to turn it on again
             db.session.info["disabled"] = True
 
-    return {"number_of_resumed_process_ids": len(resumed_process_ids), "resumed_process_ids": resumed_process_ids}
+    started_process_ids = []
+    for process_id in created_state_process_ids:
+        try:
+            process = db.session.get(ProcessTable, process_id)
+            if not process:
+                continue
+            # Workaround the commit disable function
+            db.session.info["disabled"] = False
+            processes.restart_process(process)
+            started_process_ids.append(process_id)
+        except Exception:
+            logger.exception()
+        finally:
+            # Make sure to turn it on again
+            db.session.info["disabled"] = True
+
+    return {
+        "number_of_resumed_process_ids": len(resumed_process_ids),
+        "resumed_process_ids": resumed_process_ids,
+        "number_of_started_process_ids": len(started_process_ids),
+        "started_process_ids": started_process_ids,
+    }
 
 
 @workflow("Resume all workflows that are stuck on tasks with the status 'waiting'", target=Target.SYSTEM)