Add drain timeout + DrainSummary (proposal 0010)

CompiledGraph.drain() gains an optional timeout parameter and returns
a DrainSummary frozen dataclass (undelivered_count, timeout_reached).
The timeout-fired path cancels in-flight delivery workers cleanly so
the graph remains usable for subsequent invocations.

Per-invocation dispatched/delivered counters on _InvocationContext
track undelivered events; _active_workers changes from set[Task] to
dict[Task, _InvocationContext] so drain() can read each worker's
counters at cancellation time.

Solves the slow-observer-blocks-process-exit footgun for short-lived
processes (CLIs, scripts, serverless functions).

Author: chris-colinsky

src/openarmature/graph/__init__.py

@@ -48,7 +48,7 @@
     exponential_jitter_backoff,
 )
 from .nodes import FunctionNode, Node
-from .observer import Observer, RemoveHandle, SubscribedObserver
+from .observer import DrainSummary, Observer, RemoveHandle, SubscribedObserver
 from .parallel_branches import BranchSpec, ParallelBranchesNode
 from .projection import ExplicitMapping, FieldNameMatching, ProjectionStrategy
 from .reducers import Reducer, append, last_write_wins, merge
@@ -62,6 +62,7 @@
     "ConditionalEdge",
     "ConflictingReducers",
     "DanglingEdge",
+    "DrainSummary",
     "EdgeException",
     "EndSentinel",
     "ExplicitMapping",

src/openarmature/graph/compiled.py

@@ -93,6 +93,7 @@
 from .nodes import Node
 from .observer import (
     _DRAIN_SENTINEL,
+    DrainSummary,
     Observer,
     RemoveHandle,
     SubscribedObserver,
@@ -523,10 +524,14 @@ class CompiledGraph[StateT: State]:
     # dataclass: the list reference is fixed but its contents change.
     # Parameterized factories so pyright infers the element types.
     _attached_observers: list[SubscribedObserver] = field(default_factory=list[SubscribedObserver])
-    # `set` (not list) so a per-task `add_done_callback(self._active_workers.discard)`
-    # auto-removes completed workers — long-running services that never call
-    # drain() don't accumulate completed Task references indefinitely.
-    _active_workers: set[asyncio.Task[None]] = field(default_factory=set[asyncio.Task[None]])
+    # Per-task `add_done_callback` auto-removes completed workers — long-
+    # running services that never call drain() don't accumulate completed
+    # Task references indefinitely. Values are the per-invocation
+    # `_InvocationContext` so `drain()` can read each worker's
+    # `drain_counters` to compute the undelivered-event count at timeout.
+    _active_workers: dict[asyncio.Task[None], _InvocationContext] = field(
+        default_factory=dict[asyncio.Task[None], _InvocationContext]
+    )
     # Single-element list so the frozen-dataclass binding is stable but
     # the user can swap the registered Checkpointer via
     # ``attach_checkpointer``. ``None`` when no backend is registered.
@@ -680,35 +685,71 @@ async def _migrate_record(
         )
         return migrated, summary
 
-    async def drain(self) -> None:
+    async def drain(self, timeout: float | None = None) -> DrainSummary:
         """Await delivery of every observer event produced by prior
-        invocations of this graph.
+        invocations of this graph, optionally bounded by ``timeout``.
 
         Callers running in short-lived processes (scripts, serverless
-        functions, CLIs) MUST use drain to avoid losing observer
-        events that were dispatched but not yet delivered.
+        functions, CLIs) MUST use drain to avoid losing observer events
+        that were dispatched but not yet delivered.
 
         Only events dispatched before this call are awaited; events
         from invocations started concurrently with drain may or may
         not be included. Subgraph events from active invocations are
         part of the parent invocation's worker and are covered
         automatically.
 
-        **Unbounded by design.** Drain blocks until every queued event has
-        been delivered to every subscribed observer. A slow, hung, or
-        misbehaving observer can therefore hold drain, and the calling
-        process, indefinitely. If you need a bounded wait, wrap the call
-        in `asyncio.wait_for` and accept that events still queued when the
-        deadline elapses will not be delivered::
-
-            await asyncio.wait_for(compiled.drain(), timeout=5.0)
+        ``timeout`` is a non-negative duration in seconds. If omitted
+        or ``None``, drain waits indefinitely — a slow, hung, or
+        misbehaving observer can therefore hold drain (and the calling
+        process) indefinitely. If supplied, drain returns no later
+        than ``timeout`` seconds after the call begins; any observer
+        events still queued or in-flight at that point are considered
+        undelivered. Workers are cancelled via ``Task.cancel()`` so
+        the compiled graph remains usable for subsequent invocations
+        — partial delivery state from one drain does NOT leak into
+        the next invocation.
+
+        Returns a :class:`DrainSummary` with ``undelivered_count`` and
+        ``timeout_reached`` fields. The shape is the same whether or
+        not a timeout was supplied; on the no-timeout / timeout-not-
+        fired path both fields are zero / false.
+
+        Observers SHOULD be written to be cancellation-safe
+        (idempotent writes, try/finally cleanup) so that interruption
+        by drain timeout does not leave partial side effects in an
+        inconsistent state.
         """
         if not self._active_workers:
-            return
-        # Snapshot the set: each worker's done-callback removes itself
-        # from `_active_workers`, so iterating it directly while gather
-        # awaits would mutate during iteration.
-        await asyncio.gather(*list(self._active_workers), return_exceptions=True)
+            return DrainSummary(undelivered_count=0, timeout_reached=False)
+        # Snapshot the dict: each worker's done-callback removes its
+        # entry from `_active_workers`, so iterating directly while
+        # `asyncio.wait` awaits would mutate during iteration.
+        snapshot = dict(self._active_workers)
+        workers = list(snapshot.keys())
+
+        _done, pending = await asyncio.wait(
+            workers,
+            timeout=timeout,
+            return_when=asyncio.ALL_COMPLETED,
+        )
+
+        if not pending:
+            return DrainSummary(undelivered_count=0, timeout_reached=False)
+
+        undelivered = sum(
+            snapshot[w].drain_counters.dispatched - snapshot[w].drain_counters.delivered for w in pending
+        )
+        for w in pending:
+            w.cancel()
+        # Wait for cancellations to settle so the done-callbacks fire
+        # and `_active_workers` cleans for the next invocation —
+        # load-bearing for the cross-invocation cleanliness contract.
+        # ``return_exceptions`` absorbs the ``CancelledError`` each
+        # cancelled worker raises.
+        await asyncio.gather(*pending, return_exceptions=True)
+
+        return DrainSummary(undelivered_count=undelivered, timeout_reached=True)
 
     # ------------------------------------------------------------------
     # Public invocation
@@ -893,12 +934,16 @@ async def invoke(
         # "per-invocation is OUTERMOST invoke" wording).
         correlation_token = _set_correlation_id(resolved_correlation_id)
         invocation_token = _set_invocation_id(invocation_id)
-        worker = asyncio.create_task(deliver_loop(queue))
-        self._active_workers.add(worker)
+        worker = asyncio.create_task(deliver_loop(queue, context.drain_counters))
+        self._active_workers[worker] = context
         # Auto-prune: when the worker completes (after the sentinel is
-        # processed), remove it from the active set so long-running
-        # services don't leak Task references between drain() calls.
-        worker.add_done_callback(self._active_workers.discard)
+        # processed, or after cancellation by drain() on timeout), remove
+        # it from the active set so long-running services don't leak Task
+        # references between drain() calls. ``pop(key, None)`` is the
+        # idempotent form — if a concurrent drain() removed the entry
+        # already (it shouldn't with the current design, but the no-arg
+        # form would raise KeyError), this is a safe no-op.
+        worker.add_done_callback(lambda t: self._active_workers.pop(t, None))
         # Per spec §6 cross-ref in proposal 0014: dispatch the
         # ``checkpoint_migrated`` event as soon as the delivery
         # worker is alive but before any node runs, so the OTel

src/openarmature/graph/observer.py

@@ -212,6 +212,46 @@ class _QueuedItem:
 _DRAIN_SENTINEL = None
 
 
+# Spec: realizes graph-engine §6 Drain undelivered-count bookkeeping
+# (proposal 0010). Per-invocation mutable counters; `_dispatch` bumps
+# `dispatched` after a successful `queue.put_nowait`; `deliver_loop`
+# bumps `delivered` after the per-event observer for-loop completes.
+# `undelivered = dispatched - delivered` at any point in time — and
+# specifically at `CompiledGraph.drain()` cancellation time when the
+# timeout has elapsed and pending workers' counters get summed into
+# the returned `DrainSummary`.
+@dataclass
+class _DrainCounters:
+    dispatched: int = 0
+    delivered: int = 0
+
+
+# Spec: realizes graph-engine §6 Drain summary return shape (proposal
+# 0010). The two declared fields are the spec-mandated minimum;
+# implementations MAY add richer detail in future PRs (per-observer
+# counts, sampled event metadata) without breaking the v0.19.0 shape.
+@dataclass(frozen=True)
+class DrainSummary:
+    """Outcome of a `CompiledGraph.drain()` call.
+
+    Returned from `drain()` regardless of whether a `timeout` was
+    supplied. When no timeout was supplied, or the timeout did not
+    fire, `undelivered_count == 0` and `timeout_reached is False`.
+    When the timeout fired, `undelivered_count` reports the number of
+    events that were dispatched to the delivery worker but not fully
+    delivered to every subscribed observer before cancellation, and
+    `timeout_reached is True`.
+
+    The spec-mandated minimum is these two fields. Implementations MAY
+    extend the shape with diagnostic detail (per-observer counts,
+    sampled event metadata) in subsequent versions; v0.19.0 ships the
+    minimum.
+    """
+
+    undelivered_count: int
+    timeout_reached: bool
+
+
 # Spec: realizes pipeline-utilities §10.11 per-instance progress
 # tracking in the engine. These are the MUTABLE internal-state
 # counterparts to the FROZEN public ``FanOutProgress`` /
@@ -353,6 +393,12 @@ class _InvocationContext:
     fan_out_progress_state: dict[tuple[tuple[str, ...], str], _FanOutExecutionState] = field(
         default_factory=dict[tuple[tuple[str, ...], str], _FanOutExecutionState]
     )
+    # Per spec §6 Drain (proposal 0010): shared mutable counters that
+    # the worker reads at drain-cancel time to report undelivered events
+    # in the returned ``DrainSummary``. Subgraphs share the parent's
+    # counters because subgraphs share the parent's queue + worker, so
+    # the parent context's counts naturally cover subgraph events.
+    drain_counters: _DrainCounters = field(default_factory=_DrainCounters)
 
     def full_observers(self) -> tuple[SubscribedObserver, ...]:
         """Return the ordered observer list to deliver for events from
@@ -395,6 +441,7 @@ def descend_into_subgraph(
             pending_resume_states=self.pending_resume_states,
             resume_invocation=self.resume_invocation,
             fan_out_progress_state=self.fan_out_progress_state,
+            drain_counters=self.drain_counters,
         )
 
     def descend_into_fan_out_instance(
@@ -440,6 +487,7 @@ def descend_into_fan_out_instance(
             # inner-instance node can update its own entry and so the
             # outer save sees consistent sibling state.
             fan_out_progress_state=self.fan_out_progress_state,
+            drain_counters=self.drain_counters,
         )
 
     def descend_into_parallel_branch(
@@ -485,6 +533,7 @@ def descend_into_parallel_branch(
             pending_resume_states={},
             resume_invocation=self.resume_invocation,
             fan_out_progress_state=self.fan_out_progress_state,
+            drain_counters=self.drain_counters,
         )
 
     def take_step(self) -> int:
@@ -579,9 +628,16 @@ def _dispatch(context: _InvocationContext, event: NodeEvent) -> None:
                     stacklevel=2,
                 )
     context.queue.put_nowait(_QueuedItem(event=event, observers=observers))
+    # Per spec §6 Drain (proposal 0010): increment AFTER the put so a
+    # raise from ``put_nowait`` (queue full on a bounded queue — we
+    # don't bound, but the invariant holds) doesn't desync the counter.
+    context.drain_counters.dispatched += 1
 
 
-async def deliver_loop(queue: asyncio.Queue[_QueuedItem | None]) -> None:
+async def deliver_loop(
+    queue: asyncio.Queue[_QueuedItem | None],
+    counters: _DrainCounters,
+) -> None:
     """Background worker: read queued events, deliver to observers serially.
 
     - No two observers receive the same event concurrently (we await
@@ -610,17 +666,27 @@ async def deliver_loop(queue: asyncio.Queue[_QueuedItem | None]) -> None:
                     f"observer raised {type(e).__name__}: {e}",
                     stacklevel=1,
                 )
+        # Per spec §6 Drain (proposal 0010): increment AFTER the
+        # observer for-loop completes for this event, so an event
+        # cancelled mid-for-loop is counted as undelivered
+        # (``dispatched - delivered`` includes it). The phase-filter
+        # ``continue`` above does NOT skip the increment — an event
+        # filtered out for every observer is still considered
+        # delivered (we did all the work there was to do for it).
+        counters.delivered += 1
 
 
 __all__ = [
     "ALL_PHASES",
+    "DrainSummary",
     "Observer",
     "RemoveHandle",
     "SubscribedObserver",
     # Engine-internal but listed so pyright sees them as exported (they're
     # imported by `compiled.py` and `subgraph.py`). The underscore prefix
     # is the user-facing "don't import these" signal.
     "_DRAIN_SENTINEL",
+    "_DrainCounters",
     "_FanOutExecutionState",
     "_FanOutInstanceState",
     "_InvocationContext",