beaver_project/app-instance/backend/beaver/integrations/mcp/connection.py

"""MCP connection manager."""

from __future__ import annotations

import asyncio
from contextlib import AsyncExitStack
from dataclasses import dataclass, field
from typing import Any

import httpx

from beaver.foundation.config import AuthzConfig, BackendIdentityConfig, MCPServerConfig
from beaver.integrations.authz import AuthzClient
from beaver.tools.mcp.wrapper import MCPToolWrapper
from beaver.tools.registry import ToolRegistry


@dataclass(slots=True)
class MCPConnectionReport:
    status: str = "disconnected"
    last_error: str | None = None
    tool_names: list[str] = field(default_factory=list)
    tool_count: int = 0
    transport: str = "http"

    def to_dict(self) -> dict[str, Any]:
        return {
            "status": self.status,
            "last_error": self.last_error,
            "tool_names": list(self.tool_names),
            "tool_count": self.tool_count,
            "transport": self.transport,
        }


class MCPConnectionManager:
    def __init__(
        self,
        servers: dict[str, MCPServerConfig],
        *,
        authz_config: AuthzConfig | None = None,
        backend_identity: BackendIdentityConfig | None = None,
    ) -> None:
        self.servers = servers
        self.authz_config = authz_config
        self.backend_identity = backend_identity
        self.stack = AsyncExitStack()
        self.connected = False
        self._connect_lock = asyncio.Lock()
        self.report: dict[str, MCPConnectionReport] = {}

    async def connect_all(self, registry: ToolRegistry) -> dict[str, dict[str, Any]]:
        async with self._connect_lock:
            if self.connected:
                return {key: value.to_dict() for key, value in self.report.items()}
            self.report = {}
            for server_id, cfg in self.servers.items():
                self.report[server_id] = MCPConnectionReport(transport=cfg.transport)
                try:
                    if cfg.command:
                        await self._connect_stdio(server_id, cfg, registry)
                    elif cfg.url:
                        await self._connect_http(server_id, cfg, registry)
                    else:
                        raise ValueError("MCP server requires command or url")
                    self.report[server_id].status = "connected"
                    self.report[server_id].tool_count = len(self.report[server_id].tool_names)
                except Exception as exc:
                    self.report[server_id].status = "error"
                    self.report[server_id].last_error = _describe_exception(exc, server_id=server_id, url=cfg.url or None)
            self.connected = True
            return {key: value.to_dict() for key, value in self.report.items()}

    async def close(self) -> None:
        await self.stack.aclose()
        self.connected = False

    async def _headers(self, server_id: str, cfg: MCPServerConfig) -> dict[str, str]:
        headers = dict(cfg.headers or {})
        if cfg.auth_mode != "oauth_backend_token":
            return headers
        if not (
            self.authz_config
            and self.authz_config.enabled
            and self.authz_config.base_url
            and self.backend_identity
            and self.backend_identity.client_id
            and self.backend_identity.client_secret
        ):
            raise RuntimeError("oauth_backend_token requires AuthZ and backend identity")
        audience = cfg.auth_audience or f"mcp:{server_id}"
        client = AuthzClient(self.authz_config.base_url, timeout_seconds=self.authz_config.request_timeout_seconds)
        token = await client.issue_token(
            client_id=self.backend_identity.client_id,
            client_secret=self.backend_identity.client_secret,
            audience=audience,
            scopes=list(cfg.auth_scopes),
        )
        access_token = str(token.get("access_token") or "").strip()
        if not access_token:
            raise RuntimeError("AuthZ did not return an access token")
        headers["Authorization"] = f"Bearer {access_token}"
        return headers

    async def _open_http_session(self, cfg: MCPServerConfig, headers: dict[str, str]):
        from mcp import ClientSession
        from mcp.client.streamable_http import streamable_http_client

        http_client = await self.stack.enter_async_context(
            httpx.AsyncClient(headers=headers or None, follow_redirects=True, trust_env=False)
        )
        read, write, _ = await self.stack.enter_async_context(streamable_http_client(cfg.url, http_client=http_client))
        session = await self.stack.enter_async_context(ClientSession(read, write))
        await session.initialize()
        return session

    async def _connect_http(self, server_id: str, cfg: MCPServerConfig, registry: ToolRegistry) -> None:
        headers = await self._headers(server_id, cfg)
        session = await self._open_http_session(cfg, headers)
        tools = await session.list_tools()
        for tool_def in tools.tools:
            async def call_tool(tool_name: str, args: dict[str, Any], *, _session=session) -> Any:
                return await _session.call_tool(tool_name, arguments=args)

            wrapper = MCPToolWrapper(
                server_id,
                tool_def,
                call_tool,
                cfg.tool_timeout,
                cfg.sensitive,
                cfg.kind,
                cfg.category,
                cfg.display_name,
            )
            registry.register(wrapper, replace=True)
            if wrapper.spec.name not in self.report[server_id].tool_names:
                self.report[server_id].tool_names.append(wrapper.spec.name)

    async def _connect_stdio(self, server_id: str, cfg: MCPServerConfig, registry: ToolRegistry) -> None:
        from mcp import ClientSession, StdioServerParameters
        from mcp.client.stdio import stdio_client

        params = StdioServerParameters(command=cfg.command, args=list(cfg.args), env=dict(cfg.env) or None)
        read, write = await self.stack.enter_async_context(stdio_client(params))
        session = await self.stack.enter_async_context(ClientSession(read, write))
        await session.initialize()
        tools = await session.list_tools()
        for tool_def in tools.tools:
            async def call_tool(tool_name: str, args: dict[str, Any], *, _session=session) -> Any:
                return await _session.call_tool(tool_name, arguments=args)

            wrapper = MCPToolWrapper(
                server_id,
                tool_def,
                call_tool,
                cfg.tool_timeout,
                cfg.sensitive,
                cfg.kind,
                cfg.category,
                cfg.display_name,
            )
            registry.register(wrapper, replace=True)
            if wrapper.spec.name not in self.report[server_id].tool_names:
                self.report[server_id].tool_names.append(wrapper.spec.name)


async def test_mcp_server(
    server_id: str,
    cfg: MCPServerConfig,
    *,
    authz_config: AuthzConfig | None = None,
    backend_identity: BackendIdentityConfig | None = None,
) -> dict[str, Any]:
    registry = ToolRegistry()
    manager = MCPConnectionManager({server_id: cfg}, authz_config=authz_config, backend_identity=backend_identity)
    try:
        report = await manager.connect_all(registry)
        return {"ok": report.get(server_id, {}).get("status") == "connected", "server": server_id, **report.get(server_id, {})}
    finally:
        await manager.close()


def _describe_exception(exc: BaseException, *, server_id: str, url: str | None = None) -> str:
    target = f" ({url})" if url else ""
    if isinstance(exc, httpx.TimeoutException):
        return f"MCP server '{server_id}' timed out{target}"
    if isinstance(exc, httpx.ConnectError):
        return f"MCP server '{server_id}' is unreachable{target}"
    if isinstance(exc, httpx.HTTPStatusError):
        return f"MCP server '{server_id}' returned HTTP {exc.response.status_code}{target}"
    detail = str(exc).strip() or exc.__class__.__name__
    return f"MCP server '{server_id}' failed{target}: {detail}"