graph: parameterize GraphBuilder/CompiledGraph on StateT (PEP 695 generics) (#3)

* graph: parameterize builder/compiled/subgraph on StateT (PEP 695 generics)

Carry the concrete State subclass through at type-check time.
GraphBuilder[StateT], CompiledGraph[StateT], SubgraphNode[ParentT, ChildT],
ConditionalEdge[StateT], ProjectionStrategy[ParentT, ChildT], Node[StateT],
FunctionNode[StateT] — all generic. Consumers no longer need
cast(MyState, ...) on invoke() returns, projection arguments, or edge
function parameters.

PEP 695 syntax throughout (targets Python 3.12+). SubgraphNode carries
both parent and child types; add_subgraph_node[ChildT: State] infers
ChildT from the compiled argument. FieldNameMatching[ParentT, ChildT]
is generic for Protocol conformance even though ParentT is unused in
project_in (keeps default-factory inference clean at the SubgraphNode
boundary).

_merge_partial[StateT] parameterized so type(prior).model_validate
preserves the subclass without an internal cast. Conformance adapter
pins to [State, State, State] since YAML fixtures are schema-agnostic.

Adds tests/unit/test_generics.py with assert_type to lock the typing
surface against regression. 41 tests pass, pyright strict clean.

* gitignore _docs/

Author: chris-colinsky

.gitignore

@@ -1,5 +1,9 @@
 .claude/
 
+# Local-only implementation notes (concepts walkthrough, rough edges log,
+# future-work log, langgraph comparison). Kept out of public history.
+_docs/
+
 # Python
 __pycache__/
 *.py[cod]

src/openarmature/graph/builder.py

@@ -3,6 +3,11 @@
 Per spec §2: compilation MUST fail if the graph has no declared entry,
 unreachable nodes, dangling edges, a node with more than one outgoing edge,
 or a field with more than one declared reducer.
+
+`GraphBuilder[StateT]` is parameterized on the graph's state type. Node
+functions, conditional-edge functions, and the returned `CompiledGraph[StateT]`
+all carry `StateT` forward so consumers get typed `invoke()` return values and
+a type-checked `state` parameter on every callback — without `cast(...)` calls.
 """
 
 from collections.abc import Awaitable, Callable, Mapping
@@ -24,35 +29,37 @@
 from .subgraph import SubgraphNode
 
 
-class GraphBuilder:
+class GraphBuilder[StateT: State]:
     """Mutable builder for a graph; call `compile()` to produce a `CompiledGraph`."""
 
-    def __init__(self, state_cls: type[State]) -> None:
-        self.state_cls = state_cls
-        self._nodes: dict[str, Node] = {}
-        self._edges: list[StaticEdge | ConditionalEdge] = []
+    def __init__(self, state_cls: type[StateT]) -> None:
+        self.state_cls: type[StateT] = state_cls
+        self._nodes: dict[str, Node[StateT]] = {}
+        self._edges: list[StaticEdge | ConditionalEdge[StateT]] = []
         self._entry: str | None = None
 
     def add_node(
         self,
         name: str,
-        fn: Callable[[Any], Awaitable[Mapping[str, Any]]],
+        fn: Callable[[StateT], Awaitable[Mapping[str, Any]]],
     ) -> Self:
         if name in self._nodes:
             raise ValueError(f"node {name!r} already declared")
-        self._nodes[name] = FunctionNode(name=name, fn=fn)
+        self._nodes[name] = FunctionNode[StateT](name=name, fn=fn)
         return self
 
-    def add_subgraph_node(
+    def add_subgraph_node[ChildT: State](
         self,
         name: str,
-        compiled: CompiledGraph,
-        projection: ProjectionStrategy | None = None,
+        compiled: CompiledGraph[ChildT],
+        projection: ProjectionStrategy[StateT, ChildT] | None = None,
     ) -> Self:
         if name in self._nodes:
             raise ValueError(f"node {name!r} already declared")
-        proj: ProjectionStrategy = projection if projection is not None else FieldNameMatching()
-        self._nodes[name] = SubgraphNode(name=name, compiled=compiled, projection=proj)
+        proj: ProjectionStrategy[StateT, ChildT] = (
+            projection if projection is not None else FieldNameMatching[StateT, ChildT]()
+        )
+        self._nodes[name] = SubgraphNode[StateT, ChildT](name=name, compiled=compiled, projection=proj)
         return self
 
     def add_edge(self, source: str, target: str | EndSentinel) -> Self:
@@ -62,16 +69,16 @@ def add_edge(self, source: str, target: str | EndSentinel) -> Self:
     def add_conditional_edge(
         self,
         source: str,
-        fn: Callable[[Any], str | EndSentinel],
+        fn: Callable[[StateT], str | EndSentinel],
     ) -> Self:
-        self._edges.append(ConditionalEdge(source=source, fn=fn))
+        self._edges.append(ConditionalEdge[StateT](source=source, fn=fn))
         return self
 
     def set_entry(self, name: str) -> Self:
         self._entry = name
         return self
 
-    def compile(self) -> CompiledGraph:
+    def compile(self) -> CompiledGraph[StateT]:
         # 1. ConflictingReducers — state schema check.
         per_field = field_reducers(self.state_cls)
         for fname, declared in per_field.items():
@@ -98,7 +105,7 @@ def compile(self) -> CompiledGraph:
                     raise DanglingEdge(source=edge.source, target=edge.target)
 
         # 5. MultipleOutgoingEdges + index by source for the reachability pass.
-        edges_by_source: dict[str, StaticEdge | ConditionalEdge] = {}
+        edges_by_source: dict[str, StaticEdge | ConditionalEdge[StateT]] = {}
         for edge in self._edges:
             if edge.source in edges_by_source:
                 raise MultipleOutgoingEdges(edge.source)
@@ -112,7 +119,7 @@ def compile(self) -> CompiledGraph:
             if node_name not in reachable:
                 raise UnreachableNode(node_name)
 
-        return CompiledGraph(
+        return CompiledGraph[StateT](
             state_cls=self.state_cls,
             entry=self._entry,
             nodes=dict(self._nodes),
@@ -122,7 +129,7 @@ def compile(self) -> CompiledGraph:
 
     def _reachable_nodes(
         self,
-        edges_by_source: Mapping[str, StaticEdge | ConditionalEdge],
+        edges_by_source: Mapping[str, StaticEdge | ConditionalEdge[StateT]],
     ) -> set[str]:
         assert self._entry is not None
         reachable: set[str] = {self._entry}

src/openarmature/graph/compiled.py

@@ -7,6 +7,10 @@
 
 Per spec §4 Error semantics: node, edge, reducer, and routing errors carry
 recoverable state; state validation errors do not.
+
+`CompiledGraph[StateT]` and `_merge_partial[StateT]` carry the concrete state
+subclass through to `invoke()`'s return type, so consumers don't need
+`cast(MyState, ...)` at the call site.
 """
 
 from collections.abc import Mapping
@@ -28,12 +32,12 @@
 from .state import State
 
 
-def _merge_partial(
-    prior: State,
+def _merge_partial[StateT: State](
+    prior: StateT,
     partial: Mapping[str, Any],
     reducers: Mapping[str, Reducer],
     producing_node: str,
-) -> State:
+) -> StateT:
     """Apply per-field reducers to merge a node's partial update into prior state.
 
     Re-validates the resulting state against the schema (per spec §2 SHOULD
@@ -60,6 +64,7 @@ def _merge_partial(
             ) from e
 
     try:
+        # type(prior) narrows to `type[StateT]`; model_validate returns StateT.
         return type(prior).model_validate(new_values)
     except ValidationError as e:
         offending = sorted({str(err["loc"][0]) for err in e.errors() if err["loc"]})
@@ -71,16 +76,16 @@ def _merge_partial(
 
 
 @dataclass(frozen=True)
-class CompiledGraph:
+class CompiledGraph[StateT: State]:
     """An immutable, executable graph produced by `GraphBuilder.compile()`."""
 
-    state_cls: type[State]
+    state_cls: type[StateT]
     entry: str
-    nodes: Mapping[str, Node]
-    edges: Mapping[str, StaticEdge | ConditionalEdge]
+    nodes: Mapping[str, Node[StateT]]
+    edges: Mapping[str, StaticEdge | ConditionalEdge[StateT]]
     reducers: Mapping[str, Reducer]
 
-    async def invoke(self, initial_state: State) -> State:
+    async def invoke(self, initial_state: StateT) -> StateT:
         """Run the graph from `initial_state` to END and return the final state.
 
         Raises one of the runtime error categories from spec §4 on failure.

src/openarmature/graph/edges.py

@@ -3,11 +3,17 @@
 Per spec §2 Concepts (Edge, END): edges are static or conditional; each node
 has exactly one outgoing edge. END is a distinct engine sentinel (not a
 reserved node name) used as a routing target to halt execution.
+
+`ConditionalEdge` is generic on the outer graph's state type so the routing
+function's parameter is typed against the user's `State` subclass — not
+`Any` — at type-check time.
 """
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Final
+from typing import Final
+
+from .state import State
 
 
 class EndSentinel:
@@ -29,11 +35,11 @@ class StaticEdge:
 
 
 @dataclass(frozen=True)
-class ConditionalEdge:
+class ConditionalEdge[StateT: State]:
     """Routes from `source` to whichever node `fn(state)` returns. The function
     MUST return either a declared node name or `END`; any other value raises
     `RoutingError` at runtime.
     """
 
     source: str
-    fn: Callable[[Any], str | EndSentinel]
+    fn: Callable[[StateT], str | EndSentinel]

src/openarmature/graph/nodes.py

@@ -5,29 +5,33 @@
 partial update which the engine merges via reducers.
 
 The `Node` Protocol exists so subgraphs can compose as nodes alongside
-plain function-backed nodes (see `subgraph.SubgraphNode`).
+plain function-backed nodes (see `subgraph.SubgraphNode`). Both are
+parameterized on `StateT` so the outer graph's state type flows through
+to node functions at type-check time.
 """
 
 from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass
 from typing import Any, Protocol
 
+from .state import State
 
-class Node(Protocol):
+
+class Node[StateT: State](Protocol):
     """A unit of work in a compiled graph."""
 
     @property
     def name(self) -> str: ...
 
-    async def run(self, state: Any) -> Mapping[str, Any]: ...
+    async def run(self, state: StateT) -> Mapping[str, Any]: ...
 
 
 @dataclass(frozen=True)
-class FunctionNode:
+class FunctionNode[StateT: State]:
     """A node backed by an async callable."""
 
     name: str
-    fn: Callable[[Any], Awaitable[Mapping[str, Any]]]
+    fn: Callable[[StateT], Awaitable[Mapping[str, Any]]]
 
-    async def run(self, state: Any) -> Mapping[str, Any]:
+    async def run(self, state: StateT) -> Mapping[str, Any]:
         return await self.fn(state)

src/openarmature/graph/projection.py

@@ -7,7 +7,9 @@
 
 `ProjectionStrategy` is exposed as a seam so proposal 0002 (explicit
 input/output mapping) can slot in without changes to the engine's compile or
-execute paths.
+execute paths. Parameterized on the parent and child state types so
+consumer-authored projections get typed `project_in` / `project_out`
+signatures without `cast(...)` gymnastics.
 """
 
 from collections.abc import Mapping
@@ -16,30 +18,37 @@
 from .state import State
 
 
-class ProjectionStrategy(Protocol):
+class ProjectionStrategy[ParentT: State, ChildT: State](Protocol):
     """Strategy for moving state across the parent ↔ subgraph boundary."""
 
-    def project_in(self, parent_state: State, subgraph_state_cls: type[State]) -> State: ...
+    def project_in(self, parent_state: ParentT, subgraph_state_cls: type[ChildT]) -> ChildT: ...
 
     def project_out(
         self,
-        subgraph_final_state: State,
-        parent_state: State,
-        subgraph_state_cls: type[State],
+        subgraph_final_state: ChildT,
+        parent_state: ParentT,
+        subgraph_state_cls: type[ChildT],
     ) -> Mapping[str, Any]: ...
 
 
-class FieldNameMatching:
-    """Default projection per spec v0.1.1 §2 Subgraph."""
+class FieldNameMatching[ParentT: State, ChildT: State]:
+    """Default projection per spec v0.1.1 §2 Subgraph.
 
-    def project_in(self, parent_state: State, subgraph_state_cls: type[State]) -> State:
+    Parameterized for protocol conformance under generics. `ParentT` is not
+    consumed (the default projection ignores parent state on the way in),
+    but carrying the type variable keeps the default assignable to
+    `ProjectionStrategy[ParentT, ChildT]` without type gymnastics at the
+    SubgraphNode default-factory site.
+    """
+
+    def project_in(self, parent_state: ParentT, subgraph_state_cls: type[ChildT]) -> ChildT:
         return subgraph_state_cls()
 
     def project_out(
         self,
-        subgraph_final_state: State,
-        parent_state: State,
-        subgraph_state_cls: type[State],
+        subgraph_final_state: ChildT,
+        parent_state: ParentT,
+        subgraph_state_cls: type[ChildT],
     ) -> Mapping[str, Any]:
         parent_fields = set(type(parent_state).model_fields.keys())
         sub_fields = set(subgraph_state_cls.model_fields.keys())

src/openarmature/graph/subgraph.py

@@ -4,6 +4,11 @@
 graph. The subgraph runs against its own state schema; projection between
 parent and subgraph is delegated to a `ProjectionStrategy` (default:
 `FieldNameMatching`).
+
+Parameterized on both the parent's state type (`ParentT`) and the subgraph's
+state type (`ChildT`). The outer graph only ever sees `run(state: ParentT)`
+— the `ChildT` lives on the `compiled` and `projection` fields and is
+invisible at the outer graph's node dispatch site.
 """
 
 from collections.abc import Mapping
@@ -18,14 +23,16 @@
 
 
 @dataclass(frozen=True)
-class SubgraphNode:
+class SubgraphNode[ParentT: State, ChildT: State]:
     """A node backed by a compiled subgraph."""
 
     name: str
-    compiled: "CompiledGraph"
-    projection: ProjectionStrategy = field(default_factory=FieldNameMatching)
+    compiled: "CompiledGraph[ChildT]"
+    projection: ProjectionStrategy[ParentT, ChildT] = field(
+        default_factory=FieldNameMatching[ParentT, ChildT]
+    )
 
-    async def run(self, state: State) -> Mapping[str, Any]:
+    async def run(self, state: ParentT) -> Mapping[str, Any]:
         sub_initial = self.projection.project_in(state, self.compiled.state_cls)
         sub_final = await self.compiled.invoke(sub_initial)
         return self.projection.project_out(sub_final, state, self.compiled.state_cls)

tests/conformance/adapter.py

@@ -119,9 +119,9 @@ async def fn(_state: Any) -> Mapping[str, Any]:
 
 def _make_subgraph_fn(
     node_name: str,
-    compiled: CompiledGraph,
+    compiled: CompiledGraph[State],
     trace: list[str],
-    projection: ProjectionStrategy,
+    projection: ProjectionStrategy[State, State],
 ) -> Callable[[Any], Awaitable[Mapping[str, Any]]]:
     """Outer-graph node that delegates to a compiled subgraph.
 
@@ -160,7 +160,7 @@ class BuiltGraph:
     """Result of translating a fixture into runnable engine constructs."""
 
     state_cls: type[State]
-    builder: GraphBuilder
+    builder: GraphBuilder[State]
     trace: list[str]
 
     def initial_state(self, overrides: Mapping[str, Any]) -> State:
@@ -170,7 +170,7 @@ def initial_state(self, overrides: Mapping[str, Any]) -> State:
 def build_graph(
     spec: Mapping[str, Any],
     *,
-    subgraphs: Mapping[str, CompiledGraph] | None = None,
+    subgraphs: Mapping[str, CompiledGraph[State]] | None = None,
     trace: list[str] | None = None,
     model_name: str = "FixtureState",
 ) -> BuiltGraph:
@@ -196,7 +196,7 @@ def build_graph(
             compiled = subgraphs[sub_name]
             builder.add_node(
                 node_name,
-                _make_subgraph_fn(node_name, compiled, trace, FieldNameMatching()),
+                _make_subgraph_fn(node_name, compiled, trace, FieldNameMatching[State, State]()),
             )
         elif "raises" in node_spec:
             builder.add_node(node_name, _make_raising_fn(node_name, node_spec["raises"], trace))