orchestrata_ultimate.cxp.reply_router

Reference › orchestrata_ultimate › cxp › reply_router

ReplyRouter gives each reply consumer its own queue, so that several gateways reading replies from one kernel do not compete for a shared queue (SPEC-001 US-002). servicer owns one and calls it from enqueue_reply, receive_packet and subscribe. It takes a tracer, a per-queue depth limit, an idle time-to-live, a sweep interval and a cap on consumers. See cxp.

route picks a key with _pick_key: the reply_to value in the packet metadata if present, otherwise the to_agent field, otherwise a default key. This is how the gateway's reply consumer id, stamped by the gateway cxp handler and the orchestrator handler, brings a reply back to the gateway that sent the request. route enqueues the packet and, if the queue is full, drops the oldest pending reply to make room for the newest, with a traced warning. queue_for returns the consumer's queue, creating it on demand. open_drain and close_drain take and release a lease while a consumer is reading; an entry with an open lease is never evicted. A lazy _sweep, run at most once per sweep interval from route, evicts idle entries without leases once they exceed the time-to-live, and also removes the longest-idle entries when the table exceeds the consumer cap (_evict records the reason). _QueueEntry holds the queue, the last-activity time and the lease count.

At a glance

Class diagram

Dependency diagram

API

Classes

_QueueEntry (dataclass)

Decorators: traces(spec='001', story='US-002', reqs=['REQ-001-US-002-030']), dataclass · Line 61

A per-consumer queue plus the bookkeeping the D-H10 sweep needs.

ReplyRouter

Decorators: traces(spec='001', story='US-002', reqs=['REQ-001-US-002-001', 'REQ-001-US-002-003', 'REQ-001-US-002-008', 'REQ-001-US-… · Line 88

Routes reply packets into per-consumer queues, keyed by addressee.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/cxp/reply_router.py

"""Per-consumer reply routing for the CXP servicer (SPEC-001 US-002).

Closes G3/G4/G7 from ``plans/kernel-dispatch-reply-design.md`` §4.1 D6/D7
(contract: ``plans/contracts/kernel-dispatch-reply.mmd``): before this
module, ``CXPServicer`` published every reply onto ONE unkeyed
``asyncio.Queue`` and every ``receive_packet``/``subscribe`` stream drained
that same queue regardless of the requesting ``agent_id`` -- a lottery that
measured 50/50 reply loss across concurrent gateways, plus a second
competitor (``subscribe``) that silently discarded any packet failing its
subscription filter instead of requeuing it for the consumer that actually
owned it.

``ReplyRouter`` gives every reply a home: ``route()`` keys each packet by
its addressee (``metadata.reply_to`` set by the sender, falling back to
``to_agent`` for legacy unstamped traffic, so existing callers -- notably
every consumer defaulting to ``agent_id="HUMAN"`` -- keep working
unchanged) and enqueues it onto that consumer's OWN queue. ``queue_for()``
is the read side both streaming RPCs use to find their queue.

HARDENING ADDENDUM (F7, ``plans/kernel-dispatch-reply-hardening.md`` §5,
contract ``plans/contracts/kernel-dispatch-reply-hardening.mmd`` packet
P-H7): unbounded ``_queues`` KEY growth -- every rotating gateway
(``HUMAN@<12hex>``) mints a fresh key and drain-side minting
(``receive_packet``/``subscribe``) lets ANY caller mint one -- meant a
long-lived kernel accreted dead consumer entries forever. D-H8/D-H9/D-H10
close it with idle-TTL cache eviction + drain leases, lazily swept inside
``route()`` (deliberately no background task -- F5's own lesson: one more
thing to drain at shutdown). Minting stays allowed (the shipped contract is
connect-before-first-reply); the guard is that minted-but-never-drained,
idle entries are precisely the first evicted, and total keys are hard
capped.
"""

from __future__ import annotations

import asyncio
import logging
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.tracing.tracer import Tracer

logger = logging.getLogger(__name__)

#: Key a reply lands under when the packet carries neither
#: ``metadata.reply_to`` nor a truthy ``to_agent`` -- a malformed/edge-case
#: packet, not the common path (legacy unstamped human-facing traffic
#: already carries ``to_agent="HUMAN"`` and is picked up by the second
#: branch of :meth:`ReplyRouter._pick_key`). Shares the literal every
#: existing consumer already defaults to (``gateway/server.py``
#: ``_pump_kernel_replies_once(agent_id="HUMAN")``), so even this
#: last-resort fallback stays consistent with the rest of the system.
DEFAULT_REPLY_KEY = "HUMAN"


@traces(spec="001", story="US-002", reqs=["REQ-001-US-002-030"])
@dataclass
class _QueueEntry:
    """A per-consumer queue plus the bookkeeping the D-H10 sweep needs.

    ``last_activity`` is a ``clock()`` reading (monotonic by default, see
    :meth:`ReplyRouter.__init__`), refreshed by every write
    (:meth:`ReplyRouter.route`), every read (:meth:`ReplyRouter.queue_for`),
    and every drain-lease transition (:meth:`ReplyRouter.open_drain`/
    :meth:`ReplyRouter.close_drain`). ``active_drains`` counts open drain
    leases (D-H9): while positive, this entry is eviction-immune regardless
    of how idle it looks.
    """

    queue: asyncio.Queue[dict[str, Any]]
    last_activity: float
    active_drains: int = 0


@traces(
    spec="001",
    story="US-002",
    reqs=[
        "REQ-001-US-002-001",
        "REQ-001-US-002-003",
        "REQ-001-US-002-008",
        "REQ-001-US-002-030",
    ],
)
class ReplyRouter:
    """Routes reply packets into per-consumer queues, keyed by addressee.

    Threading/concurrency note: every coroutine touching a given
    ``CXPServicer`` (``enqueue_reply``/``send_packet``, ``receive_packet``,
    ``subscribe``) runs on the SAME asyncio event loop -- one loop per
    server process, whether that's ``grpc.aio``'s server or
    ``DirectCXPServer`` (both share one ``CXPServicer`` instance; see
    ``DirectCXPServer.servicer``). ``_queues`` is therefore a plain
    ``dict``: there is no cross-thread or cross-loop access to guard
    against, so no lock is needed around create-on-demand, drain, or the
    F7 eviction sweep.
    """

    def __init__(
        self,
        tracer: Tracer,
        max_queue: int = 256,
        idle_ttl_s: float = 900.0,
        sweep_interval_s: float = 60.0,
        max_consumers: int = 1024,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        """Bind this router to *tracer* and its per-consumer eviction policy.

        Args:
            tracer: Injected :class:`Tracer` used to trace routing/overflow/
                eviction decisions (composed, not constructed here -- DI per
                the gen-skeleton contract).
            max_queue: Bounded depth for EACH per-consumer queue (a
                tunable, not an infra endpoint -- a literal default is
                permitted under RUL-003's inverse rule). Once a consumer's
                queue is at capacity, :meth:`route` drops the OLDEST
                pending reply to make room for the newest one, so a dead
                consumer's cost is bounded at ``max_queue`` dicts rather
                than growing forever.
            idle_ttl_s: How long (``clock()`` seconds) an INACTIVE entry
                (``active_drains == 0``) may sit idle before :meth:`route`'s
                lazy sweep (D-H10) evicts it -- the fix for F7's unbounded
                KEY growth. A live consumer polls well inside this window
                (default 900s vs. the gateway's ~5s poll,
                ``gateway/config.py`` ``kernel_reply_timeout_s``), so TTL
                eviction never bites a live drain; a rotated-away
                consumer's queue and its contents are eventually reclaimed
                instead of living forever.
            sweep_interval_s: The lazy sweep inside :meth:`route` runs AT
                MOST once per this many ``clock()`` seconds -- keeps the
                O(keys) sweep cost off the hot path of every single
                ``route()`` call.
            max_consumers: Hard cap on total keys in ``_queues``. If the
                TTL sweep leaves the table over this cap, the OLDEST-idle
                INACTIVE entries are evicted (also traced) until back at
                cap -- the flood backstop behind the TTL's primary
                collector, covering unauthenticated drain-side minting.
            clock: Injectable time source (default :func:`time.monotonic`)
                -- the RUL-002-compliant test seam: tests pass a
                manually-advanced forward-only clock value provider
                instead of monkeypatching ``time``.
        """
        self._tracer = tracer
        self._max_queue = max_queue
        self._idle_ttl_s = idle_ttl_s
        self._sweep_interval_s = sweep_interval_s
        self._max_consumers = max_consumers
        self._clock = clock
        self._queues: dict[str, _QueueEntry] = {}
        #: -inf guarantees the very first route() call is eligible to sweep
        #: (a no-op the first time, since nothing has had a chance to go
        #: idle yet) rather than waiting a full sweep_interval_s from boot.
        self._last_swept_at: float = float("-inf")

    def _pick_key(self, packet: dict[str, Any]) -> str:
        """Pick the consumer key a reply packet belongs under.

        Precedence: ``metadata.reply_to`` (the sender-stamped consumer id)
        wins when present and non-empty; otherwise ``to_agent`` (legacy
        unstamped traffic keeps queueing under whatever it already used,
        e.g. ``"HUMAN"`` -- backward compatible with every existing
        consumer that never stamps ``reply_to``); otherwise
        :data:`DEFAULT_REPLY_KEY`.
        """
        metadata = packet.get("metadata")
        if isinstance(metadata, dict):
            reply_to = metadata.get("reply_to")
            if reply_to:
                return str(reply_to)
        to_agent = packet.get("to_agent")
        if to_agent:
            return str(to_agent)
        return DEFAULT_REPLY_KEY

    def _get_or_create(self, consumer_id: str, now: float) -> _QueueEntry:
        """Shared create-on-demand core for :meth:`queue_for`/:meth:`open_drain`/:meth:`route`."""
        entry = self._queues.get(consumer_id)
        if entry is None:
            entry = _QueueEntry(queue=asyncio.Queue(), last_activity=now)
            self._queues[consumer_id] = entry
        else:
            entry.last_activity = now
        return entry

    @traces(
        spec="001", story="US-002", reqs=["REQ-001-US-002-001", "REQ-001-US-002-003"]
    )
    def queue_for(self, consumer_id: str) -> asyncio.Queue[dict[str, Any]]:
        """Return *consumer_id*'s queue, creating it on demand.

        The same :class:`asyncio.Queue` instance is returned on every call
        for a given ``consumer_id`` -- both :meth:`route` (write side) and
        the servicer's drain-lease methods (read side) share it, so a
        reply routed before a consumer starts draining still arrives once
        it does -- PROVIDED the consumer starts draining within
        ``idle_ttl_s`` of the entry's last activity (F7 hardening, D-H10):
        an entry idle longer than that is evicted by :meth:`route`'s lazy
        sweep, and this call then mints a fresh, empty queue in its place
        rather than resurrecting the old one and its contents.
        """
        return self._get_or_create(consumer_id, self._clock()).queue

    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-008"])
    def open_drain(self, consumer_id: str) -> asyncio.Queue[dict[str, Any]]:
        """Open an RAII drain lease on *consumer_id*'s queue (D-H9).

        Minting stays allowed here -- the shipped contract is
        connect-before-first-reply -- but while at least one lease is
        open (``active_drains > 0``), this entry is eviction-immune: the
        D-H10 lazy sweep inside :meth:`route` skips it no matter how idle
        it looks. Callers (``CXPServicer.receive_packet``/``subscribe``)
        MUST pair this with :meth:`close_drain` in a ``finally`` block.
        """
        now = self._clock()
        # Sweep on the MINT path too (reviewer finding 4, 2026-08-05): a
        # flood of rotating drain-side consumer ids never calls route(), so
        # a route()-only sweep would let the exact growth vector F7 names —
        # unauthenticated queue minting — run unbounded. Rate-limited by
        # sweep_interval_s, so the steady-state cost is a clock compare.
        self._sweep(now)
        entry = self._get_or_create(consumer_id, now)
        entry.active_drains += 1
        return entry.queue

    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-008"])
    def close_drain(self, consumer_id: str) -> None:
        """Close one drain lease opened by :meth:`open_drain` (D-H9).

        A no-op if *consumer_id* has no entry -- unreachable under correct
        pairing (an active lease is eviction-immune) but defensive against
        a caller closing twice. The lease counter never goes negative.
        """
        entry = self._queues.get(consumer_id)
        if entry is None:
            return
        entry.active_drains = max(0, entry.active_drains - 1)
        entry.last_activity = self._clock()

    @traces(
        spec="001", story="US-002", reqs=["REQ-001-US-002-001", "REQ-001-US-002-003"]
    )
    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-030"])
    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-008"])
    def route(self, packet: dict[str, Any]) -> str:
        """Key *packet* by its addressee and enqueue it, dropping-oldest on overflow.

        Returns the key the packet was queued under (mainly for
        observability/tests). When the target queue is already at
        ``max_queue`` depth, the OLDEST pending reply is discarded (never
        the new one) and a traced warning is logged -- REQ-001-US-002-030's
        "handle queue overflow gracefully": a dead/slow consumer's queue
        stays bounded instead of growing without limit, and the consumer
        still gets the freshest replies once it resumes draining.

        Also runs the D-H10 lazy idle-TTL sweep (:meth:`_sweep`, F7
        hardening) AFTER enqueuing -- at most once per
        ``sweep_interval_s`` -- so a long-lived router's ``_queues`` table
        doesn't grow without bound from rotated-away consumers and
        unauthenticated drain-side minting.
        """
        key = self._pick_key(packet)
        with self._tracer.span("ReplyRouter.route"):
            now = self._clock()
            queue = self._get_or_create(key, now).queue
            while queue.qsize() >= self._max_queue:
                dropped = queue.get_nowait()
                logger.warning(
                    "ReplyRouter: queue for consumer %r at capacity (%d); "
                    "dropping oldest reply message_id=%r",
                    key,
                    self._max_queue,
                    dropped.get("message_id"),
                )
            queue.put_nowait(packet)
            self._sweep(now)
        return key

    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-008"])
    def _sweep(self, now: float) -> None:
        """Lazily evict idle, drain-inactive queue entries (D-H10, F7).

        Runs AT MOST once per ``sweep_interval_s`` -- called
        unconditionally from :meth:`route` on every packet, but a no-op
        unless that long has elapsed since the previous sweep, so the
        O(keys) cost stays off the hot path. Two passes, both traced per
        evicted key with the dropped-packet count:

        1. **TTL pass** -- evict every entry with ``active_drains == 0``
           whose idleness (``now - last_activity``) exceeds
           ``idle_ttl_s``. The primary collector for rotated-away
           consumers.
        2. **Cap pass** -- if the table is still over ``max_consumers``
           after the TTL pass, evict the OLDEST-idle INACTIVE entries
           (smallest ``last_activity`` first) until back at cap -- the
           flood backstop for unauthenticated drain-side minting.

        An entry with an open drain lease is NEVER evicted by either
        pass, however idle it looks.

        The interval gate yields to CAP PRESSURE: while the table is over
        ``max_consumers`` the sweep runs regardless of elapsed time,
        otherwise a mint flood faster than ``sweep_interval_s`` would grow
        the table unboundedly between ticks (reviewer finding 4 follow-on,
        2026-08-05). Under pressure the cost is bounded by the flood rate
        itself, and the table can exceed the cap only within a single
        mint-to-sweep step.
        """
        over_cap = len(self._queues) > self._max_consumers
        if not over_cap and now - self._last_swept_at < self._sweep_interval_s:
            return
        self._last_swept_at = now

        stale_keys = [
            key
            for key, entry in self._queues.items()
            if entry.active_drains == 0
            and (now - entry.last_activity) > self._idle_ttl_s
        ]
        for key in stale_keys:
            self._evict(key, now, reason="idle-ttl")

        overflow = len(self._queues) - self._max_consumers
        if overflow > 0:
            oldest_idle_inactive = sorted(
                (entry.last_activity, key)
                for key, entry in self._queues.items()
                if entry.active_drains == 0
            )
            for _, key in oldest_idle_inactive[:overflow]:
                self._evict(key, now, reason="max-consumers-cap")

    def _evict(self, key: str, now: float, *, reason: str) -> None:
        """Drop *key*'s entry entirely, tracing the dropped-packet count."""
        entry = self._queues.pop(key, None)
        if entry is None:
            return
        dropped = entry.queue.qsize()
        logger.warning(
            "ReplyRouter: evicting idle consumer queue key=%r reason=%s "
            "idle_s=%.1f dropped=%d",
            key,
            reason,
            now - entry.last_activity,
            dropped,
        )