Add per-invocation observer event drain primitive (#131)

* 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.

* Tighten drain_events_for docstring + tests

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).

Author: chris-colinsky

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