From 4ceeec6ac4f901e9314c0c90a0b5771ca0565de5 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sat, 30 May 2026 20:49:44 -0700 Subject: [PATCH 1/2] Activate fixture 037 case 5 (resume re-fire) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the langfuse conformance harness for the remaining decision-tree case of proposal 0043's §8.4.1 trace.input/output sourcing fixture. The two-phase resume flow (first invoke catches NodeException → resume invoke completes) now runs end-to-end through new harness primitives: - ``flaky: {fail_first_invocation_only: true, on_success: {...}}`` compact test seam in ``_build_node_body``. - ``checkpointer: in_memory`` directive registers ``InMemoryCheckpointer`` on the graph builder. - ``returns_state_snapshot`` added to ``_TRACE_IO_HOOK_REGISTRY``. - ``_run_resume_case`` runs the two-phase flow + asserts both traces + checks the §8.4.1 invariants (distinct trace ids, shared correlation_id, first trace unchanged, hooks re-fire on resumed trace). Activation surfaced two engine bugs that PR #99 missed. The first: ``InvocationCompletedEvent.final_state`` on the failure path defaulted to ``starting_state``, but spec §8.4.1 *Resume semantics* requires the failure-path ``trace.output`` hook to receive "the partial final state captured at the failure point" (the most recent successful step's post-merge state). Adds a new ``latest_state_box`` on ``_InvocationContext`` that the engine writes after every successful step's ``state = step_result.state`` assignment; the outermost ``invoke()`` reads it in the finally-block before falling back to ``starting_state``. The second: ``latest_state_box`` MUST be per-context (unlike its sibling ``final_node_box`` which shares by reference across subgraph descents). An inner-subgraph step's success previously would overwrite the outer box with an inner-typed state; on a subsequent outer-level raise the outer ``trace.output`` hook would receive an inner state when its signature expects the outer state class. Each ``descend_into_*`` method now omits ``latest_state_box`` from the copy, so each level gets a fresh box. Four new unit-test regressions pin the bug fix across all four graph-descent shapes: flat, subgraph, fan-out instance, parallel- branches branch. Each test wires a graph where an outer node succeeds (outer_a_done=true) and a deeper raise propagates back; the ``trace_output_from_state`` hook MUST see the outer-state-typed value with the success captured. Cross-cap parser deferral for 037 stays in place — that parser still doesn't model ``langfuse_trace`` shape. Activation lives in the langfuse-specific harness only. --- CHANGELOG.md | 7 +- src/openarmature/graph/compiled.py | 35 +- src/openarmature/graph/observer.py | 26 ++ .../test_observability_langfuse.py | 226 +++++++++++- tests/unit/test_observability_langfuse.py | 335 ++++++++++++++++++ 5 files changed, 610 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 95ab3bea..63ceb1c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,11 +25,16 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). The - **`observation.metadata.detached: true` moves to the parent-side dispatching observation** (proposal 0042, observability §8.4.2). The Langfuse mapping previously emitted `detached: true` on the dispatch observation inside the detached child trace; the §8.4.2 row added by 0042 places it on the **parent-side** dispatching observation that fires the detached child (the link observation in the main trace for detached subgraphs; the parent fan-out node observation for detached fan-outs). The detached-side observation no longer carries the flag. - **`LangfuseClient.update_trace` Protocol grows `input` / `output` keyword parameters** so observer-supplied values land on the Trace's headline fields. +### Fixed + +- **`InvocationCompletedEvent.final_state` on the failure path now surfaces the partial state at failure point.** Spec §8.4.1 *Resume semantics* requires the failure-path `trace.output` hook to receive "the partial final state captured at the failure point"; the original PR #99 implementation defaulted to `starting_state`, so the hook saw pre-execution state when it should have seen post-execution-up-to-failure state. The engine now tracks the latest post-merge state via a `latest_state_box` on `_InvocationContext`, updated after every successful step and read on the failure path. Success-path behavior unchanged. +- **`latest_state_box` is per-context, not shared across subgraph descents.** Unlike the sibling `final_node_box` (which shares by reference because the spec wants the innermost failing node's name — the real culprit), `latest_state_box` must isolate per level so the outermost Langfuse trace receives outer-state-typed values. Without the isolation, a subgraph-internal step's inner-typed state would leak up to the outer trace.output hook, breaking the hook's typed contract. Each subgraph / fan-out instance / parallel-branches branch gets its own fresh box. Pinned by three regression tests covering flat, subgraph, fan-out, and parallel-branches failure paths. + ### Notes - **Pinned spec version bumped from v0.31.0 to v0.35.0.** Absorbs proposals 0042 (reserved-key extension), 0043 (Langfuse trace.input/output sourcing), and the textual additions in v0.32.0 (Gemini wire-format mapping, 0038, not yet implemented) and v0.33.0 (sessions capability, 0020, not yet implemented). - `LangfuseSDKAdapter` now applies `trace.input` / `trace.output` to the live Langfuse Trace. Input lands on the first real observation under the trace via `set_trace_io`; output uses a synthetic short-lived `openarmature.trace_io` observation as the carrier. The InMemoryLangfuseClient used by tests applies the fields directly. -- Conformance fixture `observability/conformance/037-langfuse-trace-input-output` activated for the four decision-tree cases (default stub / `disable_state_payload=False` / hooks non-null / hooks null-fallthrough). Case 5 (resume re-fire) is deferred to a follow-up — needs the langfuse harness to grow checkpointer wiring + flaky-node test seam + two-phase multi-trace assertion. +- Conformance fixture `observability/conformance/037-langfuse-trace-input-output` activated for all five cases (default stub / `disable_state_payload=False` / hooks non-null / hooks null-fallthrough / resume re-fire). The langfuse harness grew per-case `checkpointer: in_memory` wiring, a compact `flaky:` test seam, and a two-phase resume-flow assertion path. - The Langfuse v4 SDK marks `set_current_trace_io` / `Span.set_trace_io` deprecated ("removal in a future major version"). Empirical verification against Langfuse Cloud (v4.7.1, 2026-05-29) confirms it remains the **only** path that populates the Traces list view's headline `Input` / `Output` columns; `propagate_attributes(metadata=...)` does not substitute for it in the current UI. We will revisit when Langfuse publishes a concrete migration guide for v5. ## [0.10.0] — 2026-05-27 diff --git a/src/openarmature/graph/compiled.py b/src/openarmature/graph/compiled.py index 8ca7ba5a..f3f2e23f 100644 --- a/src/openarmature/graph/compiled.py +++ b/src/openarmature/graph/compiled.py @@ -1059,9 +1059,12 @@ async def invoke( # box the engine populates as nodes enter; on the failure # path that's the inner-most node that raised, on the # success path that's the last node before the END-routing - # edge. ``final_state`` is the engine's returned state on - # success and ``starting_state`` on the failure path (the - # engine doesn't expose intermediate state across raises). + # edge. ``final_state`` precedence: the engine's returned + # state on success → the most recent successful step's + # post-merge state on a mid-graph raise (per §8.4.1 + # *Resume semantics* "partial final state captured at the + # failure point") → ``starting_state`` only when no step + # ever completed. if context.final_node_box: final_node = context.final_node_box[0] else: @@ -1069,10 +1072,25 @@ async def invoke( # (e.g., resume-path validation). Fall back to the # declared entry node. final_node = self.entry + # ``latest_state_box`` is typed ``list[Any]`` on + # _InvocationContext (the context isn't parameterized on + # StateT), but at the outermost level (where this code + # runs) it always holds an outer ``StateT`` from a + # successful step's post-merge state. Cast for type + # narrowing; the per-context box-isolation pinned by + # ``test_failure_path_final_state_is_outer_type_*`` keeps + # this invariant honest. + event_final_state: StateT + if final_state is not None: + event_final_state = final_state + elif context.latest_state_box: + event_final_state = cast("StateT", context.latest_state_box[0]) + else: + event_final_state = starting_state _dispatch( context, InvocationCompletedEvent( - final_state=final_state if final_state is not None else starting_state, + final_state=event_final_state, status=status, final_node=final_node, invocation_id=invocation_id, @@ -1201,6 +1219,15 @@ async def _invoke( else: step_result = await self._step_function_node(node, current, state, context) state = step_result.state + # Proposal 0043 (post-PR-99 review): surface the most + # recent successful step's post-merge state so the + # outermost ``invoke()`` can populate + # ``InvocationCompletedEvent.final_state`` on the failure + # path with the partial state, not the bare initial state. + # Updated AFTER ``state = step_result.state`` so an + # exception inside the step bypasses this assignment and + # the previous value (or the empty box) survives. + context.latest_state_box[:] = [state] # Proposal 0043 (post-PR-99 review): restore the outer # ``current`` to the shared box after a successful step. diff --git a/src/openarmature/graph/observer.py b/src/openarmature/graph/observer.py index 0487dafd..ec8a170f 100644 --- a/src/openarmature/graph/observer.py +++ b/src/openarmature/graph/observer.py @@ -497,6 +497,23 @@ class _InvocationContext: # descents so the inner-most node's name wins on failure (the # real culprit, not the wrapper). final_node_box: list[str] = field(default_factory=list[str]) + # Per proposal 0043 (observability §8.4.1 *Resume semantics* + + # "partial final state captured at the failure point" clause). + # Tracks the most recent successful step's post-merge state at THIS + # context level so the outermost ``invoke()`` can populate + # ``InvocationCompletedEvent.final_state`` on the failure path with + # the partial outer state, not the bare ``starting_state``. On the + # success path the box is unused — the engine's return value is the + # canonical ``final_state``. **Distinct from ``final_node_box``**: + # the latest-state box is per-level (each subgraph / fan-out + # instance / parallel-branches branch gets its own fresh box), + # because the OUTER Langfuse trace cares about the outer-graph's + # state type, and an inner state has a different type. The + # ``final_node_box`` shares by reference because the spec wants the + # innermost failing node's name (the real culprit); state has the + # opposite contract — the outermost level's state is what the + # outer trace.output hook receives. + latest_state_box: list[Any] = field(default_factory=list[Any]) def full_observers(self) -> tuple[SubscribedObserver, ...]: """Return the ordered observer list to deliver for events from @@ -545,6 +562,9 @@ def descend_into_subgraph( drain_counters=self.drain_counters, state_cls=self.state_cls, final_node_box=self.final_node_box, + # latest_state_box is INTENTIONALLY NOT propagated — each + # context level tracks its own outer-state-typed latest + # successful step. See the field docstring above. ) def descend_into_fan_out_instance( @@ -596,6 +616,9 @@ def descend_into_fan_out_instance( drain_counters=self.drain_counters, state_cls=self.state_cls, final_node_box=self.final_node_box, + # latest_state_box is INTENTIONALLY NOT propagated — each + # context level tracks its own outer-state-typed latest + # successful step. See the field docstring above. ) def descend_into_parallel_branch( @@ -650,6 +673,9 @@ def descend_into_parallel_branch( drain_counters=self.drain_counters, state_cls=self.state_cls, final_node_box=self.final_node_box, + # latest_state_box is INTENTIONALLY NOT propagated — each + # context level tracks its own outer-state-typed latest + # successful step. See the field docstring above. ) def take_step(self) -> int: diff --git a/tests/conformance/test_observability_langfuse.py b/tests/conformance/test_observability_langfuse.py index b9359a68..cb2d86f5 100644 --- a/tests/conformance/test_observability_langfuse.py +++ b/tests/conformance/test_observability_langfuse.py @@ -15,6 +15,7 @@ from __future__ import annotations +import copy import json from collections.abc import Callable, Mapping, Sequence from datetime import UTC, datetime @@ -109,25 +110,16 @@ # ``(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. -# Used for proposal-0043 case 5 (resume re-fire) which needs harness -# extensions tracked separately — checkpointer wiring + flaky-node test -# seam + two-phase multi-trace assertion — landed in a follow-up PR. -_DEFERRED_CASES: frozenset[tuple[str, str]] = frozenset( - { - ( - "037-langfuse-trace-input-output", - "resume_hooks_refire_to_resumed_trace", - ), - } -) +# 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`` / # ``trace_output_from_state`` caller hooks. Each YAML hook name maps to # a Python callable matching the spec fixture's documented mock -# convention (see fixture 037's case 3 / case 4 inline comments). -# ``returns_state_snapshot`` is intentionally absent — only case 5 -# references it, and case 5 is deferred per ``_DEFERRED_CASES``. +# convention (see fixture 037's case 3 / case 4 / case 5 inline comments). def _returns_job_input_summary(_state: Any) -> dict[str, Any]: return {"summary": "job-input"} @@ -140,10 +132,20 @@ def _returns_null(_state: Any) -> None: return None +def _returns_state_snapshot(state: Any) -> dict[str, Any]: + # Fixture 037 case 5: the hook captures the state's full field set + # at hook-fire time. ``model_dump()`` returns the JSON-able + # representation; the case asserts the trace's input/output exactly + # match the values present at first-invoke entry / first-invoke + # failure-exit / resumed-invoke entry / resumed-invoke exit. + return cast("dict[str, Any]", state.model_dump()) + + _TRACE_IO_HOOK_REGISTRY: dict[str, Callable[[Any], Any]] = { "returns_job_input_summary": _returns_job_input_summary, "returns_job_output_summary": _returns_job_output_summary, "returns_null": _returns_null, + "returns_state_snapshot": _returns_state_snapshot, } @@ -665,6 +667,20 @@ async def _run_case(case: Mapping[str, Any]) -> None: target = END if target_raw == "END" else target_raw builder.add_edge(edge["from"], target) builder.set_entry(entry) + # Optional checkpointer wiring — fixture 037 case 5 needs an + # in-memory checkpointer so the first invoke's pre-failure save + # carries over to the resumed invoke. Only the literal value + # ``"in_memory"`` is recognized; other backends would need + # additional registration shimmed here. + checkpointer_spec = cast("str | None", case.get("checkpointer")) + if checkpointer_spec == "in_memory": + from openarmature.checkpoint import InMemoryCheckpointer # noqa: PLC0415 + + builder.with_checkpointer(InMemoryCheckpointer()) + elif checkpointer_spec is not None: + raise NotImplementedError( + f"langfuse harness only supports checkpointer: in_memory; got {checkpointer_spec!r}" + ) graph = builder.compile() # ``initial_state`` overrides on the case populate caller- # supplied fields; remaining fields fall back to the State @@ -713,6 +729,23 @@ async def _run_case(case: Mapping[str, Any]) -> None: caller_metadata = cast("dict[str, Any] | None", case.get("caller_metadata")) if caller_metadata is not None: invoke_kwargs["metadata"] = caller_metadata + + # Resume cases run a two-phase flow (first invoke catches expected + # error → resume invoke completes), then assert against both traces + # separately. Branch out here so the linear ``await graph.invoke`` + # below stays focused on the common case. + if "resume" in case: + await _run_resume_case( + case=case, + graph=graph, + initial_state_factory=initial_state_factory, + client=client, + invoke_kwargs=invoke_kwargs, + ) + if provider is not None: + await provider.aclose() + return + await graph.invoke(initial_state_factory(), **invoke_kwargs) await graph.drain() if provider is not None: @@ -733,6 +766,143 @@ async def _run_case(case: Mapping[str, Any]) -> None: _assert_trace(trace, expected_trace, expected_invariants=expected_invariants) +async def _run_resume_case( + *, + case: Mapping[str, Any], + graph: Any, + initial_state_factory: Callable[[], Any], + client: InMemoryLangfuseClient, + invoke_kwargs: dict[str, Any], +) -> None: + """Two-phase test flow for fixture 037 case 5. + + Phase 1 — first invoke catches the expected NodeException at the + designated node; the captured Langfuse Trace's input/output match + ``first_run_expected.langfuse_trace``. We snapshot the first trace's + headline fields immediately so the ``first_trace_unchanged`` invariant + can verify the resumed invoke leaves them untouched. + + Phase 2 — resume invoke runs the same graph with + ``resume_invocation=first_invocation_id``, completes successfully, and + the resumed Trace's input/output match ``resume.expected.langfuse_trace``. + + Phase 3 — invariants compare the two traces (distinct trace ids, + shared correlation_id, the snapshotted first trace's fields unchanged). + """ + from openarmature.graph.errors import RuntimeGraphError # noqa: PLC0415 + + # ---- Phase 1: first invoke catches expected error + first_run_expected_error = cast("dict[str, Any]", case.get("first_run_expected_error") or {}) + expected_category = cast("str", first_run_expected_error.get("category", "node_exception")) + expected_raised_from = cast("str | None", first_run_expected_error.get("raised_from")) + + # Catch the common ``RuntimeGraphError`` base so the harness handles + # any spec §4 category (node_exception / reducer_error / + # state_validation_error / edge_exception / routing_error). The + # "raised from" node attribute differs per category — check + # ``node_name`` on NodeException, ``producing_node`` on + # ReducerError, ``source_node`` on EdgeException / RoutingError — + # via a small attribute walk so we don't hardcode per-category + # accessor knowledge here. + try: + await graph.invoke(initial_state_factory(), **invoke_kwargs) + except RuntimeGraphError as exc: + assert exc.category == expected_category, ( + f"first run error category: expected {expected_category!r}, got {exc.category!r}" + ) + if expected_raised_from is not None: + actual_raised_from = ( + getattr(exc, "node_name", None) + or getattr(exc, "producing_node", None) + or getattr(exc, "source_node", None) + ) + assert actual_raised_from == expected_raised_from, ( + f"first run error raised_from: expected {expected_raised_from!r}, got {actual_raised_from!r}" + ) + else: + raise AssertionError( + f"first run expected to raise RuntimeGraphError with category={expected_category!r}; " + f"completed without error" + ) + await graph.drain() + + assert len(client.traces) == 1, ( + f"first run should produce exactly one Langfuse Trace; got {len(client.traces)}" + ) + first_invocation_id, first_trace = next(iter(client.traces.items())) + + # Snapshot the first trace's headline fields before the resume runs + # so the ``first_trace_unchanged`` invariant can compare against the + # state captured here. ``client.traces`` holds the live recorder + # objects; ``copy.deepcopy`` protects against in-place writes. + first_trace_snapshot = { + "input": copy.deepcopy(first_trace.input), + "output": copy.deepcopy(first_trace.output), + } + + first_run_expected = cast("dict[str, Any]", case["first_run_expected"]) + first_expected_trace = cast("dict[str, Any]", first_run_expected["langfuse_trace"]) + _assert_trace(first_trace, first_expected_trace, expected_invariants={}) + + # ---- Phase 2: resume invoke + resume_block = cast("dict[str, Any]", case["resume"]) + # Drop ``correlation_id`` from invoke_kwargs on resume — the engine + # restores it from the saved record per §3.1. + resume_invoke_kwargs = {k: v for k, v in invoke_kwargs.items() if k != "correlation_id"} + await graph.invoke( + initial_state_factory(), + resume_invocation=first_invocation_id, + **resume_invoke_kwargs, + ) + await graph.drain() + + # Python dicts are insertion-ordered (PEP 468; guaranteed since + # 3.7). Phase 1 added one trace; phase 2 added the resumed trace. + # Reading by position is more deterministic than scanning by + # not-equal — if a future engine change adds synthetic traces, the + # scan would silently pick the wrong key, but the position-based + # read fails the length assertion below explicitly. + trace_ids = list(client.traces.keys()) + assert len(trace_ids) == 2, ( + f"after resume there should be exactly two Langfuse Traces; got {len(trace_ids)}" + ) + assert trace_ids[0] == first_invocation_id, ( + f"first trace id changed during resume: was {first_invocation_id!r}, now {trace_ids[0]!r}" + ) + resumed_invocation_id = trace_ids[1] + resumed_trace = client.traces[resumed_invocation_id] + + resume_expected = cast("dict[str, Any]", resume_block["expected"]) + resume_expected_trace = cast("dict[str, Any]", resume_expected["langfuse_trace"]) + _assert_trace(resumed_trace, resume_expected_trace, expected_invariants={}) + + # ---- Phase 3: invariants + if resume_expected.get("first_trace_unchanged"): + assert first_trace.input == first_trace_snapshot["input"], ( + f"first_trace_unchanged failed: input was {first_trace_snapshot['input']!r}, " + f"now {first_trace.input!r}" + ) + assert first_trace.output == first_trace_snapshot["output"], ( + f"first_trace_unchanged failed: output was {first_trace_snapshot['output']!r}, " + f"now {first_trace.output!r}" + ) + + invariants = cast("dict[str, Any]", case.get("invariants") or {}) + if invariants.get("distinct_trace_ids"): + assert first_invocation_id != resumed_invocation_id, ( + f"distinct_trace_ids failed: both traces have id {first_invocation_id!r}" + ) + if invariants.get("correlation_id_consistent_across_traces"): + first_corr = first_trace.metadata.get("correlation_id") + resumed_corr = resumed_trace.metadata.get("correlation_id") + assert first_corr == resumed_corr, ( + f"correlation_id_consistent_across_traces failed: first={first_corr!r}, resumed={resumed_corr!r}" + ) + # ``hooks_refire_on_resumed_trace`` is implicit — verified by the + # ``_assert_trace`` call on the resumed trace above, which checks the + # hook-derived input/output match the resumed invocation's state. + + def _resolve_llm_model(case: Mapping[str, Any]) -> str: # Single LLM call per fixture today; pick up the per-call model if # supplied (fixture 023 explicitly sets `model: "test-model"` on @@ -801,6 +971,34 @@ async def _node_pure(_s: Any) -> dict[str, Any]: return _node_pure + # ``flaky: {fail_first_invocation_only: true, on_success: {...}}`` — + # the compact resume-fixture flaky shape (paralleling the equivalent + # form in tests/conformance/adapter.py:_make_flaky_fn). The node + # raises on its first call (a fresh ``RuntimeError`` the engine wraps + # as ``NodeException``) and returns ``on_success`` on subsequent + # calls. Used by fixture 037 case 5: the first invoke aborts at + # this node; the resumed invoke calls the same node body — the + # closure-scoped ``has_failed`` survives the resume because the + # graph (and the closure) lives for the harness's full case run, + # so the second call returns success. + flaky_spec = cast("dict[str, Any] | None", node_spec.get("flaky")) + if flaky_spec is not None: + if not flaky_spec.get("fail_first_invocation_only"): + raise NotImplementedError( + f"langfuse harness only supports the fail_first_invocation_only flaky shape; got {flaky_spec}" + ) + on_success = dict(cast("dict[str, Any]", flaky_spec.get("on_success") or {})) + has_failed = [False] + + async def _node_flaky(_s: Any) -> dict[str, Any]: + _maybe_augment() + if not has_failed[0]: + has_failed[0] = True + raise RuntimeError(f"flaky({node_name}) first-invocation failure") + return dict(on_success) + + return _node_flaky + calls_llm_spec = cast("dict[str, Any] | None", node_spec.get("calls_llm")) renders_prompt_name = cast("str | None", node_spec.get("renders_prompt")) diff --git a/tests/unit/test_observability_langfuse.py b/tests/unit/test_observability_langfuse.py index 6b11fe20..e4aa3d25 100644 --- a/tests/unit/test_observability_langfuse.py +++ b/tests/unit/test_observability_langfuse.py @@ -764,3 +764,338 @@ async def test_trace_output_status_failed_on_node_raise() -> None: trace = next(iter(client.traces.values())) assert trace.output == {"final_node": "raises", "status": "failed"} + + +class _PartialFailState(State): + a_ran: bool = False + b_ran: bool = False + + +async def _node_a_succeeds(_s: _PartialFailState) -> dict[str, Any]: + return {"a_ran": True} + + +async def _node_b_raises(_s: _PartialFailState) -> dict[str, Any]: + raise RuntimeError("node_b boom") + + +async def test_failure_path_final_state_is_state_at_failure_point() -> None: + # Spec §8.4.1 *Resume semantics* + the proposal-0043 "partial final + # state captured at the failure point" clause: a graph that + # completes node_a successfully then raises in node_b MUST surface + # the post-node-a state on the InvocationCompletedEvent so the + # ``trace_output_from_state`` hook (and the raw-state lever) see + # the partial state, not the bare initial state. Pins the engine + # fix that surfaces ``latest_state_box`` on the failure path. + + captured_output_state: list[_PartialFailState] = [] + + def output_hook(state: _PartialFailState) -> dict[str, Any]: + captured_output_state.append(state) + return {"a_ran": state.a_ran, "b_ran": state.b_ran} + + client = InMemoryLangfuseClient() + observer = LangfuseObserver(client=client, trace_output_from_state=output_hook) + graph = ( + GraphBuilder(_PartialFailState) + .add_node("node_a", _node_a_succeeds) + .add_node("node_b", _node_b_raises) + .add_edge("node_a", "node_b") + .add_edge("node_b", END) + .set_entry("node_a") + .compile() + ) + graph.attach_observer(observer) + + from openarmature.graph.errors import NodeException + + with pytest.raises(NodeException, match="node_b"): + await graph.invoke(_PartialFailState()) + await graph.drain() + + # The output hook fired with the post-node-a state (a_ran=True), + # not the initial state (a_ran=False). + assert len(captured_output_state) == 1 + assert captured_output_state[0].a_ran is True + assert captured_output_state[0].b_ran is False + trace = next(iter(client.traces.values())) + assert trace.output == {"a_ran": True, "b_ran": False} + + +class _OuterFailState(State): + outer_a_done: bool = False + sub_done: bool = False + + +class _InnerFailState(State): + inner_x_done: bool = False + + +async def _outer_node_a(_s: _OuterFailState) -> dict[str, Any]: + return {"outer_a_done": True} + + +async def _inner_node_x_succeeds(_s: _InnerFailState) -> dict[str, Any]: + return {"inner_x_done": True} + + +async def _inner_node_y_raises(_s: _InnerFailState) -> dict[str, Any]: + raise RuntimeError("inner_node_y boom") + + +async def test_failure_path_final_state_is_outer_type_when_subgraph_raises() -> None: + # Engine-bug regression: an inner-subgraph step's success previously + # overwrote the outermost ``latest_state_box`` (it was shared by + # reference across subgraph descents), so a subgraph-internal raise + # would leave the box holding an INNER state at outer ``invoke()`` + # finally time. The outer ``trace_output_from_state`` hook would + # then receive an inner-typed state when its signature expects the + # outer type — a real correctness bug. + # + # The box is now per-context: each subgraph descent gets its own + # fresh ``latest_state_box``, so the outermost level's box holds + # only outer-state-typed entries. This test exercises a graph + # where outer node_a succeeds (outer state = a_done=true), the + # subgraph step raises inside, and the outer trace.output hook + # receives the outer state with ``outer_a_done=True``, + # ``sub_done=False``. + from openarmature.graph import ExplicitMapping + + inner_graph = ( + GraphBuilder(_InnerFailState) + .add_node("inner_x", _inner_node_x_succeeds) + .add_node("inner_y", _inner_node_y_raises) + .add_edge("inner_x", "inner_y") + .add_edge("inner_y", END) + .set_entry("inner_x") + .compile() + ) + + captured_output_state: list[Any] = [] + + def output_hook(state: Any) -> dict[str, Any]: + captured_output_state.append(state) + return {"outer_a_done": state.outer_a_done, "sub_done": state.sub_done} + + client = InMemoryLangfuseClient() + observer = LangfuseObserver(client=client, trace_output_from_state=output_hook) + graph = ( + GraphBuilder(_OuterFailState) + .add_node("outer_a", _outer_node_a) + .add_subgraph_node( + "sub", + inner_graph, + projection=ExplicitMapping(inputs=None, outputs={"sub_done": "inner_x_done"}), + ) + .add_edge("outer_a", "sub") + .add_edge("sub", END) + .set_entry("outer_a") + .compile() + ) + graph.attach_observer(observer) + + from openarmature.graph.errors import NodeException + + with pytest.raises(NodeException): + await graph.invoke(_OuterFailState()) + await graph.drain() + + # The hook receives the OUTER state (with outer_a_done=True, + # sub_done=False), not the inner state — confirming the box's + # per-level isolation worked. + assert len(captured_output_state) == 1 + assert isinstance(captured_output_state[0], _OuterFailState) + assert not isinstance(captured_output_state[0], _InnerFailState) + assert captured_output_state[0].outer_a_done is True + assert captured_output_state[0].sub_done is False + trace = next(iter(client.traces.values())) + assert trace.output == {"outer_a_done": True, "sub_done": False} + + +# --------------------------------------------------------------------------- +# Per-context box isolation across fan-out + parallel-branches descents +# --------------------------------------------------------------------------- + + +class _FanOutOuterState(State): + outer_a_done: bool = False + items: list[int] = [] + results: Annotated[list[int], append] = [] + + +class _FanOutInnerState(State): + item: int = 0 + out: int = 0 + + +async def _fan_out_inner_raises(_s: _FanOutInnerState) -> dict[str, Any]: + raise RuntimeError("fan_out inner_node boom") + + +async def test_failure_path_final_state_is_outer_type_when_fan_out_inner_raises() -> None: + # Sibling to the subgraph-raise test: pins the per-context + # ``latest_state_box`` isolation across a fan-out instance descent. + # Each fan-out instance gets its own ``_InvocationContext`` + # (descend_into_fan_out_instance), so its inner step writes land on + # the instance's own box, not the outer box. When the instance + # raises, the outermost ``invoke()``'s finally-block reads the + # OUTER box — which holds outer state from ``outer_a``'s successful + # completion, not the inner instance state. + inner_graph = ( + GraphBuilder(_FanOutInnerState) + .add_node("inner_raise", _fan_out_inner_raises) + .add_edge("inner_raise", END) + .set_entry("inner_raise") + .compile() + ) + + async def _outer_a(_s: _FanOutOuterState) -> dict[str, Any]: + return {"outer_a_done": True} + + captured_output_state: list[Any] = [] + + def output_hook(state: Any) -> dict[str, Any]: + captured_output_state.append(state) + return {"outer_a_done": state.outer_a_done, "results": list(state.results)} + + client = InMemoryLangfuseClient() + observer = LangfuseObserver(client=client, trace_output_from_state=output_hook) + graph = ( + GraphBuilder(_FanOutOuterState) + .add_node("outer_a", _outer_a) + .add_fan_out_node( + "fan", + subgraph=inner_graph, + collect_field="out", + target_field="results", + items_field="items", + item_field="item", + ) + .add_edge("outer_a", "fan") + .add_edge("fan", END) + .set_entry("outer_a") + .compile() + ) + graph.attach_observer(observer) + + from openarmature.graph.errors import RuntimeGraphError + + with pytest.raises(RuntimeGraphError): + # Three fan-out instances all fail; the engine raises after the + # fan-out node completes (fail_fast default). + await graph.invoke(_FanOutOuterState(items=[1, 2, 3])) + await graph.drain() + + # The hook receives the OUTER state (FanOutOuterState), not an + # inner FanOutInnerState from the failed instance descent. + assert len(captured_output_state) == 1 + assert isinstance(captured_output_state[0], _FanOutOuterState) + assert not isinstance(captured_output_state[0], _FanOutInnerState) + assert captured_output_state[0].outer_a_done is True + # No instance succeeded, so results stays empty. + assert list(captured_output_state[0].results) == [] + + +class _ParBrOuterState(State): + outer_a_done: bool = False + branch_x_done: bool = False + branch_y_done: bool = False + + +class _ParBrBranchXState(State): + x_done: bool = False + + +class _ParBrBranchYState(State): + y_done: bool = False + + +async def _par_br_branch_x_succeeds(_s: _ParBrBranchXState) -> dict[str, Any]: + return {"x_done": True} + + +async def _par_br_branch_y_raises(_s: _ParBrBranchYState) -> dict[str, Any]: + raise RuntimeError("parallel_branches branch_y boom") + + +async def test_failure_path_final_state_is_outer_type_when_parallel_branch_raises() -> None: + # Sibling to the subgraph + fan-out tests: pins per-context + # ``latest_state_box`` isolation across a parallel-branches + # descent. Each branch's inner _invoke runs in its own + # ``_InvocationContext`` (descend_into_parallel_branch), so inner + # writes don't leak to the outer box. Even when branch_x writes + # its inner state successfully, the outermost finally-block reads + # the OUTER box on the branch_y-induced raise. + from openarmature.graph import BranchSpec + + branch_x_subgraph = ( + GraphBuilder(_ParBrBranchXState) + .add_node("succeeds", _par_br_branch_x_succeeds) + .add_edge("succeeds", END) + .set_entry("succeeds") + .compile() + ) + + branch_y_subgraph = ( + GraphBuilder(_ParBrBranchYState) + .add_node("raises", _par_br_branch_y_raises) + .add_edge("raises", END) + .set_entry("raises") + .compile() + ) + + async def _outer_a(_s: _ParBrOuterState) -> dict[str, Any]: + return {"outer_a_done": True} + + captured_output_state: list[Any] = [] + + def output_hook(state: Any) -> dict[str, Any]: + captured_output_state.append(state) + return { + "outer_a_done": state.outer_a_done, + "branch_x_done": state.branch_x_done, + "branch_y_done": state.branch_y_done, + } + + client = InMemoryLangfuseClient() + observer = LangfuseObserver(client=client, trace_output_from_state=output_hook) + graph = ( + GraphBuilder(_ParBrOuterState) + .add_node("outer_a", _outer_a) + .add_parallel_branches_node( + "dispatch", + branches={ + "branch_x": BranchSpec( + subgraph=branch_x_subgraph, + outputs={"branch_x_done": "x_done"}, + ), + "branch_y": BranchSpec( + subgraph=branch_y_subgraph, + outputs={"branch_y_done": "y_done"}, + ), + }, + ) + .add_edge("outer_a", "dispatch") + .add_edge("dispatch", END) + .set_entry("outer_a") + .compile() + ) + graph.attach_observer(observer) + + from openarmature.graph.errors import RuntimeGraphError + + with pytest.raises(RuntimeGraphError): + await graph.invoke(_ParBrOuterState()) + await graph.drain() + + # The hook receives the OUTER state (ParBrOuterState). Whether + # branch_x's success projected back into the outer state by the + # time of the raise depends on the dispatch's join semantics; + # what MUST be true is that the captured state is the OUTER + # type, not branch_x's _ParBrBranchXState or branch_y's + # _ParBrBranchYState. + assert len(captured_output_state) == 1 + assert isinstance(captured_output_state[0], _ParBrOuterState) + assert not isinstance(captured_output_state[0], _ParBrBranchXState) + assert not isinstance(captured_output_state[0], _ParBrBranchYState) + assert captured_output_state[0].outer_a_done is True From 27e3782c8205576569ac38d694ae86b9a12f0b24 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sat, 30 May 2026 21:05:09 -0700 Subject: [PATCH 2/2] Tighten fan-out regression + fix CHANGELOG count MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #102 review caught two issues: The fan-out regression test's inner subgraph contained only a raising node, so under the original shared-`latest_state_box` bug no inner step would have successfully written to the box — the test would have passed without exercising the leak it was meant to guard. The inner subgraph now has two nodes: `inner_succeeds` writes `inner_done=true` (so the descent's _invoke writes inner state to the box) followed by `inner_raises`. Confirmed by temp-reverting the descend-omit-`latest_state_box` change and observing the test fail with the typed-state-mismatch assertion. CHANGELOG said "three regression tests" but enumerated four (flat, subgraph, fan-out, parallel-branches). Bumped the count to four. --- CHANGELOG.md | 2 +- tests/unit/test_observability_langfuse.py | 30 ++++++++++++++++++++--- 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 63ceb1c1..6bca8780 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,7 +28,7 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). The ### Fixed - **`InvocationCompletedEvent.final_state` on the failure path now surfaces the partial state at failure point.** Spec §8.4.1 *Resume semantics* requires the failure-path `trace.output` hook to receive "the partial final state captured at the failure point"; the original PR #99 implementation defaulted to `starting_state`, so the hook saw pre-execution state when it should have seen post-execution-up-to-failure state. The engine now tracks the latest post-merge state via a `latest_state_box` on `_InvocationContext`, updated after every successful step and read on the failure path. Success-path behavior unchanged. -- **`latest_state_box` is per-context, not shared across subgraph descents.** Unlike the sibling `final_node_box` (which shares by reference because the spec wants the innermost failing node's name — the real culprit), `latest_state_box` must isolate per level so the outermost Langfuse trace receives outer-state-typed values. Without the isolation, a subgraph-internal step's inner-typed state would leak up to the outer trace.output hook, breaking the hook's typed contract. Each subgraph / fan-out instance / parallel-branches branch gets its own fresh box. Pinned by three regression tests covering flat, subgraph, fan-out, and parallel-branches failure paths. +- **`latest_state_box` is per-context, not shared across subgraph descents.** Unlike the sibling `final_node_box` (which shares by reference because the spec wants the innermost failing node's name — the real culprit), `latest_state_box` must isolate per level so the outermost Langfuse trace receives outer-state-typed values. Without the isolation, a subgraph-internal step's inner-typed state would leak up to the outer trace.output hook, breaking the hook's typed contract. Each subgraph / fan-out instance / parallel-branches branch gets its own fresh box. Pinned by four regression tests covering flat, subgraph, fan-out, and parallel-branches failure paths. ### Notes diff --git a/tests/unit/test_observability_langfuse.py b/tests/unit/test_observability_langfuse.py index e4aa3d25..3f722b4a 100644 --- a/tests/unit/test_observability_langfuse.py +++ b/tests/unit/test_observability_langfuse.py @@ -926,10 +926,21 @@ class _FanOutOuterState(State): class _FanOutInnerState(State): item: int = 0 out: int = 0 + inner_done: bool = False + + +async def _fan_out_inner_succeeds(_s: _FanOutInnerState) -> dict[str, Any]: + # Successful inner step — writes ``inner_done=true`` to the + # instance's _invoke ``state`` local AND to the shared + # ``latest_state_box`` (per-context, so it lands on the instance's + # OWN box). Under the original shared-box bug this write would + # leak into the outer box; under the per-context design it stays + # isolated to the instance. + return {"inner_done": True} async def _fan_out_inner_raises(_s: _FanOutInnerState) -> dict[str, Any]: - raise RuntimeError("fan_out inner_node boom") + raise RuntimeError("fan_out inner_raise boom") async def test_failure_path_final_state_is_outer_type_when_fan_out_inner_raises() -> None: @@ -941,11 +952,22 @@ async def test_failure_path_final_state_is_outer_type_when_fan_out_inner_raises( # raises, the outermost ``invoke()``'s finally-block reads the # OUTER box — which holds outer state from ``outer_a``'s successful # completion, not the inner instance state. + # + # The inner subgraph has TWO inner nodes: ``inner_succeeds`` writes + # inner state to the instance's box, then ``inner_raises`` + # propagates. Under the original shared-box bug, the box would + # end with ``_FanOutInnerState(inner_done=true)`` and the outer + # hook would receive that inner-typed value. The two-node shape + # is load-bearing — a single-node "always raise" subgraph would + # not exercise the leak because no successful inner step would + # write to the box. inner_graph = ( GraphBuilder(_FanOutInnerState) - .add_node("inner_raise", _fan_out_inner_raises) - .add_edge("inner_raise", END) - .set_entry("inner_raise") + .add_node("inner_succeeds", _fan_out_inner_succeeds) + .add_node("inner_raises", _fan_out_inner_raises) + .add_edge("inner_succeeds", "inner_raises") + .add_edge("inner_raises", END) + .set_entry("inner_succeeds") .compile() )