mcp-proxy-adapter-dev 8.11.16__py3-none-any.whl

mcp_proxy_adapter/client/jsonrpc_client/ws_job_status_parts/channel.py

@@ -0,0 +1,159 @@
+"""Bidirectional WebSocket channel of the job-status client."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any, AsyncIterator, Callable, Dict, Optional
+
+import aiohttp
+
+from mcp_proxy_adapter.client.jsonrpc_client.exceptions import classify_transport_error
+from mcp_proxy_adapter.client.jsonrpc_client.transport import JsonRpcTransport
+from mcp_proxy_adapter.client.jsonrpc_client.ws_job_status_parts.helpers import (
+    _dispatch_inbound_message_hook,
+    _raise_websocket_closed,
+    logger,
+    resolve_ws,
+)
+
+
+class BidirectionalWsChannel:
+    """
+    Bidirectional WebSocket channel: send and receive JSON messages over /ws.
+
+    Use as async context manager. Same auth (headers, TLS) as transport.
+    """
+
+    def __init__(
+        self,
+        transport: JsonRpcTransport,
+        receive_timeout: Optional[float] = None,
+        heartbeat: Optional[float] = None,
+        on_inbound_message: Optional[Callable[[Dict[str, Any]], Any]] = None,
+        close_timeout: Optional[float] = None,
+    ) -> None:
+        """
+        Initialize bidirectional WebSocket channel.
+
+        Args:
+            transport: JSON-RPC transport for connection details.
+            receive_timeout: Single-receive cap in seconds (``None`` resolves
+                through the client config and constants; 0 = no cap).
+            heartbeat: Heartbeat interval in seconds (same resolution; 0 = off).
+            close_timeout: Bound of the close handshake (same resolution).
+        """
+        self._close_timeout = close_timeout
+        self._transport = transport
+        self._receive_timeout = receive_timeout
+        self._heartbeat = heartbeat
+        self._on_inbound_message = on_inbound_message
+        self._ws: Any = None
+        self._session: Any = None
+
+    async def __aenter__(self) -> "BidirectionalWsChannel":
+        """Enter async context and establish WebSocket connection.
+
+        Returns:
+            Self for use in async with statement.
+        """
+        ws_url = self._transport.ws_url
+        headers = {
+            k: v
+            for k, v in self._transport.headers.items()
+            if k.lower() != "content-type"
+        }
+        ssl = self._transport.get_ws_ssl_context()
+        hb = resolve_ws(self._transport, "heartbeat_interval", self._heartbeat)
+        rt = resolve_ws(self._transport, "receive_timeout", self._receive_timeout)
+        if rt is not None:
+            rt += resolve_ws(self._transport, "receive_timeout_buffer", None) or 0.0
+        self._session = await self._transport.get_ws_session()
+        try:
+            self._ws = await self._session.ws_connect(
+                ws_url,
+                ssl=ssl,
+                headers=headers,
+                heartbeat=hb,
+                receive_timeout=rt,
+            )
+        except asyncio.CancelledError:
+            raise
+        except (aiohttp.ClientError, OSError) as exc:
+            # aiohttp has several provider-specific connection exceptions;
+            # the shared classifier preserves their cause and emits the
+            # documented client category.
+            raise classify_transport_error(
+                exc,
+                operation="websocket.connect",
+                url=ws_url,
+                websocket=True,
+            ) from exc
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        """Exit async context and close WebSocket connection."""
+        if self._ws:
+            close_timeout = resolve_ws(
+                self._transport, "close_timeout", self._close_timeout
+            )
+            try:
+                await asyncio.wait_for(self._ws.close(), timeout=close_timeout)
+            except asyncio.TimeoutError:
+                logger.warning(
+                    "WebSocket close timed out after %.1f seconds; "
+                    "connection may not have fully shut down",
+                    close_timeout,
+                )
+            self._ws = None
+        # Session is shared and owned by the transport; do not close it here.
+        # Just release the reference.
+        self._session = None
+
+    async def send_json(self, payload: Dict[str, Any]) -> None:
+        """Send a JSON object to the server."""
+        if self._ws is None:
+            raise RuntimeError("Channel not connected; use async with first")
+        await self._ws.send_str(json.dumps(payload))
+
+    async def receive_iter(self) -> AsyncIterator[Dict[str, Any]]:
+        """
+        Async iterator over incoming JSON messages (server -> client).
+
+        Yields each parsed JSON object. Stops when the connection closes or errors.
+        """
+        if self._ws is None:
+            raise RuntimeError("Channel not connected; use async with first")
+        while True:
+            try:
+                msg = await self._ws.receive()
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                raise classify_transport_error(
+                    exc,
+                    operation="websocket.receive",
+                    url=self._transport.ws_url,
+                    websocket=True,
+                ) from exc
+            if msg.type == aiohttp.WSMsgType.TEXT:
+                try:
+                    parsed = json.loads(msg.data)
+                    await _dispatch_inbound_message_hook(
+                        self._on_inbound_message, parsed
+                    )
+                    yield parsed
+                except json.JSONDecodeError:
+                    logger.warning("WS non-JSON: %s", msg.data[:200])
+            elif msg.type in (
+                aiohttp.WSMsgType.CLOSE,
+                aiohttp.WSMsgType.CLOSING,
+                aiohttp.WSMsgType.CLOSED,
+                aiohttp.WSMsgType.ERROR,
+            ):
+                _raise_websocket_closed(
+                    msg,
+                    message="WebSocket closed before a terminal message",
+                    operation="websocket.receive",
+                    url=self._transport.ws_url,
+                )

mcp_proxy_adapter/client/jsonrpc_client/ws_job_status_parts/command_channel.py

@@ -0,0 +1,154 @@
+"""Transport-owned command sockets of the job-status client."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from typing import Any, AsyncIterator, Dict, Optional
+
+from mcp_proxy_adapter.client.jsonrpc_client.ws_job_status_parts.channel import (
+    BidirectionalWsChannel,
+)
+from mcp_proxy_adapter.client.jsonrpc_client.ws_job_status_parts.helpers import (
+    resolve_ws,
+)
+
+
+class _CommandWsConnection:
+    """Transport-owned socket; at most one active command owns its lifecycle."""
+
+    def __init__(self, channel: BidirectionalWsChannel) -> None:
+        self.channel = channel
+        self.loop = asyncio.get_running_loop()
+        self.lock = asyncio.Lock()
+        self.leases: Dict[str, Any] = {}
+        self.reader: Optional[asyncio.Task[Any]] = None
+
+    async def acquire(self, lease: Any) -> None:
+        async with self.lock:
+            if self.reader is None or self.reader.done():
+                await self.channel.__aenter__()
+                self.reader = asyncio.create_task(self._read())
+            self.leases[lease.request_id] = lease
+            lease._ws = self.channel._ws
+
+    async def _read(self) -> None:
+        try:
+            async for message in self.channel.receive_iter():
+                rid = str(message.get("request_id") or "")
+                jid = str(message.get("job_id") or "")
+                for lease in tuple(self.leases.values()):
+                    if rid == lease.request_id or (jid and jid == lease.job_id):
+                        if jid:
+                            lease.job_id = jid
+                        if message.get("event") in {
+                            "command_completed",
+                            "command_failed",
+                            "command_cancelled",
+                            "command_error",
+                        } or (
+                            message.get("action") == "cancel_command"
+                            and message.get("success") is True
+                        ):
+                            lease.terminal = True
+                        lease.queue.put_nowait(message)
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            for lease in tuple(self.leases.values()):
+                lease.queue.put_nowait(exc)
+        finally:
+            await self.channel.__aexit__(None, None, None)
+
+    async def close(self) -> None:
+        if self.reader is not None:
+            self.reader.cancel()
+            with contextlib.suppress(asyncio.CancelledError, Exception):
+                await self.reader
+        await self.channel.__aexit__(None, None, None)
+
+
+class _CommandWsChannel(BidirectionalWsChannel):
+    """Exclusive command lease, reusing an idle socket after terminal/detach."""
+
+    def __init__(self, channel: BidirectionalWsChannel, request_id: str) -> None:
+        self._transport = channel._transport
+        self._template = channel
+        self.request_id = request_id
+        self.job_id = ""
+        self.terminal = False
+        self.detached = False
+        self.queue: asyncio.Queue[Any] = asyncio.Queue()
+        self._ws: Any = None
+        self.connection: Optional[_CommandWsConnection] = None
+
+    async def __aenter__(self) -> "_CommandWsChannel":
+        transport = self._transport
+        pools = getattr(transport, "_command_ws_connections", None)
+        if pools is None:
+            pools = transport._command_ws_connections = {}
+        # Never share across loops, authentication, TLS policy or receive caps.
+        key = (
+            asyncio.get_running_loop(),
+            transport.ws_url,
+            tuple(sorted(transport.headers.items())),
+            id(transport.get_ws_ssl_context()),
+            resolve_ws(transport, "receive_timeout", self._template._receive_timeout),
+            resolve_ws(transport, "heartbeat_interval", self._template._heartbeat),
+            resolve_ws(transport, "close_timeout", self._template._close_timeout),
+        )
+        candidates = pools.setdefault(key, [])
+        connection = next((item for item in candidates if not item.leases), None)
+        if connection is None:
+            connection = _CommandWsConnection(self._template)
+            candidates.append(connection)
+        self.connection = connection
+        # Reserve before the first await so a concurrent opener cannot claim it.
+        connection.leases[self.request_id] = self
+        try:
+            await connection.acquire(self)
+        except BaseException:
+            connection.leases.pop(self.request_id, None)
+            raise
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        connection = self.connection
+        if connection is not None:
+            if not self.terminal and not self.detached:
+                # A premature close must reach the server's existing per-socket
+                # caller-loss policy. No active sibling shares this connection.
+                await connection.close()
+            elif self.job_id and self._ws is not None and not self._ws.closed:
+                with contextlib.suppress(Exception):
+                    await self.send_json(
+                        {
+                            "action": "unsubscribe",
+                            "job_id": self.job_id,
+                            "subscription_type": "command_session",
+                        }
+                    )
+            connection.leases.pop(self.request_id, None)
+        self._ws = None
+
+    async def _close_idle_connection(self) -> None:
+        """Honor explicit close without disturbing a subsequently acquired lease."""
+        connection = self.connection
+        if connection is not None:
+            async with connection.lock:
+                if not connection.leases:
+                    await connection.close()
+
+    async def send_json(self, payload: Dict[str, Any]) -> None:
+        if self.connection is None:
+            raise RuntimeError("Channel not connected; use async with first")
+        await self.connection.channel.send_json(payload)
+        if payload.get("caller_loss") == "detach":
+            self.detached = True
+
+    async def receive_iter(self) -> AsyncIterator[Dict[str, Any]]:
+        while True:
+            item = await self.queue.get()
+            if isinstance(item, Exception):
+                raise item
+            yield item

mcp_proxy_adapter/client/jsonrpc_client/ws_job_status_parts/helpers.py

@@ -0,0 +1,76 @@
+"""Shared helpers of the WebSocket job-status client."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from typing import Any, AsyncIterator, Callable, Dict, NoReturn, Optional
+
+import aiohttp
+
+from mcp_proxy_adapter.client.jsonrpc_client.exceptions import (
+    ClientWebSocketClosedError,
+)
+from mcp_proxy_adapter.core import constants as _constants
+
+logger = logging.getLogger("mcp_proxy_adapter.client.jsonrpc_client.ws_job_status")
+
+
+@contextlib.asynccontextmanager
+async def _connect(session: Any, url: str, **kwargs: Any) -> AsyncIterator[Any]:
+    """Open ``session.ws_connect(url, **kwargs)`` whether it is awaitable or a context manager."""
+    opened = session.ws_connect(url, **kwargs)
+    if hasattr(opened, "__aenter__"):
+        async with opened as ws:
+            yield ws
+        return
+    ws = await opened
+    try:
+        yield ws
+    finally:
+        await ws.close()
+
+
+def ws_config_of(transport: Any) -> Any:
+    """Return the ``ws_config`` mapping carried by ``transport`` (or ``None``)."""
+    return getattr(transport, "ws_config", None)
+
+
+def resolve_ws(transport: Any, name: str, per_call: Any) -> Optional[float]:
+    """Resolve websocket timeout ``name`` through per-call, client config, constant."""
+    return _constants.resolve_ws_timeout(name, per_call, ws_config_of(transport))
+
+
+def _raise_websocket_closed(
+    msg: aiohttp.WSMessage,
+    *,
+    message: str,
+    operation: str,
+    url: str,
+) -> NoReturn:
+    """Raise the public close error from stable aiohttp protocol evidence."""
+    close_code = msg.data if isinstance(msg.data, int) else 1006
+    provider_cause = (
+        msg.data
+        if isinstance(msg.data, BaseException)
+        else aiohttp.WebSocketError(close_code, "WebSocket peer closed the connection")
+    )
+    raise ClientWebSocketClosedError(
+        message,
+        operation=operation,
+        url=url,
+        details={"close_type": str(msg.type), "close_code": close_code},
+    ) from provider_cause
+
+
+async def _dispatch_inbound_message_hook(
+    hook: Optional[Callable[[Dict[str, Any]], Any]],
+    data: Dict[str, Any],
+) -> None:
+    """Invoke *hook* with *data*; await if it returns a coroutine."""
+    if hook is None:
+        return
+    out = hook(data)
+    if asyncio.iscoroutine(out):
+        await out

mcp_proxy_adapter/client/proxy.py

@@ -0,0 +1,124 @@
+"""
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+
+Lightweight asynchronous client library for registry (proxy) server used in examples.
+"""
+
+from typing import Any, Dict, List, Optional
+
+import httpx
+
+
+class ProxyClient:
+    """Asynchronous client for proxy server."""
+
+    def __init__(self, base_url: str):
+        """
+        Initialize proxy client.
+
+        Args:
+            base_url: Base URL of the proxy server
+        """
+        self.base_url = base_url.rstrip("/")
+        self._client: Optional[httpx.AsyncClient] = None
+
+    async def _get_client(self) -> httpx.AsyncClient:
+        """Get or create async HTTP client."""
+        if self._client is None:
+            self._client = httpx.AsyncClient(timeout=5.0)
+        return self._client
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        if self._client:
+            await self._client.aclose()
+            self._client = None
+
+    async def _get_json(self, path: str) -> Dict[str, Any]:
+        """GET ``path`` on the proxy and return its decoded JSON body."""
+        client = await self._get_client()
+        r = await client.get(f"{self.base_url}{path}")
+        r.raise_for_status()
+        return r.json()  # type: ignore[no-any-return]
+
+    async def health(self) -> Dict[str, Any]:
+        """
+        Get proxy server health status.
+
+        Returns:
+            Health status information
+        """
+        return await self._get_json("/proxy/health")
+
+    async def register(
+        self,
+        name: str,
+        url: str,
+        capabilities: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Register a server with the proxy.
+
+        Args:
+            name: Server name
+            url: Server URL
+            capabilities: Optional list of capabilities
+            metadata: Optional metadata dictionary
+
+        Returns:
+            Registration result
+        """
+        payload = {
+            "name": name,
+            "url": url,
+            "capabilities": capabilities or [],
+            "metadata": metadata or {},
+        }
+        client = await self._get_client()
+        r = await client.post(f"{self.base_url}/register", json=payload)
+        r.raise_for_status()
+        return r.json()  # type: ignore[no-any-return]
+
+    async def unregister(self, name: str) -> Dict[str, Any]:
+        """
+        Unregister a server from the proxy.
+
+        Args:
+            name: Server name
+
+        Returns:
+            Unregistration result
+        """
+        payload = {"name": name, "url": "", "capabilities": [], "metadata": {}}
+        client = await self._get_client()
+        r = await client.post(f"{self.base_url}/unregister", json=payload)
+        r.raise_for_status()
+        return r.json()  # type: ignore[no-any-return]
+
+    async def list_servers(self) -> Dict[str, Any]:
+        """
+        List all servers registered with the proxy.
+
+        Returns:
+            List of registered servers
+        """
+        return await self._get_json("/proxy/list")
+
+    async def heartbeat(self, name: str, url: str) -> Dict[str, Any]:
+        """
+        Send heartbeat to the proxy.
+
+        Args:
+            name: Server name
+            url: Server URL
+
+        Returns:
+            Heartbeat result
+        """
+        payload = {"name": name, "url": url, "capabilities": [], "metadata": {}}
+        client = await self._get_client()
+        r = await client.post(f"{self.base_url}/proxy/heartbeat", json=payload)
+        r.raise_for_status()
+        return r.json()  # type: ignore[no-any-return]

mcp_proxy_adapter/commands/__init__.py

@@ -0,0 +1,27 @@
+"""
+Commands module initialization file.
+
+Command classes are not re-exported here: built-in commands are registered
+through ``builtin_commands.register_builtin_commands()``, and embedder
+commands through the registry/hooks. Only the framework primitives every
+consumer needs are exported.
+"""
+
+from mcp_proxy_adapter.commands.base import Command
+from mcp_proxy_adapter.commands.command_registry import CommandRegistry, registry
+from mcp_proxy_adapter.commands.dependency_container import (
+    DependencyContainer,
+    container,
+)
+from mcp_proxy_adapter.commands.result import CommandResult, ErrorResult, SuccessResult
+
+__all__ = [
+    "Command",
+    "CommandResult",
+    "SuccessResult",
+    "ErrorResult",
+    "registry",
+    "CommandRegistry",
+    "container",
+    "DependencyContainer",
+]