kiwimatecoder 0.1.0__py3-none-any.whl

kiwimatecoder/mcp/__init__.py

@@ -0,0 +1,45 @@
+"""MCP (Model Context Protocol) client support.
+
+Connects to configured MCP servers over stdio or streamable HTTP, discovers
+their tools, and registers them as ``mcp__<server>__<tool>`` in the shared tool
+registry. The ``mcp_servers`` config section holds one spec per server: a
+``command`` (+ ``args``/``env``) for stdio, or a ``url`` (+ ``headers``) for
+HTTP. Authentication is whatever headers or environment the user configures;
+interactive OAuth is not implemented.
+"""
+
+from kiwimatecoder.mcp.client import (
+    HttpMcpClient,
+    McpClient,
+    McpError,
+    StdioMcpClient,
+)
+from kiwimatecoder.mcp.manager import (
+    McpLoadResult,
+    McpManager,
+    get_manager,
+    mcp_tool_name,
+    set_manager,
+)
+from kiwimatecoder.mcp.protocol import (
+    CLIENT_INFO,
+    MCP_PROTOCOL_VERSION,
+    flatten_content,
+    parse_sse_messages,
+)
+
+__all__ = [
+    "CLIENT_INFO",
+    "HttpMcpClient",
+    "MCP_PROTOCOL_VERSION",
+    "McpClient",
+    "McpError",
+    "McpLoadResult",
+    "McpManager",
+    "StdioMcpClient",
+    "flatten_content",
+    "get_manager",
+    "mcp_tool_name",
+    "parse_sse_messages",
+    "set_manager",
+]

kiwimatecoder/mcp/client.py

@@ -0,0 +1,385 @@
+"""MCP clients over stdio subprocesses and streamable HTTP.
+
+Both transports share the JSON-RPC request/response plumbing in
+:class:`JsonRpcClient`; a transport only has to implement ``_exchange``. The
+protocol is intentionally minimal: one request in flight at a time, responses
+matched by ``id``, and bounded waits so a hung server can never wedge a
+session. Interactive OAuth is out of scope; HTTP auth is whatever headers the
+user configures.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import queue
+import subprocess
+import threading
+import time
+from collections import deque
+from collections.abc import Mapping, Sequence
+from typing import Any, Protocol
+
+import httpx
+
+from kiwimatecoder.mcp.errors import McpError
+from kiwimatecoder.mcp.protocol import (
+    CLIENT_INFO,
+    MCP_PROTOCOL_VERSION,
+    flatten_content,
+    flatten_resource_contents,
+    notification_message,
+    parse_response,
+    parse_sse_messages,
+    request_message,
+)
+
+INITIALIZE_TIMEOUT = 10.0
+CALL_TIMEOUT = 60.0
+STDERR_TAIL_LINES = 100
+
+
+class McpClient(Protocol):
+    """The surface the manager needs from any MCP transport."""
+
+    capabilities: dict[str, Any]
+    server_info: dict[str, Any]
+
+    def initialize(self) -> dict[str, Any]: ...
+    def list_tools(self) -> list[dict[str, Any]]: ...
+    def call_tool(self, name: str, arguments: dict[str, Any]) -> str: ...
+    def list_resources(self) -> list[dict[str, Any]]: ...
+    def read_resource(self, uri: str) -> str: ...
+    def close(self) -> None: ...
+
+
+class JsonRpcClient:
+    """Shared initialize/tools/resources calls over an abstract transport."""
+
+    def __init__(
+        self,
+        *,
+        initialize_timeout: float = INITIALIZE_TIMEOUT,
+        call_timeout: float = CALL_TIMEOUT,
+    ) -> None:
+        self.initialize_timeout = initialize_timeout
+        self.call_timeout = call_timeout
+        self.server_info: dict[str, Any] = {}
+        self.capabilities: dict[str, Any] = {}
+        self._next_id = 0
+
+    def _exchange(
+        self, message: dict[str, Any], timeout: float
+    ) -> dict[str, Any] | None:
+        raise NotImplementedError
+
+    def _request(
+        self, method: str, params: dict[str, Any] | None, timeout: float
+    ) -> Any:
+        self._next_id += 1
+        message = request_message(self._next_id, method, params)
+        response = self._exchange(message, timeout)
+        if response is None:
+            raise McpError(f"{method}: no response from server")
+        return parse_response(response, method)
+
+    def _notify(self, method: str, params: dict[str, Any] | None = None) -> None:
+        self._exchange(notification_message(method, params), self.initialize_timeout)
+
+    def initialize(self) -> dict[str, Any]:
+        """Perform the MCP handshake and cache server info/capabilities."""
+        result = self._request(
+            "initialize",
+            {
+                "protocolVersion": MCP_PROTOCOL_VERSION,
+                "capabilities": {},
+                "clientInfo": CLIENT_INFO,
+            },
+            self.initialize_timeout,
+        )
+        if isinstance(result, dict):
+            server_info = result.get("serverInfo")
+            capabilities = result.get("capabilities")
+            self.server_info = server_info if isinstance(server_info, dict) else {}
+            self.capabilities = capabilities if isinstance(capabilities, dict) else {}
+        else:
+            result = {}
+        self._notify("notifications/initialized")
+        return result
+
+    def list_tools(self) -> list[dict[str, Any]]:
+        result = self._request("tools/list", {}, self.call_timeout)
+        raw = result.get("tools") if isinstance(result, dict) else None
+        if not isinstance(raw, list):
+            return []
+        return [tool for tool in raw if isinstance(tool, dict)]
+
+    def call_tool(self, name: str, arguments: dict[str, Any]) -> str:
+        """Call a tool and return its flattened content.
+
+        Raises :class:`McpError` when the server reports ``isError`` so the
+        registered tool can surface a failed ``ToolResult``.
+        """
+        result = self._request(
+            "tools/call", {"name": name, "arguments": dict(arguments)}, self.call_timeout
+        )
+        payload = result if isinstance(result, dict) else {}
+        content = flatten_content(payload.get("content"))
+        if payload.get("isError"):
+            raise McpError(content or f"tool '{name}' reported an error")
+        return content
+
+    def list_resources(self) -> list[dict[str, Any]]:
+        result = self._request("resources/list", {}, self.call_timeout)
+        raw = result.get("resources") if isinstance(result, dict) else None
+        if not isinstance(raw, list):
+            return []
+        return [resource for resource in raw if isinstance(resource, dict)]
+
+    def read_resource(self, uri: str) -> str:
+        result = self._request("resources/read", {"uri": uri}, self.call_timeout)
+        payload = result if isinstance(result, dict) else {}
+        return flatten_resource_contents(payload.get("contents"))
+
+    def close(self) -> None:  # pragma: no cover - overridden by transports
+        pass
+
+
+class StdioMcpClient(JsonRpcClient):
+    """Talk to an MCP server that speaks newline-delimited JSON-RPC on stdio.
+
+    A background reader thread feeds stdout lines into a queue so requests can
+    time out instead of blocking forever, and a second thread drains stderr
+    into a bounded tail used for diagnostics (draining is what keeps a chatty
+    server from filling its pipe buffer and deadlocking).
+    """
+
+    def __init__(
+        self,
+        command: str,
+        args: Sequence[str] | None = None,
+        env: Mapping[str, str] | None = None,
+        *,
+        initialize_timeout: float = INITIALIZE_TIMEOUT,
+        call_timeout: float = CALL_TIMEOUT,
+        stderr_limit: int = STDERR_TAIL_LINES,
+    ) -> None:
+        super().__init__(
+            initialize_timeout=initialize_timeout, call_timeout=call_timeout
+        )
+        self.command = command
+        self.args = [str(arg) for arg in (args or [])]
+        self._stderr_tail: deque[str] = deque(maxlen=stderr_limit)
+        self._lines: queue.Queue[str | None] = queue.Queue()
+        self._write_lock = threading.Lock()
+        self._closed = False
+        process_env = os.environ.copy()
+        if env:
+            process_env.update({str(key): str(value) for key, value in env.items()})
+        try:
+            self._process = subprocess.Popen(
+                [command, *self.args],
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                encoding="utf-8",
+                env=process_env,
+                bufsize=1,
+            )
+        except OSError as exc:
+            raise McpError(f"could not start '{command}': {exc}") from exc
+        self._reader = threading.Thread(
+            target=self._read_stdout, name=f"mcp-stdout-{command}", daemon=True
+        )
+        self._reader.start()
+        self._stderr_reader = threading.Thread(
+            target=self._read_stderr, name="mcp-stderr", daemon=True
+        )
+        self._stderr_reader.start()
+
+    @property
+    def running(self) -> bool:
+        """Whether the server process is still alive."""
+        return self._process.poll() is None
+
+    def stderr_tail(self) -> str:
+        """The most recent stderr lines, joined (bounded)."""
+        return "\n".join(self._stderr_tail)
+
+    def _read_stdout(self) -> None:
+        stream = self._process.stdout
+        if stream is None:
+            self._lines.put(None)
+            return
+        try:
+            for line in stream:
+                self._lines.put(line)
+        except (OSError, ValueError):
+            pass
+        finally:
+            self._lines.put(None)
+
+    def _read_stderr(self) -> None:
+        stream = self._process.stderr
+        if stream is None:
+            return
+        try:
+            for line in stream:
+                self._stderr_tail.append(line.rstrip("\n"))
+        except (OSError, ValueError):
+            pass
+
+    def _exit_message(self) -> str:
+        message = "MCP server exited"
+        if self._process.returncode is not None:
+            message += f" (code {self._process.returncode})"
+        tail = self.stderr_tail()
+        if tail:
+            message += f": {tail}"
+        return message
+
+    def _exchange(
+        self, message: dict[str, Any], timeout: float
+    ) -> dict[str, Any] | None:
+        if self._closed:
+            raise McpError("MCP client is closed")
+        line = json.dumps(message) + "\n"
+        with self._write_lock:
+            stdin = self._process.stdin
+            if stdin is None or self._process.poll() is not None:
+                raise McpError(self._exit_message())
+            try:
+                stdin.write(line)
+                stdin.flush()
+            except (OSError, ValueError) as exc:
+                raise McpError(f"could not write to MCP server: {exc}") from exc
+            if "id" not in message:
+                return None
+            deadline = time.monotonic() + timeout
+            while True:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise McpError(
+                        f"timed out waiting for MCP response after {timeout:g}s"
+                    )
+                try:
+                    response_line = self._lines.get(timeout=remaining)
+                except queue.Empty:
+                    raise McpError(
+                        f"timed out waiting for MCP response after {timeout:g}s"
+                    ) from None
+                if response_line is None:
+                    raise McpError(self._exit_message())
+                try:
+                    response = json.loads(response_line)
+                except json.JSONDecodeError:
+                    continue
+                if (
+                    isinstance(response, dict)
+                    and response.get("id") == message["id"]
+                    and ("result" in response or "error" in response)
+                ):
+                    return response
+
+    def close(self) -> None:
+        """Terminate the server process, escalating to kill on timeout."""
+        if self._closed:
+            return
+        self._closed = True
+        process = self._process
+        try:
+            if process.stdin is not None:
+                try:
+                    process.stdin.close()
+                except OSError:
+                    pass
+            process.terminate()
+            try:
+                process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                try:
+                    process.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    pass
+        except OSError:
+            pass
+
+
+class HttpMcpClient(JsonRpcClient):
+    """POST JSON-RPC to a streamable-HTTP MCP endpoint.
+
+    Responses may be a plain JSON body or an SSE stream; both are accepted.
+    One :class:`httpx.Client` is reused for the life of the connection, with
+    configured headers carrying any bearer/API auth.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        headers: Mapping[str, str] | None = None,
+        *,
+        initialize_timeout: float = INITIALIZE_TIMEOUT,
+        call_timeout: float = CALL_TIMEOUT,
+        transport: httpx.BaseTransport | None = None,
+    ) -> None:
+        super().__init__(
+            initialize_timeout=initialize_timeout, call_timeout=call_timeout
+        )
+        self.url = url
+        request_headers = {
+            "Accept": "application/json, text/event-stream",
+            "Content-Type": "application/json",
+        }
+        if headers:
+            request_headers.update(
+                {str(key): str(value) for key, value in headers.items()}
+            )
+        self._client = httpx.Client(
+            headers=request_headers,
+            timeout=httpx.Timeout(call_timeout, connect=initialize_timeout),
+            transport=transport,
+            follow_redirects=True,
+        )
+
+    def _exchange(
+        self, message: dict[str, Any], timeout: float
+    ) -> dict[str, Any] | None:
+        try:
+            response = self._client.post(self.url, json=message, timeout=timeout)
+        except httpx.HTTPError as exc:
+            raise McpError(f"HTTP request to {self.url} failed: {exc}") from exc
+        if response.status_code >= 400:
+            body = response.text[:200].replace("\n", " ").strip()
+            raise McpError(
+                f"HTTP {response.status_code} from {self.url}"
+                + (f": {body}" if body else "")
+            )
+        if "id" not in message:
+            return None
+        return self._parse_body(response, message["id"])
+
+    @staticmethod
+    def _parse_body(response: httpx.Response, request_id: Any) -> dict[str, Any]:
+        text = response.text or ""
+        stripped = text.lstrip()
+        if stripped.startswith(("{", "[")):
+            try:
+                payload = json.loads(text)
+            except json.JSONDecodeError:
+                payload = None
+            if isinstance(payload, dict):
+                if "result" in payload or "error" in payload:
+                    return payload
+            elif isinstance(payload, list):
+                for item in payload:
+                    if isinstance(item, dict) and item.get("id") == request_id:
+                        return item
+        for item in parse_sse_messages(text):
+            if item.get("id") == request_id:
+                return item
+        raise McpError("MCP server returned no JSON-RPC response")
+
+    def close(self) -> None:
+        self._client.close()

kiwimatecoder/mcp/errors.py

@@ -0,0 +1,7 @@
+"""Errors raised by the MCP client layer."""
+
+from __future__ import annotations
+
+
+class McpError(Exception):
+    """Any MCP transport, protocol, or server-reported failure."""

kiwimatecoder/mcp/manager.py

@@ -0,0 +1,253 @@
+"""MCP server lifecycle: connect, register tools, and tear down cleanly.
+
+The manager reads the ``mcp_servers`` config section, connects each enabled
+server, discovers its tools, and registers them in the shared tool registry as
+``mcp__<server>__<tool>``. Servers that fail to start or initialize are
+recorded and skipped; loading never raises, so a broken MCP entry cannot stop a
+session from starting. A module-level slot lets the REPL install the session's
+manager so slash commands can inspect and reload it.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from rich.console import Console
+from rich.markup import escape
+
+from kiwimatecoder import config, events, tools
+from kiwimatecoder.mcp.client import (
+    HttpMcpClient,
+    McpClient,
+    McpError,
+    StdioMcpClient,
+)
+from kiwimatecoder.redaction import redact
+from kiwimatecoder.tools.base import FunctionTool, ToolResult
+
+MCP_SOURCE_PREFIX = "mcp:"
+
+
+def _sanitize(identifier: str) -> str:
+    """Reduce an arbitrary server/tool name to ``[A-Za-z0-9_]``."""
+    cleaned = "".join(
+        ch if (ch.isascii() and (ch.isalnum() or ch == "_")) else "_"
+        for ch in identifier
+    )
+    return cleaned or "unnamed"
+
+
+def mcp_tool_name(server: str, tool: str) -> str:
+    """Return the registry name for a tool exposed by ``server``."""
+    return f"mcp__{_sanitize(server)}__{_sanitize(tool)}"
+
+
+@dataclass
+class McpLoadResult:
+    """Outcome of one :meth:`McpManager.load` pass."""
+
+    servers: list[str] = field(default_factory=list)
+    failed: list[tuple[str, str]] = field(default_factory=list)
+    tools: list[str] = field(default_factory=list)
+
+
+class McpManager:
+    """Owns live MCP clients and the tools registered on their behalf."""
+
+    def __init__(
+        self, console: Console | None = None, bus: events.EventBus | None = None
+    ) -> None:
+        self._console = console if console is not None else Console()
+        self._bus = bus
+        self._clients: dict[str, McpClient] = {}
+        self._tools: dict[str, list[str]] = {}
+        self._failures: dict[str, str] = {}
+        self._resource_counts: dict[str, int] = {}
+
+    @property
+    def clients(self) -> dict[str, McpClient]:
+        """A copy of the connected server -> client map."""
+        return dict(self._clients)
+
+    def client(self, name: str) -> McpClient | None:
+        """Return the live client for ``name``, if connected."""
+        return self._clients.get(name)
+
+    def is_connected(self, name: str) -> bool:
+        return name in self._clients
+
+    def tools_for(self, name: str) -> list[str]:
+        """Registered tool names contributed by ``name``."""
+        return list(self._tools.get(name, []))
+
+    def failure_for(self, name: str) -> str | None:
+        """The recorded failure reason for ``name``, if it failed to load."""
+        return self._failures.get(name)
+
+    def resource_count(self, name: str) -> int | None:
+        """Number of resources discovered on ``name`` (0 when unsupported)."""
+        return self._resource_counts.get(name)
+
+    def load(
+        self, servers: dict[str, Any] | None = None
+    ) -> McpLoadResult:
+        """Connect enabled servers and register their tools. Never raises."""
+        self.shutdown()
+        if servers is None:
+            servers = config.get_mcp_servers()
+        result = McpLoadResult()
+        for name in sorted(servers):
+            spec = servers[name]
+            if not isinstance(spec, dict) or spec.get("disabled"):
+                continue
+            client: McpClient | None = None
+            try:
+                client = self._build_client(name, spec)
+                client.initialize()
+                tool_defs = client.list_tools()
+            except Exception as exc:
+                if client is not None:
+                    self._close_quietly(client)
+                self._record_failure(name, exc, result)
+                continue
+            self._clients[name] = client
+            try:
+                registered = self._register_tools(name, client, tool_defs, result)
+            except Exception as exc:
+                self._unregister_tools(name)
+                self._clients.pop(name, None)
+                self._close_quietly(client)
+                self._record_failure(name, exc, result)
+                continue
+            self._tools[name] = registered
+            self._resource_counts[name] = self._count_resources(client)
+            result.servers.append(name)
+            result.tools.extend(registered)
+        return result
+
+    def reload(self) -> McpLoadResult:
+        """Drop every connection and load the configured servers again."""
+        return self.load()
+
+    def shutdown(self) -> None:
+        """Close every client, unregister its tools, and clear state."""
+        for name in list(self._clients):
+            client = self._clients.pop(name, None)
+            if client is not None:
+                self._close_quietly(client)
+        for name in list(self._tools):
+            self._unregister_tools(name)
+        self._failures.clear()
+        self._resource_counts.clear()
+
+    def _build_client(self, name: str, spec: dict[str, Any]) -> McpClient:
+        if spec.get("command"):
+            return StdioMcpClient(
+                str(spec["command"]),
+                args=spec.get("args"),
+                env=spec.get("env"),
+            )
+        return HttpMcpClient(str(spec["url"]), headers=spec.get("headers"))
+
+    def _register_tools(
+        self,
+        server: str,
+        client: McpClient,
+        tool_defs: list[dict[str, Any]],
+        result: McpLoadResult,
+    ) -> list[str]:
+        registered: list[str] = []
+        for definition in tool_defs:
+            tool_name = str(definition.get("name") or "").strip()
+            if not tool_name:
+                continue
+            registry_name = mcp_tool_name(server, tool_name)
+            parameters = definition.get("inputSchema")
+            if not isinstance(parameters, dict):
+                parameters = {"type": "object", "properties": {}}
+            annotations = definition.get("annotations")
+            read_only = (
+                bool(annotations.get("readOnlyHint"))
+                if isinstance(annotations, dict)
+                else False
+            )
+            description = str(
+                definition.get("description")
+                or f"MCP tool '{tool_name}' from server '{server}'."
+            )
+            tool = FunctionTool(
+                name=registry_name,
+                description=description,
+                parameters=parameters,
+                func=self._make_caller(client, tool_name),
+                writes=not read_only,
+                runs=False,
+            )
+            try:
+                tools.register_tool(tool, source=f"{MCP_SOURCE_PREFIX}{server}")
+            except ValueError as exc:
+                result.failed.append((server, f"{registry_name}: {exc}"))
+                continue
+            registered.append(registry_name)
+        return registered
+
+    def _make_caller(
+        self, client: McpClient, tool_name: str
+    ) -> Callable[[dict[str, Any], Any], ToolResult]:
+        def call(arguments: dict[str, Any], session: Any) -> ToolResult:
+            try:
+                content = client.call_tool(tool_name, arguments)
+            except McpError as exc:
+                return ToolResult.error(str(exc))
+            return ToolResult(content=content or "(no output)")
+
+        return call
+
+    def _count_resources(self, client: McpClient) -> int:
+        capabilities = client.capabilities or {}
+        if "resources" not in capabilities:
+            return 0
+        try:
+            return len(client.list_resources())
+        except Exception:
+            return 0
+
+    def _unregister_tools(self, name: str) -> None:
+        for tool_name in self._tools.pop(name, []):
+            try:
+                tools.unregister_tool(tool_name)
+            except Exception:
+                pass
+
+    def _record_failure(
+        self, name: str, exc: BaseException, result: McpLoadResult
+    ) -> None:
+        reason = redact(str(exc)) or exc.__class__.__name__
+        self._failures[name] = reason
+        result.failed.append((name, reason))
+        self._console.print(
+            f"[dim]MCP server '{escape(name)}' failed: {escape(reason)}[/dim]"
+        )
+
+    @staticmethod
+    def _close_quietly(client: McpClient) -> None:
+        try:
+            client.close()
+        except Exception:
+            pass
+
+
+_MANAGER: McpManager | None = None
+
+
+def get_manager() -> McpManager | None:
+    """Return the session's MCP manager, if the REPL installed one."""
+    return _MANAGER
+
+
+def set_manager(manager: McpManager | None) -> None:
+    """Install (or clear) the session's MCP manager."""
+    global _MANAGER
+    _MANAGER = manager