beaver-gateway/src/beaver_gateway/mcp/lenient.py

"""Tolerant transports for upstream MCPs that don't strictly speak JSON-RPC.

Some MCP servers print non-JSON chatter to stdout (banners, dependency-load
messages) that the reference client forwards as exceptions, which bleed into
client UIs as warnings. This transport drops those lines instead.
"""

from __future__ import annotations

import contextlib
import logging
import sys
from typing import TYPE_CHECKING

import anyio
import anyio.lowlevel
from anyio.streams.text import TextReceiveStream
from fastmcp.client.transports.base import ClientTransport
from mcp import ClientSession, types
from mcp.client.stdio import (
    PROCESS_TERMINATION_TIMEOUT,
    StdioServerParameters,
    _create_platform_compatible_process,
    _get_executable_command,
    _terminate_process_tree,
)
from mcp.shared.message import SessionMessage

if TYPE_CHECKING:
    from collections.abc import AsyncIterator
    from typing import TextIO, Unpack

    from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
    from fastmcp.client.transports.base import SessionKwargs


_log = logging.getLogger("beaver_gateway.mcp.lenient")


class LenientStdioTransport(ClientTransport):
    """Stdio transport that tolerates non-JSON-RPC stdout noise.

    Behaves like :class:`fastmcp.client.transports.StdioTransport` from the
    consumer's perspective: one ``ClientSession`` per ``connect_session``
    block, subprocess scoped to that block. We don't replicate the upstream
    ``keep_alive`` flag because the only caller (``create_proxy``) opens
    the session lazily on first request and keeps it open for the lifetime
    of the proxy.
    """

    def __init__(
        self,
        command: str,
        args: list[str],
        env: dict[str, str] | None = None,
        cwd: str | None = None,
        log_file: TextIO | None = None,
    ) -> None:
        """``log_file`` takes an already-open ``TextIO``.

        Unlike the upstream transport, this one does not open a ``Path``
        for you.
        """
        self.command = command
        self.args = args
        self.env = env
        self.cwd = cwd
        self.log_file = log_file

    @contextlib.asynccontextmanager
    async def connect_session(
        self, **session_kwargs: Unpack[SessionKwargs]
    ) -> AsyncIterator[ClientSession]:
        errlog: TextIO = self.log_file if self.log_file is not None else sys.stderr
        async with (
            _lenient_stdio_client(
                StdioServerParameters(
                    command=self.command, args=self.args, env=self.env, cwd=self.cwd
                ),
                errlog=errlog,
            ) as (read_stream, write_stream),
            ClientSession(read_stream, write_stream, **session_kwargs) as session,
        ):
            yield session

    def __repr__(self) -> str:
        return f"<LenientStdioTransport(command={self.command!r}, args={self.args!r})>"


@contextlib.asynccontextmanager
async def _lenient_stdio_client(  # noqa: PLR0915
    server: StdioServerParameters, errlog: TextIO = sys.stderr
) -> AsyncIterator[
    tuple[
        MemoryObjectReceiveStream[SessionMessage | Exception],
        MemoryObjectSendStream[SessionMessage],
    ]
]:
    """Drop-in for ``mcp.client.stdio.stdio_client`` with a tolerant reader.

    All differences from upstream live in ``stdout_reader``: lines that fail
    ``JSONRPCMessage.model_validate_json`` are logged at DEBUG and skipped,
    never forwarded as exceptions. Shutdown still follows the MCP spec
    sequence: close stdin, wait, SIGTERM, then SIGKILL.
    """
    read_stream_writer, read_stream = anyio.create_memory_object_stream[
        SessionMessage | Exception
    ](0)
    write_stream, write_stream_reader = anyio.create_memory_object_stream[
        SessionMessage
    ](0)

    try:
        command = _get_executable_command(server.command)
        env_default = _default_inherited_env()
        process = await _create_platform_compatible_process(
            command=command,
            args=server.args,
            env=(
                {**env_default, **server.env} if server.env is not None else env_default
            ),
            errlog=errlog,
            cwd=server.cwd,
        )
    except OSError:
        await read_stream.aclose()
        await write_stream.aclose()
        await read_stream_writer.aclose()
        await write_stream_reader.aclose()
        raise

    async def stdout_reader() -> None:
        assert process.stdout, "Opened process is missing stdout"  # noqa: S101
        try:
            async with read_stream_writer:
                buffer = ""
                async for chunk in TextReceiveStream(
                    process.stdout,
                    encoding=server.encoding,
                    errors=server.encoding_error_handler,
                ):
                    lines = (buffer + chunk).split("\n")
                    buffer = lines.pop()
                    for line in lines:
                        stripped = line.strip()
                        if not stripped:
                            continue
                        try:
                            message = types.JSONRPCMessage.model_validate_json(stripped)
                        except Exception:  # noqa: BLE001
                            _log.debug(
                                "lenient stdio: dropped non-JSON line: %r",
                                stripped[:200],
                            )
                            continue
                        await read_stream_writer.send(SessionMessage(message))
        except anyio.ClosedResourceError:
            await anyio.lowlevel.checkpoint()

    async def stdin_writer() -> None:
        assert process.stdin, "Opened process is missing stdin"  # noqa: S101
        try:
            async with write_stream_reader:
                async for session_message in write_stream_reader:
                    payload = session_message.message.model_dump_json(
                        by_alias=True, exclude_none=True
                    )
                    await process.stdin.send(
                        (payload + "\n").encode(
                            encoding=server.encoding,
                            errors=server.encoding_error_handler,
                        )
                    )
        except anyio.ClosedResourceError:
            await anyio.lowlevel.checkpoint()

    async with anyio.create_task_group() as tg, process:
        tg.start_soon(stdout_reader)
        tg.start_soon(stdin_writer)
        try:
            yield read_stream, write_stream
        finally:
            if process.stdin:
                with contextlib.suppress(Exception):
                    await process.stdin.aclose()
            try:
                with anyio.fail_after(PROCESS_TERMINATION_TIMEOUT):
                    await process.wait()
            except TimeoutError:
                await _terminate_process_tree(process)
            except ProcessLookupError:
                pass
            await read_stream.aclose()
            await write_stream.aclose()
            await read_stream_writer.aclose()
            await write_stream_reader.aclose()


def _default_inherited_env() -> dict[str, str]:
    """Same env shortlist as ``mcp.client.stdio.get_default_environment``.

    Re-exported so we can compose with the user's ``env`` overrides without
    forcing a private import in this module's body.
    """
    from mcp.client.stdio import get_default_environment

    return get_default_environment()


__all__ = ["LenientStdioTransport"]