ReducerDefinition

src/wf_core/runtime/ops/merges.py

@@ -1,25 +1,58 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Mapping
-from typing import Any
+from dataclasses import dataclass
+from typing import Any, cast
 
 from wf_core.errors import WorkflowExecutionError
 from wf_core.models.reducers import ReducerRef, ReducerSpec
 from wf_core.runtime.ops.schemas import validate_payload_against_schema
 
-Reducer = Callable[[Any, Any, Mapping[str, Any]], Any]
+PlainReducer = Callable[[Any, Any], Any]
+ConfigReducer = Callable[[Any, Any, Mapping[str, Any]], Any]
+Reducer = PlainReducer | ConfigReducer
 
 
-def replace_reducer(
-    _current_value: Any, incoming_value: Any, _config: Mapping[str, Any]
+@dataclass(frozen=True, slots=True)
+class ReducerDefinition:
+    """Runtime reducer implementation paired with its inspectable spec."""
+
+    spec: ReducerSpec
+    fn: Reducer
+    accepts_config: bool = False
+
+    def apply(
+        self,
+        *,
+        reducer: ReducerRef,
+        current_value: Any,
+        incoming_value: Any,
+        destination_path: str,
     ) -> Any:
+        """Validate config, then apply the pure reducer function."""
+        validate_payload_against_schema(
+            self.spec.config_schema,
+            reducer.config,
+            f"reducer config for {self.spec.name!r}",
+        )
+        try:
+            if self.accepts_config:
+                return cast(ConfigReducer, self.fn)(
+                    current_value,
+                    incoming_value,
+                    reducer.config,
+                )
+            return cast(PlainReducer, self.fn)(current_value, incoming_value)
+        except TypeError as exc:
+            raise WorkflowExecutionError(f"{exc} at {destination_path!r}") from exc
+
+
+def replace_reducer(_current_value: Any, incoming_value: Any) -> Any:
     """Replace the current state value with the incoming value."""
     return incoming_value
 
 
-def append_reducer(
-    current_value: Any, incoming_value: Any, _config: Mapping[str, Any]
-) -> Any:
+def append_reducer(current_value: Any, incoming_value: Any) -> Any:
     """Append one value or many values into a list-valued state path."""
     if current_value is None:
         return (
@@ -34,9 +67,7 @@ def append_reducer(
     )
 
 
-def merge_object_reducer(
-    current_value: Any, incoming_value: Any, _config: Mapping[str, Any]
-) -> Any:
+def merge_object_reducer(current_value: Any, incoming_value: Any) -> Any:
     """Shallow-merge object values at one exact state path."""
     if current_value is None:
         if not isinstance(incoming_value, dict):
@@ -47,9 +78,7 @@ def merge_object_reducer(
     return current_value | incoming_value
 
 
-def set_union_reducer(
-    current_value: Any, incoming_value: Any, _config: Mapping[str, Any]
-) -> Any:
+def set_union_reducer(current_value: Any, incoming_value: Any) -> Any:
     """Merge list values while preserving stable first-seen order."""
     if current_value is None:
         current_items: list[Any] = []
@@ -68,25 +97,49 @@ def set_union_reducer(
     return merged
 
 
-def max_reducer(
-    current_value: Any, incoming_value: Any, _config: Mapping[str, Any]
-) -> Any:
+def max_reducer(current_value: Any, incoming_value: Any) -> Any:
     """Keep the larger of the current and incoming values."""
     return (
         incoming_value if current_value is None else max(current_value, incoming_value)
     )
 
 
-DEFAULT_REDUCERS: Mapping[str, Reducer] = {
-    "wf.std.replace": replace_reducer,
-    "wf.std.append": append_reducer,
-    "wf.std.merge_object": merge_object_reducer,
-    "wf.std.set_union": set_union_reducer,
-    "wf.std.max": max_reducer,
-}
-
-DEFAULT_REDUCER_SPECS: Mapping[str, ReducerSpec] = {
-    name: ReducerSpec(name=name) for name in DEFAULT_REDUCERS
+DEFAULT_REDUCER_DEFINITIONS: Mapping[str, ReducerDefinition] = {
+    "wf.std.replace": ReducerDefinition(
+        spec=ReducerSpec(
+            name="wf.std.replace",
+            description="Replace the current state value with the incoming value.",
+        ),
+        fn=replace_reducer,
+    ),
+    "wf.std.append": ReducerDefinition(
+        spec=ReducerSpec(
+            name="wf.std.append",
+            description="Append one value or many values into a list-valued state path.",
+        ),
+        fn=append_reducer,
+    ),
+    "wf.std.merge_object": ReducerDefinition(
+        spec=ReducerSpec(
+            name="wf.std.merge_object",
+            description="Shallow-merge object values at one exact state path.",
+        ),
+        fn=merge_object_reducer,
+    ),
+    "wf.std.set_union": ReducerDefinition(
+        spec=ReducerSpec(
+            name="wf.std.set_union",
+            description="Merge list values while preserving stable first-seen order.",
+        ),
+        fn=set_union_reducer,
+    ),
+    "wf.std.max": ReducerDefinition(
+        spec=ReducerSpec(
+            name="wf.std.max",
+            description="Keep the larger of the current and incoming values.",
+        ),
+        fn=max_reducer,
+    ),
 }
 
 
@@ -96,22 +149,15 @@ def apply_reducer(
     current_value: Any,
     incoming_value: Any,
     destination_path: str,
-    reducers: Mapping[str, Reducer] = DEFAULT_REDUCERS,
-    reducer_specs: Mapping[str, ReducerSpec] = DEFAULT_REDUCER_SPECS,
+    reducers: Mapping[str, ReducerDefinition] = DEFAULT_REDUCER_DEFINITIONS,
 ) -> Any:
     """Apply one named pure reducer to a state write."""
-    reducer_fn = reducers.get(reducer.name)
-    if reducer_fn is None:
+    definition = reducers.get(reducer.name)
+    if definition is None:
         raise WorkflowExecutionError(f"unknown reducer {reducer.name!r}")
-    spec = reducer_specs.get(reducer.name)
-    if spec is None:
-        raise WorkflowExecutionError(f"unknown reducer spec {reducer.name!r}")
-    validate_payload_against_schema(
-        spec.config_schema,
-        reducer.config,
-        f"reducer config for {reducer.name!r}",
+    return definition.apply(
+        reducer=reducer,
+        current_value=current_value,
+        incoming_value=incoming_value,
+        destination_path=destination_path,
     )
-    try:
-        return reducer_fn(current_value, incoming_value, reducer.config)
-    except TypeError as exc:
-        raise WorkflowExecutionError(f"{exc} at {destination_path!r}") from exc

src/wf_mcp/broker/service/builtins.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from typing import Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol
 
 from pydantic import BaseModel, Field
 
-from wf_core import ReducerSpec
+from wf_core.runtime.ops.merges import DEFAULT_REDUCER_DEFINITIONS
 from wf_authoring import (
     NodeReturn,
     NodeSpec,
@@ -35,6 +35,9 @@ from .capability_sources import (
 )
 from .specs import qualify_spec
 
+if TYPE_CHECKING:
+    from wf_core import ReducerSpec
+
 BUILTIN_CONNECTION_ID = "wf.std"
 """Internal source id for workflow standard-library node specs."""
 
@@ -112,29 +115,10 @@ def builtin_specs() -> dict[str, NodeSpec[Any, Any]]:
 
 def builtin_reducers() -> dict[str, ReducerSpec]:
     """Return built-in reducers owned by the workflow standard library."""
-    specs = (
-        ReducerSpec(
-            name="wf.std.replace",
-            description="Replace the current state value with the incoming value.",
-        ),
-        ReducerSpec(
-            name="wf.std.append",
-            description="Append one value or many values into a list-valued state path.",
-        ),
-        ReducerSpec(
-            name="wf.std.merge_object",
-            description="Shallow-merge object values at one exact state path.",
-        ),
-        ReducerSpec(
-            name="wf.std.max",
-            description="Keep the larger of the current and incoming values.",
-        ),
-        ReducerSpec(
-            name="wf.std.set_union",
-            description="Merge list values while preserving stable first-seen order.",
-        ),
-    )
-    return {spec.name: spec for spec in specs}
+    return {
+        definition.spec.name: definition.spec
+        for definition in DEFAULT_REDUCER_DEFINITIONS.values()
+    }
 
 
 def mcp_specs(service: ToolCaller) -> dict[str, NodeSpec[Any, Any]]:

tests/core/test_nested_state_paths.py

@@ -1,6 +1,14 @@
 from __future__ import annotations
 
-from wf_core import ReducerRef, SchemaRef, StateField, StateSchema, Workflow
+from wf_core import (
+    ReducerRef,
+    ReducerSpec,
+    SchemaRef,
+    StateField,
+    StateSchema,
+    Workflow,
+)
+from wf_core.runtime.ops.merges import ReducerDefinition, apply_reducer
 from wf_core.runtime.ops.state import write_state_value
 
 
@@ -48,9 +56,7 @@ def test_state_field_defaults_to_replace_reducer() -> None:
 
 
 def test_state_field_accepts_string_reducer_shorthand() -> None:
-    field = StateField.model_validate(
-        {"type": "array", "reducer": "wf.std.set_union"}
-    )
+    field = StateField.model_validate({"type": "array", "reducer": "wf.std.set_union"})
 
     assert field.reducer == ReducerRef(name="wf.std.set_union")
 
@@ -61,16 +67,15 @@ def test_state_field_accepts_configured_reducer_reference() -> None:
         reducer=ReducerRef(name="wf.std.max", config={"sample": True}),
     )
 
-    assert field.reducer.name == "wf.std.max"
-    assert field.reducer.config == {"sample": True}
+    # PYLINT!!!! what is u on ts is so clear
+    assert field.reducer.name == "wf.std.max"  # pylint: disable=no-member
+    assert field.reducer.config == {"sample": True}  # pylint: disable=no-member
 
 
 def test_unknown_state_reducer_fails_clearly() -> None:
     workflow = _workflow(
         fields={
-            "person.tags": StateField(
-                reducer=ReducerRef(name="x.nope"), type="array"
-            )
+            "person.tags": StateField(reducer=ReducerRef(name="x.nope"), type="array")
         }
     )
     state = {"person": {"tags": ["seed"]}}
@@ -133,6 +138,50 @@ def test_max_reducer_keeps_larger_value() -> None:
     assert state["best_score"] == 9
 
 
+def test_reducer_definition_can_wrap_plain_two_arg_callable() -> None:
+    definition = ReducerDefinition(
+        spec=ReducerSpec(name="test.add"),
+        fn=lambda current, incoming: (current or 0) + incoming,
+    )
+
+    result = apply_reducer(
+        reducer=ReducerRef(name="test.add"),
+        current_value=2,
+        incoming_value=3,
+        destination_path="state.total",
+        reducers={"test.add": definition},
+    )
+
+    assert result == 5
+
+
+def test_reducer_definition_can_wrap_config_aware_callable() -> None:
+    definition = ReducerDefinition(
+        spec=ReducerSpec(
+            name="test.modulo_add",
+            config_schema={
+                "type": "object",
+                "properties": {"modulus": {"type": "integer"}},
+                "required": ["modulus"],
+                "additionalProperties": False,
+            },
+        ),
+        fn=lambda current, incoming, config: ((current or 0) + incoming)
+        % config["modulus"],
+        accepts_config=True,
+    )
+
+    result = apply_reducer(
+        reducer=ReducerRef(name="test.modulo_add", config={"modulus": 10}),
+        current_value=8,
+        incoming_value=5,
+        destination_path="state.total",
+        reducers={"test.modulo_add": definition},
+    )
+
+    assert result == 3
+
+
 def _workflow(*, fields: dict[str, StateField]) -> Workflow:
     return Workflow(
         name="nested_state_paths",