refactor 3

wf_core/__init__.py

@@ -26,6 +26,7 @@ from .run_state import (
     RunState,
     RunStatus,
     RuntimeContext,
+    StepExecutionResult,
     TraceEntry,
 )
 from .tokens import END, START
@@ -53,6 +54,7 @@ __all__ = [
     "RunState",
     "RunStatus",
     "RuntimeContext",
+    "StepExecutionResult",
     "TraceEntry",
     "InterruptRequest",
     "START",

wf_core/flow_ops.py

@@ -3,7 +3,14 @@ from __future__ import annotations
 from typing import Any
 
 from .model import Workflow
-from .run_state import ExecutionFrame, FrameStatus, RunState, RunStatus, TraceEntry
+from .run_state import (
+    ExecutionFrame,
+    FrameStatus,
+    RunState,
+    RunStatus,
+    StepExecutionResult,
+    TraceEntry,
+)
 from .schema_tools import validate_payload_against_schema
 from .state_ops import project_output
 from .tokens import END
@@ -35,6 +42,28 @@ def append_trace(
     )
 
 
+def append_step_result_trace(
+    run: RunState,
+    *,
+    frame_id: str,
+    node_id: str,
+    step_type: str,
+    next_node_id: str,
+    result: StepExecutionResult,
+) -> None:
+    append_trace(
+        run,
+        frame_id=frame_id,
+        node_id=node_id,
+        step_type=step_type,
+        resolved_input=result.resolved_input,
+        outcome=result.outcome,
+        next_node_id=next_node_id,
+        output=result.output,
+        state_changes=result.state_changes,
+    )
+
+
 def advance_frame(
     run: RunState,
     frame: ExecutionFrame,

wf_core/foreach_ops.py

@@ -1,23 +1,24 @@
 from __future__ import annotations
 
-from typing import Any
-
 from .conditions import safe_resolve_path
 from .errors import WorkflowExecutionError
-from .flow_ops import advance_frame, append_trace
+from .flow_ops import advance_frame, append_step_result_trace
 from .frame_ops import frame_context_values
 from .model import ForeachNode, Workflow
-from .run_state import ExecutionFrame, FrameStatus, RunState
+from .run_state import ExecutionFrame, FrameStatus, RunState, StepExecutionResult
+from .workflow_index import WorkflowIndex
 
 
 def step_foreach(
     workflow: Workflow,
     run: RunState,
     step: ForeachNode,
-    edge_map: dict[tuple[str, str], str],
+    index: WorkflowIndex,
 ) -> RunState:
     if step.mode != "serial":
-        raise WorkflowExecutionError("parallel foreach execution is not implemented yet")
+        raise WorkflowExecutionError(
+            "parallel foreach execution is not implemented yet"
+        )
 
     frame = run.current_frame()
     progress_map = frame.metadata.setdefault("foreach_progress", {})
@@ -34,40 +35,34 @@ def step_foreach(
             f"foreach source {step.over!r} must resolve to a list"
         )
 
-    index = progress["index"]
+    loop_index = progress["index"]
-    if index >= len(iterable):
+    if loop_index >= len(iterable):
         outcome = "done"
-        next_node_id = edge_map.get((frame.node_id, outcome))
+        next_node_id = index.next_node_id(frame.node_id, outcome)
-        if next_node_id is None:
+        append_step_result_trace(
-            raise WorkflowExecutionError(
-                f"no edge found for node {frame.node_id!r} and outcome {outcome!r}"
-            )
-        append_trace(
             run,
             frame_id=frame.id,
             node_id=frame.node_id,
             step_type=step.type,
-            resolved_input={"count": len(iterable), "index": index},
-            outcome=outcome,
             next_node_id=next_node_id,
+            result=StepExecutionResult(
+                outcome=outcome,
+                resolved_input={"count": len(iterable), "index": loop_index},
                 output={},
                 state_changes={},
+            ),
         )
         advance_frame(run, frame, outcome=outcome, next_node_id=next_node_id)
         return run
 
-    loop_start = edge_map.get((frame.node_id, "loop"))
+    loop_start = index.next_node_id(frame.node_id, "loop")
-    if loop_start is None:
-        raise WorkflowExecutionError(
-            f"no edge found for foreach node {frame.node_id!r} and outcome 'loop'"
-        )
 
-    item = iterable[index]
+    item = iterable[loop_index]
-    progress["index"] = index + 1
+    progress["index"] = loop_index + 1
-    child_id = f"{frame.id}:{step.id}:{index}"
+    child_id = f"{frame.id}:{step.id}:{loop_index}"
     child_metadata = {
         "foreach_node_id": step.id,
-        "loop_index": index,
+        "loop_index": loop_index,
         "loop_item": item,
         "loop_alias": step.as_,
     }
@@ -79,16 +74,18 @@ def step_foreach(
         parent_frame_id=frame.id,
         metadata=child_metadata,
     )
-    append_trace(
+    append_step_result_trace(
         run,
         frame_id=frame.id,
         node_id=frame.node_id,
         step_type=step.type,
-        resolved_input={"item": item, "index": index},
-        outcome="loop",
         next_node_id=loop_start,
+        result=StepExecutionResult(
+            outcome="loop",
+            resolved_input={"item": item, "index": loop_index},
             output={},
             state_changes={},
+        ),
     )
     run.current_frame_id = child_id
     run.sync_from_current_frame()

wf_core/interrupt_ops.py

@@ -4,10 +4,11 @@ from typing import Any
 
 from .conditions import safe_resolve_path
 from .errors import WorkflowExecutionError
-from .flow_ops import advance_frame, append_trace
+from .flow_ops import advance_frame, append_step_result_trace
 from .model import InterruptNode, Workflow
-from .run_state import InterruptRequest, RunState
+from .run_state import InterruptRequest, RunState, StepExecutionResult
 from .state_ops import apply_mapped_state
+from .workflow_index import WorkflowIndex
 
 
 def build_interrupt_request(
@@ -40,8 +41,7 @@ def resume_interrupt(
     workflow: Workflow,
     run: RunState,
     *,
-    nodes_by_id: dict[str, Any],
+    index: WorkflowIndex,
-    edge_map: dict[tuple[str, str], str],
     resume_payload: dict[str, Any],
     resume_outcome: str,
 ) -> None:
@@ -53,7 +53,7 @@ def resume_interrupt(
         raise WorkflowExecutionError("run is interrupted but has no interrupt request")
 
     frame = run.current_frame()
-    step = nodes_by_id[frame.node_id]
+    step = index.nodes_by_id[frame.node_id]
     if not isinstance(step, InterruptNode):
         raise WorkflowExecutionError(
             f"interrupted run expected interrupt node, got {step.type!r}"
@@ -70,22 +70,19 @@ def resume_interrupt(
         run.state,
         missing_field_message="interrupt resume payload is missing required field {field}",
     )
-    next_node_id = edge_map.get((frame.node_id, resume_outcome))
+    next_node_id = index.next_node_id(frame.node_id, resume_outcome)
-    if next_node_id is None:
+    append_step_result_trace(
-        raise WorkflowExecutionError(
-            f"no edge found for interrupt node {frame.node_id!r} and outcome {resume_outcome!r}"
-        )
-
-    append_trace(
         run,
         frame_id=frame.id,
         node_id=frame.node_id,
         step_type=step.type,
-        resolved_input=resume_payload,
-        outcome=resume_outcome,
         next_node_id=next_node_id,
+        result=StepExecutionResult(
+            outcome=resume_outcome,
+            resolved_input=resume_payload,
             output=resume_payload,
             state_changes=state_changes,
+        ),
     )
     run.interrupt = None
     advance_frame(run, frame, outcome=resume_outcome, next_node_id=next_node_id)

wf_core/node_exec.py

@@ -7,7 +7,7 @@ from .conditions import safe_resolve_path
 from .errors import WorkflowExecutionError
 from .frame_ops import frame_context_values
 from .model import NodeDef, NodeResult, NodeUse, Workflow
-from .run_state import RunState, RuntimeContext
+from .run_state import RunState, RuntimeContext, StepExecutionResult
 from .schema_tools import validate_payload_against_schema
 from .state_ops import apply_output_map
 
@@ -20,7 +20,7 @@ def execute_node_use(
     node: NodeUse,
     node_def: NodeDef,
     registry: dict[str, NodeHandler],
-) -> dict[str, Any]:
+) -> StepExecutionResult:
     handler = registry.get(node.node)
     if handler is None:
         raise WorkflowExecutionError(
@@ -61,12 +61,12 @@ def execute_node_use(
         node_def.output_schema, result.output, f"node output for {node.id}"
     )
     state_changes = apply_output_map(workflow, node, result.output, run.state)
-    return {
+    return StepExecutionResult(
-        "outcome": result.outcome,
+        outcome=result.outcome,
-        "resolved_input": resolved_input,
+        resolved_input=resolved_input,
-        "output": result.output,
+        output=result.output,
-        "state_changes": state_changes,
+        state_changes=state_changes,
-    }
+    )
 
 
 def coerce_node_result(raw_result: NodeResult | dict[str, Any]) -> NodeResult:

wf_core/run_factory.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from .model import Workflow
+from .run_state import ExecutionFrame, FrameStatus, RunState, RunStatus
+
+
+def create_run_state(workflow: Workflow, workflow_input: dict[str, object]) -> RunState:
+    run = RunState(
+        workflow_name=workflow.name,
+        status=RunStatus.PENDING,
+        workflow_input=dict(workflow_input),
+        state=dict(workflow_input),
+        frames={
+            "root": ExecutionFrame(
+                id="root",
+                kind="workflow",
+                node_id=workflow.start,
+                status=FrameStatus.PENDING,
+            )
+        },
+        current_frame_id="root",
+        current_node_id=workflow.start,
+    )
+    run.sync_from_current_frame()
+    return run

wf_core/run_state.py

@@ -56,6 +56,14 @@ class TraceEntry:
     state_changes: dict[str, Any] = field(default_factory=dict)
 
 
+@dataclass(slots=True)
+class StepExecutionResult:
+    outcome: str
+    resolved_input: dict[str, Any] = field(default_factory=dict)
+    output: dict[str, Any] = field(default_factory=dict)
+    state_changes: dict[str, Any] = field(default_factory=dict)
+
+
 @dataclass(slots=True)
 class InterruptRequest:
     id: str

wf_core/runtime.py

@@ -4,7 +4,7 @@ from typing import Any
 
 from .errors import WorkflowExecutionError
 from .foreach_ops import step_foreach
-from .flow_ops import advance_frame, append_trace, finalize_run
+from .flow_ops import advance_frame, append_step_result_trace, finalize_run
 from .frame_ops import collapse_completed_frames
 from .interrupt_ops import resume_interrupt
 from .model import (
@@ -12,13 +12,12 @@ from .model import (
     ForeachNode,
     InterruptNode,
     JoinNode,
-    NodeDef,
     NodeUse,
     Workflow,
 )
 from .node_exec import NodeHandler, coerce_node_result, execute_node_use
+from .run_factory import create_run_state
 from .run_state import (
-    ExecutionFrame,
     FrameStatus,
     RunState,
     RunStatus,
@@ -30,6 +29,7 @@ from .step_handlers import (
     handle_join_step,
 )
 from .tokens import END
+from .workflow_index import WorkflowIndex, build_workflow_index
 
 __all__ = [
     "NodeHandler",
@@ -45,23 +45,7 @@ def execute_workflow(
     workflow_input: dict[str, Any],
     registry: dict[str, NodeHandler],
 ) -> RunState:
-    run = RunState(
+    run = create_run_state(workflow, workflow_input)
-        workflow_name=workflow.name,
-        status=RunStatus.PENDING,
-        workflow_input=dict(workflow_input),
-        state=dict(workflow_input),
-        frames={
-            "root": ExecutionFrame(
-                id="root",
-                kind="workflow",
-                node_id=workflow.start,
-                status=FrameStatus.PENDING,
-            )
-        },
-        current_frame_id="root",
-        current_node_id=workflow.start,
-    )
-    run.sync_from_current_frame()
 
     try:
         workflow.validate_structure().raise_for_errors()
@@ -98,9 +82,7 @@ def resume_workflow(
     if run.status == RunStatus.COMPLETED:
         return run
 
-    node_defs = {node_def.name: node_def for node_def in workflow.node_defs}
+    index = build_workflow_index(workflow)
-    nodes_by_id = {node.id: node for node in workflow.nodes}
-    edge_map = {(edge.from_, edge.outcome): edge.to for edge in workflow.edges}
 
     if run.status == RunStatus.INTERRUPTED:
         if resume_payload is None:
@@ -108,8 +90,7 @@ def resume_workflow(
         resume_interrupt(
             workflow,
             run,
-            nodes_by_id=nodes_by_id,
+            index=index,
-            edge_map=edge_map,
             resume_payload=resume_payload,
             resume_outcome=resume_outcome,
         )
@@ -129,9 +110,7 @@ def resume_workflow(
             workflow,
             run,
             registry,
-            node_defs=node_defs,
+            index=index,
-            nodes_by_id=nodes_by_id,
-            edge_map=edge_map,
         )
         if run.status == RunStatus.INTERRUPTED:
             return run
@@ -144,9 +123,7 @@ def step_workflow(
     run: RunState,
     registry: dict[str, NodeHandler],
     *,
-    node_defs: dict[str, NodeDef] | None = None,
+    index: WorkflowIndex | None = None,
-    nodes_by_id: dict[str, Any] | None = None,
-    edge_map: dict[tuple[str, str], str] | None = None,
 ) -> RunState:
     if run.current_frame_id is None:
         raise WorkflowExecutionError("run has no current frame")
@@ -161,50 +138,41 @@ def step_workflow(
         run.status = RunStatus.RUNNING
     run.error = None
 
-    node_defs = node_defs or {
+    index = index or build_workflow_index(workflow)
-        node_def.name: node_def for node_def in workflow.node_defs
-    }
-    nodes_by_id = nodes_by_id or {node.id: node for node in workflow.nodes}
-    edge_map = edge_map or {
-        (edge.from_, edge.outcome): edge.to for edge in workflow.edges
-    }
 
     frame = run.current_frame()
     if frame.status == FrameStatus.PENDING:
         frame.status = FrameStatus.RUNNING
-    step = nodes_by_id[frame.node_id]
+    step = index.nodes_by_id[frame.node_id]
 
     if isinstance(step, NodeUse):
-        node_def = node_defs[step.node]
+        node_def = index.node_defs[step.node]
         step_result = execute_node_use(workflow, run, step, node_def, registry)
-        outcome = step_result["outcome"]
     elif isinstance(step, ConditionNode):
-        outcome, step_result = handle_condition_step(run, step)
+        step_result = handle_condition_step(run, step)
     elif isinstance(step, JoinNode):
-        outcome, step_result = handle_join_step()
+        step_result = handle_join_step()
     elif isinstance(step, InterruptNode):
         return handle_interrupt_step(run, step)
     elif isinstance(step, ForeachNode):
-        return step_foreach(workflow, run, step, edge_map)
+        return step_foreach(workflow, run, step, index)
     else:
         raise WorkflowExecutionError(f"unsupported step type {step.type!r}")
 
-    next_node_id = edge_map.get((frame.node_id, outcome))
+    next_node_id = index.next_node_id(frame.node_id, step_result.outcome)
-    if next_node_id is None:
-        raise WorkflowExecutionError(
-            f"no edge found for node {frame.node_id!r} and outcome {outcome!r}"
-        )
 
-    append_trace(
+    append_step_result_trace(
         run,
         frame_id=frame.id,
         node_id=frame.node_id,
         step_type=step.type,
-        resolved_input=step_result["resolved_input"],
-        outcome=outcome,
         next_node_id=next_node_id,
-        output=step_result["output"],
+        result=step_result,
-        state_changes=step_result["state_changes"],
+    )
+    advance_frame(
+        run,
+        frame,
+        outcome=step_result.outcome,
+        next_node_id=next_node_id,
     )
-    advance_frame(run, frame, outcome=outcome, next_node_id=next_node_id)
     return run

wf_core/step_handlers.py

@@ -1,19 +1,17 @@
 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
+from .run_state import FrameStatus, RunState, RunStatus, StepExecutionResult
 
 
 def handle_condition_step(
     run: RunState,
     step: ConditionNode,
-) -> tuple[str, dict[str, Any]]:
+) -> StepExecutionResult:
     frame = run.current_frame()
     predicate = eval_condition(
         step.check,
@@ -22,19 +20,21 @@ def handle_condition_step(
         frame.prior_outcome,
     )
     outcome = "true" if predicate else "false"
-    return outcome, {
+    return StepExecutionResult(
-        "resolved_input": {},
+        outcome=outcome,
-        "output": {"predicate": predicate},
+        resolved_input={},
-        "state_changes": {},
+        output={"predicate": predicate},
-    }
+        state_changes={},
+    )
 
 
-def handle_join_step() -> tuple[str, dict[str, Any]]:
+def handle_join_step() -> StepExecutionResult:
-    return "done", {
+    return StepExecutionResult(
-        "resolved_input": {},
+        outcome="done",
-        "output": {},
+        resolved_input={},
-        "state_changes": {},
+        output={},
-    }
+        state_changes={},
+    )
 
 
 def handle_interrupt_step(

wf_core/workflow_index.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from .errors import WorkflowExecutionError
+from .model import NodeDef, Workflow
+
+
+@dataclass(slots=True)
+class WorkflowIndex:
+    node_defs: dict[str, NodeDef]
+    nodes_by_id: dict[str, Any]
+    edge_map: dict[tuple[str, str], str]
+
+    def next_node_id(self, node_id: str, outcome: str) -> str:
+        next_node_id = self.edge_map.get((node_id, outcome))
+        if next_node_id is None:
+            raise WorkflowExecutionError(
+                f"no edge found for node {node_id!r} and outcome {outcome!r}"
+            )
+        return next_node_id
+
+
+def build_workflow_index(workflow: Workflow) -> WorkflowIndex:
+    return WorkflowIndex(
+        node_defs={node_def.name: node_def for node_def in workflow.node_defs},
+        nodes_by_id={node.id: node for node in workflow.nodes},
+        edge_map={(edge.from_, edge.outcome): edge.to for edge in workflow.edges},
+    )