feat: implement mcp frontend
This commit is contained in:
@@ -0,0 +1,234 @@
|
||||
"""Tolerant transports for upstream MCPs that don't strictly speak JSON-RPC.
|
||||
|
||||
Some real-world MCP servers print non-JSON chatter to stdout before / between
|
||||
their actual JSON-RPC frames (``Processing...``, banners, dependency-load
|
||||
messages, etc.). The reference ``mcp.client.stdio.stdio_client`` parses every
|
||||
stdout line as JSON-RPC and ships any parse failure downstream as an
|
||||
exception, which the MCP ``ClientSession`` then logs as a warning that bleeds
|
||||
into client UIs (Cursor, Cline) when they connect through us.
|
||||
|
||||
``LenientStdioTransport`` re-implements the stdio-client wiring with one
|
||||
behavioural change: lines that don't parse as JSON-RPC are *silently
|
||||
dropped* (one ``DEBUG`` log entry, no exception forwarded). Downstream
|
||||
consumers see only valid messages. We keep the rest of the contract identical
|
||||
to the reference client, including the spec-mandated graceful shutdown
|
||||
sequence (close stdin → wait → SIGTERM → SIGKILL).
|
||||
|
||||
The transport plugs into ``fastmcp.server.create_proxy`` just like
|
||||
``StdioTransport`` does, so the rest of the aggregator doesn't need to
|
||||
know which flavour it got.
|
||||
"""
|
||||
|
||||
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:
|
||||
self.command = command
|
||||
self.args = args
|
||||
self.env = env
|
||||
self.cwd = cwd
|
||||
# TextIO only — pre-open Path callers themselves. The upstream
|
||||
# ``StdioTransport`` opens Path for you, but we keep this thin so
|
||||
# the type contract stays narrow and easy to validate.
|
||||
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 — mirrors mcp.client.stdio.stdio_client
|
||||
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. This is what makes warning-noisy MCPs
|
||||
quiet from the consumer's point of view.
|
||||
"""
|
||||
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 — by design, see module doc
|
||||
_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:
|
||||
# MCP spec stdio shutdown: close stdin → wait → SIGTERM → SIGKILL.
|
||||
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"]
|
||||
Reference in New Issue
Block a user