orchestrata_ultimate.gateway.mcp_transport

Reference › orchestrata_ultimate › gateway › mcp_transport

mcp_transport is the gateway's MCP transport (ADR-0099, SPEC-014 US-013): the orchestrata-mcp process that an MCP client such as Claude launches over stdio. It is not part of the gateway server. It is a client of the running gateway, so every tool call becomes one JSON-RPC request over the gateway's WebSocket API on port 8765. Nothing imports it; main is the entry point. Package narrative: gateway.

McpTransport builds its tool table from the declarative method_table: each exposed method_spec row becomes an mcp_tool_spec, and list_tools adds one extra tool, approve. call is the main flow:

• approve issues a single-use token from the mcp_approval_gate for one gated tool, and refuses any tool that is unknown or not gated.
• An unknown tool is refused. A gated tool without a redeemable token is refused with PERMISSION_DENIED before anything reaches the gateway.
• Otherwise _rpc sends the request, opening the WebSocket on first use through _connect, which authenticates with the GATEWAY_TOKEN setting and refuses to connect when it is empty. A lock serialises the exchange.
• A connection error, timeout or WebSocket error closes the socket and returns GATEWAY_UNAVAILABLE; an RPC error is returned as rpc_error.

Every outcome (approved, refused, failed, ok) is written to the mcp_ledger, and refusals and failures return the ledger row id. Host and port come from the constructor or the GATEWAY_HOST and GATEWAY_PORT settings. Settings are read first from the process environment, then from a private mapping loaded by main from the deployed environment layer, so an exported variable wins. run_stdio serves the tools through the MCP SDK server named orchestrata and closes the connection on exit.

At a glance

Class diagram

Dependency diagram

API

Classes

McpTransport

Decorators: traces(spec='014', story='US-013', reqs=['REQ-014-204', 'REQ-014-205', 'REQ-014-206']) · Line 65

Stdio MCP server whose tools forward to the running gateway.

Functions

Source

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

"""McpTransport -- the gateway's MCP transport over stdio (ADR-0099).

``orchestrata-mcp`` is a WebSocket JSON-RPC client of the *running* gateway at
``GATEWAY_HOST:GATEWAY_PORT``, authenticating with ``GATEWAY_TOKEN`` from the
deployed environment (never a flag, never written to the ledger). Its tool
table is generated at start from ``METHODS`` (``gateway/method_table.py``),
the same table ``RPCHandler`` builds its dispatch dict from -- nothing is
hand-listed, so a method added to the table appears here with no edit.
Streaming methods, legacy aliases and ``auth.authenticate`` are
``expose=False`` and are not tools.

Every call is forwarded as the JSON-RPC method it names and therefore passes
the gateway's own ``ScopeGuard``, auth and audit. A gated (mutating) tool
needs the single-use token the ``approve`` tool issues; without one the call
is refused with a ledger row id and nothing is sent to the gateway.

Author: Mohamed
"""

# Spec: 014 - WebSocket Gateway

from __future__ import annotations

import asyncio
import json
import os
from collections.abc import Callable, Mapping
from typing import Any

import websockets
from mcp import types
from mcp.server import Server
from mcp.server.stdio import stdio_server
from websockets.protocol import State

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.env_bootstrap import load_env_contract
from orchestrata_ultimate.gateway.mcp_approval_gate import McpApprovalGate
from orchestrata_ultimate.gateway.mcp_ledger import McpLedger
from orchestrata_ultimate.gateway.mcp_tool_spec import TOKEN_ARG, McpToolSpec
from orchestrata_ultimate.gateway.method_spec import MethodSpec
from orchestrata_ultimate.gateway.method_table import METHODS
from orchestrata_ultimate.gateway.models import MessageType, WSMessage

APPROVE_TOOL = "approve"
REFUSED_EXIT_CODE = 13
FAILED_EXIT_CODE = 1
DEFAULT_TIMEOUT_SECONDS = 60.0
_APPROVE_SCHEMA = {
    "type": "object",
    "properties": {
        "tool": {"type": "string", "description": "the mutating tool to approve"},
        "reason": {"type": "string", "description": "why the call is approved"},
    },
    "required": ["tool", "reason"],
    "additionalProperties": False,
}


@traces(
    spec="014",
    story="US-013",
    reqs=["REQ-014-204", "REQ-014-205", "REQ-014-206"],
)
class McpTransport:
    """Stdio MCP server whose tools forward to the running gateway."""

    def __init__(
        self,
        methods: tuple[MethodSpec, ...] = METHODS,
        *,
        ledger: McpLedger | None = None,
        gate: McpApprovalGate | None = None,
        host: str | None = None,
        port: int | None = None,
        token: str | None = None,
        timeout: float = DEFAULT_TIMEOUT_SECONDS,
        deployed_env: Mapping[str, str] | None = None,
    ) -> None:
        self._deployed: Mapping[str, str] = deployed_env or {}
        self._ledger = ledger if ledger is not None else McpLedger()
        self._gate = gate if gate is not None else McpApprovalGate()
        self._host = host
        self._port = port
        self._token = token
        self._timeout = timeout
        self.tools: dict[str, McpToolSpec] = {
            tool.name: tool
            for tool in (McpToolSpec.from_method(m) for m in methods if m.expose)
        }
        self._ws: Any = None
        self._io_lock = asyncio.Lock()

    # -- tool table -------------------------------------------------------

    def list_tools(self) -> list[types.Tool]:
        """One tool per exposed method, plus ``approve``."""
        tools = [
            types.Tool(
                name=tool.name,
                description=tool.description,
                input_schema=tool.input_schema,
            )
            for tool in self.tools.values()
        ]
        tools.append(
            types.Tool(
                name=APPROVE_TOOL,
                description="Issue a single-use approval token for one mutating tool.",
                input_schema=_APPROVE_SCHEMA,
            )
        )
        return tools

    # -- gateway connection ----------------------------------------------

    def _setting(self, name: str, default: str = "") -> str:
        """An exported variable wins; else the deployed environment layer."""
        return os.environ.get(name) or self._deployed.get(name) or default

    def _uri(self) -> str:
        host = self._host or self._setting("GATEWAY_HOST", "localhost")
        port = self._port or int(self._setting("GATEWAY_PORT", "8765"))
        return f"ws://{host}:{port}"

    async def _connect(self) -> Any:
        """Open the WebSocket and authenticate with the gateway token."""
        token = self._token or self._setting("GATEWAY_TOKEN")
        if not token:
            raise PermissionError("GATEWAY_TOKEN is not set")
        ws = await websockets.connect(self._uri(), close_timeout=self._timeout)
        try:
            frame = WSMessage(msg_type=MessageType.AUTH, payload={"token": token})
            await ws.send(json.dumps(frame.to_dict()))
            reply = json.loads(await ws.recv())
            if reply.get("msg_type") != MessageType.AUTH.value:
                raise PermissionError(
                    f"gateway refused authentication: {reply.get('payload')}"
                )
        except BaseException:
            await ws.close()
            raise
        return ws

    async def _rpc(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
        """Send one JSON-RPC request; return the matching response object."""
        async with self._io_lock:
            if self._ws is None or self._ws.state is not State.OPEN:
                self._ws = await self._connect()
            request_id = id(params) ^ int.from_bytes(os.urandom(4), "big")
            await self._ws.send(
                json.dumps(
                    {
                        "jsonrpc": "2.0",
                        "method": method,
                        "params": params,
                        "id": request_id,
                    }
                )
            )
            while True:
                data = json.loads(await self._ws.recv())
                if data.get("id") == request_id and (
                    "result" in data or "error" in data
                ):
                    return dict(data)

    async def aclose(self) -> None:
        """Close the gateway connection, if one is open."""
        if self._ws is not None:
            await self._ws.close()
            self._ws = None

    # -- call path --------------------------------------------------------

    def _refuse(
        self, ledger_tool: str, message: str, **detail: Any
    ) -> tuple[str, bool]:
        """Refuse before anything reaches the gateway; tie it to a ledger row."""
        row = self._ledger.append_row(ledger_tool, "refused", REFUSED_EXIT_CODE)
        return json.dumps({"message": message, **detail, "ledger_row": row}), True

    def _approve(self, arguments: dict[str, Any]) -> tuple[str, bool]:
        tool = str(arguments.get("tool", ""))
        spec = self.tools.get(tool)
        if spec is None or not spec.gated:
            return self._refuse(
                APPROVE_TOOL, f"{tool!r} is not an approval-gated tool", tool=tool
            )
        token = self._gate.issue(tool)
        self._ledger.append_row(tool, "approved")
        return token, False

    async def call(
        self, name: str, arguments: dict[str, Any] | None
    ) -> tuple[str, bool]:
        """Run one tool call; return ``(text, is_error)``."""
        args = dict(arguments or {})
        if name == APPROVE_TOOL:
            return self._approve(args)
        spec = self.tools.get(name)
        if spec is None:
            return self._refuse(name, f"unknown tool {name!r}", tool=name)
        token = args.pop(TOKEN_ARG, None)
        if spec.gated and not self._gate.redeem(token, name):
            return self._refuse(name, "PERMISSION_DENIED", capability=name)
        try:
            response = await asyncio.wait_for(
                self._rpc(spec.method, args), timeout=self._timeout
            )
        except (
            OSError,
            TimeoutError,
            ValueError,
            websockets.WebSocketException,
        ) as exc:
            await self.aclose()
            row = self._ledger.append_row(name, "failed", FAILED_EXIT_CODE)
            detail = {"message": "GATEWAY_UNAVAILABLE", "reason": str(exc)}
            return json.dumps({**detail, "ledger_row": row}), True
        if "error" in response:
            row = self._ledger.append_row(name, "failed", FAILED_EXIT_CODE)
            error = response["error"]
            return json.dumps({"rpc_error": error, "ledger_row": row}), True
        self._ledger.append_row(name, "ok")
        return json.dumps(response.get("result"), default=str), False

    # -- MCP SDK handlers -------------------------------------------------

    async def on_list_tools(self, ctx: Any, params: Any) -> types.ListToolsResult:
        return types.ListToolsResult(tools=self.list_tools())

    async def on_call_tool(
        self, ctx: Any, params: types.CallToolRequestParams
    ) -> types.CallToolResult:
        text, is_error = await self.call(params.name, params.arguments)
        return types.CallToolResult(
            content=[types.TextContent(text=text or "ok")], is_error=is_error
        )


async def run_stdio(
    transport_factory: Callable[[], McpTransport] = McpTransport,
) -> None:
    """Serve the generated tool table over stdio until the client closes."""
    transport = transport_factory()
    server: Server = Server(
        "orchestrata",
        on_list_tools=transport.on_list_tools,
        on_call_tool=transport.on_call_tool,
    )
    try:
        async with stdio_server() as (read_stream, write_stream):
            await server.run(
                read_stream, write_stream, server.create_initialization_options()
            )
    finally:
        await transport.aclose()


def main() -> None:
    """``orchestrata-mcp`` entry point.

    Claude launches this from a project directory with none of the gateway's
    variables exported, so the deployed environment layer (``<release>/.env``,
    the secrets file ``install-orchestrata.sh`` installs) is read here -- into
    a private mapping, never ``os.environ``, so anything exported still wins.
    """
    deployed: dict[str, str] = {}
    load_env_contract(environ=deployed)
    asyncio.run(run_stdio(lambda: McpTransport(deployed_env=deployed)))


if __name__ == "__main__":
    main()