From 7053244b3f991c8797a1242f10d37adf86142d9a Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Thu, 25 Jun 2026 18:27:19 -0700 Subject: [PATCH 1/3] Fix nested fan-out collapse under concurrency A fan-out nested inside an outer fan-out instance keyed its per-fan-out tracking entry by namespace and node name only, so the shared progress dict collided across concurrent outer instances. The second instance found the first's entry already complete and rolled its result forward, so every outer instance returned the first instance's inner result (silently wrong output) and the inner subgraph ran only once. Carry the enclosing fan-out instance lineage on the tracking key, in the in-memory dict and through the checkpoint projection, lookup, cleanup, and restore. Top-level and subgraph- or branch-nested fan-outs have an empty lineage, so their behavior, including resume, is unchanged. A fan-out nested inside an outer instance re-runs rather than skipping on resume, since the record format carries no lineage; tracked as a follow-up. --- CHANGELOG.md | 1 + src/openarmature/graph/compiled.py | 47 ++++++++++-- src/openarmature/graph/fan_out.py | 13 +++- src/openarmature/graph/observer.py | 7 +- tests/unit/test_checkpoint.py | 112 +++++++++++++++++++++++++++++ tests/unit/test_fan_out.py | 84 ++++++++++++++++++++++ 6 files changed, 254 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a90bad..6b27b52 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). The - **Langfuse per-branch dispatch-span observation** (observability §4.3 / §8.4.2, proposals 0042 / 0044). The Langfuse observer now synthesizes a per-branch Span observation under a `parallel_branches` dispatcher node, so each branch's inner observations nest under their own branch span (a three-level dispatcher / per-branch-span / inner-nodes tree) instead of parenting directly under the dispatcher. The per-branch observation carries the OA-emitted `branch_name` alongside the caller baseline metadata and any per-branch augmentation, and the Generation observation now carries `branch_name` too. The OTel observer already produced this shape (proposal 0044 shipped OTel-only in v0.11.0); this brings the Langfuse mapping into line. Callable branches (proposal 0075) are unchanged. - **Augmentation no longer lands on a shared-parent fan-out / parallel-branches node** (observability §3.4, proposal 0045). A key set via `set_invocation_metadata` inside a fan-out instance or a parallel-branches branch was incorrectly applied to the shared fan-out / dispatcher NODE span (the fork point) when the augmenting context executed at that node's own namespace, in addition to the per-instance / per-branch dispatch span where it belongs. Both the OTel and Langfuse observers now skip the shared-parent node in that case, matching the behavior already applied to strict-ancestor shared parents. The per-instance / per-branch dispatch spans and the lineage ancestors that carry the augmentation are unaffected. - **Langfuse fan-out-instance dispatch nested below the top level** (observability §5.4, proposal 0013). The Langfuse observer's per-instance dispatch synthesis and parent resolution are now prefix-general, so a fan-out node sitting inside a serial subgraph wrapper (rather than at the top namespace level) gets its per-instance dispatch observation synthesized and its inner observations parented under it. This matches the OTel observer, which already resolved across every namespace prefix. +- **Nested fan-out no longer collapses under concurrency** (engine). A fan-out nested inside an outer fan-out instance shared a single per-fan-out tracking entry across all outer instances, because that entry was keyed by namespace plus node name only. With concurrent outer instances the second instance found the first's entry already marked complete and rolled its result forward, so every outer instance returned the first instance's inner result (silently wrong output) and the inner subgraph ran only once. The tracking key now carries the enclosing fan-out instance lineage, so each outer instance gets its own inner fan-out progress and correct per-instance results. Top-level and subgraph- or branch-nested fan-outs are unaffected (their enclosing lineage is empty). Resume of a fan-out nested inside an outer fan-out instance does not yet round-trip per-outer-instance progress, so it re-runs rather than skipping on resume; tracked as a follow-up. ## [0.15.0] — 2026-06-22 diff --git a/src/openarmature/graph/compiled.py b/src/openarmature/graph/compiled.py index 5cf4209..b86f0cc 100644 --- a/src/openarmature/graph/compiled.py +++ b/src/openarmature/graph/compiled.py @@ -316,7 +316,13 @@ def _find_innermost_fan_out_instance_state( # fan-out's full key is (namespace_before_fan_out, fan_out_name) # where namespace_before_fan_out + (fan_out_name,) == prefix. for split in range(len(prefix), 0, -1): - key = (prefix[: split - 1], prefix[split - 1]) + # The fan-out at prefix[:split] registered its tracking entry keyed by + # its ENCLOSING fan-out instance lineage (the non-None fan_out_index chain + # up to its own level, prefix depth split-1). Reconstruct it from the + # current chain so a fan-out nested inside an outer instance routes to the + # right outer instance's entry. + lineage = tuple(i for i in context.fan_out_index_chain[: split - 1] if i is not None) + key = (prefix[: split - 1], prefix[split - 1], lineage) if key in state_dict: exec_state = state_dict[key] idx = context.fan_out_index @@ -326,7 +332,7 @@ def _find_innermost_fan_out_instance_state( def _project_fan_out_progress( - state_dict: Mapping[tuple[tuple[str, ...], str], _FanOutExecutionState], + state_dict: Mapping[tuple[tuple[str, ...], str, tuple[int, ...]], _FanOutExecutionState], ) -> tuple[FanOutProgress, ...]: """Project the engine-internal mutable per-fan-out state into the frozen :class:`FanOutProgress` shape on a saved record. @@ -343,7 +349,16 @@ def _project_fan_out_progress( byte-identically, which matters for backends that hash records. """ out: list[FanOutProgress] = [] - for (namespace, name), exec_state in sorted(state_dict.items()): + # The key's third element is the enclosing fan-out instance lineage; it is + # NOT projected onto the record (top-level / subgraph-nested fan-outs have an + # empty lineage, and nested-fan-out resume is a tracked limitation). One + # consequence: a fan-out nested inside an outer fan-out instance emits one + # record PER outer instance, all sharing (namespace, fan_out_node_name); + # _restore_fan_out_progress_state is last-wins on that collision (the nested + # fan-out re-runs on resume regardless). Sorting includes the lineage so + # those same-namespace entries still order deterministically (preserving the + # byte-identical-record guarantee above). + for (namespace, name, _lineage), exec_state in sorted(state_dict.items()): instances = tuple( FanOutInstanceProgress( state=inst.state, @@ -366,7 +381,7 @@ def _project_fan_out_progress( def _restore_fan_out_progress_state( saved: Sequence[FanOutProgress], -) -> dict[tuple[tuple[str, ...], str], _FanOutExecutionState]: +) -> dict[tuple[tuple[str, ...], str, tuple[int, ...]], _FanOutExecutionState]: """Inverse projection of :func:`_project_fan_out_progress`. On resume the loaded record's frozen ``fan_out_progress`` tuple gets unpacked into the mutable per-fan-out tracking dict that ``FanOutNode`` @@ -387,7 +402,7 @@ def _restore_fan_out_progress_state( the engine's canonical error-record shape, and a heuristic would misclassify them. """ - out: dict[tuple[tuple[str, ...], str], _FanOutExecutionState] = {} + out: dict[tuple[tuple[str, ...], str, tuple[int, ...]], _FanOutExecutionState] = {} for fp in saved: instances: list[_FanOutInstanceState] = [] for inst in fp.instances: @@ -400,7 +415,18 @@ def _restore_fan_out_progress_state( completed_inner_positions=list(inst.completed_inner_positions), ) ) - key = (fp.namespace, fp.fan_out_node_name) + # The enclosing fan-out instance lineage defaults to empty: the saved + # record carries no lineage, which is correct for top-level and + # subgraph/branch-nested fan-outs (all empty). A fan-out nested inside an + # outer fan-out instance does not round-trip its per-outer-instance + # progress through the current record format (it would need the lineage + # on the record): its in-memory keys carry the lineage, so the restored + # empty-lineage entry never matches and the nested fan-out RE-RUNS on + # resume. (Before the lineage fix it would instead skip, rolling forward + # the collapsed/wrong shared entry -- so re-running is the safer of two + # never-correct behaviors, and matches the spec's inner-subgraph re-entry + # model.) A full fix needs the lineage on the record; tracked separately. + key = (fp.namespace, fp.fan_out_node_name, ()) out[key] = _FanOutExecutionState( fan_out_node_name=fp.fan_out_node_name, namespace=fp.namespace, @@ -1987,7 +2013,14 @@ async def innermost(s: Any) -> Mapping[str, Any]: # raised, so subsequent saves in this invocation do not carry # stale fan-out progress and a retry middleware on the fan-out # node sees a fresh tracked state on the second attempt. - fan_out_progress_key = (context.namespace_prefix, current) + # Match the lineage-aware key FanOutNode.run registers (namespace, node + # name, enclosing fan-out instance lineage) so a nested fan-out's cleanup + # pops its own outer-instance entry, not a sibling's. + fan_out_progress_key = ( + context.namespace_prefix, + current, + tuple(i for i in context.fan_out_index_chain if i is not None), + ) try: try: try: diff --git a/src/openarmature/graph/fan_out.py b/src/openarmature/graph/fan_out.py index 3363599..4a3c461 100644 --- a/src/openarmature/graph/fan_out.py +++ b/src/openarmature/graph/fan_out.py @@ -208,7 +208,18 @@ async def run_with_context( # shared dict. Resume threads a pre-restored entry through # ``context.fan_out_progress_state``; first-run constructs a # fresh one with all instances ``not_started``. - key = (context.namespace_prefix, self.name) + # + # The key carries the ENCLOSING fan-out instance lineage (the non-None + # fan_out_index chain), not just the namespace + node name. A fan-out + # nested inside an outer fan-out instance has the same namespace for every + # outer instance, so without the lineage the shared dict collides across + # concurrent outer instances and the second instance rolls forward the + # first's "completed" results (silently wrong results). Subgraph / branch + # nesting and top-level fan-outs contribute no fan-out index, so the + # lineage is empty there -- matching the resume restore (which defaults it + # to empty), so top-level / subgraph-nested resume is unaffected. + enclosing_fan_out_lineage = tuple(i for i in context.fan_out_index_chain if i is not None) + key = (context.namespace_prefix, self.name, enclosing_fan_out_lineage) exec_state = context.fan_out_progress_state.get(key) if exec_state is None: exec_state = _FanOutExecutionState( diff --git a/src/openarmature/graph/observer.py b/src/openarmature/graph/observer.py index 383bc9f..0f583c0 100644 --- a/src/openarmature/graph/observer.py +++ b/src/openarmature/graph/observer.py @@ -525,8 +525,11 @@ class _InvocationContext: # of the fan-out so concurrent saves see consistent sibling state. # ``_maybe_save_checkpoint`` projects this into the frozen # ``FanOutProgress`` shape on the saved CheckpointRecord. - fan_out_progress_state: dict[tuple[tuple[str, ...], str], _FanOutExecutionState] = field( - default_factory=dict[tuple[tuple[str, ...], str], _FanOutExecutionState] + # Keyed by (namespace, fan_out_node_name, enclosing_fan_out_instance_lineage) + # -- the lineage (non-None outer fan_out_index chain) disambiguates a fan-out + # nested inside an outer fan-out instance across concurrent outer instances. + fan_out_progress_state: dict[tuple[tuple[str, ...], str, tuple[int, ...]], _FanOutExecutionState] = field( + default_factory=dict[tuple[tuple[str, ...], str, tuple[int, ...]], _FanOutExecutionState] ) # Per spec §6 Drain (proposal 0010): shared mutable counters that # the worker reads at drain-cancel time to report undelivered events diff --git a/tests/unit/test_checkpoint.py b/tests/unit/test_checkpoint.py index 91b58d8..d0426cd 100644 --- a/tests/unit/test_checkpoint.py +++ b/tests/unit/test_checkpoint.py @@ -745,6 +745,118 @@ async def test_nested_fan_out_records_outermost_schema_version() -> None: ) +# --------------------------------------------------------------------------- +# Nested fan-out (a fan-out inside an outer fan-out INSTANCE): the per-fan-out +# tracking key carries the enclosing instance lineage. Verify the save/restore +# round-trip projects the inner fan-out's progress and resume rolls the outer +# instances forward to the correct per-instance results. +# --------------------------------------------------------------------------- + + +class _RNestLeafState(State): + tag: str = "" + seed: str = "" + out: str = "" + + +class _RNestMidState(State): + tag: str = "" + seeds: list[str] = Field(default_factory=list[str]) + collected: list[str] = Field(default_factory=list[str]) + + +class _RNestOuterState(State): + products: list[str] = Field(default_factory=list[str]) + seeds: list[str] = Field(default_factory=list[str]) + results: list[Any] = Field(default_factory=list[Any]) + + +async def test_nested_fan_out_in_instance_resume_round_trips() -> None: + """A fan-out nested inside an outer fan-out instance round-trips through a + checkpoint save/restore: the projection emits the inner fan-out's + per-outer-instance progress, restore rebuilds the lineage-keyed dict, and + resuming from the completed record rolls the outer instances forward to the + correct per-instance results without re-running the inner leaf.""" + leaf_calls = 0 + + async def leaf(state: _RNestLeafState) -> dict[str, str]: + nonlocal leaf_calls + leaf_calls += 1 + return {"out": f"{state.tag}-{state.seed}"} + + leaf_g = ( + GraphBuilder(_RNestLeafState).add_node("ask", leaf).add_edge("ask", END).set_entry("ask").compile() + ) + mid_g = ( + GraphBuilder(_RNestMidState) + .add_fan_out_node( + "inner_fan", + subgraph=leaf_g, + items_field="seeds", + item_field="seed", + inputs={"tag": "tag"}, + collect_field="out", + target_field="collected", + ) + .add_edge("inner_fan", END) + .set_entry("inner_fan") + .compile() + ) + cp = InMemoryCheckpointer() + captured_ids: list[str] = [] + captured_records: list[CheckpointRecord] = [] + original_save = cp.save + + async def capture_save(invocation_id: str, record: CheckpointRecord) -> None: + captured_ids.append(invocation_id) + captured_records.append(record) + await original_save(invocation_id, record) + + cp.save = capture_save # type: ignore[method-assign] + + outer_g = ( + GraphBuilder(_RNestOuterState) + .add_fan_out_node( + "outer_fan", + subgraph=mid_g, + items_field="products", + item_field="tag", + inputs={"seeds": "seeds"}, + collect_field="collected", + target_field="results", + ) + .add_edge("outer_fan", END) + .set_entry("outer_fan") + .with_checkpointer(cp) + .compile() + ) + correct = [("A-x", "A-y"), ("B-x", "B-y")] + final = await outer_g.invoke(_RNestOuterState(products=["A", "B"], seeds=["x", "y"])) + assert sorted(tuple(sorted(sub)) for sub in final.results) == correct + assert leaf_calls == 4 + # The projection emitted the INNER fan-out's progress (a record whose + # fan_out_node_name is the inner node), so the lineage-keyed nested entry + # round-trips through the record format without colliding the projection's + # key destructuring. + assert any( + fp.fan_out_node_name == "inner_fan" for rec in captured_records for fp in rec.fan_out_progress + ), "expected the inner fan-out's progress in a saved record" + + # Resume from the completed invocation: the outer instances are tracked + # ``completed`` so they roll forward (no inner re-run) to the same results, + # exercising _restore_fan_out_progress_state on a record set that includes + # the nested fan-out's (collapsed-lineage) entries. + resume_id = captured_ids[-1] + leaf_calls_before_resume = leaf_calls + resumed = await outer_g.invoke( + _RNestOuterState(products=["A", "B"], seeds=["x", "y"]), resume_invocation=resume_id + ) + assert sorted(tuple(sorted(sub)) for sub in resumed.results) == correct + assert leaf_calls == leaf_calls_before_resume, ( + "completed outer instances roll forward on resume; the inner leaf must not re-run" + ) + + # --------------------------------------------------------------------------- # Resume re-entry into subgraph: parent_states populated on inner-node saves # --------------------------------------------------------------------------- diff --git a/tests/unit/test_fan_out.py b/tests/unit/test_fan_out.py index 665c244..96cccba 100644 --- a/tests/unit/test_fan_out.py +++ b/tests/unit/test_fan_out.py @@ -273,6 +273,90 @@ async def compute(state: WorkerState) -> Mapping[str, Any]: assert final.results == [11, 12, 13] +# --------------------------------------------------------------------------- +# nested fan-out (a fan-out inside an outer fan-out instance) +# --------------------------------------------------------------------------- + + +class _NestedLeafState(State): + tag: str = "" + seed: str = "" + out: str = "" + + +class _NestedMidState(State): + tag: str = "" + seeds: list[str] = Field(default_factory=list[str]) + collected: Annotated[list[str], append] = Field(default_factory=list[str]) + + +class _NestedOuterState(State): + products: list[str] = Field(default_factory=list[str]) + seeds: list[str] = Field(default_factory=list[str]) + results: Annotated[list[Any], append] = Field(default_factory=list[Any]) + + +async def test_nested_fan_out_distinct_per_outer_instance_under_concurrency() -> None: + """A fan-out nested inside an outer fan-out instance runs its inner + subgraph once per (outer, inner) pair and returns the right per-outer + result, even with the outer instances running concurrently.""" + # Regression: the per-fan-out tracking entry was keyed by (namespace, node + # name) only, so the inner fan-out's entry collided across outer instances. + # With concurrent outer instances the second found the first's entry already + # marked complete and rolled its result forward, so every outer instance + # returned the first's inner result and the inner subgraph ran only once. + leaf_calls = 0 + + async def leaf(state: _NestedLeafState) -> Mapping[str, Any]: + nonlocal leaf_calls + await asyncio.sleep(0) # yield so the concurrent outer instances interleave + leaf_calls += 1 + return {"out": f"{state.tag}-{state.seed}"} + + leaf_builder: GraphBuilder[_NestedLeafState] = GraphBuilder(_NestedLeafState) + leaf_builder.set_entry("ask") + leaf_builder.add_node("ask", leaf) + leaf_builder.add_edge("ask", END) + leaf_graph = leaf_builder.compile() + + mid_builder: GraphBuilder[_NestedMidState] = GraphBuilder(_NestedMidState) + mid_builder.set_entry("inner_fan") + mid_builder.add_fan_out_node( + "inner_fan", + subgraph=leaf_graph, + items_field="seeds", + item_field="seed", + inputs={"tag": "tag"}, + collect_field="out", + target_field="collected", + ) + mid_builder.add_edge("inner_fan", END) + mid_graph = mid_builder.compile() + + outer_builder: GraphBuilder[_NestedOuterState] = GraphBuilder(_NestedOuterState) + outer_builder.set_entry("outer_fan") + outer_builder.add_fan_out_node( + "outer_fan", + subgraph=mid_graph, + items_field="products", + item_field="tag", + inputs={"seeds": "seeds"}, + collect_field="collected", + target_field="results", + ) + outer_builder.add_edge("outer_fan", END) + outer_graph = outer_builder.compile() + + final = await outer_graph.invoke(_NestedOuterState(products=["A", "B"], seeds=["x", "y"])) + await outer_graph.drain() + # Each outer instance collected its OWN inner results; the collapse bug gave + # [["A-x", "A-y"], ["A-x", "A-y"]] (the second outer reused the first's). + got = sorted(tuple(sorted(sub)) for sub in final.results) + assert got == [("A-x", "A-y"), ("B-x", "B-y")] + # The inner leaf ran once per (outer, inner) pair, not once total. + assert leaf_calls == 4 + + # --------------------------------------------------------------------------- # concurrency # --------------------------------------------------------------------------- From c9d53f6594f12c5d0616596cb5d00e4f79db80a5 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Thu, 25 Jun 2026 18:28:14 -0700 Subject: [PATCH 2/3] Key observer dispatches by fan-out lineage A fan-out instance dispatch or a parallel-branches per-branch dispatch nested inside an outer fan-out instance was keyed by its local namespace only, so the same dispatch in different outer instances shared one key: the second instance's inner nodes reparented under the first's dispatch and an inner augmentation reached the wrong outer instance. Key dispatches in both the OTel and Langfuse observers by their full enclosing fan-out instance and branch lineage, and resolve a nested node's parent and find its dispatch node by that lineage too. The per-branch key reads the branch name from the event so callable branches still synthesize. Un-defer conformance fixture 039 cases 1 and 2 (nested-lineage augmentation) in both observer harnesses. --- CHANGELOG.md | 1 + .../observability/langfuse/observer.py | 249 +++++++++++++----- .../observability/otel/observer.py | 242 ++++++++++------- .../test_observability_langfuse.py | 182 +++++++++++-- tests/unit/test_observability_otel.py | 212 +++++++++++++++ 5 files changed, 707 insertions(+), 179 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b27b52..285acae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). The - **Langfuse per-branch dispatch-span observation** (observability §4.3 / §8.4.2, proposals 0042 / 0044). The Langfuse observer now synthesizes a per-branch Span observation under a `parallel_branches` dispatcher node, so each branch's inner observations nest under their own branch span (a three-level dispatcher / per-branch-span / inner-nodes tree) instead of parenting directly under the dispatcher. The per-branch observation carries the OA-emitted `branch_name` alongside the caller baseline metadata and any per-branch augmentation, and the Generation observation now carries `branch_name` too. The OTel observer already produced this shape (proposal 0044 shipped OTel-only in v0.11.0); this brings the Langfuse mapping into line. Callable branches (proposal 0075) are unchanged. - **Augmentation no longer lands on a shared-parent fan-out / parallel-branches node** (observability §3.4, proposal 0045). A key set via `set_invocation_metadata` inside a fan-out instance or a parallel-branches branch was incorrectly applied to the shared fan-out / dispatcher NODE span (the fork point) when the augmenting context executed at that node's own namespace, in addition to the per-instance / per-branch dispatch span where it belongs. Both the OTel and Langfuse observers now skip the shared-parent node in that case, matching the behavior already applied to strict-ancestor shared parents. The per-instance / per-branch dispatch spans and the lineage ancestors that carry the augmentation are unaffected. - **Langfuse fan-out-instance dispatch nested below the top level** (observability §5.4, proposal 0013). The Langfuse observer's per-instance dispatch synthesis and parent resolution are now prefix-general, so a fan-out node sitting inside a serial subgraph wrapper (rather than at the top namespace level) gets its per-instance dispatch observation synthesized and its inner observations parented under it. This matches the OTel observer, which already resolved across every namespace prefix. +- **Dispatch spans nested inside an outer fan-out instance no longer collide across outer instances** (observability §5.4 / §3.4, proposals 0013 / 0044 / 0045). A fan-out instance dispatch or a parallel-branches per-branch dispatch sitting inside an outer fan-out instance (a fan-out within a fan-out, or parallel-branches within a fan-out) was keyed by its local namespace only, so the same dispatch in different outer instances shared one key: the second outer instance's inner nodes reparented under the first instance's dispatch, and an inner augmentation reached the wrong outer instance's dispatch. Both the OTel and Langfuse observers now key dispatches by their full enclosing fan-out instance / branch lineage, and resolve a nested node's parent by that lineage too, so each outer instance gets its own correctly-parented dispatch subtree with isolated augmentation. Top-level and serial-nested dispatch behavior is unchanged. - **Nested fan-out no longer collapses under concurrency** (engine). A fan-out nested inside an outer fan-out instance shared a single per-fan-out tracking entry across all outer instances, because that entry was keyed by namespace plus node name only. With concurrent outer instances the second instance found the first's entry already marked complete and rolled its result forward, so every outer instance returned the first instance's inner result (silently wrong output) and the inner subgraph ran only once. The tracking key now carries the enclosing fan-out instance lineage, so each outer instance gets its own inner fan-out progress and correct per-instance results. Top-level and subgraph- or branch-nested fan-outs are unaffected (their enclosing lineage is empty). Resume of a fan-out nested inside an outer fan-out instance does not yet round-trip per-outer-instance progress, so it re-runs rather than skipping on resume; tracked as a follow-up. ## [0.15.0] — 2026-06-22 diff --git a/src/openarmature/observability/langfuse/observer.py b/src/openarmature/observability/langfuse/observer.py index 1858b2b..66636d4 100644 --- a/src/openarmature/observability/langfuse/observer.py +++ b/src/openarmature/observability/langfuse/observer.py @@ -134,6 +134,42 @@ def _observation_chain_on_path( return True +def _dispatch_key( + prefix: tuple[str, ...], + fan_out_index_chain: tuple[int | None, ...], + branch_name_chain: tuple[str | None, ...], +) -> tuple[Any, ...]: + """Lineage-aware identity key for a fan-out instance / per-branch dispatch at + namespace ``prefix``. Encodes the fan-out/pb NODE namespace plus the full + chain of fan-out instance indices / branch names along the path to it (sliced + to ``len(prefix)``). Two dispatches at the same namespace but in different + enclosing fan-out instances / branches therefore get distinct keys -- the + enclosing chain entries differ -- which is what lets a fan-out / pb nested + inside an outer fan-out instance avoid colliding across outer instances. For + a top-level or serial-nested dispatch (no enclosing fan-out/branch) the + enclosing chain entries are all None, so the key is a stable function of the + namespace plus the dispatch's own axis.""" + n = len(prefix) + return (prefix, tuple(fan_out_index_chain[:n]), tuple(branch_name_chain[:n])) + + +def _branch_dispatch_key( + prefix: tuple[str, ...], + fan_out_index_chain: tuple[int | None, ...], + branch_name_chain: tuple[str | None, ...], + branch_name: str, +) -> tuple[Any, ...]: + """Lineage-aware identity key for a per-branch dispatch at namespace + ``prefix``. The branch IDENTITY comes from ``branch_name`` explicitly (not + ``branch_name_chain[-1]``): a callable branch carries its name on the event + but never extends ``branch_name_chain`` (no subgraph descent). The key still + carries the ENCLOSING fan-out instance / branch chain (positions above this + pb node) so a pb nested inside an outer fan-out instance doesn't collide + across outer instances.""" + n = len(prefix) + return (prefix, tuple(fan_out_index_chain[:n]), tuple(branch_name_chain[: n - 1]), branch_name) + + def _empty_str_frozenset() -> frozenset[str]: """Typed empty frozenset factory for ``detached_subgraphs`` / ``detached_fan_outs`` defaults.""" @@ -476,8 +512,12 @@ def _open_started_observation(self, event: NodeEvent) -> None: # started event so synthetic per-instance dispatch observations # can attach metadata.fan_out_parent_node_name (the inner # events from inside the fan-out don't carry fan_out_config - # themselves; this cache bridges). - if event.fan_out_config is not None and event.fan_out_index is None: + # themselves; this cache bridges). fan_out_config is set only on + # the NODE's own events, so it alone identifies them -- NOT + # ``fan_out_index is None``, which would miss a fan-out node nested + # inside an outer fan-out instance (its own event carries the OUTER + # instance index), leaving the inner dispatch unsynthesized. + if event.fan_out_config is not None: inv_state.fan_out_parent_node_name[event.namespace] = event.fan_out_config.parent_node_name # Per proposal 0045: mirror cache for parallel-branches NODE @@ -536,21 +576,31 @@ def _handle_completed(self, event: NodeEvent) -> None: # spawned. Done BEFORE the regular pop so the close ordering is # children-before-parents. if event.fan_out_index is None and event.namespace and event.namespace[0] in self.detached_fan_outs: - for prefix in list(inv_state.fan_out_instance_root_prefixes): - if len(prefix) > len(event.namespace) and prefix[: len(event.namespace)] == event.namespace: + ns = event.namespace + for key in list(inv_state.fan_out_instance_root_prefixes): + anchor_ns = cast("tuple[str, ...]", key[0]) + if len(anchor_ns) >= len(ns) and anchor_ns[: len(ns)] == ns: # Detached per-instance dispatches live in # fan_out_instance_observations (same map as # non-detached); close via the matching helper. - self._close_fan_out_instance_dispatch_observation(inv_state, prefix) - inv_state.fan_out_instance_root_prefixes.discard(prefix) - inv_state.detached_traces.pop(prefix, None) + self._close_fan_out_instance_dispatch_observation(inv_state, key) + inv_state.fan_out_instance_root_prefixes.discard(key) + # detached_traces uses the top-level routing key shape; derive + # it from the lineage key's own instance index (last entry). + fi_chain = cast("tuple[int | None, ...]", key[1]) + inv_state.detached_traces.pop(anchor_ns + (str(fi_chain[-1]),), None) # Per spec proposal 0013 (v0.10.0): when the fan-out node's # own completion fires, close all per-instance dispatch # observations synthesized for it. Children-before-parents. if event.fan_out_index is None and event.fan_out_config is not None: - for prefix in list(inv_state.fan_out_instance_observations.keys()): - if len(prefix) > len(event.namespace) and prefix[: len(event.namespace)] == event.namespace: - self._close_fan_out_instance_dispatch_observation(inv_state, prefix) + ns = event.namespace + for key in list(inv_state.fan_out_instance_observations.keys()): + anchor_ns = cast("tuple[str, ...]", key[0]) + # The dispatch key is now (anchor_ns, fi_chain, bn_chain); match + # on the NODE namespace (anchor_ns) being in this completing + # node's subtree. + if len(anchor_ns) >= len(ns) and anchor_ns[: len(ns)] == ns: + self._close_fan_out_instance_dispatch_observation(inv_state, key) inv_state.fan_out_parent_node_name.pop(event.namespace, None) # Clear the detached-child-trace-ids accumulator for this # fan-out node — cyclic execution that re-enters the same @@ -562,9 +612,11 @@ def _handle_completed(self, event: NodeEvent) -> None: # per-branch dispatch observations synthesized for it (children-before- # parents) and clear the pb caches. Same shape as the fan-out cleanup. if event.parallel_branches_config is not None: - for prefix in list(inv_state.parallel_branches_branch_spans.keys()): - if len(prefix) > len(event.namespace) and prefix[: len(event.namespace)] == event.namespace: - self._close_parallel_branches_branch_dispatch_observation(inv_state, prefix) + ns = event.namespace + for key in list(inv_state.parallel_branches_branch_spans.keys()): + anchor_ns = cast("tuple[str, ...]", key[0]) + if len(anchor_ns) >= len(ns) and anchor_ns[: len(ns)] == ns: + self._close_parallel_branches_branch_dispatch_observation(inv_state, key) inv_state.parallel_branches_parent_node_name.pop(event.namespace, None) inv_state.parallel_branches_branch_names.pop(event.namespace, None) @@ -647,21 +699,17 @@ def _handle_metadata_augmentation(self, event: MetadataAugmentationEvent) -> Non if _observation_chain_on_path(observation, aug_fi_chain, aug_bn_chain): observation.handle.update(metadata=metadata_delta) - # Fan-out instance dispatch observations. + # Fan-out instance dispatch observations: on the augmenter's path iff the + # dispatch NODE namespace (key[0]) is an ancestor-or-equal of the + # augmenter AND its full lineage chain (carried on the observation) is a + # prefix of the augmenter's -- so a SIBLING outer instance's dispatch, + # whose chain differs at the enclosing position, is excluded. for key, observation in inv_state.fan_out_instance_observations.items(): - if not key: - continue - anchor_ns = key[:-1] - fi_str = key[-1] + anchor_ns = cast("tuple[str, ...]", key[0]) if not (is_strict_prefix(anchor_ns, aug_ns) or anchor_ns == aug_ns): continue - chain_pos = len(anchor_ns) - 1 - if chain_pos < 0 or chain_pos >= len(aug_fi_chain): - continue - aug_fi_at_pos = aug_fi_chain[chain_pos] - if aug_fi_at_pos is None or str(aug_fi_at_pos) != fi_str: - continue - observation.handle.update(metadata=metadata_delta) + if _observation_chain_on_path(observation, aug_fi_chain, aug_bn_chain): + observation.handle.update(metadata=metadata_delta) # Open NODE observations. Same as augmenter or strict # ancestor on the path; skip shared-parent NODE observations @@ -916,22 +964,34 @@ def _resolve_parent_observation_id(self, inv_state: _InvState, event: NodeEvent) # 3. Leaf node observation at any matching ancestor prefix, # walked longest-first. # 4. None — the Trace itself becomes the implicit parent. - # Per proposals 0044 / 0013: an inner node parents under the INNERMOST - # dispatch on its lineage -- a per-branch dispatch - # (parallel_branches_branch_spans, keyed prefix + (branch_name,)) or a - # per-instance fan-out dispatch (fan_out_instance_observations, keyed - # prefix + (str(fan_out_index),)). Walk prefixes longest-first so the - # innermost wins; this resolves arbitrary nesting (fan-out in fan-out, - # parallel-branches in fan-out, ...). Mirrors the OTel + # Per proposals 0044 / 0013 / 0045: an inner node parents under the + # INNERMOST dispatch on its lineage -- a per-branch dispatch + # (parallel_branches_branch_spans) or a per-instance fan-out dispatch + # (fan_out_instance_observations), both keyed by the lineage-aware + # _dispatch_key. Walk prefixes longest-first so the innermost wins; the + # lineage key carries the enclosing fan-out instance / branch chain, so + # this resolves arbitrary nesting (fan-out in fan-out, parallel-branches + # in fan-out, ...) to the RIGHT outer instance. Mirrors OTel # _resolve_parent_context. for prefix_len in range(len(event.namespace), 0, -1): prefix = event.namespace[:prefix_len] + fi_axis = ( + event.fan_out_index_chain[prefix_len - 1] + if prefix_len - 1 < len(event.fan_out_index_chain) + else None + ) if event.branch_name is not None: - branch_dispatch = inv_state.parallel_branches_branch_spans.get(prefix + (event.branch_name,)) + branch_dispatch = inv_state.parallel_branches_branch_spans.get( + _branch_dispatch_key( + prefix, event.fan_out_index_chain, event.branch_name_chain, event.branch_name + ) + ) if branch_dispatch is not None: return branch_dispatch.handle.id - if event.fan_out_index is not None: - dispatch = inv_state.fan_out_instance_observations.get(prefix + (str(event.fan_out_index),)) + if fi_axis is not None: + dispatch = inv_state.fan_out_instance_observations.get( + _dispatch_key(prefix, event.fan_out_index_chain, event.branch_name_chain) + ) if dispatch is not None: return dispatch.handle.id for prefix_len in range(len(event.namespace) - 1, 0, -1): @@ -1014,20 +1074,28 @@ def _sync_subgraph_observations( prefix = namespace[:depth] if prefix in inv_state.subgraph_observations: continue - # The per-instance dispatch for this event's own instance is opened - # below; skip the regular subgraph path so we don't double-open. Runs - # at ANY depth -- the fan-out may sit inside another fan-out instance - # or a subgraph wrapper -- mirroring the OTel observer. + # The fan-out instance axis at THIS depth -- the chain entry for the + # dispatch boundary into prefix -- NOT the innermost event.fan_out_index + # (which differs for an OUTER fan-out in a nested stack). Branches use + # event.branch_name directly (callable branches don't extend the chain). + fi_axis = ( + event.fan_out_index_chain[depth - 1] if depth - 1 < len(event.fan_out_index_chain) else None + ) + # The per-instance dispatch for this prefix's instance is opened + # below; skip the regular subgraph path so we don't double-open. Keyed + # by the full lineage so nested instances don't collide across outer + # ones. Runs at ANY depth (nested fan-out / subgraph wrapper). if ( - event.fan_out_index is not None - and (prefix + (str(event.fan_out_index),)) in inv_state.fan_out_instance_observations + fi_axis is not None + and _dispatch_key(prefix, event.fan_out_index_chain, event.branch_name_chain) + in inv_state.fan_out_instance_observations ): continue # Detached subgraph: kept top-level (depth == 1). _trace_id_for routes # detached events by namespace[:1], so a nested detached unit would # partially detach (its dispatch in the new Trace, inner nodes in the # main one). Nested-detached support rides with the deferred - # nested-dispatch-keying fix that generalizes _trace_id_for too. + # generalization of _trace_id_for (out of scope here). if depth == 1 and prefix[0] in self.detached_subgraphs: self._open_detached_subgraph_trace(inv_state, correlation_id, prefix, event) continue @@ -1040,11 +1108,11 @@ def _sync_subgraph_observations( continue # Non-detached fan-out: synthesize the per-instance dispatch # observation under the fan-out node observation (proposal 0013 - # v0.10.0). The fan_out_parent_node_name cache match self-gates to - # the fan-out node's namespace, so this runs at any depth (nested - # fan-out, or a fan-out inside a subgraph wrapper / branch). + # v0.10.0). Gated on the fan-out NODE at prefix and a fan-out axis at + # this depth, so it runs at any depth (nested fan-out, or a fan-out + # inside a subgraph wrapper / branch). if ( - event.fan_out_index is not None + fi_axis is not None and prefix[-1] not in self.detached_fan_outs and prefix in inv_state.fan_out_parent_node_name ): @@ -1053,8 +1121,8 @@ def _sync_subgraph_observations( # Per proposal 0044: synthesize a per-branch dispatch observation # under the pb NODE for an inner branch event, so inner branch # nodes parent under it rather than the shared pb NODE span. Mirror - # of the fan-out per-instance arm above; gated on the branch - # belonging to the pb node declared at this prefix. + # of the fan-out per-instance arm above; gated on the branch at this + # depth belonging to the pb node declared at this prefix. if ( event.branch_name is not None and prefix in inv_state.parallel_branches_parent_node_name @@ -1063,7 +1131,14 @@ def _sync_subgraph_observations( # Synthesize once per branch: _sync runs on every inner node's # started event, so guard against re-opening (a second open would # orphan the first observation and split the branch's nodes). - if prefix + (event.branch_name,) not in inv_state.parallel_branches_branch_spans: + # event.branch_name (not bn_axis) so callable branches -- which + # don't extend branch_name_chain -- still synthesize. + if ( + _branch_dispatch_key( + prefix, event.fan_out_index_chain, event.branch_name_chain, event.branch_name + ) + not in inv_state.parallel_branches_branch_spans + ): self._open_parallel_branches_branch_dispatch_observation( inv_state, correlation_id, prefix, event ) @@ -1139,18 +1214,27 @@ def _open_fan_out_instance_dispatch_observation( ) -> None: # Non-detached per-instance dispatch lives in the parent # Trace under the fan-out node's own Span observation. - fan_out_open = self._find_node_observation(inv_state, prefix) + fan_out_open = self._find_node_observation(inv_state, prefix, event) parent_observation_id = fan_out_open.handle.id if fan_out_open is not None else None parent_node_name = inv_state.fan_out_parent_node_name.get(prefix, prefix[-1]) # Per-instance dispatch is synthesized from the first inner # event inside the instance subtree; inherit scalar metadata # from that event (same pattern as ``_open_subgraph_observation``). + # The dispatch's OWN instance index is the chain entry at this depth, not + # event.fan_out_index (the innermost index of the synthesizing inner + # event) -- they differ when this is an OUTER fan-out in a nested stack. + chain_len = len(prefix) + fan_out_index = ( + event.fan_out_index_chain[chain_len - 1] + if chain_len - 1 < len(event.fan_out_index_chain) + else event.fan_out_index + ) metadata: dict[str, Any] = { "namespace": list(prefix), "step": event.step, "attempt_index": 0, "fan_out_parent_node_name": parent_node_name, - "fan_out_index": event.fan_out_index, + "fan_out_index": fan_out_index, "subgraph_name": _subgraph_identity_at(event, len(prefix)), } if correlation_id is not None: @@ -1162,9 +1246,9 @@ def _open_fan_out_instance_dispatch_observation( metadata=metadata, parent_observation_id=parent_observation_id, ) - instance_key = prefix + (str(event.fan_out_index),) - # Per proposal 0045: chain sliced to instance-dispatch depth. - chain_len = len(prefix) + # Lineage-aware key (proposal 0045): the namespace plus the full instance + # / branch chain, so nested instances don't collide across outer ones. + instance_key = _dispatch_key(prefix, event.fan_out_index_chain, event.branch_name_chain) inv_state.fan_out_instance_observations[instance_key] = _OpenObservation( handle=handle, fan_out_index_chain=event.fan_out_index_chain[:chain_len], @@ -1180,18 +1264,22 @@ def _open_parallel_branches_branch_dispatch_observation( ) -> None: # Per-branch dispatch lives under the parallel-branches NODE's own Span # observation (mirror of the fan-out per-instance dispatch). - pb_open = self._find_node_observation(inv_state, prefix) + pb_open = self._find_node_observation(inv_state, prefix, event) parent_observation_id = pb_open.handle.id if pb_open is not None else None parent_node_name = inv_state.parallel_branches_parent_node_name.get(prefix, prefix[-1]) # Synthesized from the first inner event in the branch subtree; inherit - # scalar metadata from it. branch_name is the OA-emitted §8.4.2 row; the - # caller's branchName augmentation rides in via the caller metadata. + # scalar metadata from it. branch_name is the OA-emitted §8.4.2 row (from + # event.branch_name, which a callable branch sets without extending + # branch_name_chain); the caller's branchName augmentation rides in via + # the caller metadata. + chain_len = len(prefix) + branch_name = cast("str", event.branch_name) metadata: dict[str, Any] = { "namespace": list(prefix), "step": event.step, "attempt_index": 0, "parallel_branches_parent_node_name": parent_node_name, - "branch_name": event.branch_name, + "branch_name": branch_name, "subgraph_name": _subgraph_identity_at(event, len(prefix)), } if correlation_id is not None: @@ -1199,12 +1287,16 @@ def _open_parallel_branches_branch_dispatch_observation( _apply_caller_metadata(metadata, event.caller_invocation_metadata) handle = self.client.span( trace_id=inv_state.trace_id, - name=event.branch_name, + name=branch_name, metadata=metadata, parent_observation_id=parent_observation_id, ) - branch_key = prefix + (cast("str", event.branch_name),) - chain_len = len(prefix) + # Lineage-aware key (proposal 0045): the enclosing fan-out instance / + # branch chain plus the explicit branch name, so a branch nested inside + # an outer fan-out instance doesn't collide across outer instances. + branch_key = _branch_dispatch_key( + prefix, event.fan_out_index_chain, event.branch_name_chain, branch_name + ) inv_state.parallel_branches_branch_spans[branch_key] = _OpenObservation( handle=handle, fan_out_index_chain=event.fan_out_index_chain[:chain_len], @@ -1364,7 +1456,7 @@ def _open_detached_fan_out_instance_trace( # instance's entry. ids_list = inv_state.detached_child_trace_ids.setdefault(prefix, []) ids_list.append(detached_trace_id) - fan_out_open = self._find_node_observation(inv_state, prefix) + fan_out_open = self._find_node_observation(inv_state, prefix, event) if fan_out_open is not None: # `detached: True` per §8.4.2 (proposal 0042) — the # parent-side fan-out node observation marks itself when @@ -1411,14 +1503,19 @@ def _open_detached_fan_out_instance_trace( metadata=dispatch_metadata, parent_observation_id=None, ) - instance_key = prefix + (str(event.fan_out_index),) + # Shared fan_out_instance_observations / root_prefixes use the + # lineage-aware key (consistent with resolution / close). detached_traces + # keeps the top-level routing key shape that _trace_id_for reconstructs + # (detached fan-outs are top-level; generalizing _trace_id_for is out of + # scope for this fix). + instance_key = _dispatch_key(prefix, event.fan_out_index_chain, event.branch_name_chain) chain_len = len(prefix) inv_state.fan_out_instance_observations[instance_key] = _OpenObservation( handle=handle, fan_out_index_chain=event.fan_out_index_chain[:chain_len], branch_name_chain=event.branch_name_chain[:chain_len], ) - inv_state.detached_traces[instance_key] = detached_trace_id + inv_state.detached_traces[prefix + (str(event.fan_out_index),)] = detached_trace_id inv_state.fan_out_instance_root_prefixes.add(instance_key) def _close_subgraph_observation(self, inv_state: _InvState, prefix: tuple[str, ...]) -> None: @@ -1444,18 +1541,26 @@ def _close_parallel_branches_branch_dispatch_observation( observation.handle.end() def _find_node_observation( - self, inv_state: _InvState, prefix: tuple[str, ...] + self, inv_state: _InvState, prefix: tuple[str, ...], event: NodeEvent ) -> _OpenObservation | None: # Find a NODE's own open leaf observation at ``prefix`` (the fan-out or # parallel-branches NODE, whose per-instance / per-branch dispatches - # parent under it). Scan by namespace only: the NODE may itself carry an - # outer fan_out_index / branch_name when it is nested inside another - # fan-out instance or branch, so filtering on fan_out_index is None would - # miss it. Only one entry per namespace is open at a time (retry opens - # and closes attempts serially), so the scan is unambiguous. Mirrors the - # OTel _find_fan_out_node_span. + # parent under it). Match the ENCLOSING lineage, not just the namespace: + # when the NODE is itself nested inside an outer fan-out instance / + # branch, several instances of the same NODE namespace are open at once + # under concurrency, so a namespace-only scan would bind the wrong one. + # The NODE's own event carries the instance / branch it sits in as its + # fan_out_index / branch_name (key[2] / key[3]); that equals the + # augmenting/leaf event's chain entry at the level above this NODE. + n = len(prefix) + enclosing_fi = ( + event.fan_out_index_chain[n - 2] if n >= 2 and n - 2 < len(event.fan_out_index_chain) else None + ) + enclosing_bn = ( + event.branch_name_chain[n - 2] if n >= 2 and n - 2 < len(event.branch_name_chain) else None + ) for key, observation in inv_state.open_observations.items(): - if key[0] == prefix: + if key[0] == prefix and key[2] == enclosing_fi and key[3] == enclosing_bn: return observation return None diff --git a/src/openarmature/observability/otel/observer.py b/src/openarmature/observability/otel/observer.py index 44f9fdd..a2cfead 100644 --- a/src/openarmature/observability/otel/observer.py +++ b/src/openarmature/observability/otel/observer.py @@ -299,6 +299,41 @@ def _span_chain_on_path( return True +def _dispatch_key( + prefix: tuple[str, ...], + fan_out_index_chain: tuple[int | None, ...], + branch_name_chain: tuple[str | None, ...], +) -> tuple[Any, ...]: + """Lineage-aware identity key for a fan-out instance / per-branch dispatch + span at namespace ``prefix``. Encodes the fan-out/pb NODE namespace plus the + full chain of fan-out instance indices / branch names along the path to it + (sliced to ``len(prefix)``). Two dispatches at the same namespace but in + different enclosing fan-out instances / branches get distinct keys, so a + fan-out / pb nested inside an outer fan-out instance no longer collides + across outer instances. Mirrors the LangfuseObserver helper of the same + name.""" + n = len(prefix) + return (prefix, tuple(fan_out_index_chain[:n]), tuple(branch_name_chain[:n])) + + +def _branch_dispatch_key( + prefix: tuple[str, ...], + fan_out_index_chain: tuple[int | None, ...], + branch_name_chain: tuple[str | None, ...], + branch_name: str, +) -> tuple[Any, ...]: + """Lineage-aware identity key for a per-branch dispatch span at namespace + ``prefix``. The branch IDENTITY comes from ``branch_name`` explicitly (not + ``branch_name_chain[-1]``): a callable branch carries its name on the event + but never extends ``branch_name_chain`` (no subgraph descent). The key still + carries the ENCLOSING fan-out instance / branch chain (positions above this + pb node) so a pb nested inside an outer fan-out instance doesn't collide + across outer instances. Mirrors the LangfuseObserver helper of the same + name.""" + n = len(prefix) + return (prefix, tuple(fan_out_index_chain[:n]), tuple(branch_name_chain[: n - 1]), branch_name) + + # Sorted object keys, no insignificant whitespace, UTF-8 output (per # observability §5.5.1 / §5.5.6). Within-impl determinism for identical # inputs is required; cross-impl bytewise stability is NOT required by @@ -951,8 +986,13 @@ def _handle_completed(self, event: NodeEvent) -> None: # or a pb's branch, the NODE's own completion event carries # the OUTER axis values. if event.fan_out_config is not None: + ns = event.namespace for key in list(inv_state.fan_out_instance_spans.keys()): - if len(key) > len(event.namespace) and key[: len(event.namespace)] == event.namespace: + # The dispatch key is now (anchor_ns, fi_chain, bn_chain); match + # on the NODE namespace (anchor_ns = key[0]) being in this + # completing node's subtree. + anchor_ns = cast("tuple[str, ...]", key[0]) + if len(anchor_ns) >= len(ns) and anchor_ns[: len(ns)] == ns: self._close_fan_out_instance_dispatch_span(inv_state, key) inv_state.fan_out_parent_node_name.pop(event.namespace, None) @@ -966,8 +1006,10 @@ def _handle_completed(self, event: NodeEvent) -> None: # ``branch_name is None`` (which would skip close for a pb # nested inside another pb's branch). if event.parallel_branches_config is not None: + ns = event.namespace for key in list(inv_state.parallel_branches_branch_spans.keys()): - if len(key) > len(event.namespace) and key[: len(event.namespace)] == event.namespace: + anchor_ns = cast("tuple[str, ...]", key[0]) + if len(anchor_ns) >= len(ns) and anchor_ns[: len(ns)] == ns: self._close_parallel_branches_branch_dispatch_span(inv_state, key) inv_state.parallel_branches_parent_node_name.pop(event.namespace, None) inv_state.parallel_branches_branch_names.pop(event.namespace, None) @@ -1147,40 +1189,25 @@ def _collect_augmentation_targets( if _span_chain_on_path(open_span, aug_fi_chain, aug_bn_chain): targets.append(open_span.span) - # Fan-out instance dispatch spans: keyed by anchor_ns + - # (str(fi),). The dispatch represents the descent at namespace - # position len(anchor_ns)-1 — i.e., chain position - # ``len(anchor_ns)-1`` should match the dispatch's fi. + # Fan-out instance dispatch spans: on the augmenter's path iff the + # dispatch NODE namespace (key[0]) is an ancestor-or-equal of the + # augmenter AND its full lineage chain (carried on the span) is a prefix + # of the augmenter's -- so a SIBLING outer instance's dispatch, whose + # chain differs at the enclosing position, is excluded. for key, open_span in inv_state.fan_out_instance_spans.items(): - if not key: - continue - anchor_ns = key[:-1] - fi_str = key[-1] + anchor_ns = cast("tuple[str, ...]", key[0]) if not (is_strict_prefix(anchor_ns, aug_ns) or anchor_ns == aug_ns): continue - chain_pos = len(anchor_ns) - 1 - if chain_pos < 0 or chain_pos >= len(aug_fi_chain): - continue - aug_fi_at_pos = aug_fi_chain[chain_pos] - if aug_fi_at_pos is None or str(aug_fi_at_pos) != fi_str: - continue - targets.append(open_span.span) + if _span_chain_on_path(open_span, aug_fi_chain, aug_bn_chain): + targets.append(open_span.span) - # Per-branch dispatch spans: keyed by anchor_ns + (branch_name,). - # Mirror logic to fan-out above, against the branch chain. + # Per-branch dispatch spans: same lineage match as fan-out above. for key, open_span in inv_state.parallel_branches_branch_spans.items(): - if not key: - continue - anchor_ns = key[:-1] - bn_str = key[-1] + anchor_ns = cast("tuple[str, ...]", key[0]) if not (is_strict_prefix(anchor_ns, aug_ns) or anchor_ns == aug_ns): continue - chain_pos = len(anchor_ns) - 1 - if chain_pos < 0 or chain_pos >= len(aug_bn_chain): - continue - if aug_bn_chain[chain_pos] != bn_str: - continue - targets.append(open_span.span) + if _span_chain_on_path(open_span, aug_fi_chain, aug_bn_chain): + targets.append(open_span.span) # Open NODE spans: same context (aug's own attempt span), or # strict ancestor on the augmenter's path. Skip: @@ -1803,16 +1830,32 @@ def _resolve_parent_context( # deepest matching depth is the most-immediate parent. for prefix_len in range(len(event.namespace), 0, -1): prefix = event.namespace[:prefix_len] + # Lineage-aware keys (proposal 0045): carry the enclosing fan-out + # instance / branch chain so this resolves to the RIGHT outer + # instance for arbitrary nesting. + fi_axis = ( + event.fan_out_index_chain[prefix_len - 1] + if prefix_len - 1 < len(event.fan_out_index_chain) + else None + ) if event.branch_name is not None: - branch_dispatch = inv_state.parallel_branches_branch_spans.get(prefix + (event.branch_name,)) + branch_dispatch = inv_state.parallel_branches_branch_spans.get( + _branch_dispatch_key( + prefix, event.fan_out_index_chain, event.branch_name_chain, event.branch_name + ) + ) if branch_dispatch is not None: return set_span_in_context(branch_dispatch.span) if event.fan_out_index is not None: - instance_key = prefix + (str(event.fan_out_index),) - root = inv_state.detached_roots.get(instance_key) + # Detached roots keep the top-level routing key shape (detached + # generalization is out of scope). + root = inv_state.detached_roots.get(prefix + (str(event.fan_out_index),)) if root is not None: return set_span_in_context(root.span) - instance_dispatch = inv_state.fan_out_instance_spans.get(instance_key) + if fi_axis is not None: + instance_dispatch = inv_state.fan_out_instance_spans.get( + _dispatch_key(prefix, event.fan_out_index_chain, event.branch_name_chain) + ) if instance_dispatch is not None: return set_span_in_context(instance_dispatch.span) # 1b. Detached subgraph root at any matching prefix wins @@ -1886,20 +1929,24 @@ def _sync_subgraph_spans( continue if prefix in inv_state.detached_roots: continue - # Per spec proposal 0013 (v0.10.0): for non-detached - # fan-out instances, the per-instance dispatch span lives - # at ``prefix + (str(fan_out_index),)`` in - # ``fan_out_instance_spans``. If the per-instance dispatch - # span for THIS event's instance is already open, the - # fan-out node span at ``prefix`` is already open too - # (synthesized as part of the fan-out node's started - # event), so we don't open a new subgraph span at - # ``prefix``. This dedup runs at any depth — the fan-out - # node may sit inside a subgraph wrapper, another fan-out, - # or a parallel-branches branch. + # The fan-out instance axis at THIS depth -- the chain entry for the + # dispatch boundary into prefix -- NOT the innermost event.fan_out_index + # (which differs for an OUTER fan-out in a nested stack). Branches use + # event.branch_name directly (callable branches don't extend the chain). + fi_axis = ( + event.fan_out_index_chain[depth - 1] if depth - 1 < len(event.fan_out_index_chain) else None + ) + # Per spec proposal 0013 (v0.10.0): for non-detached fan-out + # instances, the per-instance dispatch span lives in + # ``fan_out_instance_spans`` keyed by the lineage-aware _dispatch_key. + # If the dispatch span for THIS event's instance is already open, the + # fan-out node span at ``prefix`` is already open too, so we don't + # open a new subgraph span. Runs at any depth (nested fan-out / + # subgraph wrapper / branch). if ( - event.fan_out_index is not None - and (prefix + (str(event.fan_out_index),)) in inv_state.fan_out_instance_spans + fi_axis is not None + and _dispatch_key(prefix, event.fan_out_index_chain, event.branch_name_chain) + in inv_state.fan_out_instance_spans ): continue # If this prefix's first segment is configured as a @@ -1933,7 +1980,7 @@ def _sync_subgraph_spans( # detached check uses ``prefix[-1]`` (the fan-out node # name) so it remains correct at depth > 1. if ( - event.fan_out_index is not None + fi_axis is not None and prefix[-1] not in self.detached_fan_outs and prefix in inv_state.fan_out_parent_node_name ): @@ -1955,7 +2002,11 @@ def _sync_subgraph_spans( and prefix in inv_state.parallel_branches_parent_node_name and event.branch_name in inv_state.parallel_branches_branch_names.get(prefix, frozenset()) ): - branch_key = prefix + (event.branch_name,) + # event.branch_name (not bn_axis) so callable branches -- which + # don't extend branch_name_chain -- still synthesize. + branch_key = _branch_dispatch_key( + prefix, event.fan_out_index_chain, event.branch_name_chain, event.branch_name + ) if branch_key not in inv_state.parallel_branches_branch_spans: self._open_parallel_branches_branch_dispatch_span( inv_state, correlation_id, prefix, event @@ -2203,7 +2254,7 @@ def _open_detached_fan_out_instance_root( # ``fan_out_index is None`` rather than hardcoding the key. # Only one such entry is open at a time (retry opens and # closes within a single attempt's lifecycle). - fan_out_open = self._find_fan_out_node_span(inv_state, prefix) + fan_out_open = self._find_fan_out_node_span(inv_state, prefix, event) if fan_out_open is not None: fan_out_open.span.add_link(detached_sc) @@ -2272,7 +2323,7 @@ def _open_fan_out_instance_dispatch_span( """ # Find the fan-out node's open span (latest attempt under # retry) to use as parent. - fan_out_open = self._find_fan_out_node_span(inv_state, prefix) + fan_out_open = self._find_fan_out_node_span(inv_state, prefix, event) parent_ctx: object if fan_out_open is not None: parent_ctx = set_span_in_context(fan_out_open.span) @@ -2280,10 +2331,19 @@ def _open_fan_out_instance_dispatch_span( parent_ctx = otel_context.Context() parent_node_name = inv_state.fan_out_parent_node_name.get(prefix, prefix[-1]) + # The dispatch's OWN instance index is the chain entry at this depth, not + # event.fan_out_index (the innermost index of the synthesizing inner + # event) -- they differ when this is an OUTER fan-out in a nested stack. + chain_len = len(prefix) + fan_out_index = ( + event.fan_out_index_chain[chain_len - 1] + if chain_len - 1 < len(event.fan_out_index_chain) + else event.fan_out_index + ) attrs: dict[str, Any] = { "openarmature.node.name": prefix[-1], "openarmature.fan_out.parent_node_name": parent_node_name, - "openarmature.node.fan_out_index": event.fan_out_index, + "openarmature.node.fan_out_index": fan_out_index, "openarmature.subgraph.name": _subgraph_identity_at(event, len(prefix)), } if correlation_id is not None: @@ -2295,12 +2355,11 @@ def _open_fan_out_instance_dispatch_span( kind=SpanKind.INTERNAL, attributes=attrs, ) - instance_key = prefix + (str(event.fan_out_index),) - # Per proposal 0045: this dispatch span sits AT the descent - # boundary into the fan-out instance. Its chain is the slice - # of the inner event's chain up to and including the - # boundary's own position. - chain_len = len(prefix) + # Lineage-aware key (proposal 0045): the namespace plus the full instance + # / branch chain, so nested instances don't collide across outer ones. + # Its chain is the slice of the inner event's chain up to and including + # this boundary's own position. + instance_key = _dispatch_key(prefix, event.fan_out_index_chain, event.branch_name_chain) inv_state.fan_out_instance_spans[instance_key] = _OpenSpan( span=instance_span, fan_out_index_chain=event.fan_out_index_chain[:chain_len], @@ -2340,19 +2399,11 @@ def _open_parallel_branches_branch_dispatch_span( assert event.branch_name is not None, ( "parallel-branches branch dispatch synthesis requires event.branch_name" ) - # Find the parallel-branches NODE's open span (latest attempt - # under retry) to use as parent. Scan ``open_spans`` by - # namespace only — the NODE may carry an outer - # ``fan_out_index`` (if the parallel-branches node sits inside - # a fan-out instance) or ``branch_name`` (if it sits inside - # another parallel-branches branch). Only one entry per - # namespace is open at a time, so the scan is unambiguous. - node_open: _OpenSpan | None = None - for key, open_span in inv_state.open_spans.items(): - ns, _attempt, _fan_idx, _bn = key - if ns == prefix: - node_open = open_span - break + # Find the parallel-branches NODE's open span on THIS branch's lineage + # (the same enclosing-instance/branch match as the fan-out parent + # finder), so a pb nested inside an outer fan-out instance binds the + # right outer instance's NODE under concurrency. + node_open = self._find_fan_out_node_span(inv_state, prefix, event) parent_ctx: object if node_open is not None: parent_ctx = set_span_in_context(node_open.span) @@ -2360,9 +2411,13 @@ def _open_parallel_branches_branch_dispatch_span( parent_ctx = otel_context.Context() parent_node_name = inv_state.parallel_branches_parent_node_name.get(prefix, prefix[-1]) + # branch_name from event.branch_name (a callable branch sets it without + # extending branch_name_chain). The assert above narrows it to str. + chain_len = len(prefix) + branch_name = event.branch_name attrs: dict[str, Any] = { - "openarmature.node.name": event.branch_name, - "openarmature.node.branch_name": event.branch_name, + "openarmature.node.name": branch_name, + "openarmature.node.branch_name": branch_name, "openarmature.parallel_branches.parent_node_name": parent_node_name, "openarmature.subgraph.name": _subgraph_identity_at(event, len(prefix)), } @@ -2370,16 +2425,17 @@ def _open_parallel_branches_branch_dispatch_span( attrs["openarmature.correlation_id"] = correlation_id _apply_caller_metadata(attrs, event.caller_invocation_metadata) branch_span = self._tracer.start_span( - name=event.branch_name, + name=branch_name, context=cast("Any", parent_ctx), kind=SpanKind.INTERNAL, attributes=attrs, ) - branch_key = prefix + (event.branch_name,) - # Per proposal 0045: this dispatch span sits at the descent - # boundary into this branch. Its chain is the slice up to - # and including the boundary's own position. - chain_len = len(prefix) + # Lineage-aware key (proposal 0045): the enclosing fan-out instance / + # branch chain plus the explicit branch name, so a branch nested inside + # an outer fan-out instance doesn't collide across outer instances. + branch_key = _branch_dispatch_key( + prefix, event.fan_out_index_chain, event.branch_name_chain, branch_name + ) inv_state.parallel_branches_branch_spans[branch_key] = _OpenSpan( span=branch_span, fan_out_index_chain=event.fan_out_index_chain[:chain_len], @@ -2431,17 +2487,27 @@ def _drain_open_span(self, open_span: _OpenSpan) -> None: self._run_enrichers(open_span.span, None) open_span.span.end() - def _find_fan_out_node_span(self, inv_state: _InvState, prefix: tuple[str, ...]) -> _OpenSpan | None: - """Find the currently-open fan-out NODE span at ``prefix``. - Scans by namespace only — the NODE may carry an outer - ``fan_out_index`` or ``branch_name`` if the fan-out is itself - nested inside another fan-out instance or a parallel-branches - branch. Only one entry per namespace is open at a time - (retry middleware opens and closes attempts serially), so the - scan is unambiguous.""" + def _find_fan_out_node_span( + self, inv_state: _InvState, prefix: tuple[str, ...], event: NodeEvent + ) -> _OpenSpan | None: + """Find the currently-open fan-out / pb NODE span at ``prefix`` on the + given event's ENCLOSING lineage. When the NODE is itself nested inside an + outer fan-out instance / branch, several instances of the same NODE + namespace are open at once under concurrency, so a namespace-only scan + would bind the wrong one. The NODE's own event carries the instance / + branch it sits in as its fan_out_index / branch_name (key[2] / key[3]); + that equals the synthesizing event's chain entry at the level above this + NODE.""" + n = len(prefix) + enclosing_fi = ( + event.fan_out_index_chain[n - 2] if n >= 2 and n - 2 < len(event.fan_out_index_chain) else None + ) + enclosing_bn = ( + event.branch_name_chain[n - 2] if n >= 2 and n - 2 < len(event.branch_name_chain) else None + ) for key, open_span in inv_state.open_spans.items(): - ns, _attempt, _fan_idx, _bn = key - if ns == prefix: + ns, _attempt, fan_idx, bn = key + if ns == prefix and fan_idx == enclosing_fi and bn == enclosing_bn: return open_span return None diff --git a/tests/conformance/test_observability_langfuse.py b/tests/conformance/test_observability_langfuse.py index a0addf0..0f17631 100644 --- a/tests/conformance/test_observability_langfuse.py +++ b/tests/conformance/test_observability_langfuse.py @@ -27,7 +27,7 @@ import pytest import yaml -from openarmature.graph import END, ExplicitMapping, GraphBuilder +from openarmature.graph import END, BranchSpec, ExplicitMapping, GraphBuilder from openarmature.llm import OpenAIProvider from openarmature.llm.response import RuntimeConfig from openarmature.observability.langfuse import ( @@ -125,24 +125,9 @@ # ``(fixture_stem, case_name)``. The case-loop in the runner ``continue``s # past matching cases — NOT ``pytest.skip``, which would skip the whole # fixture's test invocation and hide the surrounding cases that DO run. -_DEFERRED_CASES: frozenset[tuple[str, str]] = frozenset( - { - # 039 cases 1 + 2 are TEMPORARILY deferred pending one deeper observer - # fix shared by both: dispatch keys - # (fan_out_instance_observations / parallel_branches_branch_spans) are - # namespace-local and do NOT encode the enclosing fan-out instance, so a - # dispatch INSIDE an outer fan-out instance collides across instances -- - # case 1's inner instance dispatch and case 2's per-branch dispatch both - # reparent the second outer instance's events under the first's dispatch. - # The fix (thread the enclosing fan_out_index_chain / branch_name_chain - # into the dispatch keys, across synthesis + resolution + the - # augmentation walk, in both observers) is its own focused effort + spec - # coordination. Case 3 (single fan-out level under a serial wrapper) does - # not nest dispatches, so it is wired. See _build_039_graph. - ("039-nested-lineage-augmentation", "inner_fan_out_augmenter_propagates_to_outer_dispatch_span"), - ("039-nested-lineage-augmentation", "parallel_branch_augmenter_propagates_to_outer_fan_out_instance"), - } -) +# Currently empty; the harness covers every activated case. Kept as a named hook +# so future per-case deferrals don't need to re-introduce the pattern. +_DEFERRED_CASES: frozenset[tuple[str, str]] = frozenset() # Mocks the spec fixture 037 references for ``trace_input_from_state`` / @@ -670,11 +655,170 @@ def _build_039_graph( ) -> tuple[Any, Any]: """Dispatch a 039 case to its hand-built graph; return (graph, factory).""" name = cast("str", case.get("name")) + if name == "inner_fan_out_augmenter_propagates_to_outer_dispatch_span": + return _build_039_case1(case, provider=provider, prompt_manager=prompt_manager) + if name == "parallel_branch_augmenter_propagates_to_outer_fan_out_instance": + return _build_039_case2(case, provider=provider, prompt_manager=prompt_manager) if name == "fan_out_in_serial_subgraph_augmenter_propagates_to_wrapper_span": return _build_039_case3(case, provider=provider, prompt_manager=prompt_manager) raise NotImplementedError(f"039 case not yet wired: {name!r}") +def _make_039_inner_seed_extractor(outer_item_field: str, seed_field: str, seed_subkey: str) -> Any: + # Case 1 outer fan-out instance middleware: the inner fan-out's items_field + # (inner_seed) is a SUB-FIELD of the outer product item, which the engine's + # flat item_field / inputs projection can't reach. Read the outer item from + # its slot and inject the sub-field into the per_product instance state so the + # inner fan-out can fan over it. + async def _extract(state: Any, next_: Any) -> Any: + item = getattr(state, outer_item_field, None) + if isinstance(item, Mapping): + state = state.model_copy(update={seed_field: cast("Mapping[str, Any]", item)[seed_subkey]}) + return await next_(state) + + return _extract + + +def _build_039_case1( + case: Mapping[str, Any], + *, + provider: OpenAIProvider | None, + prompt_manager: PromptManager | None, +) -> tuple[Any, Any]: + # Case 1: fan-out inside a fan-out instance. The inner fan-out + # (per_product.inner_fan_out) fans over each product's inner_seed sub-field + # and augments note=; per 0045 that note reaches the OUTER + # instance dispatch span (one level up) but NOT the shared fan-out NODEs. + assert provider is not None, "039 cases declare mock_llm, so the provider must be set" + per_category_spec = copy.deepcopy(cast("dict[str, Any]", case["inner_subgraphs"]["per_category"])) + per_category_spec.setdefault("state", {}).setdefault("fields", {}).setdefault( + "oa_fan_out_item", {"type": "dict", "default": {}} + ) + per_category = _build_inner_subgraph_with_llm( + per_category_spec, provider=provider, prompt_manager=prompt_manager, render_variables={} + ) + per_product_state = build_state_cls( + "PerProduct039C1", + { + "score": {"type": "string", "default": ""}, + "inner_seed": {"type": "list", "default": []}, + "oa_outer_item": {"type": "dict", "default": {}}, + "inner_picks": {"type": "list", "reducer": "append", "default": []}, + }, + ) + pp_builder: GraphBuilder[Any] = GraphBuilder(per_product_state) + pp_builder.set_entry("inner_fan_out") + pp_builder.add_fan_out_node( + "inner_fan_out", + subgraph=per_category, + items_field="inner_seed", + item_field="oa_fan_out_item", + collect_field="picked", + target_field="inner_picks", + instance_middleware=[_make_augment_instance_middleware({"note": "value"}, "oa_fan_out_item")], + ) + pp_builder.add_edge("inner_fan_out", END) + per_product = pp_builder.compile() + outer_state = build_state_cls( + "Outer039C1", + { + "results": {"type": "list", "reducer": "append", "default": []}, + "products": {"type": "list", "default": []}, + }, + ) + outer_builder: GraphBuilder[Any] = GraphBuilder(outer_state) + outer_builder.set_entry("outer_fan_out") + outer_builder.add_fan_out_node( + "outer_fan_out", + subgraph=per_product, + items_field="products", + item_field="oa_outer_item", + collect_field="inner_picks", + target_field="results", + # concurrency=1 while the observer's NODE-key collision under concurrent + # nested fan-out is fixed (inner nodes of different outer instances share + # a _key_for and dedup, dropping spans). The engine produces correct + # results at any concurrency; this is an observability-only constraint. + concurrency=1, + instance_middleware=[_make_039_inner_seed_extractor("oa_outer_item", "inner_seed", "inner_seed")], + ) + outer_builder.add_edge("outer_fan_out", END) + graph = outer_builder.compile() + initial = cast("dict[str, Any]", case.get("initial_state") or {}) + return graph, (lambda: outer_state(**initial)) + + +def _build_039_case2( + case: Mapping[str, Any], + *, + provider: OpenAIProvider | None, + prompt_manager: PromptManager | None, +) -> tuple[Any, Any]: + # Case 2: parallel-branches NODE inside a fan-out instance. Only the `probe` + # branch augments note= via augment_metadata_from_outer_item; + # per 0045 that note reaches the probe branch dispatch AND the OUTER fan-out + # instance dispatch, but NOT the baseline branch, the pb NODE, or the fan-out + # NODE. + assert provider is not None, "039 cases declare mock_llm, so the provider must be set" + subgraphs = cast("dict[str, Any]", case["subgraphs"]) + probe_spec = copy.deepcopy(cast("dict[str, Any]", subgraphs["probe"])) + probe_spec.setdefault("state", {}).setdefault("fields", {}).setdefault( + "oa_outer_item", {"type": "dict", "default": {}} + ) + probe = _build_inner_subgraph_with_llm( + probe_spec, provider=provider, prompt_manager=prompt_manager, render_variables={} + ) + baseline = _build_inner_subgraph_with_llm( + cast("Mapping[str, Any]", subgraphs["baseline"]), + provider=provider, + prompt_manager=prompt_manager, + render_variables={}, + ) + per_product_state = build_state_cls( + "PerProduct039C2", + { + "outcome": {"type": "string", "default": ""}, + "oa_outer_item": {"type": "dict", "default": {}}, + }, + ) + pp_builder: GraphBuilder[Any] = GraphBuilder(per_product_state) + pp_builder.set_entry("dispatcher") + pp_builder.add_parallel_branches_node( + "dispatcher", + branches={ + "probe": BranchSpec( + subgraph=probe, + inputs={"oa_outer_item": "oa_outer_item"}, + middleware=(_make_augment_instance_middleware({"note": "id"}, "oa_outer_item"),), + ), + "baseline": BranchSpec(subgraph=baseline), + }, + ) + pp_builder.add_edge("dispatcher", END) + per_product = pp_builder.compile() + outer_state = build_state_cls( + "Outer039C2", + { + "results": {"type": "list", "reducer": "append", "default": []}, + "products": {"type": "list", "default": []}, + }, + ) + outer_builder: GraphBuilder[Any] = GraphBuilder(outer_state) + outer_builder.set_entry("outer_fan_out") + outer_builder.add_fan_out_node( + "outer_fan_out", + subgraph=per_product, + items_field="products", + item_field="oa_outer_item", + collect_field="outcome", + target_field="results", + ) + outer_builder.add_edge("outer_fan_out", END) + graph = outer_builder.compile() + initial = cast("dict[str, Any]", case.get("initial_state") or {}) + return graph, (lambda: outer_state(**initial)) + + def _build_039_case3( case: Mapping[str, Any], *, diff --git a/tests/unit/test_observability_otel.py b/tests/unit/test_observability_otel.py index 748cf81..c88ddf5 100644 --- a/tests/unit/test_observability_otel.py +++ b/tests/unit/test_observability_otel.py @@ -3712,6 +3712,218 @@ async def _mid_passthrough(_s: _MidS) -> dict[str, int]: ) +async def test_nested_fan_out_in_fan_out_dispatch_lineage() -> None: + # Proposal 0045 / 0013: a fan-out nested INSIDE a fan-out instance. Each + # outer instance gets its OWN inner per-instance dispatch span (distinct + # lineage keys, no cross-instance collision -- before the fix the second + # collided with the first), and an inner leaf's augmentation reaches its own + # outer instance dispatch, not the sibling's, and not the shared NODE spans. + import asyncio + + from openarmature.observability.metadata import set_invocation_metadata + + class _LeafS(State): + tag: str = "" + seed: str = "" + out: str = "" + + class _MidS(State): + tag: str = "" + seeds: list[str] = [] + collected: Annotated[list[str], append] = [] + + class _OuterS(State): + products: list[str] = [] + seeds: list[str] = [] + results: Annotated[list[Any], append] = [] + + async def _leaf(s: _LeafS) -> dict[str, str]: + await asyncio.sleep(0) + set_invocation_metadata(note=f"{s.tag}-{s.seed}") + return {"out": f"{s.tag}-{s.seed}"} + + leaf = GraphBuilder(_LeafS).add_node("ask", _leaf).add_edge("ask", END).set_entry("ask").compile() + mid = ( + GraphBuilder(_MidS) + .add_fan_out_node( + "inner_fan", + subgraph=leaf, + items_field="seeds", + item_field="seed", + inputs={"tag": "tag"}, + collect_field="out", + target_field="collected", + ) + .add_edge("inner_fan", END) + .set_entry("inner_fan") + .compile() + ) + exporter = InMemorySpanExporter() + observer = OTelObserver(span_processor=SimpleSpanProcessor(exporter)) + g = ( + GraphBuilder(_OuterS) + .add_fan_out_node( + "outer_fan", + subgraph=mid, + items_field="products", + item_field="tag", + inputs={"seeds": "seeds"}, + collect_field="collected", + target_field="results", + # concurrency=1 while the observer's NODE-key collision under + # concurrent nested fan-out is fixed (the inner nodes of different + # outer instances share a _key_for and dedup); the engine results are + # correct at any concurrency. Tracked separately. + concurrency=1, + ) + .add_edge("outer_fan", END) + .set_entry("outer_fan") + .compile() + ) + g.attach_observer(observer) + try: + await g.invoke(_OuterS(products=["A", "B"], seeds=["x"])) + await g.drain() + finally: + observer.shutdown() + spans = exporter.get_finished_spans() + + def _attr(s: Any, k: str) -> Any: + return dict(s.attributes or {}).get(k) + + inner_dispatches = [ + s + for s in spans + if s.name == "inner_fan" and "openarmature.fan_out.parent_node_name" in dict(s.attributes or {}) + ] + assert len(inner_dispatches) == 2, ( + f"expected 2 nested inner instance dispatches, got {len(inner_dispatches)}" + ) + outer_dispatches = { + _attr(s, "openarmature.node.fan_out_index"): s + for s in spans + if s.name == "outer_fan" and "openarmature.fan_out.parent_node_name" in dict(s.attributes or {}) + } + assert outer_dispatches.keys() == {0, 1} + assert _attr(outer_dispatches[0], "openarmature.user.note") == "A-x" + assert _attr(outer_dispatches[1], "openarmature.user.note") == "B-x" + node_spans = [ + s + for s in spans + if s.name in ("outer_fan", "inner_fan") + and "openarmature.fan_out.parent_node_name" not in dict(s.attributes or {}) + ] + assert node_spans, "expected fan-out NODE spans" + assert all(_attr(s, "openarmature.user.note") is None for s in node_spans), ( + "shared fan-out NODE spans MUST NOT carry the augmentation" + ) + + +async def test_parallel_branches_in_fan_out_dispatch_lineage() -> None: + # Proposal 0045 / 0044: a parallel-branches NODE inside a fan-out instance. + # Each outer instance gets its own per-branch dispatch spans (distinct keys, + # no cross-instance collision); only the augmenting branch + its outer + # instance dispatch carry the augmentation, not the sibling branch. + import asyncio + + from openarmature.graph import BranchSpec + from openarmature.observability.metadata import set_invocation_metadata + + class _BranchS(State): + tag: str = "" + out: str = "" + + class _MidS(State): + tag: str = "" + outcome: str = "" + + class _OuterS(State): + products: list[str] = [] + results: Annotated[list[Any], append] = [] + + async def _probe(s: _BranchS) -> dict[str, str]: + await asyncio.sleep(0) + return {"out": s.tag} + + async def _baseline(_s: _BranchS) -> dict[str, str]: + await asyncio.sleep(0) + return {"out": "base"} + + probe = GraphBuilder(_BranchS).add_node("ask", _probe).add_edge("ask", END).set_entry("ask").compile() + baseline = ( + GraphBuilder(_BranchS).add_node("ask", _baseline).add_edge("ask", END).set_entry("ask").compile() + ) + + async def _augment(s: _BranchS, next_: Any) -> Any: + set_invocation_metadata(note=s.tag) + return await next_(s) + + mid = ( + GraphBuilder(_MidS) + .add_parallel_branches_node( + "dispatcher", + branches={ + "probe": BranchSpec(subgraph=probe, inputs={"tag": "tag"}, middleware=(_augment,)), + "baseline": BranchSpec(subgraph=baseline), + }, + ) + .add_edge("dispatcher", END) + .set_entry("dispatcher") + .compile() + ) + exporter = InMemorySpanExporter() + observer = OTelObserver(span_processor=SimpleSpanProcessor(exporter)) + g = ( + GraphBuilder(_OuterS) + .add_fan_out_node( + "outer_fan", + subgraph=mid, + items_field="products", + item_field="tag", + collect_field="outcome", + target_field="results", + ) + .add_edge("outer_fan", END) + .set_entry("outer_fan") + .compile() + ) + g.attach_observer(observer) + try: + await g.invoke(_OuterS(products=["A", "B"])) + await g.drain() + finally: + observer.shutdown() + spans = exporter.get_finished_spans() + + def _attr(s: Any, k: str) -> Any: + return dict(s.attributes or {}).get(k) + + probe_dispatches = [ + s + for s in spans + if s.name == "probe" and "openarmature.parallel_branches.parent_node_name" in dict(s.attributes or {}) + ] + baseline_dispatches = [ + s + for s in spans + if s.name == "baseline" + and "openarmature.parallel_branches.parent_node_name" in dict(s.attributes or {}) + ] + assert len(probe_dispatches) == 2, f"expected 2 probe dispatches, got {len(probe_dispatches)}" + assert len(baseline_dispatches) == 2, f"expected 2 baseline dispatches, got {len(baseline_dispatches)}" + outer_dispatches = { + _attr(s, "openarmature.node.fan_out_index"): s + for s in spans + if s.name == "outer_fan" and "openarmature.fan_out.parent_node_name" in dict(s.attributes or {}) + } + assert outer_dispatches.keys() == {0, 1} + assert _attr(outer_dispatches[0], "openarmature.user.note") == "A" + assert _attr(outer_dispatches[1], "openarmature.user.note") == "B" + assert all(_attr(s, "openarmature.user.note") is None for s in baseline_dispatches), ( + "the non-augmenting baseline branch MUST NOT carry the augmentation" + ) + + # --------------------------------------------------------------------------- # Callable parallel branches (proposal 0075, observability §5.7). Mirrors # spec conformance fixture 110 (otel-callable-branch-span), which is not yet From f1160c897a629854ed4491c8d8a00ae7b3f11e62 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Fri, 26 Jun 2026 08:55:38 -0700 Subject: [PATCH 3/3] Fix observer flat-key lookups, tighten key types The PR review's type-tightening surfaced two latent flat-key lookups that silently always-missed against the now-composite dispatch maps: the Langfuse LLM parent resolver and the OTel callable-branch provider-nesting lookup both still built the old flat key. Both now use the lineage-aware key, with a regression test for the Langfuse resolver (an LLM call inside a top-level fan-out instance parents under its per-instance dispatch). The tightening itself: the lineage-key helpers returned tuple[Any, ...] instead of their concrete shape. Add _DispatchKey and _BranchDispatchKey aliases for the helper returns, the dispatch maps, and the close-helper key params (all still typed tuple[str, ...] from before the lineage keys), and drop the now-unnecessary key[0] / key[1] casts. Also reword the resume comment so the nested-fan-out re-run reads as conditional on the outer instance re-entering. --- src/openarmature/graph/compiled.py | 15 +++-- .../observability/langfuse/observer.py | 50 +++++++++------- .../observability/otel/observer.py | 42 +++++++++---- tests/unit/test_observability_langfuse.py | 59 +++++++++++++++++++ 4 files changed, 128 insertions(+), 38 deletions(-) diff --git a/src/openarmature/graph/compiled.py b/src/openarmature/graph/compiled.py index b86f0cc..20098e8 100644 --- a/src/openarmature/graph/compiled.py +++ b/src/openarmature/graph/compiled.py @@ -421,11 +421,16 @@ def _restore_fan_out_progress_state( # outer fan-out instance does not round-trip its per-outer-instance # progress through the current record format (it would need the lineage # on the record): its in-memory keys carry the lineage, so the restored - # empty-lineage entry never matches and the nested fan-out RE-RUNS on - # resume. (Before the lineage fix it would instead skip, rolling forward - # the collapsed/wrong shared entry -- so re-running is the safer of two - # never-correct behaviors, and matches the spec's inner-subgraph re-entry - # model.) A full fix needs the lineage on the record; tracked separately. + # empty-lineage entry never matches. The consequence only bites when + # resume actually RE-ENTERS the nested fan-out -- i.e. its outer instance + # was in-flight at the save. A completed outer instance rolls forward and + # never re-runs its inner fan-out, so the missing inner entry is moot + # there. When the inner fan-out does re-enter, it re-runs from scratch + # rather than skipping its completed inner instances. (Before the lineage + # fix it would instead skip, rolling forward the collapsed/wrong shared + # entry -- so re-running is the safer of two never-correct behaviors, and + # matches the spec's inner-subgraph re-entry model.) A full fix needs the + # lineage on the record; tracked separately. key = (fp.namespace, fp.fan_out_node_name, ()) out[key] = _FanOutExecutionState( fan_out_node_name=fp.fan_out_node_name, diff --git a/src/openarmature/observability/langfuse/observer.py b/src/openarmature/observability/langfuse/observer.py index 66636d4..7ab8139 100644 --- a/src/openarmature/observability/langfuse/observer.py +++ b/src/openarmature/observability/langfuse/observer.py @@ -96,6 +96,13 @@ def _read_implementation_version() -> str: # Langfuse handle instead of an OTel Span. _StackKey = tuple[tuple[str, ...], int, int | None, str | None] +# Lineage-aware dispatch keys (proposal 0045): the fan-out / pb NODE namespace +# prefix plus the fan-out instance index / branch name chain slices along the +# path to it. _BranchDispatchKey carries the explicit branch name (a callable +# branch sets it without extending the chain) instead of a trailing chain entry. +_DispatchKey = tuple[tuple[str, ...], tuple[int | None, ...], tuple[str | None, ...]] +_BranchDispatchKey = tuple[tuple[str, ...], tuple[int | None, ...], tuple[str | None, ...], str] + @dataclass class _OpenObservation: @@ -138,7 +145,7 @@ def _dispatch_key( prefix: tuple[str, ...], fan_out_index_chain: tuple[int | None, ...], branch_name_chain: tuple[str | None, ...], -) -> tuple[Any, ...]: +) -> _DispatchKey: """Lineage-aware identity key for a fan-out instance / per-branch dispatch at namespace ``prefix``. Encodes the fan-out/pb NODE namespace plus the full chain of fan-out instance indices / branch names along the path to it (sliced @@ -158,7 +165,7 @@ def _branch_dispatch_key( fan_out_index_chain: tuple[int | None, ...], branch_name_chain: tuple[str | None, ...], branch_name: str, -) -> tuple[Any, ...]: +) -> _BranchDispatchKey: """Lineage-aware identity key for a per-branch dispatch at namespace ``prefix``. The branch IDENTITY comes from ``branch_name`` explicitly (not ``branch_name_chain[-1]``): a callable branch carries its name on the event @@ -262,8 +269,8 @@ class _InvState: # fan-out node's own Span observation; inner-node observations # parent under this dispatch instead of the shared fan-out node # span. Closed when the fan-out node's completed event fires. - fan_out_instance_observations: dict[tuple[str, ...], _OpenObservation] = field( - default_factory=dict[tuple[str, ...], _OpenObservation] + fan_out_instance_observations: dict[_DispatchKey, _OpenObservation] = field( + default_factory=dict[_DispatchKey, _OpenObservation] ) # Maps a namespace prefix to the detached Langfuse trace_id when # that subtree is configured detached (per the observer's @@ -276,7 +283,7 @@ class _InvState: # detached subgraph prefixes because they're closed when the # fan-out node's completed event fires, not when the namespace # cursor leaves the subtree. - fan_out_instance_root_prefixes: set[tuple[str, ...]] = field(default_factory=set[tuple[str, ...]]) + fan_out_instance_root_prefixes: set[_DispatchKey] = field(default_factory=set[_DispatchKey]) # ``parent_node_name`` cache for per-instance attribution # (spec proposal 0013 v0.10.0 — inner events from inside a # non-detached fan-out instance don't carry fan_out_config @@ -297,8 +304,8 @@ class _InvState: # (prefix = the parallel-branches NODE namespace). Inner branch-node # observations parent under this dispatch instead of the shared pb NODE # span; closed when the pb NODE's completed event fires. - parallel_branches_branch_spans: dict[tuple[str, ...], _OpenObservation] = field( - default_factory=dict[tuple[str, ...], _OpenObservation] + parallel_branches_branch_spans: dict[_BranchDispatchKey, _OpenObservation] = field( + default_factory=dict[_BranchDispatchKey, _OpenObservation] ) # Declared branch-name set per pb NODE namespace (from the NODE's # parallel_branches_config), so an inner branch event matches only the node @@ -578,7 +585,7 @@ def _handle_completed(self, event: NodeEvent) -> None: if event.fan_out_index is None and event.namespace and event.namespace[0] in self.detached_fan_outs: ns = event.namespace for key in list(inv_state.fan_out_instance_root_prefixes): - anchor_ns = cast("tuple[str, ...]", key[0]) + anchor_ns = key[0] if len(anchor_ns) >= len(ns) and anchor_ns[: len(ns)] == ns: # Detached per-instance dispatches live in # fan_out_instance_observations (same map as @@ -587,7 +594,7 @@ def _handle_completed(self, event: NodeEvent) -> None: inv_state.fan_out_instance_root_prefixes.discard(key) # detached_traces uses the top-level routing key shape; derive # it from the lineage key's own instance index (last entry). - fi_chain = cast("tuple[int | None, ...]", key[1]) + fi_chain = key[1] inv_state.detached_traces.pop(anchor_ns + (str(fi_chain[-1]),), None) # Per spec proposal 0013 (v0.10.0): when the fan-out node's # own completion fires, close all per-instance dispatch @@ -595,7 +602,7 @@ def _handle_completed(self, event: NodeEvent) -> None: if event.fan_out_index is None and event.fan_out_config is not None: ns = event.namespace for key in list(inv_state.fan_out_instance_observations.keys()): - anchor_ns = cast("tuple[str, ...]", key[0]) + anchor_ns = key[0] # The dispatch key is now (anchor_ns, fi_chain, bn_chain); match # on the NODE namespace (anchor_ns) being in this completing # node's subtree. @@ -614,7 +621,7 @@ def _handle_completed(self, event: NodeEvent) -> None: if event.parallel_branches_config is not None: ns = event.namespace for key in list(inv_state.parallel_branches_branch_spans.keys()): - anchor_ns = cast("tuple[str, ...]", key[0]) + anchor_ns = key[0] if len(anchor_ns) >= len(ns) and anchor_ns[: len(ns)] == ns: self._close_parallel_branches_branch_dispatch_observation(inv_state, key) inv_state.parallel_branches_parent_node_name.pop(event.namespace, None) @@ -705,7 +712,7 @@ def _handle_metadata_augmentation(self, event: MetadataAugmentationEvent) -> Non # prefix of the augmenter's -- so a SIBLING outer instance's dispatch, # whose chain differs at the enclosing position, is excluded. for key, observation in inv_state.fan_out_instance_observations.items(): - anchor_ns = cast("tuple[str, ...]", key[0]) + anchor_ns = key[0] if not (is_strict_prefix(anchor_ns, aug_ns) or anchor_ns == aug_ns): continue if _observation_chain_on_path(observation, aug_fi_chain, aug_bn_chain): @@ -1524,18 +1531,16 @@ def _close_subgraph_observation(self, inv_state: _InvState, prefix: tuple[str, . return observation.handle.end() - def _close_fan_out_instance_dispatch_observation( - self, inv_state: _InvState, prefix: tuple[str, ...] - ) -> None: - observation = inv_state.fan_out_instance_observations.pop(prefix, None) + def _close_fan_out_instance_dispatch_observation(self, inv_state: _InvState, key: _DispatchKey) -> None: + observation = inv_state.fan_out_instance_observations.pop(key, None) if observation is None: return observation.handle.end() def _close_parallel_branches_branch_dispatch_observation( - self, inv_state: _InvState, prefix: tuple[str, ...] + self, inv_state: _InvState, key: _BranchDispatchKey ) -> None: - observation = inv_state.parallel_branches_branch_spans.pop(prefix, None) + observation = inv_state.parallel_branches_branch_spans.pop(key, None) if observation is None: return observation.handle.end() @@ -1903,9 +1908,14 @@ def _resolve_llm_parent_observation_id( observation = inv_state.open_observations.get(key) if observation is not None: return observation.handle.id - # Per-instance fan-out dispatch. + # Per-instance fan-out dispatch. The dispatch map is keyed by the + # lineage-aware _dispatch_key; reconstruct the top-level instance's key + # (namespace[:1], the instance index, no branch axis) to match it. Only + # the innermost fan_out_index is available here, so this resolves an LLM + # call directly inside a top-level fan-out instance (the case the flat + # ``namespace[:1] + (str(index),)`` key handled before the lineage keys). if calling_fan_out_index is not None and calling_namespace_prefix: - instance_key = calling_namespace_prefix[:1] + (str(calling_fan_out_index),) + instance_key = _dispatch_key(calling_namespace_prefix[:1], (calling_fan_out_index,), (None,)) dispatch = inv_state.fan_out_instance_observations.get(instance_key) if dispatch is not None: return dispatch.handle.id diff --git a/src/openarmature/observability/otel/observer.py b/src/openarmature/observability/otel/observer.py index a2cfead..d0d701f 100644 --- a/src/openarmature/observability/otel/observer.py +++ b/src/openarmature/observability/otel/observer.py @@ -123,6 +123,14 @@ # same namespace + fan_out_index would collide on the same key. _StackKey = tuple[tuple[str, ...], int, int | None, str | None] +# Lineage-aware dispatch keys (proposal 0045): the fan-out / pb NODE namespace +# prefix plus the fan-out instance index / branch name chain slices along the +# path to it. _BranchDispatchKey carries the explicit branch name (a callable +# branch sets it without extending the chain) instead of a trailing chain entry. +# Mirrors the LangfuseObserver aliases. +_DispatchKey = tuple[tuple[str, ...], tuple[int | None, ...], tuple[str | None, ...]] +_BranchDispatchKey = tuple[tuple[str, ...], tuple[int | None, ...], tuple[str | None, ...], str] + # §5.5.5 truncation marker. The leading character is U+2026 HORIZONTAL # ELLIPSIS (3 bytes UTF-8); the marker is a fixed UTF-8 string and is @@ -303,7 +311,7 @@ def _dispatch_key( prefix: tuple[str, ...], fan_out_index_chain: tuple[int | None, ...], branch_name_chain: tuple[str | None, ...], -) -> tuple[Any, ...]: +) -> _DispatchKey: """Lineage-aware identity key for a fan-out instance / per-branch dispatch span at namespace ``prefix``. Encodes the fan-out/pb NODE namespace plus the full chain of fan-out instance indices / branch names along the path to it @@ -321,7 +329,7 @@ def _branch_dispatch_key( fan_out_index_chain: tuple[int | None, ...], branch_name_chain: tuple[str | None, ...], branch_name: str, -) -> tuple[Any, ...]: +) -> _BranchDispatchKey: """Lineage-aware identity key for a per-branch dispatch span at namespace ``prefix``. The branch IDENTITY comes from ``branch_name`` explicitly (not ``branch_name_chain[-1]``): a callable branch carries its name on the event @@ -416,6 +424,9 @@ class _InvState: # lets it OVERRIDE a prior ERROR (only a set to UNSET is ignored), so # an unconditional OK at close would erase the detached-trace ERROR. errored_detached_keys: set[tuple[str, ...]] = field(default_factory=set[tuple[str, ...]]) + # OTel keeps detached instance roots on a FLAT key (prefix + (str(index),)), + # the same shape _resolve_parent_context / detached_roots use; only the + # non-detached dispatch maps moved to the lineage-aware _DispatchKey. fan_out_instance_root_prefixes: set[tuple[str, ...]] = field(default_factory=set[tuple[str, ...]]) # Per spec §5.4 + proposal 0013 (v0.10.0): non-detached fan-outs # synthesize per-instance dispatch spans nested between the fan-out @@ -424,8 +435,8 @@ class _InvState: # ``detached_roots``. Lives in the parent trace (not a fresh # trace_id, unlike detached). Closed when the fan-out node's # ``completed`` event fires. - fan_out_instance_spans: dict[tuple[str, ...], _OpenSpan] = field( - default_factory=dict[tuple[str, ...], _OpenSpan] + fan_out_instance_spans: dict[_DispatchKey, _OpenSpan] = field( + default_factory=dict[_DispatchKey, _OpenSpan] ) # ``parent_node_name`` cache (per spec proposal 0013 correction #4): # the fan-out node's name surfaces on its own ``started`` event via @@ -442,8 +453,8 @@ class _InvState: # per-instance synthesis above). Keyed by # ``prefix + (branch_name,)``. Closed when the parallel-branches # node's ``completed`` event fires. - parallel_branches_branch_spans: dict[tuple[str, ...], _OpenSpan] = field( - default_factory=dict[tuple[str, ...], _OpenSpan] + parallel_branches_branch_spans: dict[_BranchDispatchKey, _OpenSpan] = field( + default_factory=dict[_BranchDispatchKey, _OpenSpan] ) # ``parent_node_name`` cache for parallel-branches (mirrors the # ``fan_out_parent_node_name`` cache above). Surfaces from the @@ -842,7 +853,12 @@ def prepare_sync(self, event: NodeEvent) -> None: # the branch span. if event.branch_name is not None and event.parallel_branches_config is None: dispatch = inv_state.parallel_branches_branch_spans.get( - event.namespace + (event.branch_name,) + _branch_dispatch_key( + event.namespace, + event.fan_out_index_chain, + event.branch_name_chain, + event.branch_name, + ) ) if dispatch is not None: _set_active_observer_span(dispatch.span) @@ -991,7 +1007,7 @@ def _handle_completed(self, event: NodeEvent) -> None: # The dispatch key is now (anchor_ns, fi_chain, bn_chain); match # on the NODE namespace (anchor_ns = key[0]) being in this # completing node's subtree. - anchor_ns = cast("tuple[str, ...]", key[0]) + anchor_ns = key[0] if len(anchor_ns) >= len(ns) and anchor_ns[: len(ns)] == ns: self._close_fan_out_instance_dispatch_span(inv_state, key) inv_state.fan_out_parent_node_name.pop(event.namespace, None) @@ -1008,7 +1024,7 @@ def _handle_completed(self, event: NodeEvent) -> None: if event.parallel_branches_config is not None: ns = event.namespace for key in list(inv_state.parallel_branches_branch_spans.keys()): - anchor_ns = cast("tuple[str, ...]", key[0]) + anchor_ns = key[0] if len(anchor_ns) >= len(ns) and anchor_ns[: len(ns)] == ns: self._close_parallel_branches_branch_dispatch_span(inv_state, key) inv_state.parallel_branches_parent_node_name.pop(event.namespace, None) @@ -1195,7 +1211,7 @@ def _collect_augmentation_targets( # of the augmenter's -- so a SIBLING outer instance's dispatch, whose # chain differs at the enclosing position, is excluded. for key, open_span in inv_state.fan_out_instance_spans.items(): - anchor_ns = cast("tuple[str, ...]", key[0]) + anchor_ns = key[0] if not (is_strict_prefix(anchor_ns, aug_ns) or anchor_ns == aug_ns): continue if _span_chain_on_path(open_span, aug_fi_chain, aug_bn_chain): @@ -1203,7 +1219,7 @@ def _collect_augmentation_targets( # Per-branch dispatch spans: same lineage match as fan-out above. for key, open_span in inv_state.parallel_branches_branch_spans.items(): - anchor_ns = cast("tuple[str, ...]", key[0]) + anchor_ns = key[0] if not (is_strict_prefix(anchor_ns, aug_ns) or anchor_ns == aug_ns): continue if _span_chain_on_path(open_span, aug_fi_chain, aug_bn_chain): @@ -2366,7 +2382,7 @@ def _open_fan_out_instance_dispatch_span( branch_name_chain=event.branch_name_chain[:chain_len], ) - def _close_fan_out_instance_dispatch_span(self, inv_state: _InvState, key: tuple[str, ...]) -> None: + def _close_fan_out_instance_dispatch_span(self, inv_state: _InvState, key: _DispatchKey) -> None: open_span = inv_state.fan_out_instance_spans.pop(key, None) if open_span is None: return @@ -2443,7 +2459,7 @@ def _open_parallel_branches_branch_dispatch_span( ) def _close_parallel_branches_branch_dispatch_span( - self, inv_state: _InvState, key: tuple[str, ...] + self, inv_state: _InvState, key: _BranchDispatchKey ) -> None: open_span = inv_state.parallel_branches_branch_spans.pop(key, None) if open_span is None: diff --git a/tests/unit/test_observability_langfuse.py b/tests/unit/test_observability_langfuse.py index f037d03..5c399f6 100644 --- a/tests/unit/test_observability_langfuse.py +++ b/tests/unit/test_observability_langfuse.py @@ -1573,6 +1573,65 @@ async def test_typed_failed_event_parents_under_branch_calling_node() -> None: assert error_gens[0].parent_observation_id != slow_handle.id +async def test_llm_event_parents_under_fan_out_instance_dispatch() -> None: + # Regression cover for _resolve_llm_parent_observation_id fallback #2: when an + # LLM event fires inside a top-level fan-out instance and the calling node has + # no open observation (fallback #1 misses), the Generation MUST parent under + # the per-instance fan-out dispatch observation. The dispatch map is keyed by + # the lineage-aware _dispatch_key; before the lineage keys this fallback used + # a flat namespace[:1] + (str(index),) key, which always-misses against the + # composite map and silently re-parents the Generation under the Trace. + from openarmature.observability.correlation import ( + _reset_invocation_id, + _set_invocation_id, + ) + from openarmature.observability.langfuse.observer import ( + _dispatch_key, + _InvState, + _OpenObservation, + ) + from tests._helpers.typed_event import make_failed_event + + client = InMemoryLangfuseClient() + observer = LangfuseObserver(client=client) + invocation_id = "inv-fanout-llm" + token = _set_invocation_id(invocation_id) + try: + client.trace(id=invocation_id, name="fan") + observer._inv_states[invocation_id] = _InvState(trace_id=invocation_id) # noqa: SLF001 + inv_state = observer._inv_states[invocation_id] # noqa: SLF001 + # Per-instance dispatch for top-level fan-out "fan", instance 0, keyed by + # the lineage-aware _dispatch_key. No open_observation for the calling + # node ("fan", "ask"), so the resolver must reach fallback #2. + dispatch_handle = client.span(trace_id=invocation_id, name="fan") + instance_key = _dispatch_key(("fan",), (0,), (None,)) + inv_state.fan_out_instance_observations[instance_key] = _OpenObservation(handle=dispatch_handle) + await observer( + make_failed_event( + invocation_id=invocation_id, + node_name="ask", + namespace=("fan", "ask"), + attempt_index=0, + fan_out_index=0, + branch_name=None, + model="m-test", + error_category="provider_unavailable", + error_type="ProviderUnavailable", + error_message="503 from upstream", + call_id="cc-fan", + ) + ) + finally: + _reset_invocation_id(token) + + trace = client.traces[invocation_id] + error_gens = [o for o in trace.observations if o.type == "generation" and o.level == "ERROR"] + assert len(error_gens) == 1 + assert error_gens[0].parent_observation_id == dispatch_handle.id, ( + "LLM Generation must parent under the per-instance fan-out dispatch (resolver fallback #2)" + ) + + # --------------------------------------------------------------------------- # Proposal 0063 — tool-execution Tool observation (asType "tool") # ---------------------------------------------------------------------------