smqjh-agent/smqjh-agent-runtime/app/mcp_client.py

from __future__ import annotations

import json
import urllib.error
import urllib.request
from typing import Any

from .models import McpToolDescriptor, McpToolResult


class McpClient:
    def __init__(self, endpoint: str, token: str = "") -> None:
        self.endpoint = endpoint
        self.token = token
        self._request_id = 1

    def status(self) -> Any:
        request = urllib.request.Request(self.endpoint, headers=self._headers(), method="GET")
        return self._send(request)

    def list_tools(self) -> list[McpToolDescriptor]:
        result = self._rpc("tools/list", {})
        tools = result.get("tools") if isinstance(result, dict) else None
        return tools if isinstance(tools, list) else []

    def call_tool(self, name: str, arguments: dict[str, Any]) -> McpToolResult:
        return self._rpc("tools/call", {"name": name, "arguments": arguments})

    def _rpc(self, method: str, params: dict[str, Any]) -> Any:
        payload = {
            "jsonrpc": "2.0",
            "id": self._request_id,
            "method": method,
            "params": params,
        }
        self._request_id += 1
        data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        request = urllib.request.Request(self.endpoint, data=data, headers=self._headers(), method="POST")
        response = self._send(request)
        if not isinstance(response, dict):
            raise RuntimeError("MCP returned invalid JSON-RPC payload")
        if response.get("error"):
            error = response["error"]
            if isinstance(error, dict):
                raise RuntimeError(str(error.get("message") or "MCP call failed"))
            raise RuntimeError(str(error))
        if "result" not in response:
            raise RuntimeError("MCP returned no result")
        return response["result"]

    def _send(self, request: urllib.request.Request) -> Any:
        try:
            with urllib.request.urlopen(request, timeout=45) as response:
                raw = response.read().decode("utf-8")
                return json.loads(raw) if raw else {}
        except urllib.error.HTTPError as error:
            raw = error.read().decode("utf-8", errors="replace")
            try:
                payload = json.loads(raw)
            except json.JSONDecodeError:
                payload = {}
            message = ""
            if isinstance(payload, dict):
                rpc_error = payload.get("error")
                if isinstance(rpc_error, dict):
                    message = str(rpc_error.get("message") or "")
                message = message or str(payload.get("message") or "")
            raise RuntimeError(message or f"MCP HTTP {error.code}") from error
        except urllib.error.URLError as error:
            raise RuntimeError(f"MCP connection failed: {error.reason}") from error

    def _headers(self) -> dict[str, str]:
        headers = {"Content-Type": "application/json"}
        token = self.token.strip()
        if token:
            headers["Authorization"] = f"Bearer {token}"
            headers["X-MCP-Token"] = token
        return headers