From 850a80695e0ad2a9fb7f2e9a535f083180db9d54 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Fri, 5 Jun 2026 12:01:06 -0700 Subject: [PATCH 1/2] Add per-invocation observer event drain primitive CompiledGraph.drain_events_for(invocation_id, *, timeout=5.0) gives a terminal node a way to block until every event dispatched for the in-flight invocation has been delivered to every attached observer, then proceed -- typically followed by a read against a queryable observer accumulator's per-invocation bucket whose state the drain has now caught up to. Without this, the deliver loop may still hold not-yet-dispatched events at the moment the terminal node reads, and the read silently undercounts. The implementation hangs on the existing per-invocation _DrainCounters: drain_events_for snapshots dispatched at call time, registers a (target, Future) pair on a new drain_wakers list, and awaits the Future via asyncio.wait_for. The deliver loop fulfils any waker whose target has been reached after each delivered increment. On timeout the waker is removed from the list and the partial summary is returned -- and crucially, the worker is NOT cancelled. This is the load-bearing difference from the process-wide drain: drain is shutdown semantics and cancels its workers; drain_events_for is in-flight synchronization and leaves them running so the graph keeps serving other invocations. Subgraph descents share the parent's _DrainCounters, so a drain on the outermost invocation_id covers fan-out instance events and parallel-branches branch events for free. Ten unit tests mirror the spec fixtures' case shapes: basic synchronization (028), snapshot semantic / no-deadlock on the calling node's own completed event (029), worker-not-cancelled-on-timeout (030, the load-bearing divergence), invocation-scope isolation (031), fan-out coverage (032), parallel-branches coverage (033), plus zero-timeout non-blocking check, unknown id, negative + NaN boundary rejection. Conformance manifest 0054 flips to implemented since 0.12.0. The six graph-engine fixtures (028-033) stay deferred from the cross-capability parser pending the upcoming conformance-adapter capability spec to ratify the directive vocabulary. Coord context: discuss-per-invocation-event-drain thread, 02-spec-accepted-as-0054 settling the five spec questions (section, name, default timeout, summary shape, resume interaction) and flagging the worker-cancellation divergence explicitly. --- CHANGELOG.md | 1 + conformance.toml | 53 ++- docs/concepts/observability.md | 20 +- src/openarmature/graph/compiled.py | 92 +++++ src/openarmature/graph/observer.py | 28 ++ tests/conformance/test_fixture_parsing.py | 32 +- tests/unit/test_drain.py | 394 +++++++++++++++++++++- 7 files changed, 589 insertions(+), 31 deletions(-) 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..78d171ed 100644 --- a/src/openarmature/graph/compiled.py +++ b/src/openarmature/graph/compiled.py @@ -773,6 +773,98 @@ 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 before this call returns, + 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 (including + the calling node's own ``started`` / ``completed`` pair) are + out of scope. This is what allows an in-node call to avoid + deadlocking on its own completed event. + + ``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..ace31efc 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,393 @@ 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. The graph remains usable for + # subsequent invocations after the timeout fires, and the deliver + # loop keeps processing the queue. Mirrors spec fixture 030. + deliveries: list[str] = [] + + async def very_slow_obs(event: ObserverEvent) -> None: + # 100ms per event ensures the 50ms timeout fires before any + # delivery completes (the first await in deliver_loop is + # this observer's sleep, which already exceeds the budget). + await asyncio.sleep(0.1) + if isinstance(event, NodeEvent): + deliveries.append(event.node_name) + + captured_inv_1: list[str] = [] + captured_inv_2: list[str] = [] + + async def _capture_1(_s: _S) -> Mapping[str, Any]: + captured_inv_1.append(await _capture_node_invocation_id()) + return {"v": 1} + + async def _capture_2(_s: _S) -> Mapping[str, Any]: + captured_inv_2.append(await _capture_node_invocation_id()) + return {"v": 2} + + builder: GraphBuilder[_S] = GraphBuilder(_S) + builder.set_entry("n") + builder.add_node("n", _capture_1) + builder.add_edge("n", END) + compiled_1 = builder.compile() + compiled_1.attach_observer(very_slow_obs) + await compiled_1.invoke(_S()) + + started = time.monotonic() + summary = await compiled_1.drain_events_for(captured_inv_1[0], timeout=0.05) + elapsed = time.monotonic() - started + + assert summary.timeout_reached is True + assert summary.undelivered_count > 0 + # Drain returned within the deadline (no cancellation overhead). + assert elapsed < 0.3 + # The worker is still running — a fresh invocation on the same + # graph still delivers events through the same deliver loop. + builder_2: GraphBuilder[_S] = GraphBuilder(_S) + builder_2.set_entry("m") + builder_2.add_node("m", _capture_2) + builder_2.add_edge("m", END) + compiled_2 = builder_2.compile() + compiled_2.attach_observer(very_slow_obs) + # The second invocation's drain (with a generous timeout) verifies + # the observer is still healthy and the deliver pipeline still + # processes events. + await compiled_2.invoke(_S()) + summary_2 = await compiled_2.drain_events_for(captured_inv_2[0], timeout=5.0) + assert summary_2 == DrainSummary(undelivered_count=0, timeout_reached=False) + + +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. Drains for invocation A do not wait + # on invocation B's deliveries. Mirrors spec fixture 031. + 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) + + 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) + # invocation A's events delivered before A's drain returned. + a_deliveries = [name for (inv, name) in delivery_log if inv == inv_a] + assert "n" in a_deliveries + + # Drain on an UNRELATED invocation_id (still inv_a, but the worker + # exited at the end of invocation A — so this is effectively + # querying a stale id) returns immediately. + summary_stale = await compiled.drain_events_for(inv_a, timeout=5.0) + assert summary_stale == DrainSummary(undelivered_count=0, timeout_reached=False) + + +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 From bfb5eca523464f15eaa3353ad8ae707b20ce4f81 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Fri, 5 Jun 2026 13:31:34 -0700 Subject: [PATCH 2/2] Tighten drain_events_for docstring + tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Docstring fixes: the opening sentence said the drain covers events dispatched "before this call returns", but the snapshot is taken at call entry — events dispatched between entry and exit are not in scope. Rewords to "as of this call's entry". The snapshot- semantic paragraph also claimed the calling node's started AND completed events were out of scope, but the engine dispatches started immediately before the node body runs, so started is already in the snapshot and the drain awaits its delivery normally. Only completed is guaranteed out of scope. Test fixes: - test_drain_events_for_timeout_does_not_cancel_worker now runs the follow-up invocation on the SAME compiled graph as the timed-out one, and the decisive contract check is that all NodeEvents from the originally-pending queue land in the observer's delivery list AFTER the timed-out drain returned. The previous version compiled a fresh graph for the follow-up, so the test would have passed even if the original worker had been cancelled. - test_drain_events_for_invocation_scope_isolation now actually performs two serial invocations (the previous version invoked once and drained the same id twice, testing stale-id idempotency rather than per-invocation isolation). Each drain's delivery log entries are partitioned by invocation_id; the assertion is strict (no cross-contamination, no entries outside either partition). --- src/openarmature/graph/compiled.py | 13 +-- tests/unit/test_drain.py | 134 ++++++++++++++++++----------- 2 files changed, 91 insertions(+), 56 deletions(-) diff --git a/src/openarmature/graph/compiled.py b/src/openarmature/graph/compiled.py index 78d171ed..8b65a6ee 100644 --- a/src/openarmature/graph/compiled.py +++ b/src/openarmature/graph/compiled.py @@ -791,7 +791,7 @@ async def drain_events_for( timeout: float | None = 5.0, ) -> DrainSummary: """Await delivery of every observer event tagged with - ``invocation_id`` that was dispatched before this call returns, + ``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 @@ -802,10 +802,13 @@ async def drain_events_for( every attached observer, then returns. Snapshot semantic: the drain awaits the events dispatched as - of call time. Events emitted after the call begins (including - the calling node's own ``started`` / ``completed`` pair) are - out of scope. This is what allows an in-node call to avoid - deadlocking on its own completed event. + 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 diff --git a/tests/unit/test_drain.py b/tests/unit/test_drain.py index ace31efc..27498f57 100644 --- a/tests/unit/test_drain.py +++ b/tests/unit/test_drain.py @@ -319,73 +319,90 @@ async def _capture_then_drain(_s: _S) -> Mapping[str, Any]: 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. The graph remains usable for - # subsequent invocations after the timeout fires, and the deliver - # loop keeps processing the queue. Mirrors spec fixture 030. + # 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 very_slow_obs(event: ObserverEvent) -> None: - # 100ms per event ensures the 50ms timeout fires before any - # delivery completes (the first await in deliver_loop is - # this observer's sleep, which already exceeds the budget). - await asyncio.sleep(0.1) + 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(event.node_name) + deliveries.append(f"{event.node_name}:{event.phase}") - captured_inv_1: list[str] = [] - captured_inv_2: list[str] = [] + captured_inv: list[str] = [] - async def _capture_1(_s: _S) -> Mapping[str, Any]: - captured_inv_1.append(await _capture_node_invocation_id()) + async def _capture(_s: _S) -> Mapping[str, Any]: + captured_inv.append(await _capture_node_invocation_id()) return {"v": 1} - async def _capture_2(_s: _S) -> Mapping[str, Any]: - captured_inv_2.append(await _capture_node_invocation_id()) - return {"v": 2} - builder: GraphBuilder[_S] = GraphBuilder(_S) builder.set_entry("n") - builder.add_node("n", _capture_1) + builder.add_node("n", _capture) builder.add_edge("n", END) - compiled_1 = builder.compile() - compiled_1.attach_observer(very_slow_obs) - await compiled_1.invoke(_S()) + 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 = await compiled_1.drain_events_for(captured_inv_1[0], timeout=0.05) - elapsed = time.monotonic() - started + summary_1 = await compiled.drain_events_for(captured_inv[0], timeout=0.01) + elapsed_1 = time.monotonic() - started - assert summary.timeout_reached is True - assert summary.undelivered_count > 0 + assert summary_1.timeout_reached is True + assert summary_1.undelivered_count > 0 # Drain returned within the deadline (no cancellation overhead). - assert elapsed < 0.3 - # The worker is still running — a fresh invocation on the same - # graph still delivers events through the same deliver loop. - builder_2: GraphBuilder[_S] = GraphBuilder(_S) - builder_2.set_entry("m") - builder_2.add_node("m", _capture_2) - builder_2.add_edge("m", END) - compiled_2 = builder_2.compile() - compiled_2.attach_observer(very_slow_obs) - # The second invocation's drain (with a generous timeout) verifies - # the observer is still healthy and the deliver pipeline still - # processes events. - await compiled_2.invoke(_S()) - summary_2 = await compiled_2.drain_events_for(captured_inv_2[0], timeout=5.0) + 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. Drains for invocation A do not wait - # on invocation B's deliveries. Mirrors spec fixture 031. + # 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 + # Capture (invocation_id, node_name) so we can assert # which deliveries happened under which invocation. from openarmature.observability.correlation import current_invocation_id @@ -403,19 +420,34 @@ async def _capture(_s: _S) -> Mapping[str, Any]: 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) - # invocation A's events delivered before A's drain returned. - a_deliveries = [name for (inv, name) in delivery_log if inv == inv_a] - assert "n" in a_deliveries - - # Drain on an UNRELATED invocation_id (still inv_a, but the worker - # exited at the end of invocation A — so this is effectively - # querying a stale id) returns immediately. - summary_stale = await compiled.drain_events_for(inv_a, timeout=5.0) - assert summary_stale == 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: