supporting async in the Core

wf_core/__init__.py

@@ -13,11 +13,15 @@ from .model import (
     Workflow,
 )
 from .runtime import (
+    AsyncNodeHandler,
     NodeHandler,
     WorkflowExecutionError,
     coerce_node_result,
+    execute_workflow_async,
     execute_workflow,
+    resume_workflow_async,
     resume_workflow,
+    step_workflow_async,
     step_workflow,
 )
 from .run_state import (
@@ -50,6 +54,7 @@ __all__ = [
     "SchemaRef",
     "StateField",
     "StateSchema",
+    "AsyncNodeHandler",
     "NodeHandler",
     "ExecutionFrame",
     "FrameStatus",
@@ -67,8 +72,11 @@ __all__ = [
     "Workflow",
     "WorkflowExecutionError",
     "coerce_node_result",
+    "execute_workflow_async",
     "execute_workflow",
+    "resume_workflow_async",
     "resume_workflow",
+    "step_workflow_async",
     "step_workflow",
     "validate_workflow",
 ]

wf_core/node_exec.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from collections.abc import Callable, Mapping
-from typing import Any
+from collections.abc import Awaitable, Callable, Mapping
+from typing import Any, cast
 
 from .conditions import safe_resolve_path
 from .errors import WorkflowExecutionError
@@ -12,21 +12,19 @@ from .schema_tools import validate_payload_against_schema
 from .state_ops import apply_output_map
 
 NodeHandler = Callable[[dict[str, Any], RuntimeContext], NodeResult | dict[str, Any]]
+AsyncNodeHandler = Callable[
+    [dict[str, Any], RuntimeContext],
+    Awaitable[NodeResult | dict[str, Any]] | NodeResult | dict[str, Any],
+]
 
 
-def execute_node_use(
+def _resolve_node_execution(
+    *,
     workflow: Workflow,
     run: RunState,
     node: NodeUse,
     node_def: NodeDef,
-    registry: Mapping[str, NodeHandler],
-) -> StepExecutionResult:
-    handler = registry.get(node.node)
-    if handler is None:
-        raise WorkflowExecutionError(
-            f"no handler registered for node def {node.node!r}"
-        )
-
+) -> tuple[dict[str, Any], RuntimeContext]:
     frame = run.current_frame()
     context_values = frame_context_values(frame)
     resolved_input = {
@@ -49,7 +47,18 @@ def execute_node_use(
         activated_incoming_edge=frame.activated_incoming_edge,
         metadata=dict(frame.metadata),
     )
-    raw_result = handler(resolved_input, context)
+    return resolved_input, context
+
+
+def _finalize_node_execution(
+    *,
+    workflow: Workflow,
+    run: RunState,
+    node: NodeUse,
+    node_def: NodeDef,
+    resolved_input: dict[str, Any],
+    raw_result: NodeResult | dict[str, Any],
+) -> StepExecutionResult:
     result = coerce_node_result(raw_result)
 
     if result.outcome not in node_def.outcomes:
@@ -69,6 +78,70 @@ def execute_node_use(
     )
 
 
+def execute_node_use(
+    workflow: Workflow,
+    run: RunState,
+    node: NodeUse,
+    node_def: NodeDef,
+    registry: Mapping[str, NodeHandler],
+) -> StepExecutionResult:
+    handler = registry.get(node.node)
+    if handler is None:
+        raise WorkflowExecutionError(
+            f"no handler registered for node def {node.node!r}"
+        )
+
+    resolved_input, context = _resolve_node_execution(
+        workflow=workflow,
+        run=run,
+        node=node,
+        node_def=node_def,
+    )
+    raw_result = handler(resolved_input, context)
+    return _finalize_node_execution(
+        workflow=workflow,
+        run=run,
+        node=node,
+        node_def=node_def,
+        resolved_input=resolved_input,
+        raw_result=raw_result,
+    )
+
+
+async def execute_node_use_async(
+    workflow: Workflow,
+    run: RunState,
+    node: NodeUse,
+    node_def: NodeDef,
+    registry: Mapping[str, AsyncNodeHandler],
+) -> StepExecutionResult:
+    handler = registry.get(node.node)
+    if handler is None:
+        raise WorkflowExecutionError(
+            f"no handler registered for node def {node.node!r}"
+        )
+
+    resolved_input, context = _resolve_node_execution(
+        workflow=workflow,
+        run=run,
+        node=node,
+        node_def=node_def,
+    )
+    raw_or_awaitable = handler(resolved_input, context)
+    if isinstance(raw_or_awaitable, Awaitable):
+        raw_result = await raw_or_awaitable
+    else:
+        raw_result = raw_or_awaitable
+    return _finalize_node_execution(
+        workflow=workflow,
+        run=run,
+        node=node,
+        node_def=node_def,
+        resolved_input=resolved_input,
+        raw_result=cast(NodeResult | dict[str, Any], raw_result),
+    )
+
+
 def coerce_node_result(raw_result: NodeResult | dict[str, Any]) -> NodeResult:
     if isinstance(raw_result, NodeResult):
         return raw_result

wf_core/runtime.py

@@ -16,7 +16,13 @@ from .model import (
     NodeUse,
     Workflow,
 )
-from .node_exec import NodeHandler, coerce_node_result, execute_node_use
+from .node_exec import (
+    AsyncNodeHandler,
+    NodeHandler,
+    coerce_node_result,
+    execute_node_use,
+    execute_node_use_async,
+)
 from .run_factory import create_run_state
 from .run_state import (
     FrameStatus,
@@ -33,14 +39,126 @@ from .tokens import END
 from .workflow_index import WorkflowIndex, build_workflow_index
 
 __all__ = [
+    "AsyncNodeHandler",
     "NodeHandler",
     "coerce_node_result",
+    "execute_workflow_async",
     "execute_workflow",
+    "resume_workflow_async",
     "resume_workflow",
+    "step_workflow_async",
     "step_workflow",
 ]
 
 
+def _prepare_new_run(workflow: Workflow, workflow_input: dict[str, Any]) -> RunState:
+    run = create_run_state(workflow, workflow_input)
+    workflow.validate_structure().raise_for_errors()
+    validate_payload_against_schema(
+        workflow.input_schema, workflow_input, "workflow input"
+    )
+    return run
+
+
+def _prepare_resume(
+    workflow: Workflow,
+    run: RunState,
+    *,
+    resume_payload: dict[str, Any] | None,
+    resume_outcome: str,
+) -> WorkflowIndex | None:
+    if run.workflow_name != workflow.name:
+        raise WorkflowExecutionError(
+            f"run state belongs to workflow {run.workflow_name!r}, not {workflow.name!r}"
+        )
+
+    if run.current_frame_id is None:
+        raise WorkflowExecutionError("run has no current frame")
+    collapse_completed_frames(run)
+
+    if run.current_node_id is None:
+        raise WorkflowExecutionError("run has no current node")
+
+    if run.status == RunStatus.COMPLETED:
+        return None
+
+    index = build_workflow_index(workflow)
+
+    if run.status == RunStatus.INTERRUPTED:
+        if resume_payload is None:
+            return None
+        resume_interrupt(
+            workflow,
+            run,
+            index=index,
+            resume_payload=resume_payload,
+            resume_outcome=resume_outcome,
+        )
+        collapse_completed_frames(run)
+        if run.current_node_id == END:
+            return None
+
+    run.status = RunStatus.RUNNING
+    run.error = None
+    run.current_frame().status = FrameStatus.RUNNING
+    return index
+
+
+def _prepare_step(
+    workflow: Workflow,
+    run: RunState,
+    index: WorkflowIndex | None,
+) -> tuple[WorkflowIndex, object] | None:
+    if run.current_frame_id is None:
+        raise WorkflowExecutionError("run has no current frame")
+
+    collapse_completed_frames(run)
+    if run.current_node_id is None or run.current_node_id == END:
+        return None
+    if run.status == RunStatus.INTERRUPTED:
+        return None
+
+    if run.status == RunStatus.PENDING:
+        run.status = RunStatus.RUNNING
+    run.error = None
+
+    resolved_index = index or build_workflow_index(workflow)
+    frame = run.current_frame()
+    if frame.status == FrameStatus.PENDING:
+        frame.status = FrameStatus.RUNNING
+    step = resolved_index.nodes_by_id[frame.node_id]
+    return resolved_index, step
+
+
+def _complete_step(
+    *,
+    run: RunState,
+    index: WorkflowIndex,
+    outcome: str,
+    frame_id: str,
+    node_id: str,
+    step_type: str,
+    step_result: Any,
+) -> RunState:
+    next_node_id = index.next_node_id(node_id, outcome)
+
+    append_step_result_trace(
+        run,
+        frame_id=frame_id,
+        node_id=node_id,
+        step_type=step_type,
+        next_node_id=next_node_id,
+        result=step_result,
+    )
+    advance_frame(
+        run,
+        run.frames[frame_id],
+        outcome=outcome,
+        next_node_id=next_node_id,
+    )
+    return run
+
+
 def execute_workflow(
     workflow: Workflow,
     workflow_input: dict[str, Any],
@@ -49,10 +167,7 @@ def execute_workflow(
     run = create_run_state(workflow, workflow_input)
 
     try:
-        workflow.validate_structure().raise_for_errors()
-        validate_payload_against_schema(
-            workflow.input_schema, workflow_input, "workflow input"
-        )
+        run = _prepare_new_run(workflow, workflow_input)
         return resume_workflow(workflow, run, registry)
     except Exception as exc:
         run.status = RunStatus.FAILED
@@ -60,6 +175,22 @@ def execute_workflow(
         raise
 
 
+async def execute_workflow_async(
+    workflow: Workflow,
+    workflow_input: dict[str, Any],
+    registry: Mapping[str, AsyncNodeHandler],
+) -> RunState:
+    run = create_run_state(workflow, workflow_input)
+
+    try:
+        run = _prepare_new_run(workflow, workflow_input)
+        return await resume_workflow_async(workflow, run, registry)
+    except Exception as exc:
+        run.status = RunStatus.FAILED
+        run.error = str(exc)
+        raise
+
+
 def resume_workflow(
     workflow: Workflow,
     run: RunState,
@@ -68,40 +199,16 @@ def resume_workflow(
     resume_payload: dict[str, Any] | None = None,
     resume_outcome: str = "submitted",
 ) -> RunState:
-    if run.workflow_name != workflow.name:
-        raise WorkflowExecutionError(
-            f"run state belongs to workflow {run.workflow_name!r}, not {workflow.name!r}"
-        )
-
-    if run.current_frame_id is None:
-        raise WorkflowExecutionError("run has no current frame")
-    collapse_completed_frames(run)
-
-    if run.current_node_id is None:
-        raise WorkflowExecutionError("run has no current node")
-
-    if run.status == RunStatus.COMPLETED:
-        return run
-
-    index = build_workflow_index(workflow)
-
-    if run.status == RunStatus.INTERRUPTED:
-        if resume_payload is None:
-            return run
-        resume_interrupt(
-            workflow,
-            run,
-            index=index,
-            resume_payload=resume_payload,
-            resume_outcome=resume_outcome,
-        )
-        collapse_completed_frames(run)
+    index = _prepare_resume(
+        workflow,
+        run,
+        resume_payload=resume_payload,
+        resume_outcome=resume_outcome,
+    )
+    if index is None:
         if run.current_node_id == END:
             return finalize_run(workflow, run)
-
-    run.status = RunStatus.RUNNING
-    run.error = None
-    run.current_frame().status = FrameStatus.RUNNING
+        return run
 
     while True:
         collapse_completed_frames(run)
@@ -119,6 +226,41 @@ def resume_workflow(
     return finalize_run(workflow, run)
 
 
+async def resume_workflow_async(
+    workflow: Workflow,
+    run: RunState,
+    registry: Mapping[str, AsyncNodeHandler],
+    *,
+    resume_payload: dict[str, Any] | None = None,
+    resume_outcome: str = "submitted",
+) -> RunState:
+    index = _prepare_resume(
+        workflow,
+        run,
+        resume_payload=resume_payload,
+        resume_outcome=resume_outcome,
+    )
+    if index is None:
+        if run.current_node_id == END:
+            return finalize_run(workflow, run)
+        return run
+
+    while True:
+        collapse_completed_frames(run)
+        if run.current_node_id == END:
+            break
+        await step_workflow_async(
+            workflow,
+            run,
+            registry,
+            index=index,
+        )
+        if run.status == RunStatus.INTERRUPTED:
+            return run
+
+    return finalize_run(workflow, run)
+
+
 def step_workflow(
     workflow: Workflow,
     run: RunState,
@@ -126,25 +268,11 @@ def step_workflow(
     *,
     index: WorkflowIndex | None = None,
 ) -> RunState:
-    if run.current_frame_id is None:
-        raise WorkflowExecutionError("run has no current frame")
-
-    collapse_completed_frames(run)
-    if run.current_node_id is None or run.current_node_id == END:
+    prepared = _prepare_step(workflow, run, index)
+    if prepared is None:
         return run
-    if run.status == RunStatus.INTERRUPTED:
-        return run
-
-    if run.status == RunStatus.PENDING:
-        run.status = RunStatus.RUNNING
-    run.error = None
-
-    index = index or build_workflow_index(workflow)
-
+    index, step = prepared
     frame = run.current_frame()
-    if frame.status == FrameStatus.PENDING:
-        frame.status = FrameStatus.RUNNING
-    step = index.nodes_by_id[frame.node_id]
 
     if isinstance(step, NodeUse):
         node_def = index.node_defs[step.node]
@@ -158,22 +286,58 @@ def step_workflow(
     elif isinstance(step, ForeachNode):
         return step_foreach(workflow, run, step, index)
     else:
-        raise WorkflowExecutionError(f"unsupported step type {step.type!r}")
+        raise WorkflowExecutionError(
+            f"unsupported step type {getattr(step, 'type', type(step).__name__)!r}"
+        )
 
-    next_node_id = index.next_node_id(frame.node_id, step_result.outcome)
-
-    append_step_result_trace(
-        run,
+    return _complete_step(
+        run=run,
+        index=index,
+        outcome=step_result.outcome,
         frame_id=frame.id,
         node_id=frame.node_id,
         step_type=step.type,
-        next_node_id=next_node_id,
-        result=step_result,
+        step_result=step_result,
     )
-    advance_frame(
-        run,
-        frame,
+
+
+async def step_workflow_async(
+    workflow: Workflow,
+    run: RunState,
+    registry: Mapping[str, AsyncNodeHandler],
+    *,
+    index: WorkflowIndex | None = None,
+) -> RunState:
+    prepared = _prepare_step(workflow, run, index)
+    if prepared is None:
+        return run
+    index, step = prepared
+    frame = run.current_frame()
+
+    if isinstance(step, NodeUse):
+        node_def = index.node_defs[step.node]
+        step_result = await execute_node_use_async(
+            workflow, run, step, node_def, registry
+        )
+    elif isinstance(step, ConditionNode):
+        step_result = handle_condition_step(run, step)
+    elif isinstance(step, JoinNode):
+        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, index)
+    else:
+        raise WorkflowExecutionError(
+            f"unsupported step type {getattr(step, 'type', type(step).__name__)!r}"
+        )
+
+    return _complete_step(
+        run=run,
+        index=index,
         outcome=step_result.outcome,
-        next_node_id=next_node_id,
+        frame_id=frame.id,
+        node_id=frame.node_id,
+        step_type=step.type,
+        step_result=step_result,
     )
-    return run