"""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"" @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"]