type wizards, in the folders

src/wf_authoring/__init__.py

@@ -20,7 +20,7 @@ from .ops import (
     length,
 )
 from .paths import GraphPath, context_path, graph_path, input_path, state_path
-from .spec import (
+from .nodes import (
     AsyncRegistryHandler,
     NodeReturn,
     NodeSpec,

src/wf_authoring/callables.py

@@ -0,0 +1,25 @@
+from .nodes.callables import (
+    AsyncContextNodeCallable,
+    AsyncNodeCallable,
+    AsyncPlainNodeCallable,
+    AsyncRegistryHandler,
+    ContextNodeCallable,
+    InputT,
+    NodeCallable,
+    OutputT,
+    PlainNodeCallable,
+    SyncRegistryHandler,
+)
+
+__all__ = [
+    "AsyncContextNodeCallable",
+    "AsyncNodeCallable",
+    "AsyncPlainNodeCallable",
+    "AsyncRegistryHandler",
+    "ContextNodeCallable",
+    "InputT",
+    "NodeCallable",
+    "OutputT",
+    "PlainNodeCallable",
+    "SyncRegistryHandler",
+]

src/wf_authoring/inference.py

@@ -0,0 +1,7 @@
+from .nodes.inference import accepts_context, infer_models, is_basemodel_subclass
+
+__all__ = [
+    "accepts_context",
+    "infer_models",
+    "is_basemodel_subclass",
+]

src/wf_authoring/nodes/__init__.py

@@ -0,0 +1,36 @@
+from .callables import (
+    AsyncContextNodeCallable,
+    AsyncNodeCallable,
+    AsyncPlainNodeCallable,
+    AsyncRegistryHandler,
+    ContextNodeCallable,
+    InputT,
+    NodeCallable,
+    OutputT,
+    PlainNodeCallable,
+    SyncRegistryHandler,
+)
+from .inference import accepts_context, infer_models, is_basemodel_subclass
+from .result import NodeReturn
+from .spec import NodeSpec, build_async_registry, build_registry, node
+
+__all__ = [
+    "AsyncContextNodeCallable",
+    "AsyncNodeCallable",
+    "AsyncPlainNodeCallable",
+    "AsyncRegistryHandler",
+    "ContextNodeCallable",
+    "InputT",
+    "NodeCallable",
+    "NodeReturn",
+    "NodeSpec",
+    "OutputT",
+    "PlainNodeCallable",
+    "SyncRegistryHandler",
+    "accepts_context",
+    "build_async_registry",
+    "build_registry",
+    "infer_models",
+    "is_basemodel_subclass",
+    "node",
+]

src/wf_authoring/nodes/callables.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any, Protocol, TypeVar
+
+from pydantic import BaseModel
+
+from wf_core import RuntimeContext
+
+from .result import NodeReturn
+
+InputT = TypeVar("InputT", bound=BaseModel)
+OutputT = TypeVar("OutputT", bound=BaseModel)
+
+InputT_contra = TypeVar("InputT_contra", bound=BaseModel, contravariant=True)
+OutputT_co = TypeVar("OutputT_co", bound=BaseModel, covariant=True)
+
+
+class ContextNodeCallable(Protocol[InputT_contra, OutputT_co]):
+    def __call__(
+        self,
+        payload: InputT_contra,
+        /,
+        ctx: RuntimeContext,
+    ) -> NodeReturn[OutputT_co] | OutputT_co: ...
+
+
+class PlainNodeCallable(Protocol[InputT_contra, OutputT_co]):
+    def __call__(
+        self,
+        payload: InputT_contra,
+        /,
+    ) -> NodeReturn[OutputT_co] | OutputT_co: ...
+
+
+NodeCallable = ContextNodeCallable[InputT, OutputT] | PlainNodeCallable[
+    InputT, OutputT
+]
+
+
+class AsyncContextNodeCallable(Protocol[InputT_contra, OutputT_co]):
+    def __call__(
+        self,
+        payload: InputT_contra,
+        /,
+        ctx: RuntimeContext,
+    ) -> Awaitable[NodeReturn[OutputT_co] | OutputT_co]: ...
+
+
+class AsyncPlainNodeCallable(Protocol[InputT_contra, OutputT_co]):
+    def __call__(
+        self,
+        payload: InputT_contra,
+        /,
+    ) -> Awaitable[NodeReturn[OutputT_co] | OutputT_co]: ...
+
+
+AsyncNodeCallable = AsyncContextNodeCallable[
+    InputT, OutputT
+] | AsyncPlainNodeCallable[InputT, OutputT]
+
+SyncRegistryHandler = Callable[[dict[str, Any], RuntimeContext], dict[str, Any]]
+AsyncRegistryHandler = Callable[
+    [dict[str, Any], RuntimeContext], Awaitable[dict[str, Any]]
+]

src/wf_authoring/nodes/inference.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from inspect import Parameter, signature
+from collections.abc import Callable
+from typing import cast, get_args, get_origin, get_type_hints
+
+from pydantic import BaseModel
+
+from wf_core import RuntimeContext
+
+from .result import NodeReturn
+
+
+def is_basemodel_subclass(value: object) -> bool:
+    return isinstance(value, type) and issubclass(value, BaseModel)
+
+
+def infer_models(fn: Callable[..., object]) -> tuple[type[BaseModel], type[BaseModel]]:
+    hints = get_type_hints(fn, include_extras=True)
+    params = list(signature(fn).parameters.values())
+    if len(params) not in {1, 2}:
+        raise TypeError("node function must accept (payload) or (payload, ctx)")
+
+    payload_param = params[0]
+
+    if payload_param.kind not in (
+        Parameter.POSITIONAL_ONLY,
+        Parameter.POSITIONAL_OR_KEYWORD,
+    ):
+        raise TypeError("node payload parameter must be positional")
+
+    input_model = hints.get(payload_param.name)
+    if not is_basemodel_subclass(input_model):
+        raise TypeError("node payload annotation must be a pydantic BaseModel subclass")
+
+    if len(params) == 2:
+        ctx_param = params[1]
+        if ctx_param.kind not in (
+            Parameter.POSITIONAL_ONLY,
+            Parameter.POSITIONAL_OR_KEYWORD,
+        ):
+            raise TypeError("node context parameter must be positional")
+        ctx_type = hints.get(ctx_param.name)
+        if ctx_type is not RuntimeContext:
+            raise TypeError("node context annotation must be wf_core.RuntimeContext")
+
+    return_type = hints.get("return")
+    if return_type is None:
+        raise TypeError("node function must declare a return annotation")
+
+    if is_basemodel_subclass(return_type):
+        return cast(type[BaseModel], input_model), cast(type[BaseModel], return_type)
+
+    origin = get_origin(return_type)
+    if origin is NodeReturn:
+        args = get_args(return_type)
+        if len(args) != 1 or not is_basemodel_subclass(args[0]):
+            raise TypeError(
+                "NodeReturn return annotation must wrap a BaseModel subclass"
+            )
+        return cast(type[BaseModel], input_model), cast(type[BaseModel], args[0])
+
+    raise TypeError(
+        "node return annotation must be a BaseModel subclass or NodeReturn[BaseModel]"
+    )
+
+
+def accepts_context(fn: Callable[..., object]) -> bool:
+    return len(signature(fn).parameters) >= 2

src/wf_authoring/nodes/result.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel
+
+OutputT_co = TypeVar("OutputT_co", bound=BaseModel, covariant=True)
+
+
+@dataclass(frozen=True, slots=True)
+class NodeReturn(Generic[OutputT_co]):
+    outcome: str
+    output: OutputT_co

src/wf_authoring/nodes/spec.py

@@ -0,0 +1,247 @@
+from __future__ import annotations
+
+from inspect import iscoroutinefunction
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import (
+    Any,
+    Generic,
+    Literal,
+    cast,
+    overload,
+)
+
+from pydantic import BaseModel
+
+from wf_core import NodeDef, RuntimeContext, SchemaRef
+
+from .callables import (
+    AsyncNodeCallable,
+    AsyncRegistryHandler,
+    ContextNodeCallable,
+    InputT,
+    NodeCallable,
+    OutputT,
+    PlainNodeCallable,
+    SyncRegistryHandler,
+)
+from .inference import accepts_context, infer_models
+from .result import NodeReturn
+
+
+def _schema_ref_for(model_type: type[BaseModel]) -> SchemaRef:
+    return SchemaRef.model_validate(model_type.model_json_schema())
+
+
+def _default_outcome(spec: "NodeSpec[Any, Any]") -> str:
+    return spec.outcomes[0]
+
+
+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}")
+
+
+@dataclass(slots=True)
+class NodeSpec(Generic[InputT, OutputT]):
+    name: str
+    input_model: type[InputT]
+    output_model: type[OutputT]
+    outcomes: tuple[str, ...]
+    fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]
+    description: str | None = None
+    is_async: bool = False
+    accepts_context: bool = True
+
+    def __call__(
+        self,
+        payload: InputT,
+        ctx: RuntimeContext | None = None,
+    ) -> NodeReturn[OutputT] | OutputT | Awaitable[NodeReturn[OutputT] | OutputT]:
+        if self.accepts_context:
+            if ctx is None:
+                raise TypeError(f"node {self.name!r} requires RuntimeContext")
+            return cast(ContextNodeCallable[InputT, OutputT], self.fn)(payload, ctx)
+        return cast(PlainNodeCallable[InputT, OutputT], self.fn)(payload)
+
+    def to_node_def(self) -> NodeDef:
+        return NodeDef(
+            name=self.name,
+            input_schema=_schema_ref_for(self.input_model),
+            output_schema=_schema_ref_for(self.output_model),
+            outcomes=list(self.outcomes),
+        )
+
+    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"
+            )
+
+        def handler(payload: dict[str, Any], ctx: RuntimeContext) -> dict[str, Any]:
+            parsed = self.input_model.model_validate(payload)
+            raw = self(parsed, ctx)
+            return _coerce_registry_result(
+                node_name=self.name,
+                output_model=self.output_model,
+                default_outcome=_default_outcome(self),
+                raw=cast(NodeReturn[BaseModel] | BaseModel, raw),
+            )
+
+        return handler
+
+    def to_async_registry_handler(self) -> AsyncRegistryHandler:
+        async def handler(
+            payload: dict[str, Any],
+            ctx: RuntimeContext,
+        ) -> dict[str, Any]:
+            parsed = self.input_model.model_validate(payload)
+            raw_result = self(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=_default_outcome(self),
+                raw=cast(NodeReturn[BaseModel] | BaseModel, raw),
+            )
+
+        return handler
+
+
+@overload
+def node(
+    fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT],
+    /,
+) -> NodeSpec[InputT, OutputT]: ...
+
+
+@overload
+def node(
+    fn: None = None,
+    /,
+) -> Callable[
+    [NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]],
+    NodeSpec[InputT, OutputT],
+]: ...
+
+
+@overload
+def node(
+    fn: None = None,
+    /,
+    *,
+    name: str | None = None,
+    input_model: type[InputT] | None = None,
+    output_model: type[OutputT] | None = None,
+    outcomes: tuple[str, ...] = ("ok",),
+    description: str | None = None,
+    is_async: bool | None = None,
+) -> Callable[
+    [NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]],
+    NodeSpec[InputT, OutputT],
+]: ...
+
+
+def node(
+    fn: NodeCallable[InputT, OutputT]
+    | AsyncNodeCallable[InputT, OutputT]
+    | None = None,
+    *,
+    name: str | None = None,
+    input_model: type[InputT] | None = None,
+    output_model: type[OutputT] | None = None,
+    outcomes: tuple[str, ...] = ("ok",),
+    description: str | None = None,
+    is_async: bool | None = None,
+) -> Any:
+    def decorator(
+        fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT],
+    ) -> NodeSpec[InputT, OutputT]:
+        inferred_input_model: type[BaseModel] | None = input_model
+        inferred_output_model: type[BaseModel] | None = output_model
+        if inferred_input_model is None or inferred_output_model is None:
+            inferred_input_model, inferred_output_model = infer_models(fn)
+
+        resolved_name = name or getattr(fn, "__name__", "node")
+        resolved_is_async = iscoroutinefunction(fn) if is_async is None else is_async
+        resolved_accepts_context = accepts_context(fn)
+        return cast(
+            NodeSpec[InputT, OutputT],
+            NodeSpec(
+                name=resolved_name,
+                input_model=inferred_input_model,
+                output_model=inferred_output_model,
+                outcomes=outcomes,
+                fn=cast(Any, fn),
+                description=description or fn.__doc__,
+                is_async=resolved_is_async,
+                accepts_context=resolved_accepts_context,
+            ),
+        )
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator
+
+
+def build_registry(
+    *specs: NodeSpec[Any, Any],
+) -> dict[str, SyncRegistryHandler]:
+    return _build_registry(specs, export="sync")
+
+
+def build_async_registry(
+    *specs: NodeSpec[Any, Any],
+) -> dict[str, AsyncRegistryHandler]:
+    return _build_registry(specs, export="async")
+
+
+@overload
+def _build_registry(
+    specs: tuple[NodeSpec[Any, Any], ...],
+    *,
+    export: Literal["sync"],
+) -> dict[str, SyncRegistryHandler]: ...
+
+
+@overload
+def _build_registry(
+    specs: tuple[NodeSpec[Any, Any], ...],
+    *,
+    export: Literal["async"],
+) -> dict[str, AsyncRegistryHandler]: ...
+
+
+def _build_registry(
+    specs: tuple[NodeSpec[Any, Any], ...],
+    *,
+    export: Literal["sync", "async"],
+) -> dict[str, Any]:
+    if export == "sync":
+        return {spec.name: spec.to_registry_handler() for spec in specs}
+    if export == "async":
+        return {spec.name: spec.to_async_registry_handler() for spec in specs}
+    raise ValueError(f"unknown registry export mode {export!r}")

src/wf_authoring/ops/__init__.py

@@ -0,0 +1,33 @@
+from .sequences import (
+    BoolOutput,
+    CountOutput,
+    ItemOutput,
+    MaybeItemOutput,
+    SequenceInput,
+    first_item,
+    first_item_maybe,
+    first_item_or_none,
+    is_empty,
+    last_item,
+    last_item_or_none,
+    length,
+)
+from .values import CoalesceInput, ValueOutput, coalesce
+
+__all__ = [
+    "BoolOutput",
+    "CoalesceInput",
+    "CountOutput",
+    "ItemOutput",
+    "MaybeItemOutput",
+    "SequenceInput",
+    "ValueOutput",
+    "coalesce",
+    "first_item",
+    "first_item_maybe",
+    "first_item_or_none",
+    "is_empty",
+    "last_item",
+    "last_item_or_none",
+    "length",
+]

src/wf_authoring/ops/sequences.py

@@ -4,7 +4,7 @@ from typing import Any
 
 from pydantic import BaseModel
 
-from .spec import NodeReturn, node
+from wf_authoring.nodes import NodeReturn, node
 
 
 class SequenceInput(BaseModel):
@@ -27,15 +27,6 @@ class BoolOutput(BaseModel):
     value: bool
 
 
-class CoalesceInput(BaseModel):
-    value: Any | None = None
-    fallback: Any
-
-
-class ValueOutput(BaseModel):
-    value: Any
-
-
 @node(
     name="authoring.first_item",
     input_model=SequenceInput,
@@ -63,9 +54,7 @@ def first_item_or_none(input: SequenceInput) -> ItemOutput:
     input_model=SequenceInput,
     output_model=MaybeItemOutput,
     outcomes=("found", "missing"),
-    description=(
+    description="Select the first item from a sequence, routing to found or missing.",
-        "Select the first item from a sequence, routing to found or missing."
-    ),
 )
 def first_item_maybe(input: SequenceInput) -> NodeReturn[MaybeItemOutput]:
     if not input.items:
@@ -113,13 +102,3 @@ def length(input: SequenceInput) -> CountOutput:
 )
 def is_empty(input: SequenceInput) -> BoolOutput:
     return BoolOutput(value=not input.items)
-
-
-@node(
-    name="authoring.coalesce",
-    input_model=CoalesceInput,
-    output_model=ValueOutput,
-    description="Return value when it is not None, otherwise return fallback.",
-)
-def coalesce(input: CoalesceInput) -> ValueOutput:
-    return ValueOutput(value=input.value if input.value is not None else input.fallback)

src/wf_authoring/ops/values.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel
+
+from wf_authoring.nodes import node
+
+
+class CoalesceInput(BaseModel):
+    value: Any | None = None
+    fallback: Any
+
+
+class ValueOutput(BaseModel):
+    value: Any
+
+
+@node(
+    name="authoring.coalesce",
+    input_model=CoalesceInput,
+    output_model=ValueOutput,
+    description="Return value when it is not None, otherwise return fallback.",
+)
+def coalesce(input: CoalesceInput) -> ValueOutput:
+    return ValueOutput(value=input.value if input.value is not None else input.fallback)

src/wf_authoring/spec.py

@@ -1,353 +1,35 @@
-from __future__ import annotations
+from .nodes import (
-
+    AsyncContextNodeCallable,
-from inspect import Parameter, iscoroutinefunction, signature
+    AsyncNodeCallable,
-from collections.abc import Awaitable, Callable
+    AsyncPlainNodeCallable,
-from dataclasses import dataclass
+    AsyncRegistryHandler,
-from typing import (
+    ContextNodeCallable,
-    Any,
+    InputT,
-    Generic,
+    NodeCallable,
-    Literal,
+    NodeReturn,
-    Protocol,
+    NodeSpec,
-    TypeVar,
+    OutputT,
-    cast,
+    PlainNodeCallable,
-    get_args,
+    SyncRegistryHandler,
-    get_origin,
+    build_async_registry,
-    get_type_hints,
+    build_registry,
-    overload,
+    node,
 )
 
-from pydantic import BaseModel
+__all__ = [
-
+    "AsyncContextNodeCallable",
-from wf_core import NodeDef, RuntimeContext, SchemaRef
+    "AsyncNodeCallable",
-
+    "AsyncPlainNodeCallable",
-InputT = TypeVar("InputT", bound=BaseModel)
+    "AsyncRegistryHandler",
-OutputT = TypeVar("OutputT", bound=BaseModel)
+    "ContextNodeCallable",
-
+    "InputT",
-
+    "NodeCallable",
-class ContextNodeCallable(Protocol[InputT, OutputT]):
+    "NodeReturn",
-    def __call__(
+    "NodeSpec",
-        self,
+    "OutputT",
-        payload: InputT,
+    "PlainNodeCallable",
-        /,
+    "SyncRegistryHandler",
-        ctx: RuntimeContext,
+    "build_async_registry",
-    ) -> "NodeReturn[OutputT] | OutputT": ...
+    "build_registry",
-
+    "node",
-
-class PlainNodeCallable(Protocol[InputT, OutputT]):
-    def __call__(self, payload: InputT, /) -> "NodeReturn[OutputT] | OutputT": ...
-
-
-NodeCallable = ContextNodeCallable[InputT, OutputT] | PlainNodeCallable[
-    InputT, OutputT
 ]
-
-
-class AsyncContextNodeCallable(Protocol[InputT, OutputT]):
-    def __call__(
-        self,
-        payload: InputT,
-        /,
-        ctx: RuntimeContext,
-    ) -> Awaitable["NodeReturn[OutputT] | OutputT"]: ...
-
-
-class AsyncPlainNodeCallable(Protocol[InputT, OutputT]):
-    def __call__(
-        self,
-        payload: InputT,
-        /,
-    ) -> Awaitable["NodeReturn[OutputT] | OutputT"]: ...
-
-
-AsyncNodeCallable = AsyncContextNodeCallable[
-    InputT, OutputT
-] | AsyncPlainNodeCallable[InputT, 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:
-    return SchemaRef.model_validate(model_type.model_json_schema())
-
-
-@dataclass(slots=True)
-class NodeReturn(Generic[OutputT]):
-    outcome: str
-    output: OutputT
-
-
-def _default_outcome(spec: "NodeSpec[Any, Any]") -> str:
-    return spec.outcomes[0]
-
-
-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)
-
-
-def _infer_models(
-    fn: Callable[..., object],
-) -> tuple[type[BaseModel], type[BaseModel]]:
-    hints = get_type_hints(fn, include_extras=True)
-    params = list(signature(fn).parameters.values())
-    if len(params) not in {1, 2}:
-        raise TypeError("node function must accept (payload) or (payload, ctx)")
-
-    payload_param = params[0]
-
-    if payload_param.kind not in (
-        Parameter.POSITIONAL_ONLY,
-        Parameter.POSITIONAL_OR_KEYWORD,
-    ):
-        raise TypeError("node payload parameter must be positional")
-
-    input_model = hints.get(payload_param.name)
-    if not _is_basemodel_subclass(input_model):
-        raise TypeError("node payload annotation must be a pydantic BaseModel subclass")
-
-    if len(params) == 2:
-        ctx_param = params[1]
-        if ctx_param.kind not in (
-            Parameter.POSITIONAL_ONLY,
-            Parameter.POSITIONAL_OR_KEYWORD,
-        ):
-            raise TypeError("node context parameter must be positional")
-        ctx_type = hints.get(ctx_param.name)
-        if ctx_type is not RuntimeContext:
-            raise TypeError("node context annotation must be wf_core.RuntimeContext")
-
-    return_type = hints.get("return")
-    if return_type is None:
-        raise TypeError("node function must declare a return annotation")
-
-    if _is_basemodel_subclass(return_type):
-        return cast(type[BaseModel], input_model), cast(type[BaseModel], return_type)
-
-    origin = get_origin(return_type)
-    if origin is NodeReturn:
-        args = get_args(return_type)
-        if len(args) != 1 or not _is_basemodel_subclass(args[0]):
-            raise TypeError(
-                "NodeReturn return annotation must wrap a BaseModel subclass"
-            )
-        return cast(type[BaseModel], input_model), cast(type[BaseModel], args[0])
-
-    raise TypeError(
-        "node return annotation must be a BaseModel subclass or NodeReturn[BaseModel]"
-    )
-
-
-def _accepts_context(fn: Callable[..., object]) -> bool:
-    return len(signature(fn).parameters) >= 2
-
-
-@dataclass(slots=True)
-class NodeSpec(Generic[InputT, OutputT]):
-    name: str
-    input_model: type[InputT]
-    output_model: type[OutputT]
-    outcomes: tuple[str, ...]
-    fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]
-    description: str | None = None
-    is_async: bool = False
-    accepts_context: bool = True
-
-    def __call__(
-        self,
-        payload: InputT,
-        ctx: RuntimeContext | None = None,
-    ) -> NodeReturn[OutputT] | OutputT | Awaitable[NodeReturn[OutputT] | OutputT]:
-        if self.accepts_context:
-            if ctx is None:
-                raise TypeError(f"node {self.name!r} requires RuntimeContext")
-            return cast(ContextNodeCallable[InputT, OutputT], self.fn)(payload, ctx)
-        return cast(PlainNodeCallable[InputT, OutputT], self.fn)(payload)
-
-    def to_node_def(self) -> NodeDef:
-        return NodeDef(
-            name=self.name,
-            input_schema=_schema_ref_for(self.input_model),
-            output_schema=_schema_ref_for(self.output_model),
-            outcomes=list(self.outcomes),
-        )
-
-    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"
-            )
-
-        def handler(payload: dict[str, Any], ctx: RuntimeContext) -> dict[str, Any]:
-            parsed = self.input_model.model_validate(payload)
-            raw = self(parsed, ctx)
-            return _coerce_registry_result(
-                node_name=self.name,
-                output_model=self.output_model,
-                default_outcome=_default_outcome(self),
-                raw=cast(NodeReturn[BaseModel] | BaseModel, raw),
-            )
-
-        return handler
-
-    def to_async_registry_handler(self) -> AsyncRegistryHandler:
-        async def handler(
-            payload: dict[str, Any],
-            ctx: RuntimeContext,
-        ) -> dict[str, Any]:
-            parsed = self.input_model.model_validate(payload)
-            raw_result = self(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=_default_outcome(self),
-                raw=cast(NodeReturn[BaseModel] | BaseModel, raw),
-            )
-
-        return handler
-
-
-@overload
-def node(
-    fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT],
-    /,
-) -> NodeSpec[InputT, OutputT]: ...
-
-
-@overload
-def node(
-    fn: None = None,
-    /,
-) -> Callable[
-    [NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]],
-    NodeSpec[InputT, OutputT],
-]: ...
-
-
-@overload
-def node(
-    fn: None = None,
-    /,
-    *,
-    name: str | None = None,
-    input_model: type[InputT] | None = None,
-    output_model: type[OutputT] | None = None,
-    outcomes: tuple[str, ...] = ("ok",),
-    description: str | None = None,
-    is_async: bool | None = None,
-) -> Callable[
-    [NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]],
-    NodeSpec[InputT, OutputT],
-]: ...
-
-
-def node(
-    fn: NodeCallable[InputT, OutputT]
-    | AsyncNodeCallable[InputT, OutputT]
-    | None = None,
-    *,
-    name: str | None = None,
-    input_model: type[InputT] | None = None,
-    output_model: type[OutputT] | None = None,
-    outcomes: tuple[str, ...] = ("ok",),
-    description: str | None = None,
-    is_async: bool | None = None,
-) -> Any:
-    def decorator(
-        fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT],
-    ) -> NodeSpec[InputT, OutputT]:
-        inferred_input_model: type[BaseModel] | None = input_model
-        inferred_output_model: type[BaseModel] | None = output_model
-        if inferred_input_model is None or inferred_output_model is None:
-            inferred_input_model, inferred_output_model = _infer_models(fn)
-
-        resolved_name = name or getattr(fn, "__name__", "node")
-        resolved_is_async = iscoroutinefunction(fn) if is_async is None else is_async
-        resolved_accepts_context = _accepts_context(fn)
-        return cast(
-            NodeSpec[InputT, OutputT],
-            NodeSpec(
-                name=resolved_name,
-                input_model=inferred_input_model,
-                output_model=inferred_output_model,
-                outcomes=outcomes,
-                fn=cast(Any, fn),
-                description=description or fn.__doc__,
-                is_async=resolved_is_async,
-                accepts_context=resolved_accepts_context,
-            ),
-        )
-
-    if fn is not None:
-        return decorator(fn)
-    return decorator
-
-
-def build_registry(
-    *specs: NodeSpec[Any, Any],
-) -> dict[str, SyncRegistryHandler]:
-    return _build_registry(specs, export="sync")
-
-
-def build_async_registry(
-    *specs: NodeSpec[Any, Any],
-) -> dict[str, AsyncRegistryHandler]:
-    return _build_registry(specs, export="async")
-
-
-@overload
-def _build_registry(
-    specs: tuple[NodeSpec[Any, Any], ...],
-    *,
-    export: Literal["sync"],
-) -> dict[str, SyncRegistryHandler]: ...
-
-
-@overload
-def _build_registry(
-    specs: tuple[NodeSpec[Any, Any], ...],
-    *,
-    export: Literal["async"],
-) -> dict[str, AsyncRegistryHandler]: ...
-
-
-def _build_registry(
-    specs: tuple[NodeSpec[Any, Any], ...],
-    *,
-    export: Literal["sync", "async"],
-) -> dict[str, Any]:
-    if export == "sync":
-        return {spec.name: spec.to_registry_handler() for spec in specs}
-    if export == "async":
-        return {spec.name: spec.to_async_registry_handler() for spec in specs}
-    raise ValueError(f"unknown registry export mode {export!r}")