schema validation with jsonschema, orgs

docs/schema_validation.md

@@ -1,24 +1,29 @@
 # Schema Validation Boundary
 
 `wf_core` uses `SchemaRef` to carry JSON-schema-like shapes on workflow input,
-node input, node output, and workflow output. The current runtime does not
+node input, node output, and workflow output. Runtime payload validation is
-implement full JSON Schema validation.
+delegated to the existing `jsonschema` library.
 
 ## Current Behavior
 
-`wf_core.runtime.ops.schemas.validate_payload_against_schema` currently checks:
+`wf_core.runtime.ops.schemas.validate_payload_against_schema` currently:
 
-- if `schema.type == "object"`, the payload must be a `dict`
+- converts `SchemaRef` into a JSON Schema dictionary
-- required top-level keys must be present
+- asks `jsonschema` to validate the schema itself
+- asks `jsonschema` to validate the payload
+- wraps validation failures in `WorkflowExecutionError`
 
-It does not currently check:
+This means normal JSON Schema checks such as object shape, required fields,
+property types, nested required fields, arrays, and item types are enforced by
+the library.
 
-- property value types, such as `string`, `boolean`, `array`, or nested objects
+It still does not solve:
-- array item schemas
+
-- `additionalProperties`
+- semantic compatibility between Pydantic-generated schemas and every possible
-- formats, enums, minimums, maximums, unions, discriminators, or nested required
+  external JSON Schema dialect
-  fields
+- typed Python object creation from arbitrary JSON Schema
-- whether a schema is valid JSON Schema
+- workflow state merge behavior
+- better domain-specific error payloads beyond `WorkflowExecutionError`
 
 This means schema fields are mostly contracts for authoring, planning,
 documentation, and mapping validation today. They are not yet strong runtime
@@ -26,39 +31,34 @@ guards.
 
 ## Why This Matters
 
-The engine currently looks stricter than it is. A node can return a required
+Node and workflow boundaries can now reject wrong primitive/container types when
-field with the wrong type and pass runtime validation as long as the field is
+the schema declares them. This matters before workflows are generated by an LLM
-present. That is acceptable during early design work, but it is a real product
+or backed by arbitrary MCP tools.
-boundary before workflows are generated by an LLM or backed by arbitrary MCP
-tools.
 
 ## Intended Seam
 
-The schema adapter should live behind:
+The schema adapter lives behind:
 
 ```text
 wf_core.runtime.ops.schemas.validate_payload_against_schema
 ```
 
 Callers should not choose or invoke the validation backend directly. The runtime
-should keep one small validation interface and hide whether the implementation
+keeps one small validation interface and hides the validation backend.
-uses Pydantic, `jsonschema`, a generated model cache, or a stricter custom
-adapter.
 
 ## Future Requirements
 
 - Validate workflow input, node input, node output, and final workflow output
   with the same semantics.
 - Return errors that name the failing boundary and path.
-- Avoid silently accepting unsupported schema features once schemas are
-  user/LLM-authored.
 - Keep schema validation separate from graph structure validation.
 - Keep `wf_authoring` free to generate schemas from Pydantic models without
   making the core runtime depend on authoring internals.
+- Add targeted tests as externally sourced schemas grow more complex.
 
 ## Non-Goals For Now
 
 - Do not add ad hoc type checks throughout runtime state operations.
 - Do not let each node wrapper invent separate validation behavior.
 - Do not conflate graph validation with payload validation.
-
+- Do not hand-write a general JSON Schema implementation.

docs/wf_core_architecture.md

@@ -65,9 +65,9 @@ raising at the first failure.
 ## Schema Validation
 
 Payload schema validation is intentionally isolated behind
-`wf_core.runtime.ops.schemas.validate_payload_against_schema`. That function is
+`wf_core.runtime.ops.schemas.validate_payload_against_schema` and delegated to
-not a full JSON Schema engine today; see `docs/schema_validation.md` for the
+the `jsonschema` library. See `docs/schema_validation.md` for the current
-current limits and intended adapter seam.
+limits and intended adapter seam.
 
 ## What This Cleanup Does Not Solve Yet
 
@@ -79,5 +79,5 @@ current limits and intended adapter seam.
 - Runtime errors are still ordinary exceptions plus failed run status. A richer
   error payload can be added later, but should be designed as part of trace/run
   state rather than scattered exceptions.
-- Payload schema validation is still shallow. The runtime checks object payloads
+- Payload schema validation depends on JSON Schema semantics. If external tools
-  and required top-level keys, not full JSON Schema semantics.
+  emit unusual schema dialects, add compatibility tests before adapting them.

pyproject.toml

@@ -7,6 +7,7 @@ authors = [{ name = "lda", email = "[email protected]" }]
 requires-python = ">=3.14"
 dependencies = [
     "fastmcp>=3.2.4",
+    "jsonschema>=4.26",
     "mcp[cli,rich]>=1",
     "pydantic>=2",
 ]

src/wf_core/runtime/ops/schemas.py

@@ -2,15 +2,46 @@ from __future__ import annotations
 
 from typing import Any
 
+from jsonschema import ValidationError, SchemaError, validators
+
 from wf_core.errors import WorkflowExecutionError
+from wf_core.models.schemas import SchemaRef
 
 
-def validate_payload_against_schema(schema: Any, payload: Any, label: str) -> None:
+def validate_payload_against_schema(
-    if schema.type == "object":
+    schema: SchemaRef | dict[str, Any],
-        if not isinstance(payload, dict):
+    payload: Any,
-            raise WorkflowExecutionError(f"{label} must be an object")
+    label: str,
-        for required_key in schema.required:
+) -> None:
-            if required_key not in payload:
+    """Validate a runtime payload against the schema declared at a boundary.
+
+    The runtime delegates JSON Schema semantics to `jsonschema` instead of
+    maintaining hand-written type checks. Errors are wrapped in
+    `WorkflowExecutionError` so callers keep one execution-failure surface.
+    """
+    schema_dict = _schema_dict(schema)
+    validator_cls = validators.validator_for(schema_dict)
+    try:
+        validator_cls.check_schema(schema_dict)
+        validator_cls(schema_dict).validate(payload)
+    except SchemaError as exc:
         raise WorkflowExecutionError(
-                    f"{label} is missing required field {required_key!r}"
+            f"{label} has invalid schema: {exc.message}"
-                )
+        ) from exc
+    except ValidationError as exc:
+        path = _format_error_path(exc)
+        raise WorkflowExecutionError(f"{label}{path}: {exc.message}") from exc
+
+
+def _schema_dict(schema: SchemaRef | dict[str, Any]) -> dict[str, Any]:
+    """Return a JSON-Schema-compatible dictionary for validation."""
+    if isinstance(schema, SchemaRef):
+        return schema.model_dump(exclude_none=True)
+    return schema
+
+
+def _format_error_path(exc: ValidationError) -> str:
+    """Render a compact JSON-path-like suffix for a validation error."""
+    if not exc.path:
+        return ""
+    return "".join(f"[{part!r}]" for part in exc.path)

tests/core/test_schema_validation.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import pytest
+
+from wf_core import SchemaRef, WorkflowExecutionError
+from wf_core.runtime.ops.schemas import validate_payload_against_schema
+
+
+def test_schema_validation_rejects_wrong_property_type() -> None:
+    schema = SchemaRef.model_validate(
+        {
+            "type": "object",
+            "properties": {
+                "name": {"type": "string"},
+                "count": {"type": "integer"},
+            },
+            "required": ["name", "count"],
+        }
+    )
+
+    with pytest.raises(WorkflowExecutionError, match=r"count.*not of type 'integer'"):
+        validate_payload_against_schema(
+            schema,
+            {"name": "ok", "count": "not-an-int"},
+            "node output for counter",
+        )
+
+
+def test_schema_validation_rejects_nested_missing_required_field() -> None:
+    schema = SchemaRef.model_validate(
+        {
+            "type": "object",
+            "properties": {
+                "profile": {
+                    "type": "object",
+                    "properties": {"email": {"type": "string"}},
+                    "required": ["email"],
+                }
+            },
+            "required": ["profile"],
+        }
+    )
+
+    with pytest.raises(WorkflowExecutionError, match=r"profile.*email.*required"):
+        validate_payload_against_schema(
+            schema,
+            {"profile": {}},
+            "workflow input",
+        )
+
+
+def test_schema_validation_accepts_valid_payload() -> None:
+    schema = SchemaRef.model_validate(
+        {
+            "type": "object",
+            "properties": {
+                "tags": {"type": "array", "items": {"type": "string"}},
+            },
+            "required": ["tags"],
+        }
+    )
+
+    validate_payload_against_schema(schema, {"tags": ["a", "b"]}, "node input")

tests/rewrite/__init__.py

@@ -0,0 +1 @@
+"""Rewrite/port tests for workflow authoring ergonomics."""

tests/rewrite/actions.py

@@ -0,0 +1,260 @@
+import random
+from typing import Final
+
+from tests.rewrite.models import (
+    ContextInput,
+    Countdown,
+    Counters,
+    CurrentPools,
+    CurrentRoll,
+    Entity,
+    Input,
+    PoolByCategory,
+    Rates,
+    Storage,
+    how_do_i_explain_this,
+)
+from wf_authoring import NodeReturn, node
+from wf_authoring.nodes.result import Nothing, outcome
+from wf_core.tokens import END
+
+# to the functions
+# @node(outcomes=("ok", "end")) # breaks because of input | nothing
+
+
+@node
+def init(
+    inp: Input,
+    # ) -> NodeReturn[Input | Nothing]:  # JUST doesnt work if the types are above
+) -> Input:
+    # why cant use basemodel? should we convert typeddicts to basemodels? Why cant use
+    """i hope that this is done by default, because langgraph DOESNT. why.
+
+    since input and state has the same keys
+    """
+    # if ctx.context["type"] == 'normal' and ctx.context["initial_rates"]["r_240"] != 0:
+    #     "what now" # pylint: disable=W0105
+    # return (
+    #     NodeReturn("ok", inp)
+    #     if inp.countdown > 0
+    #     else NodeReturn("end", Nothing())  # end early
+    # )
+    # if inp.pity_120_available and inp.context["type"] == "normal":
+    #    "i doesnt care"
+    return inp
+
+
+@node(outcomes=("0", "65"))
+def rate_booster(c: Counters) -> NodeReturn[Nothing]:
+    """This was an edge.
+
+    Should I implement r65 here too... i think not.
+
+    the flow was init --rate_booster-> (65, r65), (0, r0), which outputs to the same
+    rate_guarantee, which is no op.
+    """
+    if c.counter["c_80"] >= 65:
+        return outcome("65")
+    return outcome("0")
+
+
+def _popped(storage: list[Entity]) -> bool:
+    return any(c["category"] == "240" for c in reversed(storage))
+
+
+@node
+def popped(s: Storage) -> how_do_i_explain_this:
+    return how_do_i_explain_this(pity_120_available=not _popped(s.storage))
+
+
+# OHH so each of these MAY only only use whatever is needed? woah. i like
+
+
+class CountersContext(Counters, ContextInput): ...
+
+
+# this gets annoying.
+class CountersContextOutputInputAhhModelType(
+    CountersContext, Storage, how_do_i_explain_this
+): ...
+
+
+@node(outcomes=("240", "80", "10", "1"))
+def pre_roll_router(c: CountersContextOutputInputAhhModelType) -> NodeReturn[Nothing]:
+    """another edge.
+
+    the conditional router calculates which to reset to 0 (guarantee the rest)
+
+    the flow was rate_guarantee --router-> (1 -> prep), (10 -.-> RateChange.r10 --> prep) (80 -.-> r80 --> prep), (240 -.-> r240 -> prep).
+    """
+    sc, ct = c.simple_counter, c.counter
+    if c.context["type"] == "banner" and (
+        (sc == 120 and c.pity_120_available) or (sc > 0 and sc % 240 == 0)
+    ):
+        return outcome("240")
+    if ct.get("c_80", 0) % 80 == 0:
+        return outcome("80")
+    if ct.get("c_10", 0) % 10 == 0:
+        return outcome("10")
+    return outcome("1")
+
+
+# now to the weeds of it. a Class!
+
+
+class RateChange:
+    @node(name="force 6* rating")
+    @staticmethod
+    def r80(r: Rates) -> Rates:
+        return Rates.model_validate(
+            {
+                "rates": {
+                    "r_1": 0,
+                    "r_10": 0,
+                    "r_80": r.rates["r_80"],
+                    "r_240": r.rates["r_240"],
+                }
+            }
+        )
+
+    @node(name="force banner rating")
+    @staticmethod
+    def r240(_: Nothing) -> Rates:
+        return Rates.model_validate(
+            {"rates": {"r_1": 0, "r_10": 0, "r_80": 0, "r_240": 1}}
+        )
+
+    @node(name="force 5*+ rating")
+    @staticmethod
+    def r10(r: Rates) -> Rates:
+        return Rates.model_validate(
+            {
+                "rates": {
+                    "r_1": 0,
+                    "r_10": r.rates["r_10"],
+                    "r_80": r.rates["r_80"],
+                    "r_240": r.rates["r_240"],
+                }
+            }
+        )
+
+    @node(name="buff 6* rating")
+    @staticmethod
+    def r65(state: CountersContext) -> Rates:
+        c = state.counter
+        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"]
+        if state.context["type"] == "banner":
+            r240 = br["r_240"] * (1 + rpn / 2)
+            r80 = br["r_80"] * (1 + rpn / 2)
+        else:
+            r240 = 0
+            r80 = br["r_80"] * (1 + rpn)
+
+        return Rates.model_validate(
+            {
+                "rates": {
+                    "r_1": 1 - r240 - r80 - br["r_10"],
+                    "r_10": br["r_10"],
+                    "r_80": r80,
+                    "r_240": r240,
+                }
+            }
+        )
+
+    @node(name="reset rating")
+    @staticmethod
+    def r0(state: ContextInput) -> Rates:
+        return Rates.model_validate({"rates": state.context["initial_rates"].copy()})
+
+
+class CounterUp:
+    @node(name="counter 6* reset")
+    @staticmethod
+    def c80(c: Counters) -> Counters:
+        return Counters.model_validate(
+            {
+                "counter": {
+                    "c_80": 0,
+                    "c_10": 0,
+                },
+                "simple_counter": c.simple_counter,
+            }
+        )
+
+    @node(name="counter 5* reset")
+    @staticmethod
+    def c10(c: Counters) -> Counters:
+        return Counters.model_validate(
+            {
+                "counter": {"c_10": 0, "c_80": c.counter["c_80"]},  # merge with or_!
+                "simple_counter": c.simple_counter,
+            }
+        )
+
+    @node(name="counting up")
+    @staticmethod
+    def c1(state: Counters) -> Counters:
+        c = state.counter
+        return Counters(
+            simple_counter=state.simple_counter + 1,
+            counter={"c_10": (c["c_10"] + 1) % 10, "c_80": (c["c_80"] + 1) % 80},
+        )
+
+
+class RatesContextInput(Rates, ContextInput): ...
+
+
+@node(name="prepare pool")
+def prep(state: RatesContextInput) -> CurrentPools:
+    r = state.rates
+    p = state.context["pool"]
+    t: Final[tuple[tuple[str, str, str], ...]] = (  # greatest hack
+        ("1", "r_1", "n_1"),
+        ("10", "r_10", "n_10"),
+        ("80", "r_80", "n_80"),
+        ("240", "r_240", "n_240"),
+    )
+    pbc = [PoolByCategory(pool=p[pc], category=ty, rates=r[pr]) for ty, pr, pc in t]
+    # print(pbc)
+    return CurrentPools.model_validate({"current_pools": pbc})
+
+
+class ThisStorage(Storage, CurrentRoll): ...
+
+
+@node
+def roll(state: CurrentPools) -> ThisStorage:
+    r = state.current_pools
+    (t,) = random.choices(r, weights=[*map(lambda p: p["rates"], r)])
+    this = Entity(category=t["category"], name=random.choice(t["pool"]))
+    return ThisStorage.model_validate(
+        {
+            "this": this,
+            "storage": [this],  # I NEED MERGE
+        }
+    )
+
+
+@node(outcomes=("240", "80", "10", "1"))  # missed this! good job.
+def post_roll_router(
+    state: CurrentRoll,
+) -> NodeReturn[
+    Nothing
+]:  # Literal["240", "80", "10", "1"] maybe you need to encode this in node returns. Literal of strings.
+    return NodeReturn(state.this["category"], Nothing())
+
+
+@node(name="main")
+def tick(state: Countdown) -> Countdown:
+    return Countdown(countdown=state.countdown - 1)
+
+
+@node(outcomes=("tick", END))
+def keep_rolling(state: Countdown) -> NodeReturn[Nothing]:
+    return outcome("tick") if (state.countdown or 0) > 0 else outcome(END)
+
+
+# there is like no general uses for the nodes; idk tho

tests/rewrite/context.py

@@ -0,0 +1,45 @@
+from typing import Final
+
+from tests.rewrite.models import Context
+
+context: Final[Context] = {
+    "initial_rates": {
+        "r_1": 0.912,
+        "r_10": 0.08,
+        "r_80": 0.004,
+        "r_240": 0.004,
+    },
+    "pool": {
+        "n_1": ["Akekuri", "Catcher", "Flourite", "Estella", "Antal"],
+        "n_10": [
+            "Perlica",
+            "Arclight",
+            "Avywenna",
+            "Da Pan",
+            "Chen Qianyu",
+            "Wulfgard",
+            "Xaihi",
+            "Snowshine",
+            "Alesh",
+        ],
+        "n_80": [
+            "Rossi",
+            "Tangtang",
+            # "Yvonne",
+            # "Gilberta",
+            # "Laevatain",
+            "Ember",
+            "Lifeng",
+            "Ardelia",
+            "Last Rite",
+            "Pogranichnik",
+        ],
+        "n_240": [  # normal or banner / logic is hella flawed lowk ong
+            # "Rossi",
+            # "Tangtang",
+            "Zhuang Fangyi",
+        ],
+    },
+    "type": "banner",
+}
+"please can i have tang2"

tests/rewrite/models.py

@@ -0,0 +1,137 @@
+from typing import Annotated, Literal, TypedDict
+
+from pydantic import BaseModel, Field
+
+from wf_authoring.schemas import state_field
+
+
+class SophisticatedRates(TypedDict):
+    r_1: float
+    r_10: float
+    r_80: float  # 0.5 sometimes, 1
+    # not including 240 because ill force it by simple_counter, the pool split probably forces half, dealing w ts is ahh.
+    r_240: float  # i caved
+
+
+# i copy things over, not the greatest design but im not doing deep fixes.
+
+
+class SophisticatedCounter(TypedDict):
+    c_10: int  # how do i convey "add" sublevel? can we have plugins for this? should we cover this; since langgraph doesnt.
+    c_80: int
+
+
+class Counters(BaseModel):
+    # alot of state["thing"] in the og code, so i think this is the design?
+    counter: SophisticatedCounter = Field(
+        default_factory=lambda: SophisticatedCounter(c_10=0, c_80=0)
+    )  # or, that is because of langgraph limitation,
+    # id prefer counter.update with dict.update override (Overwrite, not like that ever worked)
+    simple_counter: int  # add
+
+
+class Countdown(BaseModel):
+    countdown: int  # add!
+
+
+# has to migrate from typeddict for what? for nothing.
+class Context(TypedDict):  # dataclass support? no. i mean langgraph doesnt.
+    "custom RuntimeContext[MyContext] support?"
+
+    pool: UnsophisticatedPool
+    initial_rates: SophisticatedRates
+    type: Literal["banner", "normal"]
+    "still very convoluted logic"
+
+
+# class Input(TypedDict, total=False):
+
+
+class ContextInput(BaseModel):
+    context: Context  # final!
+
+
+class how_do_i_explain_this(BaseModel):
+    pity_120_available: bool = Field(default=True)
+
+
+class Input(Counters, ContextInput, Countdown, how_do_i_explain_this):
+    "input of the graph"
+
+
+# countdown: int
+# simple_counter: int
+# counter: SophisticatedCounter
+
+# context: Context
+
+
+class Entity(TypedDict):
+    category: Literal["1", "10", "80", "240"]
+    name: str
+
+
+class PoolByCategory(TypedDict):
+    pool: list[str]
+    category: Literal["1", "10", "80", "240"]  # stricter types
+    rates: float
+
+
+# context outside? what is input?
+class UnsophisticatedPool(TypedDict):
+    n_1: list[str]
+    n_10: list[str]
+    n_80: list[str]  # normal + limited
+    n_240: list[str]  # special... have to do this
+
+
+class Storage(BaseModel):
+    "i NEED to do this?"
+
+    storage: Annotated[list[Entity], state_field(merge_strategy="append")] = Field(
+        default_factory=list
+    )  # add!
+
+
+class CurrentRoll(BaseModel):
+    this: Entity
+
+
+class PartialRates(SophisticatedRates, total=False):
+    pass
+
+
+class Rates(BaseModel):
+    rates: Annotated[PartialRates, state_field(merge_strategy="merge_object")]  # or_!
+
+
+class CurrentPools(BaseModel):
+    current_pools: list[PoolByCategory]  # replace!
+
+
+# holy refactory
+class State(
+    Counters,
+    ContextInput,
+    Countdown,
+    CurrentRoll,
+    Storage,
+    Rates,
+    CurrentPools,
+    how_do_i_explain_this,
+):
+    "this forces basemodel, i used typeddict"
+
+
+# countdown: int  # input carries here, should input have a bound like all(attr(input) in attr(state))? how tf do i even try to type that
+# simple_counter: int  # how to signal that ts adds up? we have that.
+# counter: SophisticatedCounter
+# rates: SophisticatedRates  # ts needs reworking
+# current_pools: list[PoolByCategory]
+# this: Entity
+# storage: list[
+#     Entity  # how do we convey "add" root level? Annotated again? what pydantic shit can give this thing the metadata it neeeds.
+# ]  # maybe for the 120 check we need to add some if any(entity.category = "240" for entity in storage), WHICH IS ASS btw.
+
+# fuck this yo
+# context: Context

tests/rewrite/test_this_shit_works.py

@@ -3,506 +3,19 @@
 Constraints:
     no tapping wf_core
     try to use builtin wf_authoring.ops
-
-also:
-    fix the 120 logic now that i know
 """
 
-# alr so we start with workflowbuilder.
-
 from pprint import pprint
-import random
+from typing import Any
-from typing import Annotated, Any, Final, Literal, TypedDict
 
-
+from tests.rewrite.context import context
-from wf_authoring.builder import WorkflowBuilder
+from tests.rewrite.models import (
-from pydantic import BaseModel, Field
+    Context,
-
+    Input,
-from wf_authoring import node
+    State,
-from wf_authoring import NodeReturn
+)
-from wf_authoring.dsl.conditions import expr, state
+from tests.rewrite.workflow import gacha
-from wf_authoring.nodes.result import Nothing, outcome
-from wf_authoring.schemas import state_field
 from wf_core.run_state import RunStatus
-from wf_core.tokens import END
-
-# i copy things over, not the greatest design but im not doing deep fixes.
-
-
-class SophisticatedCounter(TypedDict):
-    c_10: int  # how do i convey "add" sublevel? can we have plugins for this? should we cover this; since langgraph doesnt.
-    c_80: int
-
-
-class SophisticatedRates(TypedDict):
-    r_1: float
-    r_10: float
-    r_80: float  # 0.5 sometimes, 1
-    # not including 240 because ill force it by simple_counter, the pool split probably forces half, dealing w ts is ahh.
-    r_240: float  # i caved
-
-
-# class Input(TypedDict, total=False):
-
-
-class Counters(BaseModel):
-    # alot of state["thing"] in the og code, so i think this is the design?
-    counter: SophisticatedCounter = Field(
-        default_factory=lambda: SophisticatedCounter(c_10=0, c_80=0)
-    )  # or, that is because of langgraph limitation,
-    # id prefer counter.update with dict.update override (Overwrite, not like that ever worked)
-    simple_counter: int  # add
-
-
-class Countdown(BaseModel):
-    countdown: int  # add!
-
-
-class ContextInput(BaseModel):
-    context: Context  # final!
-
-
-class how_do_i_explain_this(BaseModel):
-    pity_120_available: bool = Field(default=True)
-
-
-class Input(Counters, ContextInput, Countdown, how_do_i_explain_this):
-    "input of the graph"
-
-    # countdown: int
-    # simple_counter: int
-    # counter: SophisticatedCounter
-
-    # context: Context
-
-
-class Entity(TypedDict):
-    category: Literal["1", "10", "80", "240"]
-    name: str
-
-
-class PoolByCategory(TypedDict):
-    pool: list[str]
-    category: Literal["1", "10", "80", "240"]  # stricter types
-    rates: float
-
-
-# context outside? what is input?
-class UnsophisticatedPool(TypedDict):
-    n_1: list[str]
-    n_10: list[str]
-    n_80: list[str]  # normal + limited
-    n_240: list[str]  # special... have to do this
-
-
-# has to migrate from typeddict for what? for nothing.
-class Context(TypedDict):  # dataclass support? no. i mean langgraph doesnt.
-    "custom RuntimeContext[MyContext] support?"
-
-    pool: UnsophisticatedPool
-    initial_rates: SophisticatedRates
-    type: Literal["banner", "normal"]
-    "still very convoluted logic"
-
-
-class Storage(BaseModel):
-    "i NEED to do this?"
-
-    storage: Annotated[list[Entity], state_field(merge_strategy="append")] = Field(
-        default_factory=list
-    )  # add!
-
-
-class CurrentRoll(BaseModel):
-    this: Entity
-
-
-class ThisStorage(Storage, CurrentRoll): ...
-
-
-class PartialRates(SophisticatedRates, total=False):
-    pass
-
-
-class Rates(BaseModel):
-    rates: Annotated[PartialRates, state_field(merge_strategy="merge_object")]  # or_!
-
-
-class CurrentPools(BaseModel):
-    current_pools: list[PoolByCategory]  # replace!
-
-
-# holy refactory
-class State(
-    Counters,
-    ContextInput,
-    Countdown,
-    CurrentRoll,
-    Storage,
-    Rates,
-    CurrentPools,
-    how_do_i_explain_this,
-):
-    "this forces basemodel, i used typeddict"
-
-    # countdown: int  # input carries here, should input have a bound like all(attr(input) in attr(state))? how tf do i even try to type that
-    # simple_counter: int  # how to signal that ts adds up? we have that.
-    # counter: SophisticatedCounter
-    # rates: SophisticatedRates  # ts needs reworking
-    # current_pools: list[PoolByCategory]
-    # this: Entity
-    # storage: list[
-    #     Entity  # how do we convey "add" root level? Annotated again? what pydantic shit can give this thing the metadata it neeeds.
-    # ]  # maybe for the 120 check we need to add some if any(entity.category = "240" for entity in storage), WHICH IS ASS btw.
-
-    # fuck this yo
-    # context: Context
-
-
-# to the functions
-# @node(outcomes=("ok", "end")) # breaks because of input | nothing
-@node
-def init(
-    inp: Input,
-    # ) -> NodeReturn[Input | Nothing]:  # JUST doesnt work if the types are above
-) -> Input:
-    # why cant use basemodel? should we convert typeddicts to basemodels? Why cant use
-    """i hope that this is done by default, because langgraph DOESNT. why.
-
-    since input and state has the same keys
-    """
-    # if ctx.context["type"] == 'normal' and ctx.context["initial_rates"]["r_240"] != 0:
-    #     "what now" # pylint: disable=W0105
-    # return (
-    #     NodeReturn("ok", inp)
-    #     if inp.countdown > 0
-    #     else NodeReturn("end", Nothing())  # end early
-    # )
-    # if inp.pity_120_available and inp.context["type"] == "normal":
-    #    "i doesnt care"
-    return inp
-
-
-# OHH so each of these MAY only only use whatever is needed? woah. i like
-
-
-class CountersContext(Counters, ContextInput): ...
-
-
-# this gets annoying.
-class CountersContextOutputInputAhhModelType(
-    CountersContext, Storage, how_do_i_explain_this
-): ...
-
-
-@node(outcomes=("0", "65"))
-def rate_booster(c: Counters) -> NodeReturn[Nothing]:
-    """This was an edge.
-
-    Should I implement r65 here too... i think not.
-
-    the flow was init --rate_booster-> (65, r65), (0, r0), which outputs to the same
-    rate_guarantee, which is no op.
-    """
-    if c.counter["c_80"] >= 65:
-        return outcome("65")
-    return outcome("0")
-
-
-def _popped(storage: list[Entity]) -> bool:
-    return any(c["category"] == "240" for c in reversed(storage))
-
-
-@node
-def popped(s: Storage) -> how_do_i_explain_this:
-    return how_do_i_explain_this(pity_120_available=not _popped(s.storage))
-
-
-@node(outcomes=("240", "80", "10", "1"))
-def pre_roll_router(c: CountersContextOutputInputAhhModelType) -> NodeReturn[Nothing]:
-    """another edge.
-
-    the conditional router calculates which to reset to 0 (guarantee the rest)
-
-    the flow was rate_guarantee --router-> (1 -> prep), (10 -.-> RateChange.r10 --> prep) (80 -.-> r80 --> prep), (240 -.-> r240 -> prep).
-    """
-    sc, ct = c.simple_counter, c.counter
-    if c.context["type"] == "banner" and (
-        (sc == 120 and c.pity_120_available) or (sc > 0 and sc % 240 == 0)
-    ):
-        return outcome("240")
-    if ct.get("c_80", 0) % 80 == 0:
-        return outcome("80")
-    if ct.get("c_10", 0) % 10 == 0:
-        return outcome("10")
-    return outcome("1")
-
-
-# now to the weeds of it. a Class!
-
-
-class RateChange:
-    @node(name="force 6* rating")
-    @staticmethod
-    def r80(r: Rates) -> Rates:
-        return Rates.model_validate(
-            {
-                "rates": {
-                    "r_1": 0,
-                    "r_10": 0,
-                    "r_80": r.rates["r_80"],
-                    "r_240": r.rates["r_240"],
-                }
-            }
-        )
-
-    @node(name="force banner rating")
-    @staticmethod
-    def r240(_: Nothing) -> Rates:
-        return Rates.model_validate(
-            {"rates": {"r_1": 0, "r_10": 0, "r_80": 0, "r_240": 1}}
-        )
-
-    @node(name="force 5*+ rating")
-    @staticmethod
-    def r10(r: Rates) -> Rates:
-        return Rates.model_validate(
-            {
-                "rates": {
-                    "r_1": 0,
-                    "r_10": r.rates["r_10"],
-                    "r_80": r.rates["r_80"],
-                    "r_240": r.rates["r_240"],
-                }
-            }
-        )
-
-    @node(name="buff 6* rating")
-    @staticmethod
-    def r65(state: CountersContext) -> Rates:
-        c = state.counter
-        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"]
-        if state.context["type"] == "banner":
-            r240 = br["r_240"] * (1 + rpn / 2)
-            r80 = br["r_80"] * (1 + rpn / 2)
-        else:
-            r240 = 0
-            r80 = br["r_80"] * (1 + rpn)
-
-        return Rates.model_validate(
-            {
-                "rates": {
-                    "r_1": 1 - r240 - r80 - br["r_10"],
-                    "r_10": br["r_10"],
-                    "r_80": r80,
-                    "r_240": r240,
-                }
-            }
-        )
-
-    @node(name="reset rating")
-    @staticmethod
-    def r0(state: ContextInput) -> Rates:
-        return Rates.model_validate({"rates": state.context["initial_rates"].copy()})
-
-
-class CounterUp:
-    @node(name="counter 6* reset")
-    @staticmethod
-    def c80(c: Counters) -> Counters:
-        return Counters.model_validate(
-            {
-                "counter": {
-                    "c_80": 0,
-                    "c_10": 0,
-                },
-                "simple_counter": c.simple_counter,
-            }
-        )
-
-    @node(name="counter 5* reset")
-    @staticmethod
-    def c10(c: Counters) -> Counters:
-        return Counters.model_validate(
-            {
-                "counter": {"c_10": 0, "c_80": c.counter["c_80"]},  # merge with or_!
-                "simple_counter": c.simple_counter,
-            }
-        )
-
-    @node(name="counting up")
-    @staticmethod
-    def c1(state: Counters) -> Counters:
-        c = state.counter
-        return Counters(
-            simple_counter=state.simple_counter + 1,
-            counter={"c_10": (c["c_10"] + 1) % 10, "c_80": (c["c_80"] + 1) % 80},
-        )
-
-
-class RatesContextInput(Rates, ContextInput): ...
-
-
-@node(name="prepare pool")
-def prep(state: RatesContextInput) -> CurrentPools:
-    r = state.rates
-    p = state.context["pool"]
-    t: Final[tuple[tuple[str, str, str], ...]] = (  # greatest hack
-        ("1", "r_1", "n_1"),
-        ("10", "r_10", "n_10"),
-        ("80", "r_80", "n_80"),
-        ("240", "r_240", "n_240"),
-    )
-    pbc = [PoolByCategory(pool=p[pc], category=ty, rates=r[pr]) for ty, pr, pc in t]
-    # print(pbc)
-    return CurrentPools.model_validate({"current_pools": pbc})
-
-
-@node
-def roll(state: CurrentPools) -> ThisStorage:
-    r = state.current_pools
-    (t,) = random.choices(r, weights=[*map(lambda p: p["rates"], r)])
-    this = Entity(category=t["category"], name=random.choice(t["pool"]))
-    return ThisStorage.model_validate(
-        {
-            "this": this,
-            "storage": [this],  # I NEED MERGE
-        }
-    )
-
-
-@node(outcomes=("240", "80", "10", "1"))  # missed this! good job.
-def post_roll_router(
-    state: CurrentRoll,
-) -> NodeReturn[
-    Nothing
-]:  # Literal["240", "80", "10", "1"] maybe you need to encode this in node returns. Literal of strings.
-    return NodeReturn(state.this["category"], Nothing())
-
-
-@node(name="main")
-def tick(state: Countdown) -> Countdown:
-    return Countdown(countdown=state.countdown - 1)
-
-
-@node(outcomes=("tick", END))
-def keep_rolling(state: Countdown) -> NodeReturn[Nothing]:
-    return outcome("tick") if (state.countdown or 0) > 0 else outcome(END)
-
-
-gacha = WorkflowBuilder(
-    name="im not hiding it no more",
-    input_schema=Input,
-    output_schema=Storage,  # could be State, since the OG doesnt care, ill probably dump out the list.
-    state_schema=State,
-)
-gacha.use(tick, id="tick")  # itd use main if we dont have id
-counter_up = gacha.use(CounterUp.c1, id="counter_up")  # 0 base to 1 base probably
-rate_up = gacha.use(RateChange.r65, id="rate_up")
-rate_same = gacha.use(RateChange.r0, id="rate_same")
-r_10 = gacha.use(RateChange.r10, id="r_g10")
-gacha.use(RateChange.r80, id="r_g80")
-gacha.use(RateChange.r240, id="r_gs")
-c_80 = gacha.use(CounterUp.c80, id="c_80")
-gacha.use(CounterUp.c10, id="c_10")
-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
-init_ref, _ = gacha.connect(init, "ok", "keep_rolling")
-gacha.connect("keep_rolling", "true", "tick")
-gacha.connect("keep_rolling", "false", END)
-
-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", "65", rate_up)
-gacha.use(pre_roll_router, id="router")
-
-gacha.connect("rate_up", "ok", "router")
-gacha.connect("rate_same", "ok", "router")
-preroll_routes = gacha.branch(
-    "router",
-    {
-        "240": "r_gs",
-        "80": "r_g80",
-        "10": r_10,
-        "1": prep,
-    },
-)
-prepare_pool = preroll_routes["1"]
-_, roll_ref = gacha.connect(prepare_pool, "ok", roll)
-gacha.connect("r_gs", "ok", "prepare_pool")
-gacha.connect(preroll_routes["80"], "ok", prepare_pool)
-gacha.connect("r_g10", "ok", prepare_pool)
-
-gacha.connect("roll", "ok", "post_roll_router")  # missed this! good job
-gacha.use(post_roll_router, id="post_roll_router")
-reset_avail = gacha.use(popped, id="reset_avail")
-gacha.connect(reset_avail, "ok", c_80)
-
-# missed this!
-gacha.branch(
-    "post_roll_router",
-    {
-        "240": "reset_avail",
-        "80": "c_80",
-        "10": "c_10",
-        "1": "keep_rolling",
-    },
-)
-
-gacha.connect(c_80, "ok", "keep_rolling")
-gacha.connect("c_10", "ok", "keep_rolling")
-gacha.set_entry_point(init_ref)
-
-# there is like no general uses for the nodes; idk tho
-
-context: Final[Context] = {
-    "initial_rates": {
-        "r_1": 0.912,
-        "r_10": 0.08,
-        "r_80": 0.004,
-        "r_240": 0.004,
-    },
-    "pool": {
-        "n_1": ["Akekuri", "Catcher", "Flourite", "Estella", "Antal"],
-        "n_10": [
-            "Perlica",
-            "Arclight",
-            "Avywenna",
-            "Da Pan",
-            "Chen Qianyu",
-            "Wulfgard",
-            "Xaihi",
-            "Snowshine",
-            "Alesh",
-        ],
-        "n_80": [
-            "Rossi",
-            "Tangtang",
-            # "Yvonne",
-            # "Gilberta",
-            # "Laevatain",
-            "Ember",
-            "Lifeng",
-            "Ardelia",
-            "Last Rite",
-            "Pogranichnik",
-        ],
-        "n_240": [  # normal or banner / logic is hella flawed lowk ong
-            # "Rossi",
-            # "Tangtang",
-            "Zhuang Fangyi",
-        ],
-    },
-    "type": "banner",
-}
 
 
 def build_input_lite(

tests/rewrite/workflow.py

@@ -0,0 +1,88 @@
+from tests.rewrite.actions import (
+    CounterUp,
+    RateChange,
+    init,
+    popped,
+    post_roll_router,
+    pre_roll_router,
+    prep,
+    rate_booster,
+    roll,
+    tick,
+)
+from tests.rewrite.models import Input, State, Storage
+from wf_authoring.builder import WorkflowBuilder
+from wf_authoring.dsl.conditions import expr, state
+from wf_core.tokens import END
+
+# alr so we start with workflowbuilder.
+
+gacha = WorkflowBuilder(
+    name="im not hiding it no more",
+    input_schema=Input,
+    output_schema=Storage,  # could be State, since the OG doesnt care, ill probably dump out the list.
+    state_schema=State,
+)
+"example workflow"
+
+gacha.use(tick, id="tick")  # itd use main if we dont have id
+counter_up = gacha.use(CounterUp.c1, id="counter_up")  # 0 base to 1 base probably
+rate_up = gacha.use(RateChange.r65, id="rate_up")
+rate_same = gacha.use(RateChange.r0, id="rate_same")
+r_10 = gacha.use(RateChange.r10, id="r_g10")
+gacha.use(RateChange.r80, id="r_g80")
+gacha.use(RateChange.r240, id="r_gs")
+c_80 = gacha.use(CounterUp.c80, id="c_80")
+gacha.use(CounterUp.c10, id="c_10")
+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
+init_ref, _ = gacha.connect(init, "ok", "keep_rolling")
+gacha.connect("keep_rolling", "true", "tick")
+gacha.connect("keep_rolling", "false", END)
+
+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", "65", rate_up)
+gacha.use(pre_roll_router, id="router")
+
+gacha.connect("rate_up", "ok", "router")
+gacha.connect("rate_same", "ok", "router")
+preroll_routes = gacha.branch(
+    "router",
+    {
+        "240": "r_gs",
+        "80": "r_g80",
+        "10": r_10,
+        "1": prep,
+    },
+)
+prepare_pool = preroll_routes["1"]
+_, roll_ref = gacha.connect(prepare_pool, "ok", roll)
+gacha.connect("r_gs", "ok", "prepare_pool")
+gacha.connect(preroll_routes["80"], "ok", prepare_pool)
+gacha.connect("r_g10", "ok", prepare_pool)
+
+gacha.connect("roll", "ok", "post_roll_router")  # missed this! good job
+gacha.use(post_roll_router, id="post_roll_router")
+reset_avail = gacha.use(popped, id="reset_avail")
+gacha.connect(reset_avail, "ok", c_80)
+
+# missed this!
+gacha.branch(
+    "post_roll_router",
+    {
+        "240": "reset_avail",
+        "80": "c_80",
+        "10": "c_10",
+        "1": "keep_rolling",
+    },
+)
+
+gacha.connect(c_80, "ok", "keep_rolling")
+gacha.connect("c_10", "ok", "keep_rolling")
+gacha.set_entry_point(init_ref)

uv.lock

@@ -522,6 +522,7 @@ version = "0.0.1"
 source = { editable = "." }
 dependencies = [
     { name = "fastmcp" },
+    { name = "jsonschema" },
     { name = "mcp", extra = ["cli", "rich"] },
     { name = "pydantic" },
 ]
@@ -534,6 +535,7 @@ dev = [
 [package.metadata]
 requires-dist = [
     { name = "fastmcp", specifier = ">=3.2.4" },
+    { name = "jsonschema", specifier = ">=4.26.0" },
     { name = "mcp", extras = ["cli", "rich"], specifier = ">=1" },
     { name = "pydantic", specifier = ">=2" },
 ]