88Per spec §4 Error semantics: node, edge, reducer, and routing errors carry
99recoverable state; state validation errors do not.
1010
11+ Per spec v0.3.0 §6 Observer hooks: between merge and edge evaluation, the
12+ engine dispatches a `NodeEvent` for the just-completed node onto the
13+ invocation's delivery queue. On node/reducer/state-validation failure, the
14+ event is dispatched (with `error` populated) before the failure propagates.
15+ Routing errors do NOT produce their own event — they arise after the
16+ preceding node's event has already been dispatched.
17+
1118`CompiledGraph[StateT]` and `_merge_partial[StateT]` carry the concrete state
1219subclass through to `invoke()`'s return type, so consumers don't need
1320`cast(MyState, ...)` at the call site.
1421"""
1522
16- from collections .abc import Mapping
17- from dataclasses import dataclass
23+ import asyncio
24+ from collections .abc import Iterable , Mapping
25+ from dataclasses import dataclass , field
1826from typing import Any
1927
2028from pydantic import ValidationError
2533 NodeException ,
2634 ReducerError ,
2735 RoutingError ,
36+ RuntimeGraphError ,
2837 StateValidationError ,
2938)
39+ from .events import NodeEvent
3040from .nodes import Node
41+ from .observer import (
42+ _DRAIN_SENTINEL ,
43+ Observer ,
44+ RemoveHandle ,
45+ _dispatch ,
46+ _InvocationContext ,
47+ _QueuedItem ,
48+ deliver_loop ,
49+ )
3150from .reducers import Reducer
3251from .state import State
52+ from .subgraph import SubgraphNode
3353
3454
3555def _merge_partial [StateT : State ](
@@ -77,38 +97,141 @@ def _merge_partial[StateT: State](
7797
7898@dataclass (frozen = True )
7999class CompiledGraph [StateT : State ]:
80- """An immutable, executable graph produced by `GraphBuilder.compile()`."""
100+ """An immutable, executable graph produced by `GraphBuilder.compile()`.
101+
102+ The compile-time topology (state class, entry, nodes, edges, reducers) is
103+ immutable. Two mutable lists ride alongside for observer plumbing —
104+ `_attached_observers` and `_active_workers` — neither of which affect the
105+ compiled topology and both of which are scoped to the same instance.
106+ """
81107
82108 state_cls : type [StateT ]
83109 entry : str
84110 nodes : Mapping [str , Node [StateT ]]
85111 edges : Mapping [str , StaticEdge | ConditionalEdge [StateT ]]
86112 reducers : Mapping [str , Reducer ]
113+ # Observer plumbing — see attach_observer/drain. Mutable on a frozen
114+ # dataclass: the list reference is fixed but its contents change.
115+ # Parameterized factories so pyright infers the element types.
116+ _attached_observers : list [Observer ] = field (default_factory = list [Observer ])
117+ _active_workers : list [asyncio .Task [None ]] = field (default_factory = list [asyncio .Task [None ]])
118+
119+ # ------------------------------------------------------------------
120+ # Observer registration (spec v0.3.0 §6)
121+ # ------------------------------------------------------------------
122+
123+ def attach_observer (self , observer : Observer ) -> RemoveHandle :
124+ """Register a graph-attached observer.
125+
126+ Per spec v0.3.0 §6: graph-attached observers fire on every invocation
127+ of this graph until removed — including when this graph runs as a
128+ subgraph inside a parent. Returns a `RemoveHandle` whose `.remove()`
129+ method detaches the observer; idempotent.
130+
131+ Per spec: changes to the registered set during a graph run do NOT
132+ take effect until the next invocation. The set of observers
133+ delivering events for an in-flight invocation is fixed at the point
134+ the invocation begins.
135+ """
136+ self ._attached_observers .append (observer )
137+ return RemoveHandle (_observers = self ._attached_observers , _observer = observer )
87138
88- async def invoke (self , initial_state : StateT ) -> StateT :
139+ async def drain (self ) -> None :
140+ """Await delivery of every observer event produced by prior
141+ invocations of this graph.
142+
143+ Per spec v0.3.0 §6: callers running in short-lived processes (scripts,
144+ serverless functions, CLIs) MUST use drain to avoid losing observer
145+ events that were dispatched but not yet delivered.
146+
147+ Only events dispatched before this call are awaited; events from
148+ invocations started concurrently with drain may or may not be
149+ included. Subgraph events from active invocations are part of the
150+ parent invocation's worker and are covered automatically.
151+ """
152+ if not self ._active_workers :
153+ return
154+ pending = list (self ._active_workers )
155+ await asyncio .gather (* pending , return_exceptions = True )
156+ for w in pending :
157+ if w in self ._active_workers :
158+ self ._active_workers .remove (w )
159+
160+ # ------------------------------------------------------------------
161+ # Public invocation
162+ # ------------------------------------------------------------------
163+
164+ async def invoke (
165+ self ,
166+ initial_state : StateT ,
167+ observers : Iterable [Observer ] | None = None ,
168+ ) -> StateT :
89169 """Run the graph from `initial_state` to END and return the final state.
90170
171+ Optional `observers` are invocation-scoped — they fire only for this
172+ run, after all graph-attached observers (including subgraph-attached
173+ ones for events originating in subgraphs) per spec v0.3.0 §6.
174+
175+ Per spec v0.3.0 §6: this method returns as soon as the graph
176+ execution loop completes, regardless of whether the observer
177+ delivery queue has finished processing every dispatched event. Use
178+ `await compiled.drain()` if you need delivery-completion guarantees.
179+
91180 Raises one of the runtime error categories from spec §4 on failure.
92181 """
93182
183+ invocation_scoped = tuple (observers ) if observers else ()
184+ queue : asyncio .Queue [_QueuedItem | None ] = asyncio .Queue ()
185+ context = _InvocationContext (
186+ queue = queue ,
187+ graph_attached = tuple (self ._attached_observers ),
188+ invocation_scoped = invocation_scoped ,
189+ )
190+ worker = asyncio .create_task (deliver_loop (queue ))
191+ self ._active_workers .append (worker )
192+ try :
193+ return await self ._invoke (initial_state , context )
194+ finally :
195+ # Sentinel terminates the worker after it processes events
196+ # already on the queue (including any error event we just
197+ # dispatched on the failure path). Drain semantics live on
198+ # `.drain()` — we do NOT await the worker here, per spec.
199+ queue .put_nowait (_DRAIN_SENTINEL )
200+
201+ # ------------------------------------------------------------------
202+ # Internal invocation (used by SubgraphNode for nested execution)
203+ # ------------------------------------------------------------------
204+
205+ async def _invoke (
206+ self ,
207+ initial_state : StateT ,
208+ context : _InvocationContext ,
209+ ) -> StateT :
210+ """Execution loop that dispatches events through the supplied context.
211+
212+ Public `invoke()` builds a fresh root context. Subgraph-as-node
213+ execution calls `_invoke` directly with a context derived from the
214+ parent's, so the queue, step counter, and observer chain thread
215+ through the boundary.
216+ """
217+
94218 state = initial_state
95219 current = self .entry
96220
97221 while True :
98222 node = self .nodes [current ]
99223
100- # Run the node. Wrap user exceptions as NodeException with
101- # recoverable_state = state at point of failure (pre-update).
102- try :
103- partial = await node . run ( state )
104- except Exception as e :
105- raise NodeException ( node_name = current , cause = e , recoverable_state = state ) from e
106-
107- # Merge partial into state via reducers (may raise ReducerError or
108- # StateValidationError; both already carry the right context).
109- state = _merge_partial ( state , partial , self . reducers , current )
224+ if isinstance ( node , SubgraphNode ):
225+ # Subgraph wrappers are transparent to the observer protocol
226+ # (per fixture 013): no event is dispatched for the wrapper
227+ # itself, the step counter does not advance for it, and any
228+ # exception bubbling up from the subgraph's _invoke is
229+ # already wrapped with the inner node's identity.
230+ partial = await node . run ( state , context = context )
231+ state = _merge_partial ( state , partial , self . reducers , current )
232+ else :
233+ state = await self . _step_function_node ( node , current , state , context )
110234
111- # Evaluate the outgoing edge against the post-update state.
112235 edge = self .edges [current ]
113236 if isinstance (edge , StaticEdge ):
114237 target : str | EndSentinel = edge .target
@@ -125,3 +248,71 @@ async def invoke(self, initial_state: StateT) -> StateT:
125248 raise RoutingError (source_node = current , returned = target , recoverable_state = state )
126249
127250 current = target
251+
252+ async def _step_function_node (
253+ self ,
254+ node : Node [StateT ],
255+ current : str ,
256+ state : StateT ,
257+ context : _InvocationContext ,
258+ ) -> StateT :
259+ """Run one function-node step: take a step, run, merge, dispatch.
260+
261+ Dispatches a `NodeEvent` exactly once per call:
262+ - On run failure (NodeException): event with error populated.
263+ - On merge failure (ReducerError or StateValidationError): event with
264+ error populated; the original error propagates unchanged after.
265+ - On success: event with post_state populated, then return.
266+ """
267+ step = context .take_step ()
268+ namespace = context .namespace_prefix + (current ,)
269+ pre_state = state
270+
271+ try :
272+ partial = await node .run (state )
273+ except Exception as e :
274+ wrapped = NodeException (node_name = current , cause = e , recoverable_state = state )
275+ self ._dispatch_failure_event (context , current , namespace , step , pre_state , wrapped )
276+ raise wrapped from e
277+
278+ try :
279+ new_state = _merge_partial (state , partial , self .reducers , current )
280+ except (ReducerError , StateValidationError ) as e :
281+ self ._dispatch_failure_event (context , current , namespace , step , pre_state , e )
282+ raise
283+
284+ _dispatch (
285+ context ,
286+ NodeEvent (
287+ node_name = current ,
288+ namespace = namespace ,
289+ step = step ,
290+ pre_state = pre_state ,
291+ post_state = new_state ,
292+ error = None ,
293+ parent_states = context .parent_states_prefix ,
294+ ),
295+ )
296+ return new_state
297+
298+ @staticmethod
299+ def _dispatch_failure_event (
300+ context : _InvocationContext ,
301+ current : str ,
302+ namespace : tuple [str , ...],
303+ step : int ,
304+ pre_state : State ,
305+ error : RuntimeGraphError ,
306+ ) -> None :
307+ _dispatch (
308+ context ,
309+ NodeEvent (
310+ node_name = current ,
311+ namespace = namespace ,
312+ step = step ,
313+ pre_state = pre_state ,
314+ post_state = None ,
315+ error = error ,
316+ parent_states = context .parent_states_prefix ,
317+ ),
318+ )
0 commit comments