graph: spec v0.4 middleware (proposal 0004) — phase 2

Implements pipeline-utilities §2-§7: middleware infrastructure plus
the two canonical middleware classes (Retry and Timing). Lands 16 of
the 18 PU middleware fixtures (001-016); fixtures 020-021 remain
skipped because they compose middleware with fan-out which arrives
in Phase 3.

Engine plumbing:
  - middleware/_core.py: Middleware Protocol, NextCall alias,
    compose_chain helper assembling outer-to-inner chains.
  - nodes.py / subgraph.py: each node carries a per-node middleware
    tuple; SubgraphNode treats its tuple as wrapping the dispatch as
    a single atomic call per §4 isolation.
  - builder.py: GraphBuilder.add_middleware for per-graph; per-node
    middleware via add_node(..., middleware=[...]) and
    add_subgraph_node(..., middleware=[...]).
  - compiled.py: _step_function_node + _step_subgraph_node build the
    runtime chain, inject the per-attempt event-dispatch innermost,
    and let raw exceptions propagate through the chain so middleware
    classifiers see the original exception's `category` attribute.
    The engine wraps any exception that escapes the chain as
    NodeException per §4.

Canonical middleware (middleware/retry.py, middleware/timing.py):
  - RetryMiddleware with configurable max_attempts, classifier
    (default matches `provider_rate_limit`, `provider_unavailable`,
    `provider_model_not_loaded` — hardcoded canonical strings,
    loose-coupled to llm-provider §7), backoff (exponential with full
    jitter; deterministic_backoff factory for tests), on_retry
    callback. CancelledError propagates via `except Exception`
    (Python's CancelledError extends BaseException).
  - TimingMiddleware with monotonic-clock duration measurement,
    per-node node_name binding, on_complete callback that receives a
    TimingRecord. Per-instance clock injection (defaulting to
    time.monotonic) lets test fixtures supply a deterministic stub
    without globally patching time.monotonic.

Conformance harness:
  - middleware_seam.py: test-only middleware (TraceRecorder,
    ShortCircuit, ErrorRecovery, ErrorRaiser, StateInspector) used
    by the basic firing / composition / short-circuit / error-
    recovery fixtures.
  - test_pipeline_utilities.py: loads PU fixtures, translates
    middleware blocks into instances, supports cases-shape fixtures
    (014, 016), determinism re-runs (011), state-aware classifiers
    (016), and the deterministic-monotonic clock_stub.
  - adapter.py: extends build_graph with graph_middleware and
    node_middleware kwargs and adds the `flaky` node directive.

Unit tests (test_middleware.py): chain composition + ordering,
short-circuit, pre/post phase symmetry, retry counting + classifier,
retry cancellation propagation, general error recovery, timing on
success and failure, subgraph isolation, default classifier
behavior across direct / via-cause / non-transient cases.

Total tests: 246 passing, 1 skipped (017 fan-out, Phase 3).

Author: chris-colinsky

src/openarmature/graph/__init__.py

@@ -26,6 +26,14 @@
     UnreachableNode,
 )
 from .events import NodeEvent
+from .middleware import (
+    Middleware,
+    RetryMiddleware,
+    TimingMiddleware,
+    TimingRecord,
+    deterministic_backoff,
+    exponential_jitter_backoff,
+)
 from .nodes import FunctionNode, Node
 from .observer import Observer, RemoveHandle, SubscribedObserver
 from .projection import ExplicitMapping, FieldNameMatching, ProjectionStrategy
@@ -48,6 +56,7 @@
     "GraphBuilder",
     "GraphError",
     "MappingReferencesUndeclaredField",
+    "Middleware",
     "MultipleOutgoingEdges",
     "Node",
     "NodeEvent",
@@ -58,15 +67,20 @@
     "Reducer",
     "ReducerError",
     "RemoveHandle",
+    "RetryMiddleware",
     "RoutingError",
     "RuntimeGraphError",
     "State",
     "StateValidationError",
     "StaticEdge",
     "SubgraphNode",
     "SubscribedObserver",
+    "TimingMiddleware",
+    "TimingRecord",
     "UnreachableNode",
     "append",
+    "deterministic_backoff",
+    "exponential_jitter_backoff",
     "last_write_wins",
     "merge",
 ]

src/openarmature/graph/builder.py

@@ -10,7 +10,7 @@
 a type-checked `state` parameter on every callback — without `cast(...)` calls.
 """
 
-from collections.abc import Awaitable, Callable, Mapping
+from collections.abc import Awaitable, Callable, Iterable, Mapping
 from typing import Any, Self, cast
 
 from .compiled import CompiledGraph
@@ -22,6 +22,7 @@
     NoDeclaredEntry,
     UnreachableNode,
 )
+from .middleware import Middleware
 from .nodes import FunctionNode, Node
 from .projection import FieldNameMatching, ProjectionStrategy
 from .reducers import Reducer
@@ -37,29 +38,56 @@ def __init__(self, state_cls: type[StateT]) -> None:
         self._nodes: dict[str, Node[StateT]] = {}
         self._edges: list[StaticEdge | ConditionalEdge[StateT]] = []
         self._entry: str | None = None
+        # Per-graph middleware in registration order (outer-to-inner).
+        # Composed OUTSIDE per-node middleware at runtime per spec §3.
+        self._middleware: list[Middleware] = []
 
     def add_node(
         self,
         name: str,
         fn: Callable[[StateT], Awaitable[Mapping[str, Any]]],
+        *,
+        middleware: Iterable[Middleware] | None = None,
     ) -> Self:
         if name in self._nodes:
             raise ValueError(f"node {name!r} already declared")
-        self._nodes[name] = FunctionNode[StateT](name=name, fn=fn)
+        self._nodes[name] = FunctionNode[StateT](
+            name=name,
+            fn=fn,
+            middleware=tuple(middleware) if middleware is not None else (),
+        )
         return self
 
     def add_subgraph_node[ChildT: State](
         self,
         name: str,
         compiled: CompiledGraph[ChildT],
         projection: ProjectionStrategy[StateT, ChildT] | None = None,
+        *,
+        middleware: Iterable[Middleware] | None = None,
     ) -> Self:
         if name in self._nodes:
             raise ValueError(f"node {name!r} already declared")
         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)
+        self._nodes[name] = SubgraphNode[StateT, ChildT](
+            name=name,
+            compiled=compiled,
+            projection=proj,
+            middleware=tuple(middleware) if middleware is not None else (),
+        )
+        return self
+
+    def add_middleware(self, middleware: Middleware) -> Self:
+        """Register a per-graph middleware applied to every node in this graph.
+
+        Per spec pipeline-utilities §3: per-graph middleware composes
+        OUTSIDE per-node middleware. Calling order is preserved
+        (outer-to-inner) — earlier ``add_middleware`` calls produce
+        outer layers in the runtime chain.
+        """
+        self._middleware.append(middleware)
         return self
 
     def add_edge(self, source: str, target: str | EndSentinel) -> Self:
@@ -142,6 +170,7 @@ def compile(self) -> CompiledGraph[StateT]:
             nodes=dict(self._nodes),
             edges=edges_by_source,
             reducers=resolved,
+            middleware=tuple(self._middleware),
         )
 
     def _reachable_nodes(

src/openarmature/graph/compiled.py

@@ -24,7 +24,7 @@
 import asyncio
 from collections.abc import Iterable, Mapping
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Any, cast
 
 from pydantic import ValidationError
 
@@ -38,6 +38,7 @@
     StateValidationError,
 )
 from .events import NodeEvent
+from .middleware import ChainCall, Middleware, compose_chain
 from .nodes import Node
 from .observer import (
     _DRAIN_SENTINEL,
@@ -113,6 +114,9 @@ class CompiledGraph[StateT: State]:
     nodes: Mapping[str, Node[StateT]]
     edges: Mapping[str, StaticEdge | ConditionalEdge[StateT]]
     reducers: Mapping[str, Reducer]
+    # Per-graph middleware in registration order (outer-to-inner). Composes
+    # OUTSIDE per-node middleware at runtime per pipeline-utilities §3.
+    middleware: tuple[Middleware, ...] = ()
     # Observer plumbing — see attach_observer/drain. Mutable on a frozen
     # dataclass: the list reference is fixed but its contents change.
     # Parameterized factories so pyright infers the element types.
@@ -263,13 +267,17 @@ async def _invoke(
                 # subgraph state-class init errors) escape the spec §4
                 # categories, so we wrap them as NodeException tagged with
                 # the wrapper's name.
-                try:
-                    partial = await node.run(state, context=context)
-                except RuntimeGraphError:
-                    raise
-                except Exception as e:
-                    raise NodeException(node_name=current, cause=e, recoverable_state=state) from e
-                state = _merge_partial(state, partial, self.reducers, current)
+                #
+                # Per pipeline-utilities §4: the parent's middleware wraps
+                # the subgraph dispatch as a single atomic call. Subgraph-
+                # internal nodes have their own middleware (from the
+                # subgraph's own CompiledGraph.middleware tuple) and do
+                # NOT see the parent's middleware. Cast erases ChildT
+                # because the dispatcher only needs to invoke `node.run`
+                # and pass the parent's chain — the inner state class
+                # lives on the subgraph's own CompiledGraph.
+                sub = cast("SubgraphNode[StateT, State]", node)
+                state = await self._step_subgraph_node(sub, current, state, context)
             else:
                 state = await self._step_function_node(node, current, state, context)
 
@@ -297,36 +305,146 @@ async def _step_function_node(
         state: StateT,
         context: _InvocationContext,
     ) -> StateT:
-        """Run one function-node step: take a step, dispatch started, run,
-        merge, dispatch completed.
-
-        Per spec v0.6.0 §6: each attempt produces a started/completed pair.
-        Both events share the same `step`. The completed event carries
-        `post_state` on success, or `error` on failure (one of run, reducer,
-        or state-validation). The completed event is dispatched before the
-        failure propagates.
+        """Run one function-node step through the middleware chain.
+
+        Per pipeline-utilities §3, the runtime chain composes:
+
+            [per_graph...] -> [per_node...] -> innermost
+
+        where ``innermost`` is the per-attempt dispatch wrapper around
+        ``node.run`` + reducer merge + observer event dispatch. Each call
+        to ``innermost`` is one attempt; middleware that calls ``next``
+        repeatedly (e.g., retry) produces multiple attempts and therefore
+        multiple started/completed event pairs from the engine, each
+        tagged with an incrementing ``attempt_index`` (graph-engine §6).
+
+        The chain is built fresh per dispatch so each step has its own
+        attempt counter. Subgraph isolation per pipeline-utilities §4 is
+        achieved by NOT including the parent's per-graph middleware when
+        the subgraph's own ``_step_function_node`` runs — each
+        CompiledGraph carries its own ``middleware`` tuple.
         """
         step = context.take_step()
         namespace = context.namespace_prefix + (current,)
-        pre_state = state
 
-        self._dispatch_started(context, current, namespace, step, pre_state)
+        # The innermost layer dispatches the per-attempt event pair around a
+        # single call to ``node.run``. Each call to this inner increments
+        # the attempt counter; middleware composes around it.
+        attempt_counter = [0]
+
+        async def innermost(s: Any) -> Mapping[str, Any]:
+            # Per pipeline-utilities §5 + graph-engine §6: per-attempt
+            # events use the wrapped §4 error type (NodeException etc.)
+            # for the observer's `error` field, but the RAW exception
+            # propagates up the chain so middleware classifiers can read
+            # the original `category` attribute (timing's
+            # exception_category, retry's classifier). The engine wraps
+            # any exception that escapes the chain, OUTSIDE this layer.
+            attempt_index = attempt_counter[0]
+            attempt_counter[0] = attempt_index + 1
+
+            self._dispatch_started(context, current, namespace, step, s, attempt_index=attempt_index)
+
+            try:
+                partial = await node.run(s)
+            except Exception as e:
+                wrapped = NodeException(node_name=current, cause=e, recoverable_state=s)
+                self._dispatch_completed(
+                    context,
+                    current,
+                    namespace,
+                    step,
+                    s,
+                    error=wrapped,
+                    attempt_index=attempt_index,
+                )
+                raise
+
+            try:
+                merged = _merge_partial(s, partial, self.reducers, current)
+            except (ReducerError, StateValidationError) as e:
+                self._dispatch_completed(
+                    context,
+                    current,
+                    namespace,
+                    step,
+                    s,
+                    error=e,
+                    attempt_index=attempt_index,
+                )
+                raise
+
+            self._dispatch_completed(
+                context,
+                current,
+                namespace,
+                step,
+                s,
+                post_state=merged,
+                attempt_index=attempt_index,
+            )
+            # Return the partial (not the merged state) so middleware sees
+            # the partial-update shape per pipeline-utilities §2. The
+            # engine's canonical merge against the original state happens
+            # below, after the chain returns.
+            return partial
+
+        chain: ChainCall = compose_chain(
+            list(self.middleware) + list(node.middleware),
+            innermost,
+        )
 
         try:
-            partial = await node.run(state)
+            final_partial = await chain(state)
+        except RuntimeGraphError:
+            raise
         except Exception as e:
-            wrapped = NodeException(node_name=current, cause=e, recoverable_state=state)
-            self._dispatch_completed(context, current, namespace, step, pre_state, error=wrapped)
-            raise wrapped from e
+            # A raw exception (node-raised or middleware-raised) escaped
+            # the chain unrecovered. Wrap as NodeException per §4.
+            raise NodeException(node_name=current, cause=e, recoverable_state=state) from e
+        # Engine's canonical merge uses the ORIGINAL state per §2: "the
+        # transformed state is passed to ``next``, NOT to the engine's
+        # merge step." If middleware transformed state mid-chain, the
+        # per-attempt completed events showed the transformed merge for
+        # observability, but the state advancing the graph loop is built
+        # from the original.
+        return _merge_partial(state, final_partial, self.reducers, current)
+
+    async def _step_subgraph_node(
+        self,
+        node: SubgraphNode[StateT, State],
+        current: str,
+        state: StateT,
+        context: _InvocationContext,
+    ) -> StateT:
+        """Run one subgraph-as-node step through the parent's middleware chain.
 
-        try:
-            new_state = _merge_partial(state, partial, self.reducers, current)
-        except (ReducerError, StateValidationError) as e:
-            self._dispatch_completed(context, current, namespace, step, pre_state, error=e)
-            raise
+        Per pipeline-utilities §4: the parent's per-graph middleware plus
+        any per-node middleware on the SubgraphNode wraps the subgraph
+        dispatch as a single atomic call. The subgraph's INTERNAL nodes
+        get their own middleware via the subgraph's own CompiledGraph;
+        parent middleware does NOT cross the boundary.
 
-        self._dispatch_completed(context, current, namespace, step, pre_state, post_state=new_state)
-        return new_state
+        No started/completed events fire for the wrapper itself; the
+        events come from the subgraph's internal node executions (per
+        fixture 013).
+        """
+
+        async def innermost(s: Any) -> Mapping[str, Any]:
+            try:
+                return await node.run(s, context=context)
+            except RuntimeGraphError:
+                raise
+            except Exception as e:
+                raise NodeException(node_name=current, cause=e, recoverable_state=s) from e
+
+        chain: ChainCall = compose_chain(
+            list(self.middleware) + list(node.middleware),
+            innermost,
+        )
+
+        final_partial = await chain(state)
+        return _merge_partial(state, final_partial, self.reducers, current)
 
     @staticmethod
     def _dispatch_started(
@@ -335,6 +453,8 @@ def _dispatch_started(
         namespace: tuple[str, ...],
         step: int,
         pre_state: State,
+        *,
+        attempt_index: int = 0,
     ) -> None:
         _dispatch(
             context,
@@ -347,6 +467,7 @@ def _dispatch_started(
                 post_state=None,
                 error=None,
                 parent_states=context.parent_states_prefix,
+                attempt_index=attempt_index,
             ),
         )
 
@@ -360,6 +481,7 @@ def _dispatch_completed(
         *,
         post_state: State | None = None,
         error: RuntimeGraphError | None = None,
+        attempt_index: int = 0,
     ) -> None:
         _dispatch(
             context,
@@ -372,5 +494,6 @@ def _dispatch_completed(
                 post_state=post_state,
                 error=error,
                 parent_states=context.parent_states_prefix,
+                attempt_index=attempt_index,
             ),
         )