refactor Even more

wf_core/flow_ops.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from typing import Any
+
+from .model import Workflow
+from .run_state import ExecutionFrame, FrameStatus, RunState, RunStatus, TraceEntry
+from .schema_tools import validate_payload_against_schema
+from .state_ops import project_output
+from .tokens import END
+
+
+def append_trace(
+    run: RunState,
+    *,
+    frame_id: str,
+    node_id: str,
+    step_type: str,
+    resolved_input: dict[str, Any],
+    outcome: str,
+    next_node_id: str,
+    output: dict[str, Any],
+    state_changes: dict[str, Any],
+) -> None:
+    run.trace.append(
+        TraceEntry(
+            frame_id=frame_id,
+            node_id=node_id,
+            step_type=step_type,
+            resolved_input=resolved_input,
+            outcome=outcome,
+            next_node_id=next_node_id,
+            output=output,
+            state_changes=state_changes,
+        )
+    )
+
+
+def advance_frame(
+    run: RunState,
+    frame: ExecutionFrame,
+    *,
+    outcome: str,
+    next_node_id: str,
+) -> None:
+    frame.prior_outcome = outcome
+    frame.activated_incoming_edge = frame.node_id
+    frame.node_id = next_node_id
+    if next_node_id == END:
+        frame.status = FrameStatus.COMPLETED
+        frame.finished_at_node_id = END
+    else:
+        frame.finished_at_node_id = None
+    run.sync_from_current_frame()
+
+
+def finalize_run(workflow: Workflow, run: RunState) -> RunState:
+    run.output = project_output(workflow, run.state)
+    validate_payload_against_schema(
+        workflow.output_schema, run.output, "workflow output"
+    )
+    run.status = RunStatus.COMPLETED
+    run.current_node_id = END
+    return run

wf_core/foreach_ops.py

@@ -4,9 +4,10 @@ from typing import Any
 
 from .conditions import safe_resolve_path
 from .errors import WorkflowExecutionError
+from .flow_ops import advance_frame, append_trace
 from .frame_ops import frame_context_values
 from .model import ForeachNode, Workflow
-from .run_state import ExecutionFrame, FrameStatus, RunState, TraceEntry
+from .run_state import ExecutionFrame, FrameStatus, RunState
 
 
 def step_foreach(
@@ -41,8 +42,8 @@ def step_foreach(
             raise WorkflowExecutionError(
                 f"no edge found for node {frame.node_id!r} and outcome {outcome!r}"
             )
-        run.trace.append(
+        append_trace(
-            TraceEntry(
+            run,
             frame_id=frame.id,
             node_id=frame.node_id,
             step_type=step.type,
@@ -52,11 +53,7 @@ def step_foreach(
             output={},
             state_changes={},
         )
-        )
+        advance_frame(run, frame, outcome=outcome, next_node_id=next_node_id)
-        frame.prior_outcome = outcome
-        frame.activated_incoming_edge = frame.node_id
-        frame.node_id = next_node_id
-        run.sync_from_current_frame()
         return run
 
     loop_start = edge_map.get((frame.node_id, "loop"))
@@ -82,8 +79,8 @@ def step_foreach(
         parent_frame_id=frame.id,
         metadata=child_metadata,
     )
-    run.trace.append(
+    append_trace(
-        TraceEntry(
+        run,
         frame_id=frame.id,
         node_id=frame.node_id,
         step_type=step.type,
@@ -93,7 +90,6 @@ def step_foreach(
         output={},
         state_changes={},
     )
-    )
     run.current_frame_id = child_id
     run.sync_from_current_frame()
     return run

wf_core/interrupt_ops.py

@@ -4,10 +4,10 @@ from typing import Any
 
 from .conditions import safe_resolve_path
 from .errors import WorkflowExecutionError
+from .flow_ops import advance_frame, append_trace
 from .model import InterruptNode, Workflow
-from .run_state import FrameStatus, InterruptRequest, RunState, TraceEntry
+from .run_state import InterruptRequest, RunState
 from .state_ops import apply_mapped_state
-from .tokens import END
 
 
 def build_interrupt_request(
@@ -76,8 +76,8 @@ def resume_interrupt(
             f"no edge found for interrupt node {frame.node_id!r} and outcome {resume_outcome!r}"
         )
 
-    run.trace.append(
+    append_trace(
-        TraceEntry(
+        run,
         frame_id=frame.id,
         node_id=frame.node_id,
         step_type=step.type,
@@ -87,11 +87,5 @@ def resume_interrupt(
         output=resume_payload,
         state_changes=state_changes,
     )
-    )
-    frame.prior_outcome = resume_outcome
-    frame.activated_incoming_edge = frame.node_id
-    frame.node_id = next_node_id
-    frame.status = FrameStatus.RUNNING if next_node_id != END else FrameStatus.COMPLETED
-    frame.finished_at_node_id = END if next_node_id == END else None
     run.interrupt = None
-    run.sync_from_current_frame()
+    advance_frame(run, frame, outcome=resume_outcome, next_node_id=next_node_id)

wf_core/runtime.py

@@ -2,11 +2,11 @@ from __future__ import annotations
 
 from typing import Any
 
-from .conditions import eval_condition
 from .errors import WorkflowExecutionError
 from .foreach_ops import step_foreach
-from .frame_ops import collapse_completed_frames, frame_context_values
+from .flow_ops import advance_frame, append_trace, finalize_run
-from .interrupt_ops import build_interrupt_request, resume_interrupt
+from .frame_ops import collapse_completed_frames
+from .interrupt_ops import resume_interrupt
 from .model import (
     ConditionNode,
     ForeachNode,
@@ -22,10 +22,13 @@ from .run_state import (
     FrameStatus,
     RunState,
     RunStatus,
-    TraceEntry,
 )
 from .schema_tools import validate_payload_against_schema
-from .state_ops import project_output
+from .step_handlers import (
+    handle_condition_step,
+    handle_interrupt_step,
+    handle_join_step,
+)
 from .tokens import END
 
 __all__ = [
@@ -112,12 +115,7 @@ def resume_workflow(
         )
         collapse_completed_frames(run)
         if run.current_node_id == END:
-            run.output = project_output(workflow, run.state)
+            return finalize_run(workflow, run)
-            validate_payload_against_schema(
-                workflow.output_schema, run.output, "workflow output"
-            )
-            run.status = RunStatus.COMPLETED
-            return run
 
     run.status = RunStatus.RUNNING
     run.error = None
@@ -138,13 +136,7 @@ def resume_workflow(
         if run.status == RunStatus.INTERRUPTED:
             return run
 
-    run.output = project_output(workflow, run.state)
+    return finalize_run(workflow, run)
-    validate_payload_against_schema(
-        workflow.output_schema, run.output, "workflow output"
-    )
-    run.status = RunStatus.COMPLETED
-    run.current_node_id = END
-    return run
 
 
 def step_workflow(
@@ -187,49 +179,11 @@ def step_workflow(
         step_result = execute_node_use(workflow, run, step, node_def, registry)
         outcome = step_result["outcome"]
     elif isinstance(step, ConditionNode):
-        predicate = eval_condition(
+        outcome, step_result = handle_condition_step(run, step)
-            step.check,
-            run.state,
-            run.workflow_input,
-            frame.prior_outcome,
-        )
-        outcome = "true" if predicate else "false"
-        step_result = {
-            "resolved_input": {},
-            "output": {"predicate": predicate},
-            "state_changes": {},
-        }
     elif isinstance(step, JoinNode):
-        outcome = "done"
+        outcome, step_result = handle_join_step()
-        step_result = {
-            "resolved_input": {},
-            "output": {},
-            "state_changes": {},
-        }
     elif isinstance(step, InterruptNode):
-        interrupt_request = build_interrupt_request(
+        return handle_interrupt_step(run, step)
-            step,
-            frame_id=frame.id,
-            state=run.state,
-            workflow_input=run.workflow_input,
-            context=frame_context_values(frame),
-        )
-        run.interrupt = interrupt_request
-        run.status = RunStatus.INTERRUPTED
-        frame.status = FrameStatus.INTERRUPTED
-        run.trace.append(
-            TraceEntry(
-                frame_id=frame.id,
-                node_id=frame.node_id,
-                step_type=step.type,
-                resolved_input=interrupt_request.payload,
-                outcome="interrupt",
-                next_node_id=frame.node_id,
-                output=interrupt_request.payload,
-                state_changes={},
-            )
-        )
-        return run
     elif isinstance(step, ForeachNode):
         return step_foreach(workflow, run, step, edge_map)
     else:
@@ -241,8 +195,8 @@ def step_workflow(
             f"no edge found for node {frame.node_id!r} and outcome {outcome!r}"
         )
 
-    run.trace.append(
+    append_trace(
-        TraceEntry(
+        run,
         frame_id=frame.id,
         node_id=frame.node_id,
         step_type=step.type,
@@ -252,13 +206,5 @@ def step_workflow(
         output=step_result["output"],
         state_changes=step_result["state_changes"],
     )
-    )
+    advance_frame(run, frame, outcome=outcome, next_node_id=next_node_id)
-
-    frame.prior_outcome = outcome
-    frame.activated_incoming_edge = frame.node_id
-    frame.node_id = next_node_id
-    if next_node_id == END:
-        frame.status = FrameStatus.COMPLETED
-        frame.finished_at_node_id = END
-    run.sync_from_current_frame()
     return run

wf_core/step_handlers.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from typing import Any
+
+from .conditions import eval_condition
+from .flow_ops import append_trace
+from .frame_ops import frame_context_values
+from .interrupt_ops import build_interrupt_request
+from .model import ConditionNode, InterruptNode, JoinNode
+from .run_state import FrameStatus, RunState, RunStatus
+
+
+def handle_condition_step(
+    run: RunState,
+    step: ConditionNode,
+) -> tuple[str, dict[str, Any]]:
+    frame = run.current_frame()
+    predicate = eval_condition(
+        step.check,
+        run.state,
+        run.workflow_input,
+        frame.prior_outcome,
+    )
+    outcome = "true" if predicate else "false"
+    return outcome, {
+        "resolved_input": {},
+        "output": {"predicate": predicate},
+        "state_changes": {},
+    }
+
+
+def handle_join_step() -> tuple[str, dict[str, Any]]:
+    return "done", {
+        "resolved_input": {},
+        "output": {},
+        "state_changes": {},
+    }
+
+
+def handle_interrupt_step(
+    run: RunState,
+    step: InterruptNode,
+) -> RunState:
+    frame = run.current_frame()
+    interrupt_request = build_interrupt_request(
+        step,
+        frame_id=frame.id,
+        state=run.state,
+        workflow_input=run.workflow_input,
+        context=frame_context_values(frame),
+    )
+    run.interrupt = interrupt_request
+    run.status = RunStatus.INTERRUPTED
+    frame.status = FrameStatus.INTERRUPTED
+    append_trace(
+        run,
+        frame_id=frame.id,
+        node_id=frame.node_id,
+        step_type=step.type,
+        resolved_input=interrupt_request.payload,
+        outcome="interrupt",
+        next_node_id=frame.node_id,
+        output=interrupt_request.payload,
+        state_changes={},
+    )
+    return run