diff --git a/examples/00-hello-world/main.py b/examples/00-hello-world/main.py index ad8206d5..1591b5ff 100644 --- a/examples/00-hello-world/main.py +++ b/examples/00-hello-world/main.py @@ -49,6 +49,7 @@ END, CompiledGraph, GraphBuilder, + MetadataAugmentationEvent, NodeEvent, State, append, @@ -194,14 +195,18 @@ def route(state: PipelineState) -> str: return state.classification.intent -async def trace(event: NodeEvent) -> None: +async def trace(event: NodeEvent | MetadataAugmentationEvent) -> None: # OpenAIProvider emits NodeEvent-shaped events for LLM-span # tracking under a sentinel namespace; those have post_state=None. + # ``set_invocation_metadata`` from within a node body emits a + # MetadataAugmentationEvent; this tracer ignores those. # Filter to events that carry a PipelineState snapshot before - # reading it. The isinstance check both narrows the type for + # reading it. The isinstance checks both narrow the type for # static checkers (post_state is typed as the base State, not - # PipelineState) and acts as a defensive guard against any + # PipelineState) and act as a defensive guard against any # foreign-state observer event the engine might dispatch. + if not isinstance(event, NodeEvent): + return if event.phase == "completed" and event.error is None and isinstance(event.post_state, PipelineState): print(f"{event.node_name}: sources={event.post_state.sources}") diff --git a/examples/03-observer-hooks/main.py b/examples/03-observer-hooks/main.py index 8c98a01e..2326369b 100644 --- a/examples/03-observer-hooks/main.py +++ b/examples/03-observer-hooks/main.py @@ -54,6 +54,7 @@ CompiledGraph, ExplicitMapping, GraphBuilder, + MetadataAugmentationEvent, NodeEvent, Observer, State, @@ -186,12 +187,18 @@ def build_review_subgraph() -> CompiledGraph[ReviewState]: # fire on every invocation of the compiled graph until removed. -async def console_tracer(event: NodeEvent) -> None: +async def console_tracer(event: NodeEvent | MetadataAugmentationEvent) -> None: """Print one structured line per node boundary to stderr. Format: `[step=N] namespace.path → fields_changed_in_this_step` On error, format flips to `... ✗ error_category`. + + Mid-invocation ``set_invocation_metadata`` augmentations also + reach observers as ``MetadataAugmentationEvent`` instances; this + tracer ignores them. """ + if isinstance(event, MetadataAugmentationEvent): + return namespace = ".".join(event.namespace) if event.error is not None: print( @@ -232,7 +239,9 @@ def __init__(self) -> None: self.errors: int = 0 self.namespaces: set[tuple[str, ...]] = set() - async def __call__(self, event: NodeEvent) -> None: + async def __call__(self, event: NodeEvent | MetadataAugmentationEvent) -> None: + if isinstance(event, MetadataAugmentationEvent): + return self.events += 1 if event.error is not None: self.errors += 1 diff --git a/examples/04-nested-subgraphs/main.py b/examples/04-nested-subgraphs/main.py index f607029f..6c34bc5a 100644 --- a/examples/04-nested-subgraphs/main.py +++ b/examples/04-nested-subgraphs/main.py @@ -49,6 +49,7 @@ CompiledGraph, ExplicitMapping, GraphBuilder, + MetadataAugmentationEvent, NodeEvent, State, append, @@ -349,7 +350,9 @@ def _fmt_state(state: Any) -> str: return " ".join(parts) if parts else "(empty)" -async def depth_observer(event: NodeEvent) -> None: +async def depth_observer(event: NodeEvent | MetadataAugmentationEvent) -> None: + if isinstance(event, MetadataAugmentationEvent): + return depth = len(event.namespace) indent = " " * (depth - 1) ns = " > ".join(event.namespace) diff --git a/examples/05-fan-out-with-retry/main.py b/examples/05-fan-out-with-retry/main.py index 9620c7f1..d2093466 100644 --- a/examples/05-fan-out-with-retry/main.py +++ b/examples/05-fan-out-with-retry/main.py @@ -78,6 +78,7 @@ END, CompiledGraph, GraphBuilder, + MetadataAugmentationEvent, NodeEvent, State, append, @@ -296,7 +297,7 @@ def build_graph(error_policy: str = "fail_fast") -> CompiledGraph[BatchState]: ) -async def fan_out_config_observer(event: NodeEvent) -> None: +async def fan_out_config_observer(event: NodeEvent | MetadataAugmentationEvent) -> None: """Print the fan-out node's resolved config when its dispatch event fires. @@ -308,6 +309,8 @@ async def fan_out_config_observer(event: NodeEvent) -> None: ``concurrency`` are callable resolvers whose value isn't visible in code. """ + if not isinstance(event, NodeEvent): + return if event.fan_out_config is None: return if event.phase != "started": diff --git a/examples/06-parallel-branches/main.py b/examples/06-parallel-branches/main.py index 78a2ed13..edce6500 100644 --- a/examples/06-parallel-branches/main.py +++ b/examples/06-parallel-branches/main.py @@ -70,6 +70,7 @@ BranchSpec, CompiledGraph, GraphBuilder, + MetadataAugmentationEvent, NodeEvent, State, append, @@ -240,7 +241,7 @@ async def present(s: ArticleState) -> Mapping[str, Any]: return {"trace": ["present"]} -async def branch_attribution_observer(event: NodeEvent) -> None: +async def branch_attribution_observer(event: NodeEvent | MetadataAugmentationEvent) -> None: """Print which branch each inner-node event came from. NodeEvent carries ``branch_name`` on events from nodes that @@ -250,6 +251,8 @@ async def branch_attribution_observer(event: NodeEvent) -> None: observer skips events with no branch attribution and prints ``(branch=…) node_name`` for the rest. """ + if not isinstance(event, NodeEvent): + return if event.branch_name is None or event.phase != "started": return print(f" [observer] (branch={event.branch_name}) inner node {event.node_name!r} started") diff --git a/src/openarmature/graph/__init__.py b/src/openarmature/graph/__init__.py index 49742598..bd45a8b5 100644 --- a/src/openarmature/graph/__init__.py +++ b/src/openarmature/graph/__init__.py @@ -35,7 +35,7 @@ StateValidationError, UnreachableNode, ) -from .events import NodeEvent +from .events import MetadataAugmentationEvent, NodeEvent from .fan_out import FanOutConfig, FanOutNode from .middleware import ( Middleware, @@ -78,6 +78,7 @@ "GraphBuilder", "GraphError", "MappingReferencesUndeclaredField", + "MetadataAugmentationEvent", "Middleware", "MultipleOutgoingEdges", "NextCall", diff --git a/src/openarmature/graph/events.py b/src/openarmature/graph/events.py index 9fd84489..c2cb50bd 100644 --- a/src/openarmature/graph/events.py +++ b/src/openarmature/graph/events.py @@ -231,4 +231,46 @@ class NodeEvent: caller_invocation_metadata: Mapping[str, AttributeValue] = field(default_factory=lambda: _EMPTY_METADATA) -__all__ = ["FanOutEventConfig", "NodeEvent"] +# Spec: realizes observability §3.4 + graph-engine §6 augmentation +# event mechanism (proposal 0040). Emitted by +# ``set_invocation_metadata`` when called mid-invocation; carries the +# delta + the augmenting context's lineage identity so observers can +# resolve which of their open observations belong to the augmenting +# context's subtree and apply the entries in place. +@dataclass(frozen=True) +class MetadataAugmentationEvent: + """A metadata-augmentation event delivered to observers. + + Emitted by :func:`openarmature.observability.metadata.set_invocation_metadata` + when called mid-invocation. Carries: + + - ``entries``: the delta merged into the per-async-context + invocation metadata mapping by the call. Read-only view. + - ``namespace`` / ``attempt_index`` / ``fan_out_index`` / + ``branch_name``: the four lineage fields that jointly identify + the augmenting execution context (the calling node's identity + tuple). When ``set_invocation_metadata`` is called from outside + a node body, ``namespace`` is the empty tuple, ``attempt_index`` + is ``0``, and both ``fan_out_index`` and ``branch_name`` are + ``None`` — the invocation-level identity. + + Distinct from :class:`NodeEvent` because there is no node phase, + no pre/post state, and no error: this event reports a side-channel + augmentation, not a node-attempt boundary. Per graph-engine §6 the + event is NOT subject to the observer ``phases`` filter (which only + governs ``NodeEvent`` phases); the delivery worker forwards it to + every subscribed observer. Observers that handle it iterate their + open observations whose lineage is an ancestor of (or equal to) + the augmenting context's lineage and apply the entries as + ``openarmature.user.`` (OTel, §5.6) / + ``metadata.`` (Langfuse, §8.4.1+§8.4.2). + """ + + entries: Mapping[str, AttributeValue] + namespace: tuple[str, ...] + attempt_index: int = 0 + fan_out_index: int | None = None + branch_name: str | None = None + + +__all__ = ["FanOutEventConfig", "MetadataAugmentationEvent", "NodeEvent"] diff --git a/src/openarmature/graph/observer.py b/src/openarmature/graph/observer.py index efaae55a..47ee619f 100644 --- a/src/openarmature/graph/observer.py +++ b/src/openarmature/graph/observer.py @@ -34,19 +34,20 @@ from dataclasses import dataclass, field from typing import Any, Literal, Protocol -from .events import NodeEvent +from .events import MetadataAugmentationEvent, NodeEvent from .state import State class Observer(Protocol): - """The shape of a callable that receives node-boundary events. + """The shape of a callable that receives observer events. `Observer` is a structural Protocol; any async callable matching the signature qualifies, no subclass required. Plain functions, bound methods, and class instances with `__call__` all work:: - async def log_observer(event: NodeEvent) -> None: - print(event.node_name, event.phase) + async def log_observer(event: NodeEvent | MetadataAugmentationEvent) -> None: + if isinstance(event, NodeEvent): + print(event.node_name, event.phase) compiled.attach_observer(log_observer) @@ -63,6 +64,27 @@ async def log_observer(event: NodeEvent) -> None: conformance doesn't pin you to that name; any of `event`, `_event`, `e`, etc. matches. + Two event variants reach observers (graph-engine §6 + proposal + 0040). The signature is the union; observers ``isinstance``-narrow + on the first line and choose which variants they handle. + + - :class:`NodeEvent` — the started/completed/checkpoint phase + events. Subject to the ``phases`` filter on + :class:`SubscribedObserver`; observers whose phase set excludes + ``event.phase`` do NOT receive it. + - :class:`MetadataAugmentationEvent` — emitted by + :func:`openarmature.observability.metadata.set_invocation_metadata` + when called mid-invocation. Carries the augmenting context's + lineage tuple (``namespace``, ``attempt_index``, + ``fan_out_index``, ``branch_name``) so rich backends can update + their open observations in place + (``span.set_attribute(openarmature.user., v)`` for OTel, + ``observation.update(metadata=...)`` for Langfuse). Per spec §6 + this variant is NOT subject to the ``phases`` filter — every + subscribed observer sees it and isinstance-narrows to decide + whether to act. Simple user observers typically early-return + after ``isinstance(event, NodeEvent)`` checks. + Optional ``prepare_sync`` extension ----------------------------------- An observer MAY additionally define a synchronous method:: @@ -81,9 +103,13 @@ def prepare_sync(self, event: NodeEvent, /) -> None: ... the synchronous prep entirely; observers that do define it run only for ``"started"``-phase events, with errors warned-not- propagated (same isolation contract as the async path). + ``prepare_sync`` is never invoked for + :class:`MetadataAugmentationEvent` (the synchronous-prep contract + is anchored on the ``started`` phase, which only ``NodeEvent`` + carries). """ - async def __call__(self, event: NodeEvent, /) -> None: ... + async def __call__(self, event: NodeEvent | MetadataAugmentationEvent, /) -> None: ... # Per spec v0.6.0 §6: the two valid phase strings. Used as the default @@ -200,15 +226,22 @@ class _QueuedItem: receive it. The list is computed at dispatch time so events from different depths in nested subgraphs carry the correct observer chain without the worker needing to know the graph topology. + + ``event`` is the union of ``NodeEvent`` (started / completed / + checkpoint phases) and ``MetadataAugmentationEvent`` (proposal + 0040, side-channel augmentation). The delivery worker branches by + type to apply the right delivery contract (phase-filter for + ``NodeEvent``, no filter for the augmentation event). """ - event: NodeEvent + event: NodeEvent | MetadataAugmentationEvent observers: tuple[SubscribedObserver, ...] # A sentinel value the engine puts on the queue to signal the worker to # return after draining the events ahead of it. None is unambiguous — -# observers receive `NodeEvent` instances, never None. +# the queue carries `NodeEvent` and `MetadataAugmentationEvent` instances +# wrapped in `_QueuedItem`, never None. _DRAIN_SENTINEL = None @@ -587,16 +620,29 @@ def take_step(self) -> int: return n -def _dispatch(context: _InvocationContext, event: NodeEvent) -> None: - """Enqueue a node event for the delivery worker. - - For ``"started"``-phase events, also call any subscribed observer's - optional ``prepare_sync(event)`` synchronously — in the engine task, - BEFORE queueing — so observers that need to publish per-event state - the engine itself reads in the same engine-task scope (e.g., the - OTel observer setting ``current_active_observer_span`` for the - engine to attach into the OTel context) can do so before the node - body runs. +def _dispatch( + context: _InvocationContext, + event: NodeEvent | MetadataAugmentationEvent, +) -> None: + """Enqueue an event for the delivery worker. + + Handles two event variants: + + - :class:`NodeEvent`: the started/completed/checkpoint pair model. + For ``"started"``-phase events, also calls any subscribed + observer's optional ``prepare_sync(event)`` synchronously — in + the engine task, BEFORE queueing — so observers that need to + publish per-event state the engine itself reads in the same + engine-task scope (e.g., the OTel observer setting + ``current_active_observer_span`` for the engine to attach into + the OTel context) can do so before the node body runs. + - :class:`MetadataAugmentationEvent` (proposal 0040): a side- + channel augmentation event emitted by + ``set_invocation_metadata`` mid-invocation. Bypasses the + ``prepare_sync`` branch entirely — the sync-prep contract is + anchored on ``"started"``, which only ``NodeEvent`` carries. + Queued onto the same serial worker so observers see it in + strict order with the surrounding node events. Phase-gated forwarding: ``prepare_sync`` only fires when ``"started"`` is in the subscribed observer's ``phases`` set, mirroring how the @@ -616,7 +662,7 @@ def _dispatch(context: _InvocationContext, event: NodeEvent) -> None: observers = context.full_observers() if not observers: return - if event.phase == "started": + if isinstance(event, NodeEvent) and event.phase == "started": for subscribed in observers: if "started" not in subscribed.phases: continue @@ -686,9 +732,15 @@ async def deliver_loop( each). - No observer receives event N+1 until everyone has finished N (the loop processes one item fully before pulling the next). - - Observers whose ``phases`` set excludes the event's phase do - NOT receive it. Phase filter applies at delivery, not dispatch; - the engine still produces both events for every attempt. + - For :class:`NodeEvent`, observers whose ``phases`` set excludes + the event's phase do NOT receive it. Phase filter applies at + delivery, not dispatch; the engine still produces both events + for every attempt. + - For :class:`MetadataAugmentationEvent` (proposal 0040), the + ``phases`` filter is bypassed entirely — the event isn't a + node-phase event, so every subscribed observer receives it + regardless of ``phases``. Observers ``isinstance``-narrow on + the first line and choose whether to act. - Observer exceptions don't propagate, don't break siblings, don't block subsequent events. Reported via ``warnings.warn``. @@ -698,11 +750,12 @@ async def deliver_loop( item = await queue.get() if item is None: return + event = item.event for subscribed in item.observers: - if item.event.phase not in subscribed.phases: + if isinstance(event, NodeEvent) and event.phase not in subscribed.phases: continue try: - await subscribed.observer(item.event) + await subscribed.observer(event) except Exception as e: warnings.warn( f"observer raised {type(e).__name__}: {e}", diff --git a/src/openarmature/llm/providers/openai.py b/src/openarmature/llm/providers/openai.py index d01fff12..73fa25eb 100644 --- a/src/openarmature/llm/providers/openai.py +++ b/src/openarmature/llm/providers/openai.py @@ -55,6 +55,7 @@ from openarmature.graph.events import NodeEvent from openarmature.observability.correlation import ( current_attempt_index, + current_branch_name, current_dispatch, current_fan_out_index, current_namespace_prefix, @@ -1256,6 +1257,7 @@ def _make_llm_event( calling_namespace_prefix=current_namespace_prefix(), calling_attempt_index=current_attempt_index(), calling_fan_out_index=current_fan_out_index(), + calling_branch_name=current_branch_name(), active_prompt=active_prompt, active_prompt_group=active_prompt_group, input_messages=input_messages, diff --git a/src/openarmature/observability/correlation.py b/src/openarmature/observability/correlation.py index 6bdcb41b..53f139b3 100644 --- a/src/openarmature/observability/correlation.py +++ b/src/openarmature/observability/correlation.py @@ -36,7 +36,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from openarmature.graph.events import NodeEvent + from openarmature.graph.events import MetadataAugmentationEvent, NodeEvent from openarmature.graph.observer import SubscribedObserver @@ -208,34 +208,35 @@ def _reset_active_observers(token: Token[tuple[SubscribedObserver, ...]]) -> Non # --------------------------------------------------------------------------- -_active_dispatch_var: ContextVar[Callable[[NodeEvent], None] | None] = ContextVar( +_active_dispatch_var: ContextVar[Callable[[NodeEvent | MetadataAugmentationEvent], None] | None] = ContextVar( "openarmature.active_dispatch", default=None ) -def current_dispatch() -> Callable[[NodeEvent], None] | None: +def current_dispatch() -> Callable[[NodeEvent | MetadataAugmentationEvent], None] | None: """Return the engine's dispatch callable for the current invocation, or ``None`` outside any invocation. Capability code emitting observer events from inside a node body - calls this to put a ``NodeEvent``-shaped record on the engine's - delivery queue. The queue's serial worker preserves - per-invocation event ordering across all event sources (engine, - checkpoint, LLM provider, future backends). + calls this to put a ``NodeEvent``-shaped record (or a proposal- + 0040 ``MetadataAugmentationEvent``) on the engine's delivery + queue. The queue's serial worker preserves per-invocation event + ordering across all event sources (engine, checkpoint, LLM + provider, mid-invocation metadata augmentation, future backends). """ return _active_dispatch_var.get() def _set_active_dispatch( - dispatch: Callable[[NodeEvent], None], -) -> Token[Callable[[NodeEvent], None] | None]: + dispatch: Callable[[NodeEvent | MetadataAugmentationEvent], None], +) -> Token[Callable[[NodeEvent | MetadataAugmentationEvent], None] | None]: """Set the engine's dispatch callable in scope. Internal — engine-only.""" return _active_dispatch_var.set(dispatch) def _reset_active_dispatch( - token: Token[Callable[[NodeEvent], None] | None], + token: Token[Callable[[NodeEvent | MetadataAugmentationEvent], None] | None], ) -> None: _active_dispatch_var.reset(token) diff --git a/src/openarmature/observability/langfuse/observer.py b/src/openarmature/observability/langfuse/observer.py index da31acba..5d408292 100644 --- a/src/openarmature/observability/langfuse/observer.py +++ b/src/openarmature/observability/langfuse/observer.py @@ -26,8 +26,10 @@ import uuid from collections.abc import Mapping from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, cast +from typing import Any, cast +from openarmature.graph.events import MetadataAugmentationEvent, NodeEvent +from openarmature.observability.lineage import is_prefix_or_equal, is_strict_prefix from openarmature.observability.llm_event import LLM_NAMESPACE, LlmEventPayload from .client import ( @@ -37,10 +39,6 @@ LangfuseUsage, ) -if TYPE_CHECKING: - from openarmature.graph.events import NodeEvent - - # §5.5.5 / §8.7 truncation: when the serialized payload exceeds the # configured cap, the marker below is appended and the unparseable # JSON serves as the "this was truncated" signal in Langfuse's input @@ -62,10 +60,14 @@ def _read_spec_version() -> str: # In-flight Span observation handle, keyed by the standard span-stack -# key (namespace, attempt_index, fan_out_index). Mirrors the OTel -# observer's _OpenSpan shape but holds a Langfuse handle instead of an -# OTel Span. -_StackKey = tuple[tuple[str, ...], int, int | None] +# key (namespace, attempt_index, fan_out_index, branch_name). +# ``branch_name`` discriminates concurrent same-named inner nodes +# across sibling parallel-branches branches (pipeline-utilities §11); +# without it the two inner ``ask`` nodes of two branches with the +# same namespace + fan_out_index would collide on the same key. +# Mirrors the OTel observer's ``_StackKey`` shape but holds a +# Langfuse handle instead of an OTel Span. +_StackKey = tuple[tuple[str, ...], int, int | None, str | None] @dataclass @@ -250,7 +252,10 @@ def __post_init__(self) -> None: f"minimum of {_PAYLOAD_MIN_BYTES} bytes" ) - async def __call__(self, event: NodeEvent) -> None: + async def __call__(self, event: NodeEvent | MetadataAugmentationEvent) -> None: + if isinstance(event, MetadataAugmentationEvent): + self._handle_metadata_augmentation(event) + return # LLM provider events use a sentinel namespace per §5.5; route # them to the dedicated Generation path. if event.namespace == LLM_NAMESPACE: @@ -372,6 +377,89 @@ def _handle_completed(self, event: NodeEvent) -> None: # detached_traces entry so a subsequent re-entry mints fresh. inv_state.detached_traces.pop(event.namespace, None) + # ------------------------------------------------------------------ + # Metadata augmentation (proposal 0040 §3.4 + §6) + # ------------------------------------------------------------------ + + def _handle_metadata_augmentation(self, event: MetadataAugmentationEvent) -> None: + # Spec proposal 0040 §3.4 MUST: open observations whose lineage + # ancestor-or-equals the augmenting context get the entries + # applied in place via the Langfuse handle's + # ``update(metadata=...)`` method. Sibling instances / branches + # and ancestors above the containment are skipped (same scoping + # rule as the OTel mapping — see + # ``OTelObserver._handle_metadata_augmentation`` for the algebra). + # + # For an outermost-serial augmenter (FI=None, BN=None), the + # invocation's Trace itself is updated via + # ``client.update_trace`` so the augmented keys land on + # ``trace.metadata.`` for §8.4-style top-level filtering. + # Inside a fan-out instance / parallel-branches branch the + # Trace is OUT of scope (it's shared with siblings); only the + # innermost containment + the augmenter's own subtree update. + # + # Per-instance / per-branch isolation: + # ``set_invocation_metadata`` runs in the calling node's task + # whose Context already carries the per-async-context COW + # mapping (proposal 0034 §3.4). The augmentation event's + # ``entries`` are that delta only — applying them to matching + # open observations preserves the per-async-context isolation + # 029 / 030 encode. + from openarmature.observability.correlation import current_invocation_id + + invocation_id = current_invocation_id() + if invocation_id is None or not event.entries: + return + inv_state = self._inv_states.get(invocation_id) + aug_fi = event.fan_out_index + aug_bn = event.branch_name + aug_ns = event.namespace + metadata_delta = dict(event.entries) + + # Trace.metadata: only for outermost-serial. The Trace is + # shared across siblings, so a fan-out instance's per-item + # productId would leak across siblings if it updated the + # Trace — 029's fixture explicitly rules that out. + if aug_fi is None and aug_bn is None: + self.client.update_trace(id=invocation_id, metadata=metadata_delta) + + if inv_state is None: + return + + # Subgraph wrapper observations on the ancestor path + # (outermost-serial only — inside a fan-out instance the + # subgraph wrapper above the fan-out node is sibling-shared). + if aug_fi is None and aug_bn is None: + for prefix, observation in inv_state.subgraph_observations.items(): + if is_strict_prefix(prefix, aug_ns): + observation.handle.update(metadata=metadata_delta) + + # Fan-out instance dispatch observation(s) when the augmenter + # is inside a fan-out instance. Keys are anchor_ns + + # (str(fan_out_index),) per + # ``_open_fan_out_instance_dispatch_observation``. + if aug_fi is not None: + fi_str = str(aug_fi) + for key, observation in inv_state.fan_out_instance_observations.items(): + if not key or key[-1] != fi_str: + continue + anchor_ns = key[:-1] + if is_strict_prefix(anchor_ns, aug_ns) or anchor_ns == aug_ns: + observation.handle.update(metadata=metadata_delta) + + # Open node observations on the augmenter's call stack. Match + # on ``fan_out_index`` AND ``branch_name`` to skip sibling + # instances / branches; namespace must prefix (or equal) the + # augmenter's. + for key, observation in inv_state.open_observations.items(): + ns, _ai, fi, bn = key + if fi != aug_fi: + continue + if bn != aug_bn: + continue + if is_prefix_or_equal(ns, aug_ns): + observation.handle.update(metadata=metadata_delta) + def _open_trace(self, invocation_id: str, correlation_id: str | None, event: NodeEvent) -> None: # ``entry_node`` and the trace name MUST identify the outer-graph # entry, not whichever node fired first. Subgraph wrappers do not @@ -399,7 +487,7 @@ def _open_trace(self, invocation_id: str, correlation_id: str | None, event: Nod self._inv_states[invocation_id] = _InvState(trace_id=invocation_id) def _key_for(self, event: NodeEvent) -> _StackKey: - return (event.namespace, event.attempt_index, event.fan_out_index) + return (event.namespace, event.attempt_index, event.fan_out_index, event.branch_name) def _resolve_parent_observation_id(self, inv_state: _InvState, event: NodeEvent) -> str | None: # Parent precedence (innermost wins): @@ -990,6 +1078,7 @@ def _resolve_llm_parent_observation_id( payload.calling_namespace_prefix, payload.calling_attempt_index, payload.calling_fan_out_index, + payload.calling_branch_name, ) observation = inv_state.open_observations.get(key) if observation is not None: diff --git a/src/openarmature/observability/lineage.py b/src/openarmature/observability/lineage.py new file mode 100644 index 00000000..696e558c --- /dev/null +++ b/src/openarmature/observability/lineage.py @@ -0,0 +1,23 @@ +# Spec: cross-cutting helpers for the observer-side lineage match +# introduced by proposal 0040 (metadata-augmentation event open-span +# update, observability §3.4 + §6). + +"""Tuple-prefix predicates used by the OTel + Langfuse observers to +match an augmentation event's namespace against the namespaces of +open spans / observations. Shared so both observers express the +ancestor-or-equal rule identically. +""" + +from __future__ import annotations + +__all__ = ["is_prefix_or_equal", "is_strict_prefix"] + + +def is_strict_prefix(prefix: tuple[str, ...], full: tuple[str, ...]) -> bool: + """True iff ``prefix`` is a strict prefix of ``full`` (NOT equal).""" + return len(prefix) < len(full) and full[: len(prefix)] == prefix + + +def is_prefix_or_equal(prefix: tuple[str, ...], full: tuple[str, ...]) -> bool: + """True iff ``prefix`` is a prefix of (or equal to) ``full``.""" + return len(prefix) <= len(full) and full[: len(prefix)] == prefix diff --git a/src/openarmature/observability/llm_event.py b/src/openarmature/observability/llm_event.py index 134f4968..e7c53c27 100644 --- a/src/openarmature/observability/llm_event.py +++ b/src/openarmature/observability/llm_event.py @@ -95,6 +95,12 @@ class LlmEventPayload(BaseModel): calling_namespace_prefix: tuple[str, ...] = () calling_attempt_index: int = 0 calling_fan_out_index: int | None = None + # Calling-node branch_name (pipeline-utilities §11). Mirrors the + # other ``calling_*`` fields; the OTel observer's open-span key + # widening (``_StackKey`` now includes ``branch_name``) needs this + # to disambiguate concurrent same-named inner nodes across sibling + # branches. + calling_branch_name: str | None = None # Prompt-context snapshot captured at dispatch time. ``Any`` # because the prompts package imports State indirectly; the typed # shapes are PromptResult / PromptGroup from openarmature.prompts. diff --git a/src/openarmature/observability/metadata.py b/src/openarmature/observability/metadata.py index 1e037805..0ce6ac75 100644 --- a/src/openarmature/observability/metadata.py +++ b/src/openarmature/observability/metadata.py @@ -119,22 +119,27 @@ def set_invocation_metadata(**entries: AttributeValue) -> None: overwritten; other keys are preserved. Per spec §3.4: affects spans / observations emitted AFTER the - call returns; spans already closed are NOT retroactively updated. - Implementations MAY update open root-level surfaces (e.g., the - Langfuse Trace's metadata) where the backend SDK supports it; - Langfuse's ``trace.update`` is the canonical example. The - framework's helper here just maintains the ContextVar; per- - backend update propagation is the observer's concern. + call returns. Open observations whose lineage covers the calling + context ARE updated in place per proposal 0040 — implementations + enqueue a :class:`~openarmature.graph.events.MetadataAugmentationEvent` + on the engine's serial observer-delivery queue carrying the + delta + the calling context's lineage tuple (namespace, + attempt_index, fan_out_index, branch_name); observers correlate + the lineage with their open observations and apply + ``observation.update(metadata=...)`` / ``span.set_attribute(...)`` + in place. Spans already CLOSED at call time are NOT retroactively + updated. Raises :class:`ValueError` if any key violates the reserved- namespace rule or any value is not OTel-attribute-compatible. Outside any active invocation, this still updates the ContextVar (a fresh per-context override), but the value will - not be observed by any backend since no observer is in scope. - The empty-invocation case is supported for symmetry; users - typically call this from inside a node body, middleware, or - observer where an invocation is already in flight. + not be observed by any backend since no observer is in scope: + :func:`current_dispatch` returns ``None`` and no augmentation + event is emitted. The empty-invocation case is supported for + symmetry; users typically call this from inside a node body, + middleware, or observer where an invocation is already in flight. """ if not entries: return @@ -144,6 +149,37 @@ def set_invocation_metadata(**entries: AttributeValue) -> None: merged: dict[str, AttributeValue] = dict(_invocation_metadata_var.get()) merged.update(entries) _invocation_metadata_var.set(MappingProxyType(merged)) + # Proposal 0040: emit a MetadataAugmentationEvent so observers can + # update their open observations in place. Local imports break the + # observability -> graph -> observability cycle (events.py imports + # AttributeValue from this module; observer.py imports NodeEvent; + # correlation.py forward-declares both under TYPE_CHECKING). + # ``current_dispatch`` is ``None`` outside an invocation (boundary + # ``invoke()`` hasn't installed a dispatch closure yet) — we still + # mutated the ContextVar above so a node body called later in the + # same async context sees the entries; we just don't enqueue a + # delivery event for observers that don't exist. + from openarmature.graph.events import MetadataAugmentationEvent + + from .correlation import ( + current_attempt_index, + current_branch_name, + current_dispatch, + current_fan_out_index, + current_namespace_prefix, + ) + + dispatch = current_dispatch() + if dispatch is None: + return + event = MetadataAugmentationEvent( + entries=MappingProxyType(dict(entries)), + namespace=current_namespace_prefix(), + attempt_index=current_attempt_index(), + fan_out_index=current_fan_out_index(), + branch_name=current_branch_name(), + ) + dispatch(event) def validate_invocation_metadata(mapping: object) -> MappingProxyType[str, AttributeValue]: diff --git a/src/openarmature/observability/otel/observer.py b/src/openarmature/observability/otel/observer.py index 0ed63ec2..c598aa97 100644 --- a/src/openarmature/observability/otel/observer.py +++ b/src/openarmature/observability/otel/observer.py @@ -77,7 +77,7 @@ import json from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, cast +from typing import Any, cast from opentelemetry import context as otel_context from opentelemetry import trace as otel_trace @@ -96,16 +96,18 @@ ) from opentelemetry.trace.propagation import set_span_in_context +from openarmature.graph.events import MetadataAugmentationEvent, NodeEvent +from openarmature.observability.lineage import is_prefix_or_equal, is_strict_prefix from openarmature.observability.llm_event import LLM_NAMESPACE, LlmEventPayload -if TYPE_CHECKING: - from openarmature.graph.events import NodeEvent - - -# Span-stack key shape: ``(namespace, attempt_index, fan_out_index)`` -# — these three fields uniquely identify any node attempt within an -# invocation. -_StackKey = tuple[tuple[str, ...], int, int | None] +# Span-stack key shape: +# ``(namespace, attempt_index, fan_out_index, branch_name)`` — these +# four fields jointly identify any node attempt within an invocation. +# ``branch_name`` discriminates concurrent same-named inner nodes +# across sibling parallel-branches branches (pipeline-utilities §11); +# without it the two inner ``ask`` nodes of two branches with the +# same namespace + fan_out_index would collide on the same key. +_StackKey = tuple[tuple[str, ...], int, int | None, str | None] # Re-export the LLM-event namespace sentinel under the same name the @@ -448,10 +450,14 @@ def _inv_state_for(self, invocation_id: str) -> _InvState: return state # ------------------------------------------------------------------ - # Observer protocol — async callable accepting a NodeEvent + # Observer protocol — async callable accepting node events + the + # proposal-0040 metadata-augmentation event variant. # ------------------------------------------------------------------ - async def __call__(self, event: NodeEvent) -> None: + async def __call__(self, event: NodeEvent | MetadataAugmentationEvent) -> None: + if isinstance(event, MetadataAugmentationEvent): + self._handle_metadata_augmentation(event) + return # LLM provider events use a sentinel namespace so we can route # them to the dedicated §5.5 span path. if event.namespace == _LLM_NAMESPACE: @@ -632,6 +638,94 @@ def _handle_completed(self, event: NodeEvent) -> None: # subsequent re-entry mints a fresh trace. inv_state.detached_roots.pop(event.namespace, None) + # ------------------------------------------------------------------ + # Metadata augmentation (proposal 0040 §3.4 + §6) + # ------------------------------------------------------------------ + + def _handle_metadata_augmentation(self, event: MetadataAugmentationEvent) -> None: + # Spec proposal 0040: spans whose lineage ancestor-or-equals the + # augmenting context (within the same fan-out instance / + # parallel-branch boundary) get ``openarmature.user.`` + # applied in place. Sibling instances / branches and ancestors + # ABOVE the boundary are skipped. + # + # Match rule (using the augmentation event's lineage tuple + # ``(NS, AI, FI, BN)``): + # - Invocation span: included iff ``FI is None and BN is None`` + # (outermost-serial context). The shared fan-out node span and + # the invocation span are explicitly out of scope when + # augmenting from inside a fan-out instance or branch. + # - Subgraph wrapper spans: included on the outermost-serial + # path when their namespace is a strict prefix of NS. + # - Fan-out instance dispatch spans: included iff the dispatch + # span's FI suffix matches ``str(FI)`` and the anchor namespace + # is a strict prefix of NS. + # - Per-attempt node spans (``open_spans``): included iff the + # span's FI equals the augmenter's FI and its namespace is a + # prefix of (or equal to) NS. + from openarmature.observability.correlation import current_invocation_id + + invocation_id = current_invocation_id() + if invocation_id is None: + return + if not event.entries: + return + targets = self._collect_augmentation_targets(invocation_id, event) + for span in targets: + for key, value in event.entries.items(): + # OTel forbids None as an attribute value; the metadata + # validator at the engine boundary rejects None already, + # so we can pass through directly. + span.set_attribute(f"openarmature.user.{key}", value) + + def _collect_augmentation_targets( + self, invocation_id: str, event: MetadataAugmentationEvent + ) -> list[Span]: + targets: list[Span] = [] + aug_fi = event.fan_out_index + aug_bn = event.branch_name + aug_ns = event.namespace + inv_state = self._inv_states.get(invocation_id) + if aug_fi is None and aug_bn is None: + # Outermost-serial context: the invocation span is in scope. + inv_open = self._invocation_span.get(invocation_id) + if inv_open is not None: + targets.append(inv_open.span) + if inv_state is not None: + # Subgraph wrapper spans on the ancestor path. + for prefix, open_span in inv_state.subgraph_spans.items(): + if is_strict_prefix(prefix, aug_ns): + targets.append(open_span.span) + if inv_state is None: + return targets + if aug_fi is not None: + # Fan-out instance dispatch span(s) on the ancestor path. + # Keys are anchor_ns + (str(fan_out_index),) per + # ``_open_fan_out_instance_dispatch_span``. + fi_str = str(aug_fi) + for key, open_span in inv_state.fan_out_instance_spans.items(): + if not key or key[-1] != fi_str: + continue + anchor_ns = key[:-1] + if is_strict_prefix(anchor_ns, aug_ns) or anchor_ns == aug_ns: + targets.append(open_span.span) + # Open node spans on the augmenter's call stack. Match on + # ``fan_out_index`` AND ``branch_name`` to skip sibling + # instances / branches; namespace must prefix (or equal) the + # augmenter's. The BN discriminator is what keeps two + # concurrent same-named inner nodes across sibling + # parallel-branches branches from leaking augmentation to + # each other. + for key, open_span in inv_state.open_spans.items(): + ns, _ai, fi, bn = key + if fi != aug_fi: + continue + if bn != aug_bn: + continue + if is_prefix_or_equal(ns, aug_ns): + targets.append(open_span.span) + return targets + # ------------------------------------------------------------------ # Special-event paths # ------------------------------------------------------------------ @@ -902,6 +996,7 @@ def _resolve_llm_parent( payload.calling_namespace_prefix, payload.calling_attempt_index, payload.calling_fan_out_index, + payload.calling_branch_name, ) calling = inv_state.open_spans.get(calling_key) if calling is not None: @@ -954,7 +1049,7 @@ def _open_invocation_span( self._invocation_span[invocation_id] = _OpenSpan(span=span) def _key_for(self, event: NodeEvent) -> _StackKey: - return (event.namespace, event.attempt_index, event.fan_out_index) + return (event.namespace, event.attempt_index, event.fan_out_index, event.branch_name) def _resolve_parent_context( self, @@ -1358,7 +1453,7 @@ def _find_fan_out_node_span(self, inv_state: _InvState, prefix: tuple[str, ...]) closes within each attempt's lifecycle), so a scan finds it unambiguously.""" for key, open_span in inv_state.open_spans.items(): - ns, _attempt, fan_idx = key + ns, _attempt, fan_idx, _bn = key if ns == prefix and fan_idx is None: return open_span return None diff --git a/tests/conformance/adapter.py b/tests/conformance/adapter.py index d12c18a1..efc05c21 100644 --- a/tests/conformance/adapter.py +++ b/tests/conformance/adapter.py @@ -38,7 +38,7 @@ merge, merge_all, ) -from openarmature.graph.events import NodeEvent +from openarmature.graph.events import MetadataAugmentationEvent, NodeEvent from openarmature.graph.observer import Observer if TYPE_CHECKING: @@ -855,7 +855,9 @@ def make_observer_fn( the event unrecorded and the counter shows it as undelivered. """ - async def observer(event: NodeEvent) -> None: + async def observer(event: NodeEvent | MetadataAugmentationEvent) -> None: + if isinstance(event, MetadataAugmentationEvent): + return sleep_ms = _resolve_sleep_ms(fixture) if sleep_ms > 0: await asyncio.sleep(sleep_ms / 1000.0) diff --git a/tests/conformance/test_conformance.py b/tests/conformance/test_conformance.py index ef8f8648..ef6962f5 100644 --- a/tests/conformance/test_conformance.py +++ b/tests/conformance/test_conformance.py @@ -27,7 +27,7 @@ State, SubscribedObserver, ) -from openarmature.graph.events import NodeEvent +from openarmature.graph.events import MetadataAugmentationEvent, NodeEvent from openarmature.graph.observer import Observer from .adapter import ( @@ -605,7 +605,8 @@ class FixtureState(State): received: list[NodeEvent] = [] - async def observer(event: NodeEvent) -> None: + async def observer(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append(event) async def node_a(_state: Any) -> dict[str, Any]: diff --git a/tests/conformance/test_fixture_parsing.py b/tests/conformance/test_fixture_parsing.py index 9740f97b..5fbab19e 100644 --- a/tests/conformance/test_fixture_parsing.py +++ b/tests/conformance/test_fixture_parsing.py @@ -92,17 +92,22 @@ def _id(case: tuple[str, Path]) -> str: "observability/028-caller-metadata-namespace-rejection": ( "Rejection invariants live in the dedicated _run_fixture_028 runner" ), - # Proposal 0034 fan-out / parallel-branches caller-metadata - # fixtures need the harness primitive - # ``augment_metadata_from_field`` (per-instance / per-branch - # ``set_invocation_metadata`` calls). The 026/027/028 fixtures - # (cross-cutting + boundary rejection) shipped with PR 4; the - # augmentation primitive lands in a follow-up. + # Proposal 0034 / 0040 fan-out / parallel-branches caller-metadata + # fixtures. The augmentation MECHANISM is implemented in v0.11.0 + # (#22) and covered end-to-end by unit tests + # (test_observability_otel.py + test_observability_langfuse.py). + # The CONFORMANCE FIXTURES stay deferred for harness-shape gaps: + # - 029 omits ``collect_field`` / ``target_field`` on the fan-out + # cfg AND a ``state:`` block on the inner subgraph (both + # required by the cross-cap adapter). + # - 030 expects per-branch dispatch spans in the Langfuse tree; + # the spec direction for that span layer is pending in coord + # thread ``discuss-otel-parallel-branches-dispatch-span``. "observability/029-caller-metadata-fan-out-per-instance": ( - "Per-instance augmentation harness primitive lands in a follow-up" + "Fixture-shape gaps (no collect_field/target_field/state); mechanism covered by unit tests" ), "observability/030-caller-metadata-parallel-branches-per-branch": ( - "Per-branch augmentation harness primitive lands in a follow-up" + "Per-branch dispatch span shape pending spec; mechanism covered by unit tests" ), # proposal 0033 added typed directive shapes (`secondary_manager`, # `label_resolver`, `cases`) the canonical parser doesn't model. @@ -153,8 +158,13 @@ def _id(case: tuple[str, Path]) -> str: ), # Proposal 0040 (open-span metadata update) — task #22 implements # the §6 augmentation-event mechanism + un-defers 029/030 + 034. + # Fixture 034 lands in the Langfuse-specific harness directly + # via the augment_metadata directive (see + # ``tests/conformance/test_observability_langfuse.py``); the + # cross-capability parser still doesn't model langfuse_trace, so + # defer the parser-side activation per the 022-024 pattern. "observability/034-caller-metadata-open-span-update-serial": ( - "Open-span augmentation-event mechanism lands with #22 (0040 not-yet)" + "Langfuse shape models live in the dedicated test_observability_langfuse harness" ), # Proposal 0039 (caller-supplied invocation_id) Langfuse trace.id # derivation fixtures use the langfuse_trace expected shape the diff --git a/tests/conformance/test_observability.py b/tests/conformance/test_observability.py index 21588dc5..ed0af901 100644 --- a/tests/conformance/test_observability.py +++ b/tests/conformance/test_observability.py @@ -884,25 +884,6 @@ async def _run_fixture_028(spec: Mapping[str, Any]) -> None: cases = cast("list[dict[str, Any]]", spec["cases"]) for case in cases: case_name = cast("str", case["name"]) - # Cases using the `augment_metadata` directive exercise §3.4 - # mid-invocation rejection at set_invocation_metadata. The - # augment_metadata harness primitive (per fixture 034) lands - # with proposal 0040 / task #22; surface the deferral via - # warnings.warn so pytest's end-of-run summary lists it (rather - # than silently passing) and continue to the other cases. - nodes_check = cast("dict[str, Any]", case.get("nodes", {})) - if any( - isinstance(n, dict) and "augment_metadata" in cast("dict[str, Any]", n) - for n in nodes_check.values() - ): - import warnings # noqa: PLC0415 - - warnings.warn( - f"028 case {case_name!r} deferred: augment_metadata harness primitive " - f"lands with proposal 0040 / #22", - stacklevel=2, - ) - continue try: # Build a minimal graph from the case's nodes/edges. The # fixture's node is a noop update — we never expect it to @@ -916,14 +897,29 @@ async def _run_fixture_028(spec: Mapping[str, Any]) -> None: for node_name, node_spec in nodes_spec.items(): node_dict = cast("dict[str, Any]", node_spec) update_block = cast("dict[str, Any]", node_dict["update"]) + augment_block = cast("dict[str, Any] | None", node_dict.get("augment_metadata")) + + def _make_body( + payload: dict[str, Any], + augment: dict[str, Any] | None, + ) -> Any: + # Per spec §3.4 + proposal 0040: the augment_metadata + # primitive injects a ``set_invocation_metadata(**augment)`` + # call at the top of the node body. Used by 028's + # mid-invocation-rejection case (reserved name `step`) + # and by 034 for the open-span update demonstration. + from openarmature.observability.metadata import ( # noqa: PLC0415 + set_invocation_metadata, + ) - def _make_body(payload: dict[str, Any]) -> Any: async def _body(_s: Any) -> dict[str, Any]: + if augment is not None: + set_invocation_metadata(**augment) return dict(payload) return _body - builder.add_node(node_name, _make_body(update_block)) + builder.add_node(node_name, _make_body(update_block, augment_block)) for edge in cast("list[dict[str, str]]", case["edges"]): target_raw = edge["to"] target = END if target_raw == "END" else target_raw @@ -939,20 +935,46 @@ async def _body(_s: Any) -> dict[str, Any]: graph.attach_observer(langfuse_observer) caller_metadata = cast("dict[str, Any]", case["caller_metadata"]) + expected = cast("dict[str, Any]", case["expected"]) + expects_boundary_rejection = expected.get("invoke_rejects_at_api_boundary", False) + expects_call_site_rejection = expected.get("augment_rejects_at_call_site", False) try: - # Covers both rejection paths: the prefix-namespace - # rejection (openarmature.* / gen_ai.*, from 0034) and - # the exact-key-name rejection (0041's §8.4 reserved - # set). Both error messages contain "reserved". - with pytest.raises(ValueError, match="reserved"): - await graph.invoke(state_cls(), metadata=caller_metadata) + if expects_boundary_rejection: + # Boundary-rejection path: invoke()'s caller_metadata + # validator rejects before any work begins. Covers + # both the prefix-namespace rejection (openarmature.* + # / gen_ai.*, from 0034) and the exact-key-name + # rejection (0041's §8.4 reserved set). Both error + # messages contain "reserved". + with pytest.raises(ValueError, match="reserved"): + await graph.invoke(state_cls(), metadata=caller_metadata) + elif expects_call_site_rejection: + # Mid-invocation rejection path: caller_metadata + # passes the boundary; the node body's + # ``set_invocation_metadata(**augment)`` raises a + # ValueError at the call site. The engine wraps the + # node-body raise in NodeException whose + # ``__cause__`` is the ValueError. The §3.4 contract + # is that the helper raises at the call site — the + # reserved key MUST NOT reach any emission, hence + # no spans / no Langfuse observations afterward. + from openarmature.graph import NodeException # noqa: PLC0415 + + with pytest.raises(NodeException) as exc_info: + await graph.invoke(state_cls(), metadata=caller_metadata) + cause = exc_info.value.__cause__ + assert isinstance(cause, ValueError), ( + f"expected NodeException.__cause__ to be ValueError; got {type(cause).__name__}" + ) + assert "reserved" in str(cause), f"expected 'reserved' in cause message; got {cause!s}" + else: + raise AssertionError( + "case has neither invoke_rejects_at_api_boundary nor augment_rejects_at_call_site set" + ) + await graph.drain() finally: otel_observer.shutdown() - expected = cast("dict[str, Any]", case["expected"]) - if expected.get("invoke_rejects_at_api_boundary"): - # Already verified above via pytest.raises. - pass if expected.get("no_spans_emitted"): spans = exporter.get_finished_spans() assert len(spans) == 0, f"expected zero spans, got {[s.name for s in spans]}" @@ -1140,6 +1162,20 @@ async def _run_fixture_005_case(case: Mapping[str, Any]) -> None: global_exporter = InMemorySpanExporter() global_provider = TracerProvider() global_provider.add_span_processor(SimpleSpanProcessor(global_exporter)) + # OTel SDK 1.x's ``set_tracer_provider`` is guarded by a + # ``_TRACER_PROVIDER_SET_ONCE`` primitive — once a non-default + # provider is set, subsequent calls are silent no-ops (with a + # WARNING log "Overriding of current TracerProvider is not + # allowed"). If a prior test in the suite-run order left a + # non-default provider behind, the call below would no-op and + # this case's ``global_exporter`` would receive 0 spans. Reset + # both the value AND the Once explicitly so this case's set + # always wins. The finally block below restores ``prior_global`` + # via the same direct reset so the next test starts clean. + once = otel_trace._TRACER_PROVIDER_SET_ONCE # type: ignore[attr-defined] + with once._lock: # pyright: ignore[reportPrivateUsage] + otel_trace._TRACER_PROVIDER = None # type: ignore[attr-defined] + once._done = False # pyright: ignore[reportPrivateUsage] otel_trace.set_tracer_provider(global_provider) try: diff --git a/tests/conformance/test_observability_langfuse.py b/tests/conformance/test_observability_langfuse.py index 1f4506f4..aed7b98b 100644 --- a/tests/conformance/test_observability_langfuse.py +++ b/tests/conformance/test_observability_langfuse.py @@ -68,10 +68,151 @@ "031-langfuse-subgraph-span-hierarchy", "032-langfuse-fan-out-per-instance-spans", "033-langfuse-detached-trace-mode", + # 034 — proposal 0040 outermost-serial open-span update. + # Single-node graph; the ``augment_metadata`` directive on + # the node body injects a ``set_invocation_metadata`` call + # before the LLM call, exercising the §3.4 MUST that open + # spans in the augmenting context's lineage update in place. + "034-caller-metadata-open-span-update-serial", + # 029 + 030 stay deferred in v0.11.0: + # - 029 (fan-out per-instance): fixture omits ``collect_field`` + # and ``target_field`` on the fan_out cfg, plus the inner + # subgraph omits a ``state:`` block — both are required by + # the cross-cap adapter. The augmentation behavior IS + # verified end-to-end by the unit test + # ``test_observability_langfuse.py::test_metadata_augmentation_in_fan_out_isolates_per_instance`` + # plus the OTel counterpart. + # - 030 (parallel-branches per-branch): the expected trace + # requires a per-branch dispatch span the Langfuse mapping + # doesn't synthesize today; the spec direction is in + # ``discuss-otel-parallel-branches-dispatch-span``. + # Sibling-skip behavior IS verified by the OTel unit test + # ``test_metadata_augmentation_in_parallel_branches_skips_sibling``. + # Both fixtures land once spec settles the dispatch-span + # shape AND the adapter learns to infer fan-out aggregation + # defaults from inner subgraphs. } ) +def _normalize_fan_out_subgraph_keys(spec: dict[str, Any]) -> None: + """In-place rename of fan-out config keys that fixture 029 uses + but the cross-capability adapter doesn't: + + - ``inner_subgraph`` → ``subgraph`` (within each ``fan_out`` block) + - top-level ``inner_subgraphs`` → ``subgraphs`` + + The directive intent is identical; only the key naming differs + across the spec fixture style and the cross-cap adapter's + parser. Keep the original keys intact in the source spec; this + function mutates a deepcopy in the harness wrapper. + """ + if "inner_subgraphs" in spec and "subgraphs" not in spec: + spec["subgraphs"] = spec.pop("inner_subgraphs") + for node_spec in cast("dict[str, Any]", spec.get("nodes") or {}).values(): + if not isinstance(node_spec, dict): + continue + node_dict = cast("dict[str, Any]", node_spec) + fan_out_cfg = cast("dict[str, Any] | None", node_dict.get("fan_out")) + if fan_out_cfg is None: + continue + if "inner_subgraph" in fan_out_cfg and "subgraph" not in fan_out_cfg: + fan_out_cfg["subgraph"] = fan_out_cfg.pop("inner_subgraph") + + +def _build_augment_middlewares( + case: Mapping[str, Any], +) -> tuple[ + dict[str, list[Any]], # fan_out_instance_middleware: node_name -> [Middleware] + dict[str, dict[str, list[Any]]], # parallel_branches_branch_middleware: node -> branch -> [Middleware] +]: + """Detect proposal-0040 augment directives in the case spec and + synthesize the middlewares that drive them via the adapter's + ``fan_out_instance_middleware`` / ``parallel_branches_branch_middleware`` + hooks. + + - Fan-out ``augment_metadata_from_field: {key: field_path}`` → + one instance middleware that reads ``current_fan_out_index()``, + indexes into the parent's ``items_field`` list captured at + fixture-build time, and calls ``set_invocation_metadata(**entries)`` + where entries are pulled from the per-instance item via field_path. + - Parallel-branches ``branches..augment_metadata: {key: value}`` + → per-branch middleware that calls + ``set_invocation_metadata(**entries)`` once at branch entry. + """ + fan_out_mw: dict[str, list[Any]] = {} + branch_mw: dict[str, dict[str, list[Any]]] = {} + initial_state = cast("dict[str, Any]", case.get("initial_state") or {}) + + for node_name, node_spec_any in cast("dict[str, Any]", case.get("nodes") or {}).items(): + if not isinstance(node_spec_any, dict): + continue + node_spec = cast("dict[str, Any]", node_spec_any) + fan_out_cfg = cast("dict[str, Any] | None", node_spec.get("fan_out")) + if fan_out_cfg is not None: + augment_field_map = cast("dict[str, str] | None", fan_out_cfg.get("augment_metadata_from_field")) + if augment_field_map: + items_field = cast("str | None", fan_out_cfg.get("items_field")) + items_list = ( + cast("list[dict[str, Any]]", initial_state.get(items_field, [])) if items_field else [] + ) + fan_out_mw[node_name] = [_make_augment_instance_middleware(augment_field_map, items_list)] + pb_cfg = cast("dict[str, Any] | None", node_spec.get("parallel_branches")) + if pb_cfg is not None: + branches_cfg = cast("dict[str, dict[str, Any]]", pb_cfg.get("branches") or {}) + per_branch: dict[str, list[Any]] = {} + for branch_name, branch_cfg in branches_cfg.items(): + augment_entries = cast("dict[str, Any] | None", branch_cfg.get("augment_metadata")) + if augment_entries: + per_branch[branch_name] = [_make_augment_branch_middleware(augment_entries)] + if per_branch: + branch_mw[node_name] = per_branch + return fan_out_mw, branch_mw + + +def _make_augment_instance_middleware(field_map: dict[str, str], items: list[dict[str, Any]]) -> Any: + """Per-instance fan-out middleware that calls + ``set_invocation_metadata`` with per-item entries pulled from + ``items[current_fan_out_index()][field_path]``. Captures ``items`` + at fixture-build time so each instance reads the same list.""" + + class _AugmentInstanceMW: + async def __call__(self, state: Any, next_: Any, /) -> Any: + from openarmature.observability.correlation import ( # noqa: PLC0415 + current_fan_out_index, + ) + from openarmature.observability.metadata import ( # noqa: PLC0415 + set_invocation_metadata, + ) + + idx = current_fan_out_index() + if idx is not None and 0 <= idx < len(items): + item = items[idx] + entries = {key: item[field_path] for key, field_path in field_map.items()} + set_invocation_metadata(**entries) + return await next_(state) + + return _AugmentInstanceMW() + + +def _make_augment_branch_middleware(entries: dict[str, Any]) -> Any: + """Per-branch middleware that calls ``set_invocation_metadata`` + once at branch entry. Captures ``entries`` at fixture-build + time so the call inside the middleware doesn't need to read + the case spec at runtime.""" + + class _AugmentBranchMW: + async def __call__(self, state: Any, next_: Any, /) -> Any: + from openarmature.observability.metadata import ( # noqa: PLC0415 + set_invocation_metadata, + ) + + set_invocation_metadata(**entries) + return await next_(state) + + return _AugmentBranchMW() + + def _fixture_paths() -> list[Path]: return sorted(p for p in CONFORMANCE_DIR.glob("[0-9][0-9][0-9]-*.yaml") if p.stem in _LANGFUSE_FIXTURES) @@ -154,7 +295,18 @@ async def fetch(self, name: str, label: str = "production") -> Prompt: async def test_langfuse_fixture(fixture_path: Path) -> None: spec = _load(fixture_path) if "cases" in spec: + # Fold fixture-level ``subgraphs`` / ``inner_subgraphs`` into + # each case so the per-case runner sees them locally. Fixture + # 030 declares its branch subgraphs at fixture-level (alongside + # ``cases:``); without this fold the per-case build can't + # resolve ``branches.fraud_check.subgraph: fraud_check``. + fixture_subgraphs = cast("dict[str, Any] | None", spec.get("subgraphs")) + fixture_inner_subgraphs = cast("dict[str, Any] | None", spec.get("inner_subgraphs")) for case in cast("list[dict[str, Any]]", spec["cases"]): + if fixture_subgraphs is not None and "subgraphs" not in case: + case["subgraphs"] = fixture_subgraphs + if fixture_inner_subgraphs is not None and "inner_subgraphs" not in case: + case["inner_subgraphs"] = fixture_inner_subgraphs try: await _run_case(case) except AssertionError as e: @@ -207,10 +359,25 @@ def patch_nodes(graph_block: Mapping[str, Any] | None) -> None: patch_nodes(cast("Mapping[str, Any]", sub)) -def _compile_subgraphs(spec: Mapping[str, Any]) -> dict[str, Any]: +def _compile_subgraphs( + spec: Mapping[str, Any], + *, + provider: OpenAIProvider | None = None, + prompt_manager: PromptManager | None = None, + render_variables: dict[str, Any] | None = None, +) -> dict[str, Any]: """Build any subgraphs declared by the fixture and return a name→compiled-graph registry the adapter consumes. Mirrors the - OTel-side helper in ``test_observability.py``.""" + OTel-side helper in ``test_observability.py``. + + When ``provider`` is supplied, inner subgraph nodes carrying the + ``calls_llm:`` / ``renders_prompt:`` directives (fixtures 029 / + 030) are built via the langfuse-specific + :func:`_build_node_body` rather than the cross-cap adapter — the + adapter doesn't model LLM directives. Outer subgraphs without + LLM directives still resolve through ``build_graph`` so the + existing 031/032/033 wiring is unchanged. + """ subgraph_specs: dict[str, Any] = {} if "subgraph" in spec: single = cast("Mapping[str, Any]", spec["subgraph"]) @@ -221,11 +388,97 @@ def _compile_subgraphs(spec: Mapping[str, Any]) -> dict[str, Any]: subgraph_specs[k] = v compiled_subgraphs: dict[str, Any] = {} for name, sub_spec in subgraph_specs.items(): - sub_built = build_graph(sub_spec, trace=[]) - compiled_subgraphs[name] = sub_built.builder.compile() + if provider is not None and _has_llm_nodes(sub_spec): + compiled_subgraphs[name] = _build_inner_subgraph_with_llm( + sub_spec, + provider=provider, + prompt_manager=prompt_manager, + render_variables=render_variables or {}, + ) + else: + sub_built = build_graph(sub_spec, trace=[]) + compiled_subgraphs[name] = sub_built.builder.compile() return compiled_subgraphs +def _has_llm_nodes(spec: Mapping[str, Any]) -> bool: + """True iff any node in the subgraph spec declares an LLM + directive (``calls_llm`` / ``renders_prompt``) — those need the + langfuse-specific node body builder rather than the cross-cap + adapter.""" + nodes_spec = cast("dict[str, Any]", spec.get("nodes") or {}) + for node_spec in nodes_spec.values(): + if not isinstance(node_spec, dict): + continue + node_dict = cast("dict[str, Any]", node_spec) + if "calls_llm" in node_dict or "renders_prompt" in node_dict: + return True + return False + + +def _infer_state_fields_from_nodes(nodes_spec: dict[str, Any]) -> dict[str, dict[str, Any]]: + """Build a minimal state-fields block from nodes' partial-update + targets so an inner subgraph without an explicit ``state:`` block + (fixture 029) still compiles. Walks ``stores_response_in`` + directives on ``calls_llm`` blocks; defaults each inferred field + to ``string``.""" + fields: dict[str, dict[str, Any]] = {} + for node_spec_any in nodes_spec.values(): + if not isinstance(node_spec_any, dict): + continue + node_spec = cast("dict[str, Any]", node_spec_any) + calls_llm = cast("dict[str, Any] | None", node_spec.get("calls_llm")) + if calls_llm is None: + continue + stores_in = cast("str | None", calls_llm.get("stores_response_in")) + if stores_in is not None and stores_in not in fields: + fields[stores_in] = {"type": "string", "default": ""} + return fields + + +def _build_inner_subgraph_with_llm( + spec: Mapping[str, Any], + *, + provider: OpenAIProvider, + prompt_manager: PromptManager | None, + render_variables: dict[str, Any], +) -> Any: + """Compile an inner subgraph spec into a CompiledGraph using the + langfuse-specific node body builder so ``calls_llm`` / ``renders_prompt`` + directives resolve correctly. Used by fixtures 029 / 030 whose + branch / per-instance subgraphs each make an LLM call.""" + # Some inner-subgraph specs (fixture 029) omit a ``state:`` block. + # Synthesize one from ``stores_response_in`` directives so the + # partial update each node returns has a corresponding field on + # the state class. Default field type is ``string`` with empty + # default, matching the canonical fixture convention. + state_block = cast("dict[str, Any] | None", spec.get("state")) + if state_block is not None: + state_fields = cast("dict[str, dict[str, Any]]", state_block["fields"]) + else: + state_fields = _infer_state_fields_from_nodes(cast("dict[str, Any]", spec.get("nodes") or {})) + state_cls = build_state_cls("InnerSubgraphState", state_fields) + nodes_spec = cast("dict[str, Any]", spec["nodes"]) + entry = cast("str", spec["entry"]) + edges = cast("list[dict[str, str]]", spec["edges"]) + builder = GraphBuilder(state_cls) + for node_name, node_spec in nodes_spec.items(): + node_body = _build_node_body( + node_name=node_name, + node_spec=cast("dict[str, Any]", node_spec), + provider=provider, + prompt_manager=prompt_manager, + render_variables=render_variables, + ) + builder.add_node(node_name, node_body) + for edge in edges: + target_raw = edge["to"] + target = END if target_raw == "END" else target_raw + builder.add_edge(edge["from"], target) + builder.set_entry(entry) + return builder.compile() + + def _resolve_detached_wrapper_names(case: Mapping[str, Any]) -> frozenset[str]: """Translate fixture-level ``detached_subgraphs`` (a list of SUBGRAPH IDENTITY names) into the set of WRAPPER NODE names the observer keys @@ -287,8 +540,32 @@ async def _run_case(case: Mapping[str, Any]) -> None: # no-op so the graph is runnable, mirroring the OTel harness's # ``_patch_unsupported_directives``. _patch_unsupported_directives(case) - subgraphs = _compile_subgraphs(case) - built = build_graph(case, subgraphs=subgraphs, trace=[]) + # Per proposal 0040 fixture 029: rename ``inner_subgraph(s)`` → + # ``subgraph(s)`` so the cross-cap adapter resolves the + # references. Pure key normalization; semantics unchanged. + if isinstance(case, dict): + _normalize_fan_out_subgraph_keys(case) + # Per proposal 0040 fixtures 029 / 030: synthesize the + # augmentation middlewares that drive the per-instance / + # per-branch ``set_invocation_metadata`` calls. Both flow into + # ``build_graph`` via the adapter's standard middleware hooks; + # the augmentation event then fires through the engine and + # the LangfuseObserver handles it via + # ``_handle_metadata_augmentation``. + fan_out_instance_mw, branch_mw = _build_augment_middlewares(case) + subgraphs = _compile_subgraphs( + case, + provider=provider, + prompt_manager=prompt_manager, + render_variables=cast("dict[str, Any]", case.get("render_variables") or {}), + ) + built = build_graph( + case, + subgraphs=subgraphs, + trace=[], + fan_out_instance_middleware=fan_out_instance_mw or None, + parallel_branches_branch_middleware=branch_mw or None, + ) graph = built.builder.compile() initial_state_factory = lambda: built.initial_state(case.get("initial_state", {})) # noqa: E731 else: @@ -394,10 +671,26 @@ def _build_node_body( # named prompt, then call the LLM under `with_active_prompt` # so the Generation's prompt-linkage metadata + entity link # populate per §8.4.4 (024). + # Per proposal 0040 fixture 034 the ``augment_metadata`` directive + # MAY wrap any of the above shapes: at body entry, the harness + # calls ``set_invocation_metadata(**augment)``. Open spans + # outermost-serial (the invocation span / the calling node span) + # MUST then carry the augmented keys in place. + augment_spec = cast("dict[str, Any] | None", node_spec.get("augment_metadata")) + + def _maybe_augment() -> None: + if augment_spec is not None: + from openarmature.observability.metadata import ( # noqa: PLC0415 + set_invocation_metadata, + ) + + set_invocation_metadata(**augment_spec) + update_spec = cast("dict[str, Any] | None", node_spec.get("update")) if update_spec is not None: async def _node(_s: Any) -> dict[str, Any]: + _maybe_augment() return dict(update_spec) return _node @@ -407,6 +700,7 @@ async def _node(_s: Any) -> dict[str, Any]: async def _llm_node(_s: Any) -> dict[str, Any]: assert provider is not None, f"node {node_name!r} has calls_llm but no mock_llm responses" + _maybe_augment() messages_spec = cast( "list[dict[str, Any]] | None", (calls_llm_spec or {}).get("messages"), diff --git a/tests/unit/test_drain.py b/tests/unit/test_drain.py index 25d3df76..49288d77 100644 --- a/tests/unit/test_drain.py +++ b/tests/unit/test_drain.py @@ -26,7 +26,7 @@ GraphBuilder, State, ) -from openarmature.graph.events import NodeEvent +from openarmature.graph.events import MetadataAugmentationEvent, NodeEvent class _S(State): @@ -75,10 +75,11 @@ async def test_drain_without_timeout_waits_for_all_events() -> None: # DrainSummary with the consistent shape. received: list[str] = [] - async def slow_obs(event: NodeEvent) -> None: + async def slow_obs(event: NodeEvent | MetadataAugmentationEvent) -> None: # ~50ms per event; the 3-node graph fires 6 events # (3 nodes × started + completed) so ~300ms of work total. await asyncio.sleep(0.05) + assert isinstance(event, NodeEvent) received.append(event.node_name) compiled = _build_compiled() @@ -102,7 +103,8 @@ async def test_drain_with_timeout_not_reached_for_fast_observers() -> None: # fast. Summary reports clean delivery. received: list[str] = [] - async def fast_obs(event: NodeEvent) -> None: + async def fast_obs(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append(event.node_name) compiled = _build_compiled() @@ -121,10 +123,11 @@ async def test_drain_with_timeout_fires_reports_undelivered() -> None: # generous slack for cancellation settlement). received: list[str] = [] - async def slow_obs(event: NodeEvent) -> None: + async def slow_obs(event: NodeEvent | MetadataAugmentationEvent) -> None: # 200ms per event vs 100ms drain timeout; at most 0-1 events # complete before the deadline fires. await asyncio.sleep(0.2) + assert isinstance(event, NodeEvent) received.append(event.node_name) compiled = _build_compiled() @@ -153,7 +156,7 @@ async def test_drain_after_timeout_leaves_graph_usable() -> None: call_count = [0] received_invocation_two: list[str] = [] - async def obs(event: NodeEvent) -> None: + async def obs(event: NodeEvent | MetadataAugmentationEvent) -> None: # First invocation: slow enough to force the timeout to # fire. Second invocation: fast, so drain completes cleanly. # The mode is controlled by `call_count[0]`: we bump it @@ -161,6 +164,7 @@ async def obs(event: NodeEvent) -> None: if call_count[0] == 0: await asyncio.sleep(0.1) else: + assert isinstance(event, NodeEvent) received_invocation_two.append(event.node_name) compiled = _build_compiled() @@ -201,8 +205,9 @@ async def test_drain_with_zero_timeout_fires_immediately() -> None: # immediately with whatever the worker hasn't gotten to yet. received: list[str] = [] - async def obs(event: NodeEvent) -> None: + async def obs(event: NodeEvent | MetadataAugmentationEvent) -> None: await asyncio.sleep(0.05) + assert isinstance(event, NodeEvent) received.append(event.node_name) compiled = _build_compiled() diff --git a/tests/unit/test_observability_langfuse.py b/tests/unit/test_observability_langfuse.py index 0a7249be..e20c885b 100644 --- a/tests/unit/test_observability_langfuse.py +++ b/tests/unit/test_observability_langfuse.py @@ -427,3 +427,155 @@ async def test_subgraph_dispatch_observation_ended_on_invocation_close() -> None trace = next(iter(client.traces.values())) for obs in trace.observations: assert obs.ended, f"observation {obs.name!r} not ended after shutdown()" + + +# --------------------------------------------------------------------------- +# §3.4 mid-invocation augmentation (proposal 0040) +# --------------------------------------------------------------------------- + + +class _AugmentState(State): + answer: str = "" + + +async def test_metadata_augmentation_updates_trace_and_node_for_outermost() -> None: + # Spec §3.4 MUST + proposal 0040 §6: an outermost-serial + # ``set_invocation_metadata`` call MUST update both the open Trace + # (via client.update_trace, surfacing the entries on + # trace.metadata. for §8.4 top-level filtering) AND the + # calling node's open observation (via handle.update(metadata=)). + # Mirrors fixture 034's Langfuse expectations. + from openarmature.observability.metadata import set_invocation_metadata + + async def node_augments(_s: _AugmentState) -> dict[str, str]: + set_invocation_metadata(request_id="req-xyz") + return {"answer": "ok"} + + g = ( + GraphBuilder(_AugmentState) + .add_node("ask", node_augments) + .add_edge("ask", END) + .set_entry("ask") + .compile() + ) + graph, client, observer = _attach(g) + try: + await graph.invoke(_AugmentState()) + await graph.drain() + finally: + observer.shutdown() + + trace = next(iter(client.traces.values())) + # Trace metadata: augmented key landed on the open Trace. + assert trace.metadata.get("request_id") == "req-xyz" + # Calling node's observation: augmented key landed via in-place + # update before the observation closed. + ask_obs = _find_observation(trace, "ask") + assert ask_obs.metadata.get("request_id") == "req-xyz" + + +async def test_metadata_augmentation_in_fan_out_isolates_per_instance() -> None: + # Fixture 029-shaped: each fan-out instance augments metadata with + # its own product_id. The Trace MUST NOT carry any product_id + # (it's shared across siblings); the per-instance dispatch + # observation AND the inner ask observation for each instance + # MUST carry that instance's OWN product_id. + import asyncio + + from openarmature.observability.correlation import current_fan_out_index + from openarmature.observability.metadata import set_invocation_metadata + + class _ParentState(State): + products: list[dict[str, str]] = [] + results: list[str] = [] + + class _ChildState(State): + product: dict[str, str] = {} + out: str = "" + + async def _ask(s: _ChildState) -> dict[str, str]: + await asyncio.sleep(0) + idx = current_fan_out_index() + assert idx is not None + product_id = s.product["id"] + set_invocation_metadata(product_id=product_id) + return {"out": f"ok-{product_id}"} + + inner = ( + GraphBuilder(_ChildState) + .add_node("inner_ask", _ask) + .add_edge("inner_ask", END) + .set_entry("inner_ask") + .compile() + ) + parent = ( + GraphBuilder(_ParentState) + .add_fan_out_node( + "fan", + subgraph=inner, + collect_field="out", + target_field="results", + items_field="products", + item_field="product", + concurrency=3, + ) + .add_edge("fan", END) + .set_entry("fan") + .compile() + ) + graph, client, observer = _attach(parent) + try: + products = [{"id": "prod-A"}, {"id": "prod-B"}, {"id": "prod-C"}] + await graph.invoke(_ParentState(products=products)) + await graph.drain() + finally: + observer.shutdown() + + trace = next(iter(client.traces.values())) + # Trace metadata MUST NOT carry per-instance product_id (sibling + # isolation — fixture 029's central invariant). + assert "product_id" not in trace.metadata, ( + f"per-instance augmentation leaked onto Trace metadata: {trace.metadata}" + ) + # Each per-instance dispatch observation carries ITS OWN product_id. + instance_obs = [ + obs for obs in trace.observations if obs.name == "fan" and "fan_out_index" in obs.metadata + ] + assert len(instance_obs) == 3 + seen_dispatch: dict[int, str] = {} + for obs in instance_obs: + fan_idx_value = obs.metadata.get("fan_out_index") + product_id_value = obs.metadata.get("product_id") + assert isinstance(fan_idx_value, int) + assert isinstance(product_id_value, str) + seen_dispatch[fan_idx_value] = product_id_value + assert seen_dispatch == {0: "prod-A", 1: "prod-B", 2: "prod-C"} + + +async def test_metadata_augmentation_outside_invocation_is_silent() -> None: + # Plumbing safety: no invocation in scope means no dispatch and no + # observer event — set_invocation_metadata is a Context-only + # mutation. The Langfuse handler is never called in this path so + # no client / no Trace state is created. + from openarmature.observability.metadata import set_invocation_metadata + + set_invocation_metadata(local_key="local_value") + + +async def test_metadata_augmentation_no_op_when_no_entries() -> None: + # Direct-call safety: an augmentation event with empty entries + # should be a no-op on the observer side. + from openarmature.graph.events import MetadataAugmentationEvent + + client = InMemoryLangfuseClient() + observer = LangfuseObserver(client=client) + event = MetadataAugmentationEvent( + entries={}, + namespace=("ask",), + attempt_index=0, + fan_out_index=None, + branch_name=None, + ) + observer._handle_metadata_augmentation(event) # noqa: SLF001 + # No Trace was opened (no invocation in scope) and no exception. + assert client.traces == {} diff --git a/tests/unit/test_observability_otel.py b/tests/unit/test_observability_otel.py index 8f18cfd2..9fc58054 100644 --- a/tests/unit/test_observability_otel.py +++ b/tests/unit/test_observability_otel.py @@ -1395,3 +1395,311 @@ def test_force_flush_delegates_to_provider() -> None: assert observer.force_flush(timeout_ms=1000) is True finally: observer.shutdown() + + +# --------------------------------------------------------------------------- +# §3.4 mid-invocation augmentation (proposal 0040) +# --------------------------------------------------------------------------- + + +class _AugmentState(State): + answer: str = "" + + +async def test_metadata_augmentation_updates_outermost_open_spans() -> None: + # Spec §3.4 MUST + proposal 0040 §6: when a node body calls + # ``set_invocation_metadata`` mid-invocation, every open span whose + # lineage ancestor-or-equals the calling context's MUST be updated + # in place to carry the augmented entries. In a single-node + # outermost-serial graph, that's the invocation root span AND the + # calling node's span. + from openarmature.observability.metadata import set_invocation_metadata + + captured: dict[str, str] = {} + + async def node_augments(_s: _AugmentState) -> dict[str, str]: + set_invocation_metadata(request_id="req-xyz") + captured["seen"] = "yes" + return {"answer": "ok"} + + exporter = InMemorySpanExporter() + observer = OTelObserver(span_processor=SimpleSpanProcessor(exporter)) + g = ( + GraphBuilder(_AugmentState) + .add_node("ask", node_augments) + .add_edge("ask", END) + .set_entry("ask") + .compile() + ) + g.attach_observer(observer) + try: + await g.invoke(_AugmentState()) + await g.drain() + finally: + observer.shutdown() + + spans = exporter.get_finished_spans() + invocation_spans = [s for s in spans if s.name == "openarmature.invocation"] + ask_spans = [s for s in spans if s.name == "ask"] + assert len(invocation_spans) == 1 + assert len(ask_spans) == 1 + inv_attrs = dict(invocation_spans[0].attributes or {}) + ask_attrs = dict(ask_spans[0].attributes or {}) + # Augmentation reached both the invocation span (open at the call) + # and the calling node's span (the augmenter itself). + assert inv_attrs.get("openarmature.user.request_id") == "req-xyz" + assert ask_attrs.get("openarmature.user.request_id") == "req-xyz" + + +async def test_metadata_augmentation_outside_invocation_is_silent() -> None: + # Plumbing safety: ``set_invocation_metadata`` outside any active + # invocation updates the ContextVar but emits no augmentation event + # (no dispatch is in scope). The observer never sees an event so + # no observer-side error surfaces. + from openarmature.observability.metadata import set_invocation_metadata + + # No graph, no observer attached — should not raise. + set_invocation_metadata(local_only="value") + + +async def test_metadata_augmentation_no_op_when_no_entries() -> None: + # Empty entries dict is a no-op at the public API (the helper + # short-circuits before validating or dispatching). The observer + # still must tolerate the case in any future direct test path. + from openarmature.graph.events import MetadataAugmentationEvent + + observer = OTelObserver(span_processor=SimpleSpanProcessor(InMemorySpanExporter())) + try: + # Direct call to the handler bypasses the engine so we can + # confirm an empty-entries augmentation is silently dropped. + event = MetadataAugmentationEvent( + entries={}, + namespace=("ask",), + attempt_index=0, + fan_out_index=None, + branch_name=None, + ) + observer._handle_metadata_augmentation(event) # noqa: SLF001 + finally: + observer.shutdown() + + +async def test_metadata_augmentation_in_fan_out_isolates_per_instance() -> None: + # Spec §3.4 + proposal 0040 scoping rule: a fan-out instance + # augmenting metadata MUST update its own instance dispatch span + # and its own inner-node span, but NOT the shared fan_out_node + # parent span, NOT the invocation span, and NOT sibling instances' + # spans. Each ``inner_ask`` span ends up tagged with its own + # ``product_id`` only. + import asyncio + + from openarmature.observability.correlation import current_fan_out_index + from openarmature.observability.metadata import set_invocation_metadata + + class _ParentState(State): + products: list[dict[str, str]] = Field(default_factory=list[dict[str, str]]) + results: list[str] = Field(default_factory=list[str]) + + class _ChildState(State): + product: dict[str, str] = Field(default_factory=dict[str, str]) + out: str = "" + + async def _ask(s: _ChildState) -> dict[str, str]: + # Yield once so concurrent instances interleave their + # augmentation events on the observer queue. + await asyncio.sleep(0) + idx = current_fan_out_index() + assert idx is not None + product_id = s.product["id"] + set_invocation_metadata(product_id=product_id) + return {"out": f"ok-{product_id}"} + + inner = ( + GraphBuilder(_ChildState) + .add_node("inner_ask", _ask) + .add_edge("inner_ask", END) + .set_entry("inner_ask") + .compile() + ) + parent = ( + GraphBuilder(_ParentState) + .add_fan_out_node( + "fan", + subgraph=inner, + collect_field="out", + target_field="results", + items_field="products", + item_field="product", + concurrency=3, + ) + .add_edge("fan", END) + .set_entry("fan") + ) + exporter = InMemorySpanExporter() + observer = OTelObserver(span_processor=SimpleSpanProcessor(exporter)) + compiled = parent.compile() + compiled.attach_observer(observer) + try: + products = [ + {"id": "prod-A"}, + {"id": "prod-B"}, + {"id": "prod-C"}, + ] + await compiled.invoke(_ParentState(products=products)) + await compiled.drain() + finally: + observer.shutdown() + + spans = exporter.get_finished_spans() + inner_spans = [s for s in spans if s.name == "inner_ask"] + assert len(inner_spans) == 3 + seen: dict[str, str] = {} + for span in inner_spans: + attrs = dict(span.attributes or {}) + product_id = attrs.get("openarmature.user.product_id") + fan_out_idx = attrs.get("openarmature.node.fan_out_index") + assert isinstance(product_id, str), f"missing per-instance augmentation on {span.name}" + assert isinstance(fan_out_idx, int) + seen[str(fan_out_idx)] = product_id + # Each instance carries its OWN product_id; no sibling leakage. + assert seen == {"0": "prod-A", "1": "prod-B", "2": "prod-C"} + + # The shared fan-out parent node span and the invocation span MUST + # NOT carry any per-instance product_id. The PER-INSTANCE dispatch + # spans (synthesized for non-detached fan-outs per §5.4 + proposal + # 0013) are IN scope, so each one SHOULD carry its own product_id. + invocation_spans = [s for s in spans if s.name == "openarmature.invocation"] + fan_spans = [s for s in spans if s.name == "fan"] + assert len(invocation_spans) == 1 + # The shared fan-out parent has ``openarmature.fan_out.item_count`` + # set; per-instance dispatch spans don't. + parent_fan_spans = [s for s in fan_spans if "openarmature.fan_out.item_count" in dict(s.attributes or {})] + instance_fan_spans = [ + s for s in fan_spans if "openarmature.fan_out.item_count" not in dict(s.attributes or {}) + ] + assert len(parent_fan_spans) == 1 + assert len(instance_fan_spans) == 3 + # Parent + invocation: no per-instance product_id leakage. + for span in (*parent_fan_spans, *invocation_spans): + attrs = dict(span.attributes or {}) + assert "openarmature.user.product_id" not in attrs, ( + f"per-instance augmentation leaked onto {span.name} span" + ) + # Per-instance dispatch spans: each one carries its own product_id. + seen_dispatch: dict[int, str] = {} + for span in instance_fan_spans: + attrs = dict(span.attributes or {}) + idx_value = attrs.get("openarmature.node.fan_out_index") + product_value = attrs.get("openarmature.user.product_id") + assert isinstance(idx_value, int) + assert isinstance(product_value, str), f"per-instance dispatch span missing product_id; attrs={attrs}" + seen_dispatch[idx_value] = product_value + assert seen_dispatch == {0: "prod-A", 1: "prod-B", 2: "prod-C"} + + +async def test_metadata_augmentation_in_parallel_branches_skips_sibling() -> None: + # Sibling-skip for parallel-branches: two concurrent branches each + # augment metadata with their own branch identifier. Each branch's + # inner-node span carries ONLY its own ``branch_label``; no + # cross-branch leakage. This also implicitly verifies that the + # OTel observer's open-span key disambiguates concurrent same- + # named inner nodes across sibling branches (pre-fix, both + # branches' ``ask`` opens collided on the same _StackKey). + import asyncio + + from openarmature.graph import BranchSpec + from openarmature.observability.metadata import set_invocation_metadata + + class _DispatchState(State): + fraud_result: str = "" + audit_result: str = "" + + class _FraudState(State): + score: str = "" + + class _AuditState(State): + summary: str = "" + + async def _fraud_ask(_s: _FraudState) -> dict[str, str]: + await asyncio.sleep(0) + set_invocation_metadata(branch_label="fraud_check") + return {"score": "low"} + + async def _audit_ask(_s: _AuditState) -> dict[str, str]: + await asyncio.sleep(0) + set_invocation_metadata(branch_label="policy_audit") + return {"summary": "compliant"} + + fraud_subgraph = ( + GraphBuilder(_FraudState).add_node("ask", _fraud_ask).add_edge("ask", END).set_entry("ask").compile() + ) + audit_subgraph = ( + GraphBuilder(_AuditState).add_node("ask", _audit_ask).add_edge("ask", END).set_entry("ask").compile() + ) + exporter = InMemorySpanExporter() + observer = OTelObserver(span_processor=SimpleSpanProcessor(exporter)) + g = ( + GraphBuilder(_DispatchState) + .add_parallel_branches_node( + "dispatcher", + branches={ + "fraud_check": BranchSpec( + subgraph=fraud_subgraph, + outputs={"fraud_result": "score"}, + ), + "policy_audit": BranchSpec( + subgraph=audit_subgraph, + outputs={"audit_result": "summary"}, + ), + }, + ) + .add_edge("dispatcher", END) + .set_entry("dispatcher") + .compile() + ) + g.attach_observer(observer) + try: + await g.invoke(_DispatchState()) + await g.drain() + finally: + observer.shutdown() + + spans = exporter.get_finished_spans() + # Pre-fix: two concurrent ``ask`` spans would collide on the + # _StackKey, so only ONE ask span would land. Post-fix: both + # branches' ask spans land, each tagged with its own branch_name. + ask_spans = [s for s in spans if s.name == "ask"] + assert len(ask_spans) == 2 + by_branch: dict[str, dict[str, Any]] = {} + for span in ask_spans: + attrs = dict(span.attributes or {}) + bn = attrs.get("openarmature.branch_name") + assert isinstance(bn, str) + by_branch[bn] = attrs + # Each branch's ask carries its OWN branch_label augmentation. + assert by_branch["fraud_check"].get("openarmature.user.branch_label") == "fraud_check" + assert by_branch["policy_audit"].get("openarmature.user.branch_label") == "policy_audit" + # No cross-branch leakage: fraud's ask does NOT carry policy_audit's + # label and vice versa. The branch_label key is the same name; what + # matters is each span shows ONLY its own value. + assert by_branch["fraud_check"].get("openarmature.user.branch_label") != "policy_audit" + assert by_branch["policy_audit"].get("openarmature.user.branch_label") != "fraud_check" + + # The parallel-branches NODE span(s) and the invocation span MUST + # NOT carry either branch's branch_label (per-async-context + # isolation). Note: the current OTel mapping synthesizes a + # subgraph wrapper at the parallel-branches NODE's namespace in + # addition to the NODE's own span — that's a pre-existing + # divergence from fixture 030's expected Langfuse shape that + # `discuss-otel-parallel-branches-dispatch-span` is asking spec + # to settle. For this test both dispatcher-named spans MUST be + # augmentation-clean. + dispatcher_spans = [s for s in spans if s.name == "dispatcher"] + invocation_spans = [s for s in spans if s.name == "openarmature.invocation"] + assert len(invocation_spans) == 1 + assert len(dispatcher_spans) >= 1 + for span in (*dispatcher_spans, *invocation_spans): + attrs = dict(span.attributes or {}) + assert "openarmature.user.branch_label" not in attrs, ( + f"per-branch augmentation leaked onto {span.name} span" + ) diff --git a/tests/unit/test_observer.py b/tests/unit/test_observer.py index e4ab6760..d37946e1 100644 --- a/tests/unit/test_observer.py +++ b/tests/unit/test_observer.py @@ -9,10 +9,11 @@ import asyncio import warnings +from types import MappingProxyType from typing import Literal from openarmature.graph import Observer, State, SubscribedObserver -from openarmature.graph.events import NodeEvent +from openarmature.graph.events import MetadataAugmentationEvent, NodeEvent from openarmature.graph.observer import ( _DRAIN_SENTINEL, RemoveHandle, @@ -22,6 +23,7 @@ _QueuedItem, deliver_loop, ) +from openarmature.observability.metadata import set_invocation_metadata class DummyState(State): @@ -62,7 +64,8 @@ async def _drain(queue: asyncio.Queue[_QueuedItem | None], worker: asyncio.Task[ async def test_events_delivered_in_queue_order() -> None: received: list[str] = [] - async def observer(event: NodeEvent) -> None: + async def observer(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append(event.node_name) queue: asyncio.Queue[_QueuedItem | None] = asyncio.Queue() @@ -78,10 +81,12 @@ async def observer(event: NodeEvent) -> None: async def test_multiple_observers_fire_in_registration_order() -> None: received: list[str] = [] - async def obs1(event: NodeEvent) -> None: + async def obs1(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append(f"obs1:{event.node_name}") - async def obs2(event: NodeEvent) -> None: + async def obs2(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append(f"obs2:{event.node_name}") queue: asyncio.Queue[_QueuedItem | None] = asyncio.Queue() @@ -101,7 +106,7 @@ async def obs2(event: NodeEvent) -> None: async def test_observer_exception_does_not_propagate_to_caller() -> None: - async def boom(_event: NodeEvent) -> None: + async def boom(_event: NodeEvent | MetadataAugmentationEvent) -> None: raise RuntimeError("nope") queue: asyncio.Queue[_QueuedItem | None] = asyncio.Queue() @@ -119,10 +124,11 @@ async def boom(_event: NodeEvent) -> None: async def test_raising_observer_does_not_block_siblings_on_same_event() -> None: received: list[str] = [] - async def obs1(_event: NodeEvent) -> None: + async def obs1(_event: NodeEvent | MetadataAugmentationEvent) -> None: raise RuntimeError("obs1 boom") - async def obs2(event: NodeEvent) -> None: + async def obs2(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append(event.node_name) queue: asyncio.Queue[_QueuedItem | None] = asyncio.Queue() @@ -139,10 +145,11 @@ async def obs2(event: NodeEvent) -> None: async def test_raising_observer_does_not_block_subsequent_events() -> None: received: list[str] = [] - async def always_raises(_event: NodeEvent) -> None: + async def always_raises(_event: NodeEvent | MetadataAugmentationEvent) -> None: raise RuntimeError("always boom") - async def silent(event: NodeEvent) -> None: + async def silent(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append(event.node_name) queue: asyncio.Queue[_QueuedItem | None] = asyncio.Queue() @@ -164,7 +171,8 @@ async def silent(event: NodeEvent) -> None: async def test_phase_filter_skips_unsubscribed_phase() -> None: received: list[tuple[str, str]] = [] - async def obs(event: NodeEvent) -> None: + async def obs(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append((event.node_name, event.phase)) queue: asyncio.Queue[_QueuedItem | None] = asyncio.Queue() @@ -180,7 +188,7 @@ async def obs(event: NodeEvent) -> None: async def test_subscribed_observer_rejects_empty_phases() -> None: - async def obs(_event: NodeEvent) -> None: + async def obs(_event: NodeEvent | MetadataAugmentationEvent) -> None: pass try: @@ -191,7 +199,7 @@ async def obs(_event: NodeEvent) -> None: async def test_subscribed_observer_rejects_unknown_phase() -> None: - async def obs(_event: NodeEvent) -> None: + async def obs(_event: NodeEvent | MetadataAugmentationEvent) -> None: pass try: @@ -207,7 +215,8 @@ async def obs(_event: NodeEvent) -> None: async def test_sentinel_terminates_worker_after_processing_queued_events() -> None: received: list[str] = [] - async def observer(event: NodeEvent) -> None: + async def observer(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append(event.node_name) queue: asyncio.Queue[_QueuedItem | None] = asyncio.Queue() @@ -235,10 +244,10 @@ async def test_dispatch_skips_when_no_observers_for_depth() -> None: async def test_dispatch_enqueues_with_full_observer_chain_in_order() -> None: - async def graph_obs(_event: NodeEvent) -> None: + async def graph_obs(_event: NodeEvent | MetadataAugmentationEvent) -> None: pass - async def invocation_obs(_event: NodeEvent) -> None: + async def invocation_obs(_event: NodeEvent | MetadataAugmentationEvent) -> None: pass graph_subscribed = _wrap(graph_obs) @@ -263,13 +272,13 @@ async def invocation_obs(_event: NodeEvent) -> None: async def test_descend_extends_chain_namespace_and_parent_states() -> None: - async def outer_obs(_event: NodeEvent) -> None: + async def outer_obs(_event: NodeEvent | MetadataAugmentationEvent) -> None: pass - async def sub_obs(_event: NodeEvent) -> None: + async def sub_obs(_event: NodeEvent | MetadataAugmentationEvent) -> None: pass - async def invocation_obs(_event: NodeEvent) -> None: + async def invocation_obs(_event: NodeEvent | MetadataAugmentationEvent) -> None: pass outer_subscribed = _wrap(outer_obs) @@ -315,7 +324,7 @@ async def test_take_step_shares_counter_across_descended_contexts() -> None: def test_remove_handle_detaches_observer() -> None: - async def obs(_event: NodeEvent) -> None: + async def obs(_event: NodeEvent | MetadataAugmentationEvent) -> None: pass subscribed = _wrap(obs) @@ -328,7 +337,7 @@ async def obs(_event: NodeEvent) -> None: def test_remove_handle_is_idempotent() -> None: - async def obs(_event: NodeEvent) -> None: + async def obs(_event: NodeEvent | MetadataAugmentationEvent) -> None: pass subscribed = _wrap(obs) @@ -338,3 +347,125 @@ async def obs(_event: NodeEvent) -> None: handle.remove() handle.remove() # second call is a no-op, doesn't raise assert subscribed not in observers + + +# ===== Metadata-augmentation event delivery (proposal 0040) ===== + + +async def test_metadata_augmentation_event_bypasses_phase_filter() -> None: + """Augmentation events flow through ``__call__`` on the union-typed + Observer Protocol and ignore the per-observer ``phases`` set + entirely (they aren't phase events). Observers that only care + about NodeEvent ``isinstance``-narrow and early-return. + """ + augment_received: list[MetadataAugmentationEvent] = [] + node_received: list[NodeEvent] = [] + + async def observer(event: NodeEvent | MetadataAugmentationEvent) -> None: + if isinstance(event, MetadataAugmentationEvent): + augment_received.append(event) + else: + node_received.append(event) + + queue: asyncio.Queue[_QueuedItem | None] = asyncio.Queue() + worker = asyncio.create_task(deliver_loop(queue, _DrainCounters())) + # Subscribe to ``completed`` only — to prove the augmentation event + # bypasses the phase filter (it has no phase). + completed_only = (SubscribedObserver(observer=observer, phases=frozenset({"completed"})),) + augmentation = MetadataAugmentationEvent( + entries=MappingProxyType({"region": "us-east-1"}), + namespace=("router", "classify"), + attempt_index=0, + fan_out_index=None, + branch_name=None, + ) + queue.put_nowait(_QueuedItem(event=augmentation, observers=completed_only)) + await _drain(queue, worker) + + assert augment_received == [augmentation] + assert node_received == [] + + +async def test_metadata_augmentation_observer_exception_is_isolated() -> None: + """A raise on the augmentation event follows the same isolation + contract as a raise on a NodeEvent — warned, sibling observers + still run, the worker keeps draining.""" + sibling_received: list[MetadataAugmentationEvent] = [] + + async def boom(_event: NodeEvent | MetadataAugmentationEvent) -> None: + raise RuntimeError("boom") + + async def good(event: NodeEvent | MetadataAugmentationEvent) -> None: + if isinstance(event, MetadataAugmentationEvent): + sibling_received.append(event) + + queue: asyncio.Queue[_QueuedItem | None] = asyncio.Queue() + worker = asyncio.create_task(deliver_loop(queue, _DrainCounters())) + augmentation = MetadataAugmentationEvent( + entries=MappingProxyType({"k": "v"}), + namespace=(), + ) + subscribed = (_wrap(boom), _wrap(good)) + with warnings.catch_warnings(record=True) as captured: + warnings.simplefilter("always") + queue.put_nowait(_QueuedItem(event=augmentation, observers=subscribed)) + await _drain(queue, worker) + + assert sibling_received == [augmentation] + assert any("observer raised RuntimeError" in str(w.message) for w in captured) + + +async def test_set_invocation_metadata_emits_augmentation_event_via_dispatch() -> None: + """``set_invocation_metadata`` reads the current_dispatch closure + (engine-installed in real runs) and constructs a + MetadataAugmentationEvent carrying the delta + lineage from the + correlation ContextVars.""" + from openarmature.observability.correlation import ( + _reset_active_dispatch, + _reset_attempt_index, + _reset_branch_name, + _reset_fan_out_index, + _reset_namespace_prefix, + _set_active_dispatch, + _set_attempt_index, + _set_branch_name, + _set_fan_out_index, + _set_namespace_prefix, + ) + + captured: list[NodeEvent | MetadataAugmentationEvent] = [] + + def dispatch(event: NodeEvent | MetadataAugmentationEvent) -> None: + captured.append(event) + + dispatch_token = _set_active_dispatch(dispatch) + namespace_token = _set_namespace_prefix(("outer", "leaf")) + fan_out_token = _set_fan_out_index(2) + branch_token = _set_branch_name("primary") + attempt_token = _set_attempt_index(1) + try: + set_invocation_metadata(region="us-east-1", retries=3) + finally: + _reset_attempt_index(attempt_token) + _reset_branch_name(branch_token) + _reset_fan_out_index(fan_out_token) + _reset_namespace_prefix(namespace_token) + _reset_active_dispatch(dispatch_token) + + assert len(captured) == 1 + event = captured[0] + assert isinstance(event, MetadataAugmentationEvent) + assert dict(event.entries) == {"region": "us-east-1", "retries": 3} + assert event.namespace == ("outer", "leaf") + assert event.attempt_index == 1 + assert event.fan_out_index == 2 + assert event.branch_name == "primary" + + +def test_set_invocation_metadata_outside_invocation_skips_dispatch() -> None: + """Without a current_dispatch installed (no engine in scope), + ``set_invocation_metadata`` still updates the ContextVar but + does NOT raise and does NOT attempt to enqueue an event.""" + # Sanity: by default outside any engine the dispatch ContextVar is + # None, so the call should be a no-op on the queue side. + set_invocation_metadata(local_key="local_value") diff --git a/tests/unit/test_runtime_errors.py b/tests/unit/test_runtime_errors.py index acd0c8bf..e08c14ec 100644 --- a/tests/unit/test_runtime_errors.py +++ b/tests/unit/test_runtime_errors.py @@ -163,11 +163,12 @@ async def test_routing_error_lands_on_preceding_node_completed_event() -> None: preceding node's `completed` event with `error` populated, NOT in a separate event pair. The downstream node never fires events.""" from openarmature.graph import RoutingError - from openarmature.graph.events import NodeEvent + from openarmature.graph.events import MetadataAugmentationEvent, NodeEvent received: list[NodeEvent] = [] - async def observer(event: NodeEvent) -> None: + async def observer(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append(event) async def node_a(_state: Any) -> dict[str, Any]: @@ -217,11 +218,12 @@ async def test_edge_exception_lands_on_preceding_node_completed_event() -> None: conditional edge function raising lands on the preceding node's `completed` event with `error` populated, NOT in a separate event pair. The downstream node never fires events.""" - from openarmature.graph.events import NodeEvent + from openarmature.graph.events import MetadataAugmentationEvent, NodeEvent received: list[NodeEvent] = [] - async def observer(event: NodeEvent) -> None: + async def observer(event: NodeEvent | MetadataAugmentationEvent) -> None: + assert isinstance(event, NodeEvent) received.append(event) async def node_a(_state: Any) -> dict[str, Any]: