orchestrata_ultimate.cxp.direct_cxp_adapter

Reference › orchestrata_ultimate › cxp › direct_cxp_adapter

direct_cxp_adapter is the CXP transport implementation: it puts the shared servicer behind the cxpt direct transport from direct_transport. transport returns its two classes from create_cxp_client and create_cxp_server. The packet conversion is done by the codec in proto_codec, using the messages from proto. See cxp.

_resolve_target builds the cxpt target: a Unix socket path if CXP_DIRECT_UDS_PATH is set, otherwise tcp://<host>:<port>.

DirectCXPServer owns a CXPServicer and a method table. Four unary methods (SendPacket, HealthCheck, GetStatus, Acknowledge) and two streaming methods (ReceivePacket, Subscribe) are mapped to servicer coroutines. _dispatch runs on a cxpt handler thread, decodes the request, runs the servicer coroutine on the server's event loop, and returns the encoded response. It reports "unimplemented" for an unknown method and "unavailable" when the loop is gone. _drive_stream turns an async generator into an iterator and honours cancellation. start, stop and serve_forever manage the lifecycle.

DirectCXPClient implements the client protocol. connect creates the client in a worker thread, because the underlying connect blocks. Every call is bounded by CXPClientConfig.timeout_seconds through _call_bounded. send_packet converts the packet, retries on transport failure up to max_retries with an increasing delay, and returns a dictionary with success, message id, error and latency. Receiving, subscribing, acknowledging, health_check and get_status follow the same pattern.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

_StreamCancelled

Bases: Exception · Line 90

Internal sentinel: :meth:DirectCXPServer._next_or_cancelled saw the cancel flag before agen.__anext__() settled.

DirectCXPServer

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

The CXP server, over the cxpt direct transport.

DirectCXPClient

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

The CXP client, over the cxpt direct transport.

Functions

Source

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

"""The CXP transport: cxpt direct transport over the shared servicer.

``DirectCXPServer``/``DirectCXPClient`` are THE CXP server and client. They
were written as drop-in-shaped mirrors of the ``grpc.aio`` pair
(``DirectCXPServer``/``DirectCXPClient``) so a call site could switch transport by
construction alone; that migration completed on 2026-08-07 when grpc.aio was
retired (SPEC-023: its EventEngine threads re-enable the GIL, and its
vendored free-threaded wheel could not be resolved by a fresh clone). The
names keep their ``Direct``/``GRPC`` prefixes for now purely to avoid
churning call sites in the same change that removed the transport; see
``plans/grpc-retirement-design.md``.

The six CXP RPCs (``SendPacket``, ``HealthCheck``, ``GetStatus``,
``ReceivePacket``, ``Subscribe``, ``Acknowledge``) dispatch to
:class:`~orchestrata_ultimate.cxp.servicer.CXPServicer`, the transport-neutral
protocol core, over :mod:`orchestrata_ultimate.cxp.direct_transport`.

The wire format is still the CXP protos: request/response payloads are the
protos' ``SerializeToString()`` bytes, framed by cxpt rather than HTTP/2.
protobuf is a codec and was never coupled to the RPC runtime, so it stays.

Spec: 004-cxp / 001-kernel.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import os
from collections.abc import (
    AsyncGenerator,
    Awaitable,
    Callable,
    Iterator,
    Mapping,
)
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.packet import CXPPacket as CXPPacketModel
from orchestrata_ultimate.cxp.direct_transport import (
    STATUS_OK,
    STATUS_UNAVAILABLE,
    STATUS_UNIMPLEMENTED,
    AsyncDirectClient,
    DirectServer,
    RpcError,
)
from orchestrata_ultimate.cxp.proto.cxp_pb2 import (
    AckStatus,
    CXPAck,
    CXPHealthRequest,
    CXPHealthResponse,
    CXPPacket,
    CXPReceipt,
    CXPReceiveRequest,
    CXPSubscription,
    StatusReport,
    StatusRequest,
)
from orchestrata_ultimate.cxp.proto_codec import CXPClientConfig, CXPPacketCodec
from orchestrata_ultimate.cxp.servicer import CXPServerConfig, CXPServicer

logger = logging.getLogger(__name__)

# The env var that opts a Direct*-transport peer into a Unix domain socket
# (co-located peers on the same Incus host) instead of TCP. Mirrors the
# host/port fields CXPServerConfig/CXPClientConfig already carry -- when
# unset, the target is derived from those fields exactly as
# the CXP server/client derive their TCP address.
_UDS_PATH_ENV_VAR = "CXP_DIRECT_UDS_PATH"

# CXPServicer's RPC methods never read their `context` argument (verified:
# zero `context.` uses in the servicer), so a bare placeholder is a
# sufficient duck-typed stand-in. That is precisely what made retiring
# grpc.aio a deletion rather than a rewrite: the servicer never depended on
# it. Typed Any to satisfy every RPC method's `context: Any` parameter.
_DISPATCH_CONTEXT: Any = object()

# How often DirectCXPServer._next_or_cancelled polls `cancelled` while an
# agen.__anext__() pull is in flight (backlog #23). Bounds the worst-case
# extra latency between a CANCEL frame arriving and the servicer's async
# generator actually being torn down -- independent of, and far shorter
# than, CXPServicer.receive_packet/subscribe's own internal queue-wait
# timeouts (seconds to 60s).
_CANCEL_POLL_INTERVAL_S = 0.05


class _StreamCancelled(Exception):
    """Internal sentinel: :meth:`DirectCXPServer._next_or_cancelled` saw the
    cancel flag before ``agen.__anext__()`` settled. Never escapes this
    module -- :meth:`DirectCXPServer._drive_stream` catches it and returns.
    """


def _resolve_target(host: str, port: int, env: Mapping[str, str] | None) -> str:
    """Resolve a cxpt ``target`` string from host/port + an optional env override.

    ``CXP_DIRECT_UDS_PATH`` (if set) selects a Unix domain socket; otherwise
    falls back to ``tcp://<host>:<port>``.

    Args:
        host: The gRPC-mirrored host field (``CXPServerConfig.host`` /
            ``CXPClientConfig.host``).
        port: The gRPC-mirrored port field.
        env: Optional explicit environment mapping (tests inject one instead
            of mutating ``os.environ``); defaults to the process environment.
    """
    environ = os.environ if env is None else env
    uds_path = environ.get(_UDS_PATH_ENV_VAR)
    if uds_path:
        return f"unix:{uds_path}"
    return f"tcp://{host}:{port}"


@traces(spec="001", story="US-002", reqs=["REQ-001-US-002-012"])
class DirectCXPServer:
    """The CXP server, over the cxpt direct transport.

    Constructs one :class:`~orchestrata_ultimate.cxp.servicer.CXPServicer`
    and bridges each cxpt handler call onto it via
    :func:`asyncio.run_coroutine_threadsafe` against the loop captured in
    :meth:`start`.
    """

    def __init__(
        self,
        config: CXPServerConfig | None = None,
        packet_handler: (
            Callable[[dict[str, Any]], dict[str, Any] | Awaitable[dict[str, Any]]]
            | None
        ) = None,
        *,
        target: str | None = None,
        env: Mapping[str, str] | None = None,
    ) -> None:
        self._config = config or CXPServerConfig()
        self._servicer = CXPServicer(packet_handler)
        self._target = target or _resolve_target(
            self._config.host, self._config.port, env
        )
        self._server: DirectServer | None = None
        self._loop: asyncio.AbstractEventLoop | None = None
        self._stop_event = asyncio.Event()

        # Unary RPC -> (bound servicer coroutine method, request proto class).
        self._unary: dict[bytes, tuple[Any, Any]] = {
            b"SendPacket": (self._servicer.send_packet, CXPPacket),
            b"HealthCheck": (self._servicer.health_check, CXPHealthRequest),
            b"GetStatus": (self._servicer.get_status, StatusRequest),
            b"Acknowledge": (self._servicer.acknowledge, CXPAck),
        }
        # Server-streaming RPC -> (bound servicer async-generator method,
        # request proto class).
        self._streams: dict[bytes, tuple[Any, Any]] = {
            b"ReceivePacket": (self._servicer.receive_packet, CXPReceiveRequest),
            b"Subscribe": (self._servicer.subscribe, CXPSubscription),
        }

    @property
    def servicer(self) -> CXPServicer:
        """The underlying CXPServicer -- exposed for parity/introspection."""
        return self._servicer

    async def start(self) -> None:
        """Capture the running loop and bring up the cxpt server."""
        self._loop = asyncio.get_running_loop()
        self._server = DirectServer(
            self._target, self._dispatch, workers=self._config.max_workers or 8
        )
        logger.info("CXP direct-transport server started on %s", self._target)

    async def stop(self, grace: int = 5) -> None:
        """Stop the cxpt server.

        Args:
            grace: Accepted for signature parity with ``CXPServer.stop``;
                unused because ``DirectServer.close()`` joins its handler
                threads with its own fixed timeout.

        ``DirectServer.close()`` is blocking (it joins every handler thread
        with a 2s timeout each). Calling it synchronously here -- on the
        event-loop thread, since ``stop()`` is a coroutine awaited on the
        loop -- would deadlock: a handler thread can be parked inside
        ``_dispatch`` at ``asyncio.run_coroutine_threadsafe(coro,
        self._loop).result()``, which needs THIS loop to keep running to
        ever resolve. With the join blocking the loop thread, that call
        never completes, the join times out, and the handler thread survives
        as a zombie that can go on to answer a later test's connection.
        Running ``close()`` in an executor keeps the loop servicing those
        in-flight ``run_coroutine_threadsafe`` calls while the join waits on
        them, so handler threads drain for real and no zombies survive.
        """
        if self._server is not None:
            server = self._server
            self._server = None
            await asyncio.get_running_loop().run_in_executor(None, server.close)
        self._stop_event.set()
        logger.info("CXP direct-transport server stopped")

    async def serve_forever(self) -> None:
        """Block until :meth:`stop` is called."""
        if self._server is None:
            raise RuntimeError(
                "DirectCXPServer.start() must be called before serve_forever()"
            )
        await self._stop_event.wait()

    def _dispatch(
        self, method: bytes, payload: bytes, cancelled: Callable[[], bool]
    ) -> tuple[int, bytes] | Iterator[tuple[int, bytes]]:
        """Runs on a cxpt handler thread; bridges to the async servicer.

        Unary methods return a ``(status, bytes)`` tuple; the two
        server-streaming methods return a generator DirectServer drains into
        ``STREAM_ITEM``/``STREAM_END`` frames. Any servicer exception simply
        propagates -- DirectServer turns it into ``STATUS_INTERNAL``, which is
        an acceptable default (CXPServicer signals errors via ``CXPReceipt
        .success=False`` rather than raising, so this path is effectively
        best-effort/defensive).

        ``cancelled`` (DirectServer.Handler's third argument, backlog #23) is
        threaded straight through to ``_drive_stream`` for the streaming
        branch; unary calls have no use for it.

        Defense-in-depth: a closed loop means ``stop()`` (or the surrounding
        ``asyncio.run()``) has already torn down the event loop this handler
        thread depends on, so a ``run_coroutine_threadsafe(...).result()``
        call here would either hang (loop gone but not yet reflected) or
        raise out of a worker thread. Reporting ``STATUS_UNAVAILABLE``
        up front is cheap and keeps every dispatch path off a dead loop.
        """
        loop = self._loop
        if loop is None:
            return STATUS_UNIMPLEMENTED, b""
        if loop.is_closed():
            return STATUS_UNAVAILABLE, b""
        unary = self._unary.get(method)
        if unary is not None:
            handler, request_cls = unary
            request = request_cls.FromString(payload)
            response = asyncio.run_coroutine_threadsafe(
                handler(request, _DISPATCH_CONTEXT), loop
            ).result()
            return STATUS_OK, response.SerializeToString()
        stream = self._streams.get(method)
        if stream is not None:
            handler, request_cls = stream
            request = request_cls.FromString(payload)
            agen = handler(request, _DISPATCH_CONTEXT)
            return self._drive_stream(agen, loop, cancelled)
        return STATUS_UNIMPLEMENTED, b""

    @staticmethod
    def _drive_stream(
        agen: AsyncGenerator[Any],
        loop: asyncio.AbstractEventLoop,
        cancelled: Callable[[], bool],
    ) -> Iterator[tuple[int, bytes]]:
        """Pull a servicer async generator's items onto the handler thread.

        Re-checks ``loop.is_closed()`` on every item (not just once at
        ``_dispatch`` entry): the loop can close mid-stream, between one
        ``__anext__()`` and the next. On a closed loop the stream just ends
        cleanly here -- DirectServer's caller (``_serve_stream``) still sends
        the terminating ``STREAM_END`` frame, so no hang, no raise out of the
        handler thread.

        Backlog #23: ``DirectServer._serve_stream`` closes the generator this
        method returns as soon as it notices a client CANCEL (or a dead
        connection) -- via ``.close()``, which raises ``GeneratorExit`` here at
        the current ``yield``. The same happens on ordinary garbage collection
        once nothing references this generator. Either way the ``finally``
        below closes the underlying servicer ``agen`` on the loop thread (the
        only thread allowed to touch it) so it releases whatever it holds --
        e.g. ``CXPServicer.subscribe``'s queue wait -- instead of continuing to
        run for a caller that no longer exists.

        That covers cancellation observed BETWEEN items. It is not enough on
        its own: a single ``agen.__anext__()`` can block far longer than one
        item's worth of work -- ``CXPServicer.subscribe``/``receive_packet``
        each await their packet queue with a many-second internal timeout --
        and nothing checks ``cancelled`` while that single call is still in
        flight, so a CANCEL landing mid-wait would otherwise sit unnoticed for
        the rest of that timeout (a live, reproduced regression, not a
        theoretical one). ``_next_or_cancelled`` closes that gap by racing
        the pull against a ``cancelled`` poll on the SAME loop.
        """
        try:
            while True:
                if loop.is_closed():
                    return
                try:
                    item = asyncio.run_coroutine_threadsafe(
                        DirectCXPServer._next_or_cancelled(agen, cancelled), loop
                    ).result()
                except StopAsyncIteration:
                    return
                except _StreamCancelled:
                    return
                yield STATUS_OK, item.SerializeToString()
        finally:
            if not loop.is_closed():
                with contextlib.suppress(Exception):
                    asyncio.run_coroutine_threadsafe(agen.aclose(), loop).result()

    @staticmethod
    async def _next_or_cancelled(
        agen: AsyncGenerator[Any], cancelled: Callable[[], bool]
    ) -> Any:
        """Await ``agen.__anext__()``, polling ``cancelled`` concurrently.

        Runs entirely on the loop thread (scheduled via
        ``run_coroutine_threadsafe`` by the caller). If ``cancelled()`` flips
        true before the pull settles, the in-flight ``__anext__()`` task is
        cancelled -- interrupting even a many-second internal wait inside the
        servicer -- and :class:`_StreamCancelled` is raised so
        ``_drive_stream`` stops promptly instead of only noticing between
        items.
        """
        anext_task = asyncio.ensure_future(agen.__anext__())
        try:
            while not anext_task.done():
                if cancelled():
                    anext_task.cancel()
                    with contextlib.suppress(BaseException):
                        await anext_task
                    raise _StreamCancelled
                await asyncio.wait({anext_task}, timeout=_CANCEL_POLL_INTERVAL_S)
            return anext_task.result()
        finally:
            if not anext_task.done():
                anext_task.cancel()


@traces(spec="001", story="US-002", reqs=["REQ-001-US-002-012"])
class DirectCXPClient:
    """The CXP client, over the cxpt direct transport.

    Proto encode/decode is delegated to
    :class:`~orchestrata_ultimate.cxp.proto_codec.CXPPacketCodec`. Until
    2026-08-07 that meant constructing an unconnected ``GRPCClient`` purely to
    borrow its private ``_packet_to_proto``/``_proto_to_packet`` -- safe only
    because that constructor performed no I/O. The codec is now a real,
    stateless API, so the borrowed-private dependency is gone.
    """

    def __init__(
        self,
        config: CXPClientConfig | None = None,
        *,
        target: str | None = None,
        nworkers: int = 4,
        env: Mapping[str, str] | None = None,
    ) -> None:
        self._config = config or CXPClientConfig()
        self._target = target or _resolve_target(
            self._config.host, self._config.port, env
        )
        self._nworkers = nworkers
        self._client: AsyncDirectClient | None = None
        self._connected = False
        self._codec = CXPPacketCodec()

    async def connect(self) -> None:
        """Connect to the cxpt direct-transport server.

        ``cxpt.Client(target, nworkers)`` does a blocking ``connect(2)`` on the
        calling thread, so construction runs in a thread — otherwise no
        ``asyncio.wait_for`` around this coroutine can bound it.
        ``AsyncDirectClient.__init__`` binds no loop state (loop is bound
        lazily on first call), so off-loop construction is safe.
        """
        self._client = await asyncio.to_thread(
            AsyncDirectClient, self._target, self._nworkers
        )
        self._connected = True
        logger.info("Connected to CXP direct-transport server at %s", self._target)

    async def disconnect(self) -> None:
        """Disconnect from the cxpt direct-transport server."""
        if self._client is not None:
            client, self._client = self._client, None
            await client.aclose()
        self._connected = False
        logger.info("Disconnected from CXP direct-transport server")

    def _require_client(self) -> AsyncDirectClient:
        if self._client is None:
            raise RuntimeError("direct-transport client not initialized")
        return self._client

    async def _call_bounded(self, client: Any, method: str, payload: bytes) -> bytes:
        """Issue a unary call bounded by the configured ``timeout_seconds``.

        ``CXPClientConfig.timeout_seconds`` was dead configuration until
        2026-08-07: nothing read it, and the transport awaited its reply
        future with no deadline. A kernel that died mid-request therefore hung
        the caller forever -- the gateway's "fall back when the kernel is
        unreachable" path could never run, because it never regained control.
        """
        result: bytes = await client.call(
            method, payload, timeout=self._config.timeout_seconds
        )
        return result

    async def send_packet(self, packet: CXPPacketModel) -> dict[str, Any]:
        """Send a CXP packet via the direct transport, with bounded retries.

        Retries on transport failure, honouring ``CXPClientConfig.max_retries``
        with the same escalating backoff the retired grpc client used. Until
        2026-08-07 this method called once and never retried, so ``max_retries``
        was dead configuration: the gateway and dispatcher both construct
        clients expecting the default of 3 attempts and were silently getting
        one, turning every transient blip into a propagated failure. The
        transport swap must not quietly cost reliability the callers already
        assume.

        ``max_retries=0`` means the loop never runs and no RPC is attempted --
        the caller asked for zero attempts and gets the default failure.
        """
        if not self._connected:
            await self.connect()
        client = self._require_client()
        proto_packet = self._codec.packet_to_proto(packet)
        wire = proto_packet.SerializeToString()

        for attempt in range(self._config.max_retries):
            try:
                response = await self._call_bounded(client, "SendPacket", wire)
            except (RpcError, TimeoutError, OSError) as exc:
                logger.warning("Attempt %d failed: %s", attempt + 1, exc)
                if attempt == self._config.max_retries - 1:
                    raise
                await asyncio.sleep(0.5 * (attempt + 1))
                continue
            receipt = CXPReceipt.FromString(response)
            return {
                "success": receipt.success,
                "message_id": receipt.message_id,
                "error": receipt.error_message,
                "latency_ms": receipt.delivery_latency_ms,
            }

        return {"success": False, "error": "Max retries exceeded"}

    async def send_packets(self, packets: list[CXPPacketModel]) -> list[dict[str, Any]]:
        """Send multiple CXP packets.

        The direct transport has no client-streaming RPC equivalent to gRPC's
        ``send_packet_stream``, so (per design) this loops ``send_packet`` --
        one receipt dict per packet, unlike ``GRPCClient.send_packets``'s
        single batch-summary dict.
        """
        return [await self.send_packet(packet) for packet in packets]

    async def receive_packets(
        self,
        agent_id: str,
        timeout_seconds: int = 10,
        max_messages: int = 100,
    ) -> list[dict[str, Any]]:
        """Receive packets queued for an agent."""
        if not self._connected:
            await self.connect()
        client = self._require_client()
        request = CXPReceiveRequest(
            agent_id=agent_id,
            timeout_seconds=timeout_seconds,
            max_messages=max_messages,
        )
        packets: list[dict[str, Any]] = []
        async for chunk in client.call_stream(
            "ReceivePacket", request.SerializeToString()
        ):
            packets.append(self._codec.proto_to_packet(CXPPacket.FromString(chunk)))
            if len(packets) >= max_messages:
                break
        return packets

    async def subscribe(
        self,
        agent_id: str,
        subjects: list[str] | None = None,
        from_agents: list[str] | None = None,
    ) -> list[dict[str, Any]]:
        """Subscribe to packets matching criteria."""
        if not self._connected:
            await self.connect()
        client = self._require_client()
        subscription = CXPSubscription(
            agent_id=agent_id,
            subjects=subjects or [],
            from_agents=from_agents or [],
        )
        return [
            self._codec.proto_to_packet(CXPPacket.FromString(chunk))
            async for chunk in client.call_stream(
                "Subscribe", subscription.SerializeToString()
            )
        ]

    async def acknowledge(
        self,
        message_id: str,
        agent_id: str,
        status: str = "success",
    ) -> bool:
        """Acknowledge receipt of a packet."""
        if not self._connected:
            await self.connect()
        client = self._require_client()
        ack_status = (
            AckStatus.ACK_STATUS_SUCCESS
            if status == "success"
            else AckStatus.ACK_STATUS_FAILURE
        )
        ack = CXPAck(message_id=message_id, agent_id=agent_id, status=ack_status)
        response = await self._call_bounded(
            client, "Acknowledge", ack.SerializeToString()
        )
        receipt = CXPReceipt.FromString(response)
        return bool(receipt.success)

    async def health_check(self) -> dict[str, Any]:
        """Check server health."""
        if not self._connected:
            await self.connect()
        client = self._require_client()
        request = CXPHealthRequest(agent_id="client")
        response = await self._call_bounded(
            client, "HealthCheck", request.SerializeToString()
        )
        health = CXPHealthResponse.FromString(response)
        return {
            "healthy": health.healthy,
            "status": health.status,
            "uptime_seconds": health.uptime_seconds,
            "packets_sent": health.packets_sent,
            "packets_received": health.packets_received,
        }

    async def get_status(self, requester: str = "gateway") -> list[dict[str, Any]]:
        """Fetch the live volk roster from the kernel (synchronous status query)."""
        if not self._connected:
            await self.connect()
        client = self._require_client()
        request = StatusRequest(requester=requester)
        response = await self._call_bounded(
            client, "GetStatus", request.SerializeToString()
        )
        report = StatusReport.FromString(response)
        return [
            {
                "alias": v.alias,
                "role": v.role,
                "status": v.status,
                "summary": v.summary,
                "current_thread": v.current_thread,
                "busy": v.busy,
            }
            for v in report.volks
        ]