This thing where

tests/test_demo_workflow.py

@@ -13,7 +13,7 @@ from wf_core import (
 )
 from wf_core.demo_workflow import build_demo_registry, build_demo_workflow
 from wf_core.run_factory import create_run_state
-from wf_authoring import WorkflowBuilder, node
+from wf_authoring import NodeReturn, WorkflowBuilder, build_registry, node
 
 
 class DriveListFilesInput(BaseModel):
@@ -106,11 +106,11 @@ def combine_summaries_spec(
 def send_email_spec(
     payload: SendEmailInput,
     ctx: RuntimeContext,
-) -> dict[str, object]:
-    return {
-        "outcome": "sent",
-        "output": {"email_status": f"sent: {payload.summary}"},
-    }
+) -> NodeReturn[SendEmailOutput]:
+    return NodeReturn(
+        outcome="sent",
+        output=SendEmailOutput(email_status=f"sent: {payload.summary}"),
+    )
 
 
 @node(
@@ -210,13 +210,13 @@ def build_authoring_demo_workflow():
     builder.connect("send_email", "sent", END)
     builder.connect("skip_email", "ok", END)
 
-    registry = {
-        drive_list_files_spec.name: drive_list_files_spec.to_registry_handler(),
-        summarize_document_spec.name: summarize_document_spec.to_registry_handler(),
-        combine_summaries_spec.name: combine_summaries_spec.to_registry_handler(),
-        send_email_spec.name: send_email_spec.to_registry_handler(),
-        mark_email_skipped_spec.name: mark_email_skipped_spec.to_registry_handler(),
-    }
+    registry = build_registry(
+        drive_list_files_spec,
+        summarize_document_spec,
+        combine_summaries_spec,
+        send_email_spec,
+        mark_email_skipped_spec,
+    )
     return builder.compile(), registry
 
 
@@ -368,3 +368,30 @@ def test_builder_compiled_workflow_executes_like_declared_demo() -> None:
     assert built_run.output == declared_run.output
     assert built_run.state == declared_run.state
     assert built_run.current_node_id == declared_run.current_node_id
+
+
+def test_async_node_spec_cannot_export_sync_registry_handler() -> None:
+    class AsyncInput(BaseModel):
+        value: str
+
+    class AsyncOutput(BaseModel):
+        echoed: str
+
+    @node(
+        name="async_echo",
+        input_model=AsyncInput,
+        output_model=AsyncOutput,
+        is_async=True,
+    )
+    async def async_echo(
+        payload: AsyncInput,
+        ctx: RuntimeContext,
+    ) -> AsyncOutput:
+        return AsyncOutput(echoed=payload.value)
+
+    try:
+        async_echo.to_registry_handler()
+    except TypeError as exc:
+        assert "async" in str(exc)
+    else:
+        raise AssertionError("expected async node export to fail for sync registry")

wf_authoring/__init__.py

@@ -1,11 +1,13 @@
 from .builder import WorkflowBuilder
 from .catalog import NodeCatalog, NodeCatalogEntry
-from .spec import NodeSpec, node
+from .spec import NodeReturn, NodeSpec, build_registry, node
 
 __all__ = [
     "NodeCatalog",
     "NodeCatalogEntry",
+    "NodeReturn",
     "NodeSpec",
     "WorkflowBuilder",
+    "build_registry",
     "node",
 ]

wf_authoring/spec.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from typing import Any, Generic, TypeVar
 
@@ -10,19 +10,29 @@ from wf_core import NodeDef, RuntimeContext, SchemaRef
 
 InputT = TypeVar("InputT", bound=BaseModel)
 OutputT = TypeVar("OutputT", bound=BaseModel)
+NodeCallable = Callable[[InputT, RuntimeContext], "NodeReturn[OutputT] | OutputT"]
+AsyncNodeCallable = Callable[
+    [InputT, RuntimeContext], Awaitable["NodeReturn[OutputT] | OutputT"]
+]
 
 
 def _schema_ref_for(model_type: type[BaseModel]) -> SchemaRef:
     return SchemaRef.model_validate(model_type.model_json_schema())
 
 
+@dataclass(slots=True)
+class NodeReturn(Generic[OutputT]):
+    outcome: str
+    output: OutputT
+
+
 @dataclass(slots=True)
 class NodeSpec(Generic[InputT, OutputT]):
     name: str
     input_model: type[InputT]
     output_model: type[OutputT]
     outcomes: tuple[str, ...]
-    fn: Callable[[InputT, RuntimeContext], OutputT | dict[str, Any]]
+    fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]
     description: str | None = None
     is_async: bool = False
 
@@ -30,7 +40,7 @@ class NodeSpec(Generic[InputT, OutputT]):
         self,
         payload: InputT,
         ctx: RuntimeContext,
-    ) -> OutputT | dict[str, Any]:
+    ) -> NodeReturn[OutputT] | OutputT | Awaitable[NodeReturn[OutputT] | OutputT]:
         return self.fn(payload, ctx)
 
     def to_node_def(self) -> NodeDef:
@@ -42,13 +52,26 @@ class NodeSpec(Generic[InputT, OutputT]):
         )
 
     def to_registry_handler(self) -> Callable[[dict[str, Any], RuntimeContext], dict[str, Any]]:
+        if self.is_async:
+            raise TypeError(
+                f"node {self.name!r} is async and cannot be exported to the sync registry"
+            )
+
         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):
+                if not isinstance(raw.output, self.output_model):
+                    raise TypeError(
+                        f"node {self.name!r} returned NodeReturn with unsupported output "
+                        f"{type(raw.output)!r}"
+                    )
+                return {
+                    "outcome": raw.outcome,
+                    "output": raw.output.model_dump(),
+                }
             if isinstance(raw, self.output_model):
                 return {"outcome": self.outcomes[0], "output": raw.model_dump()}
-            if isinstance(raw, dict):
-                return raw
             raise TypeError(
                 f"node {self.name!r} returned unsupported value {type(raw)!r}"
             )
@@ -63,12 +86,13 @@ def node(
     output_model: type[OutputT],
     outcomes: tuple[str, ...] = ("ok",),
     description: str | None = None,
+    is_async: bool = False,
 ) -> Callable[
-    [Callable[[InputT, RuntimeContext], OutputT | dict[str, Any]]],
+    [NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]],
     NodeSpec[InputT, OutputT],
 ]:
     def decorator(
-        fn: Callable[[InputT, RuntimeContext], OutputT | dict[str, Any]],
+        fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT],
     ) -> NodeSpec[InputT, OutputT]:
         resolved_name = name or getattr(fn, "__name__", "node")
         return NodeSpec(
@@ -78,7 +102,13 @@ def node(
             outcomes=outcomes,
             fn=fn,
             description=description or fn.__doc__,
-            is_async=False,
+            is_async=is_async,
         )
 
     return decorator
+
+
+def build_registry(
+    *specs: NodeSpec[Any, Any],
+) -> dict[str, Callable[[dict[str, Any], RuntimeContext], dict[str, Any]]]:
+    return {spec.name: spec.to_registry_handler() for spec in specs}