async ver, starting high level

tests/test_demo_workflow.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import asyncio
+
 from pydantic import BaseModel
 
 from wf_core import (
@@ -17,6 +19,7 @@ from wf_authoring import (
     NodeReturn,
     WorkflowBuilder,
     bind_fields,
+    build_async_registry,
     bind_state,
     build_registry,
     expr,
@@ -431,6 +434,34 @@ def test_async_node_spec_cannot_export_sync_registry_handler() -> None:
         raise AssertionError("expected async node export to fail for sync registry")
 
 
+def test_async_registry_accepts_sync_and_async_specs() -> None:
+    @node()
+    def sync_echo(
+        payload: InferredEchoInput,
+        ctx: RuntimeContext,
+    ) -> InferredEchoOutput:
+        return InferredEchoOutput(echoed=payload.value)
+
+    @node()
+    async def async_echo(
+        payload: InferredAsyncInput,
+        ctx: RuntimeContext,
+    ) -> InferredAsyncOutput:
+        return InferredAsyncOutput(echoed=f"async:{payload.value}")
+
+    registry = build_async_registry(sync_echo, async_echo)
+    ctx = RuntimeContext(current_node_id="demo")
+
+    async def run_handler(name: str, value: str) -> dict[str, object]:
+        return await registry[name]({"value": value}, ctx)
+
+    sync_result = asyncio.run(run_handler("sync_echo", "hello"))
+    async_result = asyncio.run(run_handler("async_echo", "world"))
+
+    assert sync_result == {"outcome": "ok", "output": {"echoed": "hello"}}
+    assert async_result == {"outcome": "ok", "output": {"echoed": "async:world"}}
+
+
 @node()
 def inferred_echo(
     payload: InferredEchoInput,

wf_authoring/__init__.py

@@ -3,7 +3,15 @@ from .catalog import NodeCatalog, NodeCatalogEntry
 from .conditions import context, exists, expr, input, state
 from .mapping import bind_fields, bind_state, merge_maps
 from .paths import GraphPath, context_path, graph_path, input_path, state_path
-from .spec import NodeReturn, NodeSpec, build_registry, node
+from .spec import (
+    AsyncRegistryHandler,
+    NodeReturn,
+    NodeSpec,
+    SyncRegistryHandler,
+    build_async_registry,
+    build_registry,
+    node,
+)
 from .subgraph import subgraph_node
 
 __all__ = [
@@ -12,8 +20,11 @@ __all__ = [
     "GraphPath",
     "NodeReturn",
     "NodeSpec",
+    "AsyncRegistryHandler",
+    "SyncRegistryHandler",
     "WorkflowBuilder",
     "bind_fields",
+    "build_async_registry",
     "build_registry",
     "bind_state",
     "merge_maps",

wf_authoring/spec.py

@@ -15,6 +15,10 @@ NodeCallable = Callable[[InputT, RuntimeContext], "NodeReturn[OutputT] | OutputT
 AsyncNodeCallable = Callable[
     [InputT, RuntimeContext], Awaitable["NodeReturn[OutputT] | OutputT"]
 ]
+SyncRegistryHandler = Callable[[dict[str, Any], RuntimeContext], dict[str, Any]]
+AsyncRegistryHandler = Callable[
+    [dict[str, Any], RuntimeContext], Awaitable[dict[str, Any]]
+]
 
 
 def _schema_ref_for(model_type: type[BaseModel]) -> SchemaRef:
@@ -27,6 +31,28 @@ class NodeReturn(Generic[OutputT]):
     output: OutputT
 
 
+def _coerce_registry_result(
+    *,
+    node_name: str,
+    output_model: type[BaseModel],
+    default_outcome: str,
+    raw: NodeReturn[BaseModel] | BaseModel,
+) -> dict[str, Any]:
+    if isinstance(raw, NodeReturn):
+        if not isinstance(raw.output, output_model):
+            raise TypeError(
+                f"node {node_name!r} returned NodeReturn with unsupported output "
+                f"{type(raw.output)!r}"
+            )
+        return {
+            "outcome": raw.outcome,
+            "output": raw.output.model_dump(),
+        }
+    if isinstance(raw, output_model):
+        return {"outcome": default_outcome, "output": raw.model_dump()}
+    raise TypeError(f"node {node_name!r} returned unsupported value {type(raw)!r}")
+
+
 def _is_basemodel_subclass(value: object) -> bool:
     return isinstance(value, type) and issubclass(value, BaseModel)
 
@@ -113,7 +139,7 @@ class NodeSpec(Generic[InputT, OutputT]):
             outcomes=list(self.outcomes),
         )
 
-    def to_registry_handler(self) -> Callable[[dict[str, Any], RuntimeContext], dict[str, Any]]:
+    def to_registry_handler(self) -> SyncRegistryHandler:
         if self.is_async:
             raise TypeError(
                 f"node {self.name!r} is async and cannot be exported to the sync registry"
@@ -122,20 +148,34 @@ class NodeSpec(Generic[InputT, OutputT]):
         def handler(payload: dict[str, Any], ctx: RuntimeContext) -> dict[str, Any]:
             parsed = self.input_model.model_validate(payload)
             raw = self.fn(parsed, ctx)
-            if isinstance(raw, NodeReturn):
+            return _coerce_registry_result(
-                if not isinstance(raw.output, self.output_model):
+                node_name=self.name,
-                    raise TypeError(
+                output_model=self.output_model,
-                        f"node {self.name!r} returned NodeReturn with unsupported output "
+                default_outcome=self.outcomes[0],
-                        f"{type(raw.output)!r}"
+                raw=cast(NodeReturn[BaseModel] | BaseModel, raw),
             )
-                return {
+
-                    "outcome": raw.outcome,
+        return handler
-                    "output": raw.output.model_dump(),
+
-                }
+    def to_async_registry_handler(self) -> AsyncRegistryHandler:
-            if isinstance(raw, self.output_model):
+        async def handler(
-                return {"outcome": self.outcomes[0], "output": raw.model_dump()}
+            payload: dict[str, Any],
-            raise TypeError(
+            ctx: RuntimeContext,
-                f"node {self.name!r} returned unsupported value {type(raw)!r}"
+        ) -> dict[str, Any]:
+            parsed = self.input_model.model_validate(payload)
+            raw_result = self.fn(parsed, ctx)
+            if self.is_async:
+                raw = await cast(
+                    Awaitable[NodeReturn[OutputT] | OutputT],
+                    raw_result,
+                )
+            else:
+                raw = cast(NodeReturn[OutputT] | OutputT, raw_result)
+            return _coerce_registry_result(
+                node_name=self.name,
+                output_model=self.output_model,
+                default_outcome=self.outcomes[0],
+                raw=cast(NodeReturn[BaseModel] | BaseModel, raw),
             )
 
         return handler
@@ -217,5 +257,11 @@ def node(
 
 def build_registry(
     *specs: NodeSpec[Any, Any],
-) -> dict[str, Callable[[dict[str, Any], RuntimeContext], dict[str, Any]]]:
+) -> dict[str, SyncRegistryHandler]:
     return {spec.name: spec.to_registry_handler() for spec in specs}
+
+
+def build_async_registry(
+    *specs: NodeSpec[Any, Any],
+) -> dict[str, AsyncRegistryHandler]:
+    return {spec.name: spec.to_async_registry_handler() for spec in specs}