This thing where

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}