and more good stuff, fmt

src/wf_authoring/builder.py

@@ -3,7 +3,8 @@ from __future__ import annotations
 from collections.abc import Mapping
 from dataclasses import dataclass, field
 import re
-from typing import Any, Literal, TypeAlias, cast
+from typing import Any, Literal, TypeAlias, TypeGuard, cast
+import warnings
 
 from wf_core import (
     ConditionNode,
@@ -14,6 +15,8 @@ from wf_core import (
     SchemaRef,
     StateSchema,
     Workflow,
+    RunState,
+    execute_workflow,
 )
 from wf_core.errors import WorkflowExecutionError
 from wf_core.model import Condition as CoreCondition
@@ -25,6 +28,7 @@ from .schemas import SchemaLike, StateSchemaLike, schema_ref_from, state_schema_
 from .spec import NodeSpec
 
 StepRef: TypeAlias = str | NodeUse | ConditionNode | ForeachNode | InterruptNode
+BranchRef: TypeAlias = StepRef | NodeSpec[Any, Any]
 MapArg: TypeAlias = Mapping[Any, Any]
 
 
@@ -53,6 +57,10 @@ def _step_id(ref: StepRef) -> str:
     return ref.id
 
 
+def _is_node_spec(ref: object) -> TypeGuard[NodeSpec[Any, Any]]:
+    return isinstance(ref, NodeSpec)
+
+
 def _slug_id(value: str) -> str:
     slug = re.sub(r"[^0-9A-Za-z_]+", "_", value).strip("_").lower()
     return slug or "step"
@@ -65,7 +73,9 @@ def _auto_input_map(
     state_schema: StateSchema,
 ) -> dict[str, str]:
     return {
-        _auto_source_path(field, input_schema=input_schema, state_schema=state_schema): field
+        _auto_source_path(
+            field, input_schema=input_schema, state_schema=state_schema
+        ): field
         for field in spec.input_model.model_json_schema().get("properties", {})
     }
 
@@ -165,6 +175,15 @@ class WorkflowBuilder:
         """Export handlers for all node specs used by this builder."""
         return build_registry(*self.node_specs.values())
 
+    def execute(self, workflow_input: dict[str, Any]) -> RunState:
+        """Compile and execute this workflow with its used node registry.
+
+        This is intended for tests, examples, and local authoring loops. Production
+        callers that need custom registries, persistence, or resume behavior should
+        call wf_core execution functions directly.
+        """
+        return execute_workflow(self.compile(), workflow_input, self.registry())
+
     def condition(self, *, id: str, check: CoreCondition | Expr) -> ConditionNode:
         node = ConditionNode(
             id=id,
@@ -223,6 +242,33 @@ class WorkflowBuilder:
             )
         )
 
+    def branch(
+        self,
+        from_: BranchRef,
+        branches: Mapping[str, BranchRef],
+    ) -> dict[str, StepRef]:
+        """Connect multiple outcomes from one branch source.
+
+        Passing a NodeSpec creates a node use with auto-mapping and an auto id.
+        Passing an existing step or id only wires edges. Empty branch maps are
+        allowed but warn because they usually indicate an unfinished router.
+        """
+        if not branches:
+            warnings.warn(
+                "WorkflowBuilder.branch called with no branches",
+                UserWarning,
+                stacklevel=2,
+            )
+            return {}
+
+        source = self.use(from_) if _is_node_spec(from_) else from_
+        resolved_targets: dict[str, StepRef] = {}
+        for outcome, target in branches.items():
+            resolved = self.use(target) if _is_node_spec(target) else target
+            self.connect(cast(StepRef, source), outcome, cast(StepRef, resolved))
+            resolved_targets[outcome] = cast(StepRef, resolved)
+        return resolved_targets
+
     def compile(self) -> Workflow:
         if self.start is None:
             raise WorkflowExecutionError(

src/wf_authoring/nodes/callables.py

@@ -33,9 +33,7 @@ class PlainNodeCallable(Protocol[InputT_contra, OutputT_co]):
     ) -> NodeReturn[OutputT_co] | OutputT_co: ...
 
 
-NodeCallable = ContextNodeCallable[InputT, OutputT] | PlainNodeCallable[
+NodeCallable = ContextNodeCallable[InputT, OutputT] | PlainNodeCallable[InputT, OutputT]
-    InputT, OutputT
-]
 
 
 class AsyncContextNodeCallable(Protocol[InputT_contra, OutputT_co]):
@@ -55,9 +53,9 @@ class AsyncPlainNodeCallable(Protocol[InputT_contra, OutputT_co]):
     ) -> Awaitable[NodeReturn[OutputT_co] | OutputT_co]: ...
 
 
-AsyncNodeCallable = AsyncContextNodeCallable[
+AsyncNodeCallable = (
-    InputT, OutputT
+    AsyncContextNodeCallable[InputT, OutputT] | AsyncPlainNodeCallable[InputT, OutputT]
-] | AsyncPlainNodeCallable[InputT, OutputT]
+)
 
 SyncRegistryHandler = Callable[[dict[str, Any], RuntimeContext], dict[str, Any]]
 AsyncRegistryHandler = Callable[

src/wf_mcp/names.py

@@ -43,6 +43,7 @@ def parse_namespaced_tool_name(
 def is_admin_tool_name(proxy_name: str) -> bool:
     return proxy_name.startswith(f"{ADMIN_NAMESPACE}_")
 
+
 class LdaNamespace(Namespace):
     def __init__(self, prefix: str) -> None:
         super().__init__(prefix)

src/wf_mcp/transparent_proxy.py

@@ -60,7 +60,9 @@ class TransparentProxyRuntime:
         if prompts_as_tools:
             self.server.add_transform(PromptsAsTools(self.server))
         if search_tools:
-            self.server.add_transform(BM25SearchTransform(always_visible=_ADMIN_TOOL_NAMES))
+            self.server.add_transform(
+                BM25SearchTransform(always_visible=_ADMIN_TOOL_NAMES)
+            )
 
     def current_config(self) -> BrokerConfig:
         if self.manager is None:
@@ -143,9 +145,7 @@ class TransparentProxyRuntime:
     ) -> dict[str, Any]:
         tools = await self._list_proxy_tools()
         if connection_id is not None:
-            tools = [
+            tools = [tool for tool in tools if tool["connection_id"] == connection_id]
-                tool for tool in tools if tool["connection_id"] == connection_id
-            ]
         if query:
             needle = query.casefold()
             tools = [

tests/rewrite/test_this_shit_works.py

@@ -298,7 +298,7 @@ class RateChange:
     @staticmethod
     def r65(state: CountersContext) -> Rates:
         c = state.counter
-        assert c["c_80"] >= 65, f"routed wrongly, 80 pity currently at {c["c_80"]}"
+        assert c["c_80"] >= 65, f"routed wrongly, 80 pity currently at {c['c_80']}"
         n = c["c_80"] - 64
         rpn = n * 0.05
         br = state.context["initial_rates"]
@@ -401,7 +401,7 @@ def post_roll_router(
 
 
 @node(name="main")
-def tick(state: Countdown) -> Countdown:  # type: ignore
+def tick(state: Countdown) -> Countdown:
     return Countdown(countdown=state.countdown - 1)
 
 
@@ -412,23 +412,19 @@ def keep_rolling(state: Countdown) -> NodeReturn[Nothing]:
 
 # could be @graph.(something combining node and use)...
 
-gacha.use(
+gacha.use(init)
-    init,
+gacha.use(tick, id="tick") # itd use main
-    id="init",  # use the node name if not defined?
+counter_up = gacha.use(CounterUp.c1, id="counter_up")  # 0 base to 1 base probably
-    # should i HAVE to declare in / out maps every time?
+rate_up = gacha.use(RateChange.r65, id="rate_up")
-)
+rate_same = gacha.use(RateChange.r0, id="rate_same")
-gacha.use(tick, id="tick")
+r_10 = gacha.use(RateChange.r10, id="r_g10")
-gacha.use(CounterUp.c1, id="counter_up")  # 0 base to 1 base probably
-gacha.use(RateChange.r65, id="rate_up")
-gacha.use(RateChange.r0, id="rate_same")
-gacha.use(RateChange.r10, id="r_g10")
 gacha.use(RateChange.r80, id="r_g80")
 gacha.use(RateChange.r240, id="r_gs")
-gacha.use(prep, id="prep")
+prepare_pool = gacha.use(prep)
-gacha.use(roll, id="roll")
+gacha.use(roll)
-gacha.use(CounterUp.c80, id="c_80")
+c_80 = gacha.use(CounterUp.c80, id="c_80")
 gacha.use(CounterUp.c10, id="c_10")
-gacha.condition(
+gacha.condition(  # condition dont ignore id; you can...
     id="keep_rolling", check=expr(state("countdown")) > 0
 )  # replaces keep_rolling
 # Outcome is currently hidden from the docs (there is none), outcome_map is insane, should we have it
@@ -440,40 +436,44 @@ gacha.connect("tick", "ok", "counter_up")
 gacha.use(rate_booster, id="rate_booster")
 
 gacha.connect("counter_up", "ok", "rate_booster")
-gacha.connect("rate_booster", "0", "rate_same")
+gacha.connect("rate_booster", "0", rate_same)
-gacha.connect("rate_booster", "65", "rate_up")
+gacha.connect("rate_booster", "65", rate_up)
 gacha.use(pre_roll_router, id="router")
 
 gacha.connect("rate_up", "ok", "router")
 gacha.connect("rate_same", "ok", "router")
-for outcome, node_id in {
+preroll_routes = gacha.branch(
+    "router",
+    {
         "240": "r_gs",
         "80": "r_g80",
-    "10": "r_g10",
+        "10": r_10,
-    "1": "prep",
+        "1": "prepare_pool",
-}.items():
+    },
-    gacha.connect("router", outcome, node_id)
+)
-
+print(preroll_routes)
-
+gacha.connect(prepare_pool, "ok", "roll")
-gacha.connect("prep", "ok", "roll")
+gacha.connect("r_gs", "ok", "prepare_pool")
-gacha.connect("r_gs", "ok", "prep")
+gacha.connect(preroll_routes["80"], "ok", prepare_pool)
-gacha.connect("r_g80", "ok", "prep")
+gacha.connect("r_g10", "ok", prepare_pool)
-gacha.connect("r_g10", "ok", "prep")
 
 gacha.connect("roll", "ok", "post_roll_router")  # missed this! good job
 gacha.use(post_roll_router, id="post_roll_router")
-gacha.use(popped, id="reset_avail")
+reset_avail = gacha.use(popped, id="reset_avail")
-gacha.connect("reset_avail", "ok", "c_80")
+gacha.connect(reset_avail, "ok", c_80)
 
-for outcome, node_id in {
+# missed this!
+gacha.branch(
+    "post_roll_router",
+    {
         "240": "reset_avail",
         "80": "c_80",
         "10": "c_10",
         "1": "keep_rolling",
-}.items():  # missed this!
+    },
-    gacha.connect("post_roll_router", outcome, node_id)
+)
 
-gacha.connect("c_80", "ok", "keep_rolling")
+gacha.connect(c_80, "ok", "keep_rolling")
 gacha.connect("c_10", "ok", "keep_rolling")
 
 # there is like no general uses for the nodes; idk tho
@@ -560,19 +560,24 @@ def execute(graph: WorkflowBuilder, input: Input):
     c = graph.compile()
     r = build_registry(*(graph.node_specs.values()))
     i = input.model_dump()
-    # pprint(i)
+    pprint(i)
-    # pprint(c)
+    pprint(c)
-    # pprint(r)
+    pprint(r)
     return execute_workflow(c, i, r)
 
 
 # twice in a row! it took 100+ and a miss tho
 
+
 def test():
     d = execute(
         gacha,
         build_input(context)(
-            20, rolled_previously=240 - 135, until_5=5, until_6=73, good_stuff=False # lets pretend
+            20,
+            rolled_previously=240 - 135,
+            until_5=5,
+            until_6=73,
+            good_stuff=False,  # lets pretend
         ),
     )
     assert d.status == RunStatus.COMPLETED, "oops"
@@ -580,5 +585,6 @@ def test():
     assert any(i["name"] in context["pool"]["n_240"] for i in state.storage)
     pprint(state.storage)
 
+
 if __name__ == "__main__":
     test()

tests/test_authoring_schemas.py

@@ -7,7 +7,7 @@ from pydantic import BaseModel, Field
 import pytest
 
 from wf_authoring import WorkflowBuilder, node, state_field
-from wf_core import RunStatus, WorkflowExecutionError, execute_workflow
+from wf_core import RunStatus, WorkflowExecutionError
 
 
 class WorkflowInput(BaseModel):
@@ -61,6 +61,12 @@ def auto_bind_node(input: AutoBindInput) -> AutoBindOutput:
     return AutoBindOutput(text=input.text.upper(), count=input.count + 1)
 
 
+@node(name="test.branch_router", outcomes=("left", "right"))
+def branch_router(input: AutoBindInput) -> AutoBindOutput:
+    """Route to left or right while preserving state shape."""
+    return AutoBindOutput(text=input.text, count=input.count)
+
+
 def test_builder_accepts_basemodel_classes_for_workflow_schemas() -> None:
     builder = WorkflowBuilder(
         name="model_schema_demo",
@@ -136,11 +142,8 @@ def test_builder_auto_binds_matching_node_inputs_and_outputs_to_state() -> None:
     step = builder.use(auto_bind_node, id="update")
     builder.connect(step, "ok", "__end__")
 
-    workflow = builder.compile()
+    run = builder.execute(
-    run = execute_workflow(
-        workflow,
         {"text": "hello", "count": 1},
-        builder.registry(),
     )
 
     assert step.in_map == {
@@ -210,3 +213,74 @@ def test_builder_registry_exports_used_node_specs() -> None:
     builder.use(auto_bind_node)
 
     assert set(builder.registry()) == {"test.auto_bind"}
+
+
+def test_builder_execute_compiles_and_runs_with_used_registry() -> None:
+    builder = WorkflowBuilder(
+        name="execute_demo",
+        input_schema=AutoBindInput,
+        state_schema=AutoBindState,
+        output_schema=AutoBindOutput,
+    )
+    step = builder.use(auto_bind_node)
+    builder.set_entry_point(step)
+    builder.connect(step, "ok", "__end__")
+
+    run = builder.execute({"text": "hello", "count": 1})
+
+    assert run.status == RunStatus.COMPLETED
+    assert run.state["text"] == "HELLO"
+    assert run.state["count"] == 2
+
+
+def test_builder_branch_connects_existing_steps() -> None:
+    builder = WorkflowBuilder(
+        name="branch_existing_demo",
+        input_schema=AutoBindInput,
+        state_schema=AutoBindState,
+        output_schema=AutoBindOutput,
+    )
+    router = builder.use(branch_router)
+    left = builder.use(auto_bind_node, id="left")
+    right = builder.use(auto_bind_node, id="right")
+
+    builder.branch(router, {"left": left, "right": right})
+
+    assert [(edge.from_, edge.outcome, edge.to) for edge in builder.edges] == [
+        ("test_branch_router", "left", "left"),
+        ("test_branch_router", "right", "right"),
+    ]
+
+
+def test_builder_branch_can_use_node_specs_as_targets() -> None:
+    builder = WorkflowBuilder(
+        name="branch_specs_demo",
+        input_schema=AutoBindInput,
+        state_schema=AutoBindState,
+        output_schema=AutoBindOutput,
+    )
+    router = builder.use(branch_router)
+
+    targets = builder.branch(router, {"left": auto_bind_node})
+
+    target = targets["left"]
+    assert not isinstance(target, str)
+    assert target.id == "test_auto_bind"
+    assert builder.edges[0].from_ == "test_branch_router"
+    assert builder.edges[0].outcome == "left"
+    assert builder.edges[0].to == "test_auto_bind"
+
+
+def test_builder_branch_warns_on_empty_branch_map() -> None:
+    builder = WorkflowBuilder(
+        name="branch_empty_demo",
+        input_schema=AutoBindInput,
+        state_schema=AutoBindState,
+        output_schema=AutoBindOutput,
+    )
+    router = builder.use(branch_router)
+
+    with pytest.warns(UserWarning, match="no branches"):
+        targets = builder.branch(router, {})
+
+    assert targets == {}

tests/test_wf_mcp_transparent_proxy.py

@@ -265,9 +265,10 @@ def test_transparent_proxy_proxy_tool_listing_supports_filters_and_cursor() -> N
             )
             second_page = _structured(second_page_result)
             assert len(second_page["tools"]) == 1
-            assert second_page["tools"][0]["proxy_name"] != first_page["tools"][0][
+            assert (
-                "proxy_name"
+                second_page["tools"][0]["proxy_name"]
-            ]
+                != first_page["tools"][0]["proxy_name"]
+            )
 
             filtered_result = await client.call_tool(
                 "wf.mcp_list_proxy_tools",