mcp_client.py

from contextlib import asynccontextmanager

from mcp.client.streamable_http import streamable_http_client
from mcp.shared._httpx_utils import create_mcp_http_client
from mcp import ClientSession
from config import get_settings


class RemoteMCPClient:
    """
    Async context-manager wrapper around an MCP streamable-HTTP session.

    Usage (all work must happen inside the `async with` block so that every
    anyio cancel scope is entered and exited in the *same* task):

        async with RemoteMCPClient(url, token) as client:
            tools = await client.list_tools()
    """

    def __init__(self, server_url: str, api_key: str):
        self.server_url = server_url
        self.api_key = api_key
        self.session: ClientSession | None = None
        self._exit_stack = None

    # ── async context manager ─────────────────────────────────────────────────

    async def __aenter__(self):
        from contextlib import AsyncExitStack
        self._exit_stack = AsyncExitStack()
        await self._exit_stack.__aenter__()

        headers = {"Authorization": f"Bearer {self.api_key}"}
        http_client = create_mcp_http_client(headers=headers)
        await self._exit_stack.enter_async_context(http_client)

        # streamable_http_client yields (read_stream, write_stream) — 2 values
        read_stream, write_stream = await self._exit_stack.enter_async_context(
            streamable_http_client(self.server_url, http_client=http_client)
        )

        self.session = await self._exit_stack.enter_async_context(
            ClientSession(read_stream, write_stream)
        )
        await self.session.initialize()
        return self

    async def __aexit__(self, *exc_info):
        await self._exit_stack.__aexit__(*exc_info)
        self.session = None

    # ── public API ────────────────────────────────────────────────────────────

    async def list_tools(self) -> list[dict]:
        result = await self.session.list_tools()
        return [{"name": t.name, "description": t.description} for t in result.tools]

    async def call_tool(self, tool_name: str, params: dict) -> dict:
        try:
            result = await self.session.call_tool(tool_name, params)
            return {"status": "success", "data": result.content}
        except Exception as e:
            return {"status": "error", "message": str(e)}


def get_github_mcp_client() -> RemoteMCPClient:
    settings = get_settings()
    return RemoteMCPClient(settings.github_mcp_url, settings.github_mcp_token)