LLeMbas/src/lembas/services/mcp/client.py

"""Talking to a remote MCP server over streamable HTTP.

One session per call, deliberately. A cached session would need an owner, a
lifetime, eviction, a lock -- a round runs its tools concurrently -- and a
shutdown hook, and the server may expire it underneath all of that anyway.
`ToolContext` is a session-free snapshot precisely so that nothing inside a tool
holds live state. The cost is one extra POST in front of a call that is already
a network round trip inside a reply taking seconds; the upgrade, if it is ever
worth it, is a dict in this module and invisible to everything else.

Redirects are followed by hand and every hop is re-checked, for the reason
`services/fetch.py` gives: an administrator can point this at any URL, and a
name resolving to 127.0.0.1 walks past any check that only reads the text.
"""

from __future__ import annotations

import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager, suppress
from dataclasses import dataclass, field
from typing import Any
from urllib.parse import quote, urlparse

import httpx

from lembas.db.models import SECRET_BEARER, SECRET_HEADER, SECRET_QUERY, McpServer
from lembas.services import fetch as fetch_service
from lembas.services.crypto import decrypt
from lembas.services.mcp import protocol
from lembas.services.mcp.protocol import McpError

log = logging.getLogger(__name__)

# A JSON-RPC message is small; anything this large is a server misbehaving, and
# reading it into memory before parsing is the failure to avoid.
MAX_MESSAGE_BYTES = 4 * 1024 * 1024

# tools/list is paged. Both bounds exist because the whole list goes into every
# request as schema.
MAX_PAGES = 10
MAX_TOOLS = 100

MIN_TIMEOUT, MAX_TIMEOUT = 1, 120
MIN_CHARS, MAX_CHARS = 200, 40_000


@dataclass(frozen=True)
class McpSpec:
    """One server, read while the session was open. See `custom_tools.HttpSpec`."""

    slug: str
    name: str
    url: str
    headers: dict[str, str] = field(default_factory=dict)
    secret: str = ""
    secret_placement: str = SECRET_BEARER
    secret_name: str = "Authorization"
    timeout: int = 30
    max_chars: int = 8000
    allow_private: bool = False


def spec_from(row: McpServer) -> McpSpec:
    return McpSpec(
        slug=row.slug,
        name=row.name or row.slug,
        url=row.url or "",
        headers=dict(row.headers_json or {}),
        secret=decrypt(row.secret_encrypted),
        secret_placement=row.secret_placement,
        secret_name=row.secret_name or "Authorization",
        timeout=min(max(int(row.timeout or 0), MIN_TIMEOUT), MAX_TIMEOUT),
        max_chars=min(max(int(row.max_chars or 0), MIN_CHARS), MAX_CHARS),
        allow_private=bool(row.allow_private),
    )


def _headers(spec: McpSpec) -> dict[str, str]:
    headers = {
        "User-Agent": fetch_service.USER_AGENT,
        "Content-Type": "application/json",
        # Both, because a server may answer either for the same request.
        "Accept": "application/json, text/event-stream",
        **{str(k): str(v) for k, v in spec.headers.items()},
    }
    if spec.secret:
        if spec.secret_placement == SECRET_BEARER:
            headers[spec.secret_name or "Authorization"] = f"Bearer {spec.secret}"
        elif spec.secret_placement == SECRET_HEADER:
            headers[spec.secret_name or "Authorization"] = spec.secret
    return headers


def _endpoint(spec: McpSpec) -> str:
    url = fetch_service.check_url(spec.url, allow_private=spec.allow_private)
    if spec.secret and spec.secret_placement == SECRET_QUERY:
        joiner = "&" if urlparse(url).query else "?"
        url = f"{url}{joiner}{quote(spec.secret_name)}={quote(spec.secret, safe='')}"
    return url


class Session:
    """One initialised conversation with a server."""

    def __init__(self, spec: McpSpec, client: httpx.AsyncClient) -> None:
        self.spec = spec
        self._client = client
        self._url = ""
        self._headers = _headers(spec)
        self._session_id = ""
        self._next_id = 0
        self.protocol_version = ""
        self.server_info: dict[str, Any] = {}

    # --- Transport -----------------------------------------------------------
    async def _post(self, message: dict[str, Any]) -> httpx.Response:
        current = self._url
        headers = dict(self._headers)
        if self._session_id:
            headers["Mcp-Session-Id"] = self._session_id
        if self.protocol_version:
            headers["MCP-Protocol-Version"] = self.protocol_version

        origin = (urlparse(current).scheme, urlparse(current).netloc)
        for _ in range(fetch_service.MAX_REDIRECTS + 1):
            try:
                response = await self._client.post(current, json=message, headers=headers)
            except httpx.RequestError as exc:
                raise McpError(f"Could not reach {self.spec.name}: {exc}") from exc

            if not response.is_redirect:
                return response

            location = response.headers.get("location", "")
            if not location:
                raise McpError(f"{self.spec.name} redirected to nowhere.")
            if response.status_code not in (307, 308):
                # 301, 302 and 303 turn a POST into a GET, which means nothing
                # to a JSON-RPC endpoint. Refused rather than guessed at.
                raise McpError(
                    f"{self.spec.name} answered {response.status_code}, which would "
                    "turn the request into a GET. Point the URL at the endpoint itself."
                )

            current = fetch_service.check_url(
                str(response.url.join(location)), allow_private=self.spec.allow_private
            )
            if (urlparse(current).scheme, urlparse(current).netloc) != origin:
                headers.pop(self.spec.secret_name or "Authorization", None)
                origin = (urlparse(current).scheme, urlparse(current).netloc)

        raise McpError(f"{self.spec.name} redirected too many times.")

    def _read(self, response: httpx.Response, *, request_id: int) -> dict[str, Any]:
        if response.status_code >= 400:
            raise McpError(f"{self.spec.name} returned HTTP {response.status_code}.")
        return protocol.result_of(
            response.content[:MAX_MESSAGE_BYTES],
            response.headers.get("content-type", ""),
            request_id=request_id,
        )

    # --- Lifecycle -----------------------------------------------------------
    async def open(self) -> None:
        self._url = _endpoint(self.spec)
        self._next_id += 1
        request_id = self._next_id

        response = await self._post(
            protocol.request(
                "initialize",
                {
                    "protocolVersion": protocol.PROTOCOL_VERSION,
                    "capabilities": {},
                    "clientInfo": protocol.CLIENT_INFO,
                },
                request_id=request_id,
            )
        )
        # Captured before the body is read: a server that issues one expects it
        # on everything after this, including the initialized notification.
        self._session_id = response.headers.get("mcp-session-id", "")
        result = self._read(response, request_id=request_id)

        self.protocol_version = str(result.get("protocolVersion") or protocol.PROTOCOL_VERSION)
        info = result.get("serverInfo")
        self.server_info = info if isinstance(info, dict) else {}
        if self.protocol_version != protocol.PROTOCOL_VERSION:
            log.info(
                "%s speaks MCP %s, we asked for %s",
                self.spec.name,
                self.protocol_version,
                protocol.PROTOCOL_VERSION,
            )

        await self._post(protocol.notification("notifications/initialized"))

    async def call(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
        """One request, re-initialising once if the session has expired."""
        self._next_id += 1
        request_id = self._next_id
        response = await self._post(protocol.request(method, params, request_id=request_id))

        if response.status_code == 404 and self._session_id:
            # The server dropped the session. One retry, then it is a failure
            # like any other -- a loop here would be a loop against a server
            # that has decided to forget us.
            log.info("%s expired its session; re-initialising", self.spec.name)
            self._session_id = ""
            await self.open()
            self._next_id += 1
            request_id = self._next_id
            response = await self._post(protocol.request(method, params, request_id=request_id))

        return self._read(response, request_id=request_id)

    async def close(self) -> None:
        """Best effort. A server with nothing to clean up answers 405."""
        if not self._session_id:
            return
        headers = {**self._headers, "Mcp-Session-Id": self._session_id}
        if self.protocol_version:
            headers["MCP-Protocol-Version"] = self.protocol_version
        with suppress(httpx.RequestError, McpError):
            await self._client.delete(self._url, headers=headers)


@asynccontextmanager
async def session_for(spec: McpSpec) -> AsyncIterator[Session]:
    """An initialised session, closed afterwards whatever happened."""
    client = httpx.AsyncClient(timeout=spec.timeout, follow_redirects=False)
    session = Session(spec, client)
    try:
        await session.open()
        yield session
    finally:
        await session.close()
        await client.aclose()


async def list_tools(spec: McpSpec) -> tuple[list[dict[str, Any]], dict[str, Any], str]:
    """Every tool a server advertises, plus what it said about itself.

    Returns (tools, serverInfo, protocolVersion). Bounded at MAX_TOOLS: the list
    is sent as schema on every request, so a server offering two hundred is a
    server that would fill the window before anything was asked.
    """
    tools: list[dict[str, Any]] = []
    async with session_for(spec) as session:
        cursor = ""
        for _ in range(MAX_PAGES):
            result = await session.call("tools/list", {"cursor": cursor} if cursor else {})
            for entry in result.get("tools") or []:
                cleaned = protocol.clean_tool(entry)
                if cleaned is None:
                    log.info("%s advertised an unusable tool entry", spec.name)
                elif len(tools) < MAX_TOOLS:
                    tools.append(cleaned)
            cursor = str(result.get("nextCursor") or "")
            if not cursor or len(tools) >= MAX_TOOLS:
                break
        return tools, session.server_info, session.protocol_version


async def call_tool(spec: McpSpec, name: str, arguments: dict[str, Any]) -> tuple[str, bool]:
    """Run one tool. Returns (text, is_error)."""
    async with session_for(spec) as session:
        result = await session.call("tools/call", {"name": name, "arguments": arguments})
        return protocol.content_to_text(result), bool(result.get("isError"))


__all__ = ["McpError", "McpSpec", "Session", "call_tool", "list_tools", "spec_from"]