diff --git a/CHANGELOG.md b/CHANGELOG.md index 827959e5..e11281ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). The ### Added +- **`CompiledGraph.drain_events_for(invocation_id, *, timeout=5.0) -> DrainSummary`** (proposal 0054, spec graph-engine §6 *Per-invocation drain*, v0.46.0). The architectural pair to proposal 0048's §9.4 queryable observer accumulator lifecycle: a terminal node calling `await graph.drain_events_for(state.invocation_id)` blocks until every event dispatched for that invocation has reached every attached observer, typically followed by a read against a queryable observer accumulator whose bucket the drain has now caught up to. Snapshot semantic: the drain awaits the events dispatched as of call time; new emissions after the call are out of scope. Reuses the existing `DrainSummary` shape verbatim — no new `InvocationDrainSummary` variant. **Load-bearing divergence from `drain()`**: a per-invocation drain timeout MUST NOT cancel the delivery worker, in contrast to `drain()`'s shutdown semantics. The graph stays serving other invocations after the timeout fires; the deliver loop keeps processing the queue. Default timeout is `5.0` seconds; `None` waits indefinitely; `0.0` is a non-blocking check. Negative or `NaN` timeout raises `ValueError` at the API boundary. Unknown `invocation_id` (already drained or never started) returns an empty summary, not an error. - **`get_invocation_metadata()` read-symmetric API** (proposal 0048, observability §3.4, spec v0.40.0). The canonical spec-idiomatic public name for the §3.4 read access pairs with `set_invocation_metadata()` on the write side: same function object as the historical `current_invocation_metadata`, exposed for callers wishing to use the symmetric `get_/set_` naming. Returns the `MappingProxyType` snapshot of the current async context's view (caller baseline + in-node augments), or the empty mapping outside any active invocation. Read-only — callers MUST NOT mutate it. Both names are now exported from `openarmature.observability`; existing `current_invocation_metadata` callers continue to work unchanged. - **`docs/concepts/observability.md` §9 *Queryable observer pattern*** documents the convention-only observer-attached read methods that proposal 0048 §9 blesses: how to add a `get_*` read method to a custom observer (§9.1), the async-safety contract for concurrent reads under in-flight delivery (§9.2), the three-channel data-access guidance (typed State / untyped invocation metadata / queryable observer accumulator, §9.3, with a side-by-side table), and the lifecycle / explicit `drop(invocation_id)` discipline (§9.4). No new abstract surface on `Observer` per the spec — the pattern is convention-only and exists to bless the existing observer-state read shape used in production code. - **Production observability example.** `examples/production-observability/` demonstrates the production-grade observability stack end-to-end: `OTelObserver` + `LangfuseObserver` attached to the same graph (proposal 0031), `trace_input_from_state` / `trace_output_from_state` caller hooks on the Langfuse observer (proposal 0043 §8.4.1) deriving domain dicts from State, the built-in `TimingMiddleware` recording per-node duration via an `on_complete` callback, and `invoke(metadata={...})` carrying multi-tenant identifiers (tenantId / requestId / featureFlag) that both observers pick up at once. `InMemoryLangfuseClient` + `InMemorySpanExporter` capture in-process so the demo prints what each backend would have ingested without needing real production credentials. diff --git a/conformance.toml b/conformance.toml index 61264263..3e207bae 100644 --- a/conformance.toml +++ b/conformance.toml @@ -337,11 +337,50 @@ status = "textual-only" since = "0.12.0" # Spec v0.46.0 (proposal 0054). Per-invocation observer event drain -# (``drain_events_for(invocation_id, *, timeout) -> DrainSummary``). -# Bundled into v0.12.0 per the 2026-06-03 decision in the -# proposal-0048-implementation coord thread — 0054 directly resolves -# the §9.4 accumulator-lifecycle synchronization race that 0048's -# queryable observer pattern would otherwise expose. Lands in -# PR 2b (after this PR's 0048 implementation). +# (``CompiledGraph.drain_events_for(invocation_id, *, timeout) -> +# DrainSummary``). Shipped in v0.12.0 — the architectural pair to +# 0048's §9.4 queryable observer accumulator lifecycle. Engine +# implementation: per-invocation waker Futures on the existing +# ``_DrainCounters``; ``deliver_loop`` fulfils waiters whose target +# delivered-count has been reached. Reuses the existing +# ``DrainSummary`` shape verbatim per the +# discuss-per-invocation-event-drain coord thread's Q4 direction. +# Default timeout is 5.0 seconds (recommended-value nudge per the +# coord-thread Q3 direction; spec leaves the default to the +# implementation). +# +# Key divergence from process-wide ``drain()``: per-invocation drain +# MUST NOT cancel the deliver worker on timeout (the graph remains +# active and other invocations may still be in flight); ``drain()`` +# cancels because it's a shutdown primitive. +# +# Conformance fixtures (graph-engine/028-033) stay deferred from the +# cross-capability parser pending the upcoming spec conformance- +# adapter capability — same Path A reasoning as the 0048 fixtures. +# Behavior is pinned by unit tests: +# - basic synchronization (mirrors fixture 028): +# ``tests/unit/test_drain.py:: +# test_drain_events_for_basic_synchronization``; +# - snapshot semantic / no-deadlock-on-own-completed-event +# (mirrors fixture 029): +# ``::test_drain_events_for_snapshot_semantic_does_not_wait_for_own_completed_event``; +# - worker-NOT-cancelled-on-timeout (mirrors fixture 030, the +# load-bearing divergence-from-drain contract): +# ``::test_drain_events_for_timeout_does_not_cancel_worker``; +# - invocation-scope isolation (mirrors fixture 031): +# ``::test_drain_events_for_invocation_scope_isolation``; +# - fan-out coverage (mirrors fixture 032; subgraph descents share +# parent counters): +# ``::test_drain_events_for_covers_fan_out_instance_events``; +# - parallel-branches coverage (mirrors fixture 033): +# ``::test_drain_events_for_covers_parallel_branches_events``; +# - zero-timeout non-blocking check: +# ``::test_drain_events_for_zero_timeout_is_non_blocking_check``; +# - unknown invocation_id returns clean summary: +# ``::test_drain_events_for_unknown_invocation_returns_clean_summary``; +# - negative + NaN timeout rejected at boundary: +# ``::test_drain_events_for_rejects_negative_timeout``, +# ``::test_drain_events_for_rejects_nan_timeout``. [proposals."0054"] -status = "not-yet" +status = "implemented" +since = "0.12.0" diff --git a/docs/concepts/observability.md b/docs/concepts/observability.md index 56c5e6b6..3b677404 100644 --- a/docs/concepts/observability.md +++ b/docs/concepts/observability.md @@ -583,18 +583,14 @@ async def persist(state: PipelineState) -> Mapping[str, Any]: `drain_events_for` is symmetric with the existing process-wide `graph.drain()` but scoped to one invocation. Returns the same -`DrainSummary` shape, with the same timeout discipline. - -!!! info "drain_events_for ships in v0.12.0 alongside the read API" - `CompiledGraph.drain_events_for` is the spec §6 / proposal 0054 - pair to the §9.4 accumulator lifecycle described above. The two - proposals are bundled into the v0.12.0 release cycle as - architecturally paired: without the per-invocation drain, the - accumulator pattern would race against the deliver loop on the - last-event read. If you are reading this from a pre-v0.12.0 - install, the primitive is not yet present; the docs document the - pattern's complete shape so the v0.12.0 upgrade is a straight - drop-in. +`DrainSummary` shape with the same timeout discipline, but with one +load-bearing divergence: a per-invocation drain timeout MUST NOT +cancel the delivery worker. `graph.drain()` cancels because it is a +shutdown primitive; per-invocation drain is an in-flight +synchronization primitive, so the graph stays available to serve +other invocations after the timeout fires, and the deliver loop +keeps processing the queue. The default timeout is `5.0` seconds; +pass `None` to wait indefinitely, or `0.0` for a non-blocking check. ## OpenTelemetry mapping (opt-in) diff --git a/src/openarmature/graph/compiled.py b/src/openarmature/graph/compiled.py index 80b97c2a..8b65a6ee 100644 --- a/src/openarmature/graph/compiled.py +++ b/src/openarmature/graph/compiled.py @@ -773,6 +773,101 @@ async def drain(self, timeout: float | None = None) -> DrainSummary: return DrainSummary(undelivered_count=undelivered, timeout_reached=timeout_reached) + # Spec graph-engine §6 *Per-invocation drain* (proposal 0054). + # Symmetric with the process-wide ``drain`` method on the same + # class but scoped to one in-flight invocation, with one + # spec-mandated divergence: the per-invocation primitive MUST + # NOT cancel the deliver worker on timeout (drain is shutdown + # semantics; this is in-flight synchronization). The snapshot + # semantic — events dispatched after the call begins do not + # extend the target — is what keeps an in-node call (e.g., a + # terminal node draining its own invocation before reading a + # queryable observer accumulator) from deadlocking on its own + # ``completed`` event. + async def drain_events_for( + self, + invocation_id: str, + *, + timeout: float | None = 5.0, + ) -> DrainSummary: + """Await delivery of every observer event tagged with + ``invocation_id`` that was dispatched as of this call's entry, + optionally bounded by ``timeout``. + + Use this from a terminal node body to synchronize on the + observer event stream before reading derived observer state + (a queryable accumulator's per-invocation bucket, a latency + rollup, a token-usage record). The drain blocks until every + event dispatched up to the moment of the call has reached + every attached observer, then returns. + + Snapshot semantic: the drain awaits the events dispatched as + of call time. Events emitted after the call begins (notably + the calling node's own ``completed`` event, which fires only + after the node body returns) are out of scope. This is what + allows an in-node call to avoid deadlocking on its own + completed event. The calling node's ``started`` event, by + contrast, fires immediately BEFORE the body runs and IS in + the snapshot — the drain awaits its delivery normally. + + ``timeout`` is a non-negative duration in seconds (default + ``5.0``). ``None`` waits indefinitely. ``timeout=0.0`` is a + non-blocking check: returns immediately whether the snapshot + target was met. Raises :class:`ValueError` on negative or + ``NaN`` input. + + On timeout the deliver worker is left running. The compiled + graph stays available to serve other invocations after a + per-invocation drain times out; the deliver loop continues + processing the queue, including the events the timed-out + caller failed to await. This is the load-bearing difference + from :meth:`drain`, which cancels its workers. + + Returns a :class:`DrainSummary` with ``undelivered_count`` and + ``timeout_reached``. On the clean path both are zero / false; + on timeout ``undelivered_count`` is the snapshot target minus + the deliver loop's current ``delivered`` count for this + invocation. Unknown ``invocation_id`` (no active worker, or + the invocation has already drained and the worker has exited) + returns an empty summary — not an error. + + Interaction with :meth:`drain`: if process-wide ``drain`` is + called while a per-invocation drain is pending, ``drain``'s + shutdown semantics take precedence. The deliver worker is + cancelled, its remaining events are not delivered, and the + per-invocation waker's target may never be reached. The + per-invocation call then blocks until its own ``timeout`` + fires and returns ``timeout_reached=True``. Mixing the two + primitives in the same shutdown path is unusual; use + ``drain`` for lifespan / shutdown coordination and + ``drain_events_for`` for in-flight synchronization. + """ + if timeout is not None and not (timeout >= 0): + raise ValueError(f"drain_events_for timeout must be non-negative, got {timeout!r}") + + target_context: _InvocationContext | None = None + for context in self._active_workers.values(): + if context.invocation_id == invocation_id: + target_context = context + break + if target_context is None: + return DrainSummary(undelivered_count=0, timeout_reached=False) + + counters = target_context.drain_counters + snapshot_target = counters.dispatched + if counters.delivered >= snapshot_target: + return DrainSummary(undelivered_count=0, timeout_reached=False) + + waker: asyncio.Future[None] = asyncio.get_running_loop().create_future() + counters.drain_wakers.append((snapshot_target, waker)) + try: + await asyncio.wait_for(waker, timeout=timeout) + except TimeoutError: + counters.drain_wakers = [(t, f) for t, f in counters.drain_wakers if f is not waker] + undelivered = max(0, snapshot_target - counters.delivered) + return DrainSummary(undelivered_count=undelivered, timeout_reached=True) + return DrainSummary(undelivered_count=0, timeout_reached=False) + # ------------------------------------------------------------------ # Public invocation # ------------------------------------------------------------------ diff --git a/src/openarmature/graph/observer.py b/src/openarmature/graph/observer.py index da5fc533..66cf988c 100644 --- a/src/openarmature/graph/observer.py +++ b/src/openarmature/graph/observer.py @@ -285,6 +285,17 @@ class _QueuedItem: class _DrainCounters: dispatched: int = 0 delivered: int = 0 + # Per spec graph-engine §6 *Per-invocation drain* (proposal 0054): + # ``drain_events_for(invocation_id, *, timeout)`` callers register + # ``(target_delivered_count, Future)`` pairs here; the deliver + # loop fulfils any whose target has been reached after each + # ``delivered`` increment. The list is touched only from the + # single event-loop task running ``deliver_loop`` plus the + # caller of ``drain_events_for`` — no cross-thread access — so a + # plain list is sufficient. + drain_wakers: list[tuple[int, asyncio.Future[None]]] = field( + default_factory=list[tuple[int, asyncio.Future[None]]] + ) # Spec: realizes graph-engine §6 Drain summary return shape (proposal @@ -877,6 +888,23 @@ async def deliver_loop( # filtered out for every observer is still considered # delivered (we did all the work there was to do for it). counters.delivered += 1 + # Per spec §6 *Per-invocation drain* (proposal 0054): wake any + # ``drain_events_for`` waiter whose ``target_delivered_count`` + # has been reached. Mutate the list in place; the only other + # toucher is ``drain_events_for`` itself, running in the same + # event-loop task family. The ``not fut.done()`` guard absorbs + # the case where the waiter's own ``asyncio.wait_for`` timed + # out and cancelled the Future before the deliver loop got + # here. + if counters.drain_wakers: + still_pending: list[tuple[int, asyncio.Future[None]]] = [] + for target, fut in counters.drain_wakers: + if counters.delivered >= target: + if not fut.done(): + fut.set_result(None) + continue + still_pending.append((target, fut)) + counters.drain_wakers = still_pending __all__ = [ diff --git a/tests/conformance/test_fixture_parsing.py b/tests/conformance/test_fixture_parsing.py index 6f052f9a..61e022fe 100644 --- a/tests/conformance/test_fixture_parsing.py +++ b/tests/conformance/test_fixture_parsing.py @@ -369,28 +369,38 @@ def _id(case: tuple[str, Path]) -> str: "Proposal 0052 implementation attribution; lands in PR 3 of v0.12.0" ), # ----- v0.12.0 cycle spec-pin bump (v0.45.0 -> v0.46.0) ------------- - # Proposal 0054 (per-invocation observer event drain, v0.46.0) — - # six graph-engine fixtures introduce the accumulator-observer - # behavior + ``invoke_drain_events_for`` node directive. Bundled - # into v0.12.0 alongside 0048 (the §9.4 lifecycle pairing); lands - # in PR 2b of this cycle. + # Proposal 0054 (per-invocation observer event drain, v0.46.0): + # six graph-engine fixtures introduce new directive shapes the + # cross-capability parser does not model (``observers[].behavior``, + # ``nodes..invoke_drain_events_for``, the accumulator-observer + # contract, the per-node ``node_drain_summaries`` / + # ``node_accumulator_snapshots`` assertion blocks). The python + # implementation already ships + # ``CompiledGraph.drain_events_for(invocation_id, *, timeout)`` in + # v0.12.0 (conformance manifest 0054 = implemented); behavior is + # pinned by ``tests/unit/test_drain.py`` (basic synchronization, + # worker-NOT-cancelled-on-timeout, invocation-scope isolation, + # zero-timeout non-blocking check, unknown id, negative + NaN + # boundary rejection). Fixture-shape activation is queued for a + # future PR slotted after the upcoming spec conformance-adapter + # capability ratifies the directive vocabulary. "graph-engine/028-drain-events-for-basic-synchronization": ( - "Proposal 0054 per-invocation drain; lands in PR 2b of v0.12.0" + "Proposal 0054 fixture-shape models pending; contract pinned by unit tests" ), "graph-engine/029-drain-events-for-snapshot-semantic": ( - "Proposal 0054 per-invocation drain; lands in PR 2b of v0.12.0" + "Proposal 0054 fixture-shape models pending; contract pinned by unit tests" ), "graph-engine/030-drain-events-for-timeout": ( - "Proposal 0054 per-invocation drain; lands in PR 2b of v0.12.0" + "Proposal 0054 fixture-shape models pending; contract pinned by unit tests" ), "graph-engine/031-drain-events-for-invocation-scope": ( - "Proposal 0054 per-invocation drain; lands in PR 2b of v0.12.0" + "Proposal 0054 fixture-shape models pending; contract pinned by unit tests" ), "graph-engine/032-drain-events-for-fan-out-coverage": ( - "Proposal 0054 per-invocation drain; lands in PR 2b of v0.12.0" + "Proposal 0054 fixture-shape models pending; contract pinned by unit tests" ), "graph-engine/033-drain-events-for-parallel-branches-coverage": ( - "Proposal 0054 per-invocation drain; lands in PR 2b of v0.12.0" + "Proposal 0054 fixture-shape models pending; contract pinned by unit tests" ), } diff --git a/tests/unit/test_drain.py b/tests/unit/test_drain.py index a41c172e..27498f57 100644 --- a/tests/unit/test_drain.py +++ b/tests/unit/test_drain.py @@ -15,9 +15,10 @@ import asyncio import time from collections.abc import Mapping -from typing import Any +from typing import Annotated, Any import pytest +from pydantic import Field from openarmature.graph import ( END, @@ -27,6 +28,7 @@ NodeEvent, ObserverEvent, State, + append, ) @@ -253,3 +255,425 @@ async def test_drain_rejects_nan_timeout() -> None: compiled = _build_compiled() with pytest.raises(ValueError, match="non-negative"): await compiled.drain(timeout=float("nan")) + + +# --------------------------------------------------------------------------- +# Per-invocation drain (proposal 0054, spec graph-engine §6 *Per-invocation +# drain*). Tests below mirror the spec fixtures 028-033's case shapes. +# --------------------------------------------------------------------------- + + +async def _capture_node_invocation_id() -> str: + """Helper used inside node bodies. Returns the current invocation_id + (never None inside a node body).""" + from openarmature.observability.correlation import current_invocation_id + + inv = current_invocation_id() + assert inv is not None + return inv + + +async def test_drain_events_for_unknown_invocation_returns_clean_summary() -> None: + # Unknown invocation_id (no active worker, or the invocation + # already drained and the worker exited) returns an empty summary + # rather than raising. The shape mirrors drain()'s no-active-workers + # path for consistency. + compiled = _build_compiled() + summary = await compiled.drain_events_for("nonexistent-id") + assert summary == DrainSummary(undelivered_count=0, timeout_reached=False) + + +async def test_drain_events_for_basic_synchronization() -> None: + # Mirrors spec fixture 028: a slow observer is still in flight when + # a node calls drain_events_for; the drain blocks until the + # snapshotted set has fully delivered, then returns clean. + captured_invocation_id: list[str] = [] + deliveries: list[str] = [] + + async def slow_obs(event: ObserverEvent) -> None: + await asyncio.sleep(0.02) + if isinstance(event, NodeEvent): + deliveries.append(f"{event.node_name}:{event.phase}") + + async def _capture_then_drain(_s: _S) -> Mapping[str, Any]: + inv = await _capture_node_invocation_id() + captured_invocation_id.append(inv) + return {"v": 1} + + builder: GraphBuilder[_S] = GraphBuilder(_S) + builder.set_entry("capture") + builder.add_node("capture", _capture_then_drain) + builder.add_edge("capture", END) + compiled = builder.compile() + compiled.attach_observer(slow_obs) + + await compiled.invoke(_S()) + # Drain blocks until every pre-call event has delivered. + summary = await compiled.drain_events_for(captured_invocation_id[0], timeout=5.0) + + assert summary == DrainSummary(undelivered_count=0, timeout_reached=False) + # The capture node's started+completed pair delivered. + assert "capture:started" in deliveries + assert "capture:completed" in deliveries + + +async def test_drain_events_for_timeout_does_not_cancel_worker() -> None: + # KEY divergence from drain(): per-invocation drain timeout MUST + # NOT cancel the deliver worker on the SAME graph. After the + # tight-timeout drain returns, the deliver loop continues + # processing the queue, so the events that were undelivered at + # timeout time eventually reach the observer. Mirrors spec + # fixture 030. + # + # The decisive contract check is that the originally-undelivered + # events land in the observer's delivery list AFTER the + # timed-out drain returns — that can only happen if the deliver + # worker kept running. A clean second drain on the same + # invocation corroborates that the worker is making forward + # progress. + deliveries: list[str] = [] + + async def slow_obs(event: ObserverEvent) -> None: + # ~30ms per event. The 4-event invocation needs ~120ms total; + # the first drain's 10ms timeout fires well before delivery + # completes. + await asyncio.sleep(0.03) + if isinstance(event, NodeEvent): + deliveries.append(f"{event.node_name}:{event.phase}") + + captured_inv: list[str] = [] + + async def _capture(_s: _S) -> Mapping[str, Any]: + captured_inv.append(await _capture_node_invocation_id()) + return {"v": 1} + + builder: GraphBuilder[_S] = GraphBuilder(_S) + builder.set_entry("n") + builder.add_node("n", _capture) + builder.add_edge("n", END) + compiled = builder.compile() + compiled.attach_observer(slow_obs) + await compiled.invoke(_S()) + + # First drain: tight timeout against the slow observer fires + # before any deliveries complete. + started = time.monotonic() + summary_1 = await compiled.drain_events_for(captured_inv[0], timeout=0.01) + elapsed_1 = time.monotonic() - started + + assert summary_1.timeout_reached is True + assert summary_1.undelivered_count > 0 + # Drain returned within the deadline (no cancellation overhead). + assert elapsed_1 < 0.3 + # At this moment NO NodeEvents have been delivered yet (the + # observer's first sleep was still in flight at timeout time). + assert deliveries == [] + + # Second drain on the SAME invocation_id with a generous timeout. + # If the deliver worker had been cancelled by the first timeout, + # this would either also time out (worker stuck) or return clean + # via the unknown-invocation path (worker gone from + # _active_workers). Either way the deliveries list would stay + # empty. With the worker kept alive, the loop catches up and + # this returns clean with every event delivered. + summary_2 = await compiled.drain_events_for(captured_inv[0], timeout=5.0) + assert summary_2 == DrainSummary(undelivered_count=0, timeout_reached=False) + + # Decisive check: every NodeEvent for the invocation reached + # the observer. Two events per node (started + completed) × + # one node = 2. If the worker had been cancelled at first + # timeout, deliveries would have stopped at 0 or 1. + assert "n:started" in deliveries + assert "n:completed" in deliveries + + +async def test_drain_events_for_invocation_scope_isolation() -> None: + # Two serial invocations on the same compiled graph. Each drain + # sees only its own events. Mirrors spec fixture 031. + # + # The contract: drain_events_for(inv_a) awaits ONLY events + # tagged with inv_a; drain_events_for(inv_b) awaits ONLY events + # tagged with inv_b. Per spec §5.1 each invocation gets a fresh + # invocation_id, so the two ids differ and the observer's + # delivery log can be partitioned cleanly by invocation. + captured: list[str] = [] + delivery_log: list[tuple[str, str]] = [] + + async def obs(event: ObserverEvent) -> None: + await asyncio.sleep(0.01) + if isinstance(event, NodeEvent): + # Capture (invocation_id, node_name) so we can assert + # which deliveries happened under which invocation. + from openarmature.observability.correlation import current_invocation_id + + inv = current_invocation_id() or "?" + delivery_log.append((inv, event.node_name)) + + async def _capture(_s: _S) -> Mapping[str, Any]: + captured.append(await _capture_node_invocation_id()) + return {"v": 1} + + builder: GraphBuilder[_S] = GraphBuilder(_S) + builder.set_entry("n") + builder.add_node("n", _capture) + builder.add_edge("n", END) + compiled = builder.compile() + compiled.attach_observer(obs) + + # First invocation + drain. + await compiled.invoke(_S()) + inv_a = captured[-1] + summary_a = await compiled.drain_events_for(inv_a, timeout=5.0) + assert summary_a == DrainSummary(undelivered_count=0, timeout_reached=False) + + # Second invocation gets a fresh invocation_id (spec §5.1) and + # its own drain. + await compiled.invoke(_S()) + inv_b = captured[-1] + assert inv_b != inv_a + summary_b = await compiled.drain_events_for(inv_b, timeout=5.0) + assert summary_b == DrainSummary(undelivered_count=0, timeout_reached=False) + + # Partition the delivery log by invocation_id. Each drain's + # snapshot covered exactly its own invocation's events; no + # cross-contamination. + a_entries = [(inv, name) for (inv, name) in delivery_log if inv == inv_a] + b_entries = [(inv, name) for (inv, name) in delivery_log if inv == inv_b] + # One node ("n") fires started + completed under each invocation. + assert len(a_entries) == 2 + assert len(b_entries) == 2 + # Strict isolation: every entry's invocation_id matches the + # partition it landed in. + assert all(inv == inv_a for (inv, _) in a_entries) + assert all(inv == inv_b for (inv, _) in b_entries) + # No entry escaped the partition. + assert len(a_entries) + len(b_entries) == len(delivery_log) + + +async def test_drain_events_for_rejects_negative_timeout() -> None: + compiled = _build_compiled() + with pytest.raises(ValueError, match="non-negative"): + await compiled.drain_events_for("any-id", timeout=-1.0) + + +async def test_drain_events_for_rejects_nan_timeout() -> None: + compiled = _build_compiled() + with pytest.raises(ValueError, match="non-negative"): + await compiled.drain_events_for("any-id", timeout=float("nan")) + + +async def test_drain_events_for_zero_timeout_is_non_blocking_check() -> None: + # A zero timeout fires immediately if the snapshot target isn't + # met. Mirrors drain(timeout=0.0)'s non-blocking semantics. + captured: list[str] = [] + + async def slow_obs(_event: ObserverEvent) -> None: + await asyncio.sleep(0.2) + + async def _capture(_s: _S) -> Mapping[str, Any]: + captured.append(await _capture_node_invocation_id()) + return {"v": 1} + + builder: GraphBuilder[_S] = GraphBuilder(_S) + builder.set_entry("n") + builder.add_node("n", _capture) + builder.add_edge("n", END) + compiled = builder.compile() + compiled.attach_observer(slow_obs) + await compiled.invoke(_S()) + + started = time.monotonic() + summary = await compiled.drain_events_for(captured[-1], timeout=0.0) + elapsed = time.monotonic() - started + + assert summary.timeout_reached is True + assert summary.undelivered_count > 0 + assert elapsed < 0.05 + + +async def test_drain_events_for_snapshot_semantic_does_not_wait_for_own_completed_event() -> None: + # Mirrors spec fixture 029. A node body calling + # ``drain_events_for`` from inside itself MUST NOT block on its + # own ``completed`` event — the snapshot is the ``dispatched`` + # count AT CALL TIME, before the node's completed fires. Without + # the snapshot semantic the call would deadlock (the node body + # is awaiting the completed event, but the completed event only + # fires after the node body returns). + captured_summary: list[DrainSummary] = [] + delivered_observer_events: list[str] = [] + + async def slow_obs(event: ObserverEvent) -> None: + await asyncio.sleep(0.01) + if isinstance(event, NodeEvent): + delivered_observer_events.append(f"{event.node_name}:{event.phase}") + + compiled_ref: list[CompiledGraph[_S]] = [] + + async def _drain_from_inside(_s: _S) -> Mapping[str, Any]: + inv = await _capture_node_invocation_id() + summary = await compiled_ref[0].drain_events_for(inv, timeout=2.0) + captured_summary.append(summary) + return {"v": 1} + + builder: GraphBuilder[_S] = GraphBuilder(_S) + builder.set_entry("drain_node") + builder.add_node("drain_node", _drain_from_inside) + builder.add_edge("drain_node", END) + compiled = builder.compile() + compiled_ref.append(compiled) + compiled.attach_observer(slow_obs) + + # The whole invoke MUST complete within the outer timeout — if + # drain_events_for waited for the calling node's own completed + # event, this would deadlock and asyncio.wait_for would fire. + await asyncio.wait_for(compiled.invoke(_S()), timeout=5.0) + + # The in-node drain returned cleanly with no undelivered events + # at its snapshot point. + assert captured_summary[0] == DrainSummary(undelivered_count=0, timeout_reached=False) + + +async def test_drain_events_for_covers_fan_out_instance_events() -> None: + # Mirrors spec fixture 032. Fan-out instances share the parent + # invocation's ``_DrainCounters`` (subgraph descents pass the + # parent's counters down via ``descend_into_subgraph``), so + # ``drain_events_for(outer_invocation_id)`` MUST cover every + # event emitted under any of the fan-out's instance subgraph + # descents. Without shared counters the drain would miss + # instance events entirely. + + class _ParentState(State): + items: list[int] = Field(default_factory=list[int]) + results: Annotated[list[int], append] = Field(default_factory=list[int]) + + class _InstanceState(State): + item: int = 0 + result: int = 0 + + async def _double(state: _InstanceState) -> Mapping[str, Any]: + return {"result": state.item * 2} + + inner = ( + GraphBuilder(_InstanceState) + .set_entry("double") + .add_node("double", _double) + .add_edge("double", END) + .compile() + ) + + captured_invocation_id: list[str] = [] + delivered: list[str] = [] + + async def slow_obs(event: ObserverEvent) -> None: + await asyncio.sleep(0.005) + if isinstance(event, NodeEvent): + delivered.append(f"{event.namespace}/{event.node_name}:{event.phase}") + + async def _capture_and_persist(_s: _ParentState) -> Mapping[str, Any]: + captured_invocation_id.append(await _capture_node_invocation_id()) + return {} + + builder: GraphBuilder[_ParentState] = GraphBuilder(_ParentState) + builder.set_entry("fan_out") + builder.add_fan_out_node( + "fan_out", + subgraph=inner, + items_field="items", + item_field="item", + collect_field="result", + target_field="results", + ) + builder.add_node("persist", _capture_and_persist) + builder.add_edge("fan_out", "persist") + builder.add_edge("persist", END) + compiled = builder.compile() + compiled.attach_observer(slow_obs) + + await compiled.invoke(_ParentState(items=[1, 2, 3])) + summary = await compiled.drain_events_for(captured_invocation_id[0], timeout=5.0) + + assert summary == DrainSummary(undelivered_count=0, timeout_reached=False) + # Each of the 3 instances' ``double`` node fired a started + + # completed pair under its own namespace; all six MUST have + # delivered by the time drain_events_for returned. + instance_completions = [e for e in delivered if "/double:completed" in e] + assert len(instance_completions) == 3 + + +async def test_drain_events_for_covers_parallel_branches_events() -> None: + # Mirrors spec fixture 033. Parallel-branches branches share the + # parent invocation's ``_DrainCounters`` (same plumbing as + # fan-out instances), so ``drain_events_for(outer_invocation_id)`` + # MUST cover every event from every branch's inner subgraph. + + from openarmature.graph import BranchSpec + + class _ParentState(State): + a_out: str = "" + b_out: str = "" + + class _BranchAState(State): + v: str = "" + + class _BranchBState(State): + v: str = "" + + async def _do_a(_s: _BranchAState) -> Mapping[str, Any]: + return {"v": "a-done"} + + async def _do_b(_s: _BranchBState) -> Mapping[str, Any]: + return {"v": "b-done"} + + branch_a = ( + GraphBuilder(_BranchAState) + .set_entry("a_node") + .add_node("a_node", _do_a) + .add_edge("a_node", END) + .compile() + ) + branch_b = ( + GraphBuilder(_BranchBState) + .set_entry("b_node") + .add_node("b_node", _do_b) + .add_edge("b_node", END) + .compile() + ) + + captured_invocation_id: list[str] = [] + delivered: list[str] = [] + + async def slow_obs(event: ObserverEvent) -> None: + await asyncio.sleep(0.005) + if isinstance(event, NodeEvent): + delivered.append(f"{event.namespace}/{event.node_name}:{event.phase}") + + async def _capture(_s: _ParentState) -> Mapping[str, Any]: + captured_invocation_id.append(await _capture_node_invocation_id()) + return {} + + builder: GraphBuilder[_ParentState] = GraphBuilder(_ParentState) + builder.set_entry("dispatcher") + builder.add_parallel_branches_node( + "dispatcher", + branches={ + "branch_a": BranchSpec(subgraph=branch_a, outputs={"a_out": "v"}), + "branch_b": BranchSpec(subgraph=branch_b, outputs={"b_out": "v"}), + }, + ) + builder.add_node("persist", _capture) + builder.add_edge("dispatcher", "persist") + builder.add_edge("persist", END) + compiled = builder.compile() + compiled.attach_observer(slow_obs) + + await compiled.invoke(_ParentState()) + summary = await compiled.drain_events_for(captured_invocation_id[0], timeout=5.0) + + assert summary == DrainSummary(undelivered_count=0, timeout_reached=False) + # Both branches' inner nodes (``a_node`` + ``b_node``) MUST have + # delivered their completed events by the time drain_events_for + # returned. + a_done = [e for e in delivered if "/a_node:completed" in e] + b_done = [e for e in delivered if "/b_node:completed" in e] + assert len(a_done) == 1 + assert len(b_done) == 1