orchestrata_ultimate.gateway.rpc

Reference › orchestrata_ultimate › gateway › rpc

rpc is the thin JSON-RPC dispatcher of the gateway. RPCHandler takes the connection manager (session), the pub/sub manager (pubsub), the ring buffer (ring_buffer) and the audit writer (audit), and it constructs one handler object per domain from the handlers package. protocol and server are its importers. See gateway.

_build_method_registry turns the declarative method_table into a dict from method name to bound handler coroutine. It indexes the handler instances by class path, binds each method_spec row to its instance, substitutes an "unavailable" answer for owners that are not wired, and raises LookupError if a row names an owner with no instance. Nothing is listed by hand, so a new table row becomes dispatchable without editing this function.

dispatch looks up the method and raises RPCError with code -32601 if it is unknown. It then writes an audit event named rpc.<method>, runs the handler, and records the call latency in metrics whether the handler succeeded or raised. The set_* methods inject late-bound collaborators into the handlers, including the processor, graph repository, memory engine, test runner, kernel client and kernel status, pipeline client, policy engine, intent evaluator, admin token and the runtime_config store; server calls them during start-up.

At a glance

Class diagram

Dependency diagram

Truncated at 25 nodes: showing 22 of 31 imports and 2 of 2 importers.

API

Constants

Classes

RPCHandler

Line 66

Thin dispatcher that routes RPC calls to handler classes.

Functions

Source

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

# mypy: ignore-errors
"""Thin RPC dispatcher for the WebSocket gateway.

Routes all SPEC-014 RPC methods to focused handler classes:
CXP (FR-015), Volk (FR-016), Session (FR-017), Governance (FR-018),
Knowledge (FR-019), Gateway (FR-020), State sync, and Auth.
Plus SPEC-018 Fireflies MCP bridge methods.
"""

import logging
import time
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.gateway.audit import AuditWriter
from orchestrata_ultimate.gateway.exceptions import RPCError
from orchestrata_ultimate.gateway.handlers._policy_gate import PolicyGate
from orchestrata_ultimate.gateway.handlers._scope_guard import ScopeGuard
from orchestrata_ultimate.gateway.handlers.auth import AuthHandler
from orchestrata_ultimate.gateway.handlers.config import ConfigHandler
from orchestrata_ultimate.gateway.handlers.cxp import CxpHandler
from orchestrata_ultimate.gateway.handlers.fireflies import FirefliesHandler
from orchestrata_ultimate.gateway.handlers.gateway import (
    GATEWAY_VERSION,
    GatewayHandler,
)
from orchestrata_ultimate.gateway.handlers.governance import (
    GovernanceHandler,
)
from orchestrata_ultimate.gateway.handlers.knowledge import (
    KnowledgeHandler,
)
from orchestrata_ultimate.gateway.handlers.memory import MemoryHandler
from orchestrata_ultimate.gateway.handlers.orchestrator import (
    OrchestratorHandler,
)
from orchestrata_ultimate.gateway.handlers.projects import (
    ProjectsHandler,
)
from orchestrata_ultimate.gateway.handlers.session import SessionHandler
from orchestrata_ultimate.gateway.handlers.state import StateHandler
from orchestrata_ultimate.gateway.handlers.test_metrics import TestMetricsHandler
from orchestrata_ultimate.gateway.handlers.ultimatetui import UltimateTUIHandler
from orchestrata_ultimate.gateway.handlers.volk import (
    VolkHandler,
    load_volk_registry,
)
from orchestrata_ultimate.gateway.handlers.workflow import WorkflowHandler
from orchestrata_ultimate.gateway.method_spec import MethodSpec
from orchestrata_ultimate.gateway.method_table import METHODS
from orchestrata_ultimate.gateway.metrics import rpc_latency
from orchestrata_ultimate.gateway.pubsub import PubSubManager
from orchestrata_ultimate.gateway.ring_buffer import RedisRingBuffer
from orchestrata_ultimate.gateway.runtime_config import SecretRedactor
from orchestrata_ultimate.gateway.session import ConnectionManager

__all__ = [
    "GATEWAY_VERSION",
    "RPCHandler",
    "load_volk_registry",
]

logger = logging.getLogger(__name__)


class RPCHandler:
    """Thin dispatcher that routes RPC calls to handler classes."""

    def __init__(
        self,
        connections: ConnectionManager,
        pubsub: PubSubManager,
        ring_buffer: RedisRingBuffer,
        audit: AuditWriter,
        processor: Any = None,
        session_root: str | None = None,
    ) -> None:
        """Wire dispatcher to gateway subsystems.

        Args:
            connections: Connection lifecycle manager.
            pubsub: Pub/sub subscription manager.
            ring_buffer: Redis ring buffer for event history.
            audit: Graph audit writer.
            processor: Message processor for broadcast/push.
            session_root: Explicit session-persistence root, sourced from
                the caller's ``GatewayConfig.session_root`` (never a
                literal) -- forwarded to ``SessionHandler``, which itself
                falls back to the ``SESSION_ROOT`` environment variable and
                raises if neither is available (fix/session-root, RUL-011).
        """
        self._audit = audit
        self._connections = connections

        # Build focused handler instances
        self._gateway = GatewayHandler(connections, pubsub, ring_buffer, processor)
        self._cxp = CxpHandler(pubsub, processor)
        self._volk = VolkHandler(pubsub)
        self._session = SessionHandler(pubsub, session_root=session_root)
        self._governance = GovernanceHandler(pubsub)
        self._knowledge = KnowledgeHandler()
        self._projects = ProjectsHandler()
        self._orchestrator = OrchestratorHandler()
        self._tui = UltimateTUIHandler()
        self._state = StateHandler(ring_buffer)
        self._auth = AuthHandler()
        self._fireflies = FirefliesHandler()
        self._memory = MemoryHandler()
        self._test_metrics = TestMetricsHandler()

        # SPEC-014 resurrection (US-014-005): workflow control + runtime
        # config service. ScopeGuard is shared by both -- one seam for the
        # per-connection scope every resurrected method attaches to.
        self._scope_guard = ScopeGuard(connections)
        self._workflow = WorkflowHandler(PolicyGate(), self._scope_guard)
        self._config = ConfigHandler(SecretRedactor(), self._scope_guard, pubsub)

        # Wire pub/sub into the orchestrator unconditionally so the UI is
        # subscribed to pipeline.<thread_id> before dispatch even when the graph is
        # absent. set_graph_repository() also does this (via
        # CxpHandler.set_orchestrator_handler), but that path only fires when a
        # graph repository is configured; this call is idempotent with it.
        self._orchestrator.set_pubsub(pubsub)

        # SPEC-003 ingress reconciliation: front the NL ingress with the shared,
        # extensible IntentEvaluator instead of the coarse heuristic parser. Uses
        # the vocabulary tier by default (no infra needed); a monolithic
        # deployment can override with the boot broker's Prolog-backed evaluator
        # via set_intent_evaluator(). Non-fatal — the handler falls back to the
        # heuristic if construction fails.
        try:
            from orchestrata_ultimate.orchestrator.intent_classifier import (
                IntentEvaluator,
            )

            self._orchestrator.set_evaluator(IntentEvaluator())
        except Exception:
            logger.warning("IntentEvaluator unavailable; ingress uses heuristic parser")

        self._methods = _build_method_registry(
            self._gateway,
            self._cxp,
            self._volk,
            self._session,
            self._governance,
            self._knowledge,
            self._projects,
            self._orchestrator,
            self._state,
            self._auth,
            self._fireflies,
            self._memory,
            self._test_metrics,
            tui=self._tui,
            workflow=self._workflow,
            config=self._config,
        )

    def set_processor(self, processor: Any) -> None:
        """Inject the message processor after construction.

        Args:
            processor: MessageProcessor instance.
        """
        self._gateway.set_processor(processor)
        self._cxp.set_processor(processor)
        self._volk.set_processor(processor)
        self._session.set_processor(processor)
        self._governance.set_processor(processor)

    def set_graph_repository(self, graph_repository: Any) -> None:
        """Inject the operational-graph repository for knowledge queries.

        Args:
            graph_repository: An AsyncGraphRepository instance.
        """
        self._knowledge.set_graph_repository(graph_repository)
        try:
            from orchestrata_ultimate.knowledge_graph.graph_client import GraphClient
            from orchestrata_ultimate.knowledge_graph.query_engine import QueryEngine

            self._knowledge.set_query_engine(
                QueryEngine(GraphClient(driver=graph_repository.driver))
            )
        except Exception:
            logger.warning("QueryEngine unavailable for knowledge work-context RPCs")
        self._projects.set_graph_repository(graph_repository)
        self._session.set_graph_repository(graph_repository)
        self._governance.set_graph_repository(graph_repository)
        self._test_metrics.set_graph_repository(graph_repository)
        self._fireflies.set_graph_repository(graph_repository)
        self._orchestrator.set_projects_handler(self._projects)
        self._cxp.set_orchestrator_handler(self._orchestrator)

    @traces(spec="010", story="US-002", reqs=["REQ-010-032"])
    def set_memory_engine(self, engine: Any) -> None:
        """Inject the gateway's shared memory engine singleton.

        Wires the one DI-registered ``MemoryEngine`` into both the memory RPC
        handler (so ``memory.*`` uses it) and the session handler (so
        ``session.close`` migrates that session's hot memories to cold storage).
        Both then operate on the same instance (SPEC-010 FR-001/FR-008).

        Args:
            engine: The gateway's singleton ``MemoryEngine``.
        """
        self._memory.set_engine(engine)
        self._session.set_memory_engine(engine)

    def set_test_runner(self, runner: Any) -> None:
        """Inject the test runner for executing pytest.

        Args:
            runner: TestRunner instance.
        """
        self._test_metrics._runner = runner

    def set_kernel_client(self, client: Any) -> None:
        """Inject the CXP gRPC client (the live volk kernel) into the handlers
        that dispatch user intents to it.

        Args:
            client: A connected CXP client to the agent-runtime CXP kernel.
        """
        self._orchestrator.set_kernel_client(client)
        self._cxp.set_kernel_client(client)
        self._volk.set_kernel_client(client)

    def set_kernel_status(
        self,
        mode: str,
        endpoint: str | None,
        probe: Any = None,
    ) -> None:
        """Record which path ``WebSocketGateway._init_kernel`` took, so
        ``gateway.info`` can report it (the "kernel" block).

        Forwards to the gateway handler -- mirrors the ``set_kernel_client``/
        ``set_pipeline_client`` injection shape: ``WebSocketGateway`` owns the
        connection attempt, ``RPCHandler`` is the thin pass-through, and the
        handler that answers ``gateway.info`` is the one that stores it.

        Args:
            mode: ``"grpc"`` | ``"in-process"`` | ``"unavailable"``.
            endpoint: ``"host:port"`` when ``mode == "grpc"``, else ``None``.
            probe: Optional awaitable returning whether the kernel answers
                right now; keeps the reported block from going stale.
        """
        self._gateway.set_kernel_status(mode, endpoint, probe)

    @traces(
        spec="001",
        story="US-002",
        reqs=["REQ-001-US-002-001", "REQ-001-US-002-003"],
    )
    def set_reply_consumer_id(self, consumer_id: str) -> None:
        """Inject this gateway's minted reply-consumer id (design D6).

        Mirrors :meth:`set_kernel_client`'s pass-through chain shape: both
        the CXP/chat handler (P-B3) and the orchestrator/intent handler
        (P-B6) stamp their outbound kernel-bound packets with it.

        Args:
            consumer_id: This gateway's minted reply-consumer id.
        """
        self._cxp.set_reply_consumer_id(consumer_id)
        self._orchestrator.set_reply_consumer_id(consumer_id)

    def set_intent_evaluator(self, evaluator: Any) -> None:
        """Override the ingress NL evaluator (e.g. with a monolithic
        deployment's boot broker evaluator, which carries the Prolog DCG-grammar
        tier). By default the dispatcher already wires a vocabulary-tier
        ``IntentEvaluator`` at construction.

        Args:
            evaluator: An ``IntentEvaluator`` instance.
        """
        self._orchestrator.set_evaluator(evaluator)

    def set_pipeline_client(self, client: Any) -> None:
        """Inject the in-process pipeline kernel client used as a fallback when
        no live volk kernel is reachable.

        Wired into the *orchestrator* handler (orchestrator.intent ->
        pipeline stage events), NOT the CXP/chat handler: the pipeline client
        cannot round-trip a conversational Concierge reply on ``cxp.traffic``,
        so chat (``cxp.send``) stays kernel-less and reports the explicit
        assistant-unavailable notice (SPEC-032 G243) instead of a fabricated
        reply. Also wired as the workflow handler's control port (SPEC-014
        US-014-005): workflow.start/pause/resume/cancel are only operable
        against this in-process client (D3) -- a pure gRPC-kernel deployment
        never calls this setter, so ``WorkflowHandler`` answers those verbs
        with a structured refusal instead of a silent no-op.

        Args:
            client: A connected PipelineKernelClient.
        """
        self._orchestrator.set_kernel_client(client)
        self._workflow.set_control(client)

    def set_policy_engine(self, engine: Any) -> None:
        """Inject a pre-built governance PolicyEngine into the governance handler.

        Args:
            engine: A PolicyEngine for dry-run evaluation + active rules.
        """
        self._governance.set_policy_engine(engine)

    def set_policy_engine_factory(self, factory: Any) -> None:
        """Register a deferred PolicyEngine factory on the governance handler
        (built lazily on first ``governance.evaluate``), and feed the same
        factory to the workflow PAM++ gate (SPEC-014 REQ-014-080) -- ONE
        real Prolog rule base now enforces both ``governance.evaluate`` and
        ``workflow.*``.

        Args:
            factory: ``() -> PolicyEngine``.
        """
        self._governance.set_policy_engine_factory(factory)
        self._workflow.set_policy_engine_factory(factory)

    def set_admin_token(self, token: Any) -> None:
        """Inject the admin token gating volk.register/deregister.

        Args:
            token: Shared admin secret (or None to leave admin ops denied).
        """
        self._volk.set_admin_token(token)

    def set_runtime_config_store(self, store: Any) -> None:
        """Inject the runtime-config store (SPEC-014 US-014-005) built by
        ``WebSocketGateway._init_config_service``.

        Args:
            store: A connected ``RuntimeConfigStore``.
        """
        self._config.set_store(store)

    @traces(spec="014", story="US-011", reqs=["REQ-014-178"])
    async def dispatch(
        self,
        method: str,
        connection_id: str,
        params: dict[str, Any],
    ) -> dict[str, Any]:
        """Route an RPC call to the appropriate handler.

        Args:
            method: RPC method name.
            connection_id: Caller connection ID.
            params: Method parameters.

        Returns:
            Result dictionary from the handler.

        Raises:
            RPCError: If the method is unknown (-32601).
        """
        handler_fn = self._methods.get(method)
        if not handler_fn:
            raise RPCError(
                f"Method not found: {method}",
                rpc_code=-32601,
            )
        await self._audit.log_event(f"rpc.{method}", connection_id, params)
        # REQ-014-146: dispatch latency by method, recorded around the
        # handler call itself (not the audit write) -- covers both the
        # success and error-raising paths.
        start = time.monotonic()
        try:
            return await handler_fn(connection_id, params)
        finally:
            rpc_latency(method, time.monotonic() - start)


def _class_path(cls: type) -> str:
    """Dotted ``<module>.<Class>`` path, the key a ``MethodSpec`` owner uses."""
    return f"{cls.__module__}.{cls.__qualname__}"


#: Handlers ``RPCHandler`` may be built without; a method owned by one that is
#: absent answers with this structured "unavailable" payload instead of
#: failing dispatch (the behaviour the hand-written lambdas had).
_UNAVAILABLE: dict[str, str] = {
    _class_path(UltimateTUIHandler): "TUI unavailable",
    _class_path(TestMetricsHandler): "Test metrics unavailable",
    _class_path(WorkflowHandler): "Workflow control unavailable",
    _class_path(ConfigHandler): "Config service unavailable",
}


@traces(spec="014", story="US-013", reqs=["REQ-014-203"])
def _build_method_registry(
    gateway: GatewayHandler,
    cxp: CxpHandler,
    volk: VolkHandler,
    session: SessionHandler,
    governance: GovernanceHandler,
    knowledge: KnowledgeHandler,
    projects: ProjectsHandler,
    orchestrator: OrchestratorHandler,
    state: StateHandler,
    auth: AuthHandler,
    fireflies: FirefliesHandler,
    memory: MemoryHandler,
    test_metrics: TestMetricsHandler | None = None,
    tui: UltimateTUIHandler | None = None,
    workflow: WorkflowHandler | None = None,
    config: ConfigHandler | None = None,
    methods: tuple[MethodSpec, ...] = METHODS,
) -> dict[str, Any]:
    """Build the method name to handler function mapping from ``METHODS``.

    The table (``gateway/method_table.py``) names each handler by dotted
    path; this function binds every row to the matching handler instance
    passed in. Nothing is hand-listed here, so a row added to the table is
    dispatchable with no edit to this function.

    Args:
        gateway: Gateway introspection handler.
        cxp: CXP protocol handler.
        volk: Volk registry handler.
        session: Session lifecycle handler.
        governance: Governance query handler.
        knowledge: Knowledge graph handler.
        projects: Projects and tasks handler.
        orchestrator: Orchestrator intent handler.
        state: State sync handler.
        auth: Authentication handler.
        fireflies: Fireflies MCP bridge handler.
        memory: Memory handler.
        test_metrics: Test metrics handler (optional).
        tui: UltimateTUI bridge handler (optional).
        workflow: Workflow control handler (SPEC-014 US-014-005).
        config: Runtime config service handler (SPEC-014 US-014-005).
        methods: The method table to bind (default ``METHODS``).

    Returns:
        Dict mapping RPC method names to async handler functions.

    Raises:
        LookupError: A row names an owner no passed handler instance has.
    """
    owners: dict[str, Any] = {
        _class_path(type(handler)): handler
        for handler in (
            gateway,
            cxp,
            volk,
            session,
            governance,
            knowledge,
            projects,
            orchestrator,
            state,
            auth,
            fireflies,
            memory,
            test_metrics,
            tui,
            workflow,
            config,
        )
        if handler is not None
    }
    registry: dict[str, Any] = {}
    for spec in methods:
        instance = owners.get(spec.owner)
        if instance is not None:
            registry[spec.name] = getattr(instance, spec.attribute)
        elif spec.owner in _UNAVAILABLE:
            registry[spec.name] = _unavailable(_UNAVAILABLE[spec.owner])
        else:
            raise LookupError(
                f"method {spec.name!r}: no handler instance for {spec.owner}"
            )
    return registry


def _unavailable(message: str) -> Any:
    """Handler standing in for an absent optional handler."""

    async def answer(_connection_id: str, _params: dict[str, Any]) -> dict[str, Any]:
        return {"error": message}

    return answer