Add per-invocation observer event drain primitive

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.

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"

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)
 

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
     # ------------------------------------------------------------------

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__ = [

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.<name>.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"
     ),
 }