|
| 1 | +import logging |
| 2 | +import sys |
| 3 | +from io import TextIOWrapper |
| 4 | +from typing import Generic |
| 5 | + |
| 6 | +import anyio |
| 7 | +import anyio.lowlevel |
| 8 | + |
| 9 | +from mcp import types |
| 10 | +from mcp.server.minimcp.exceptions import InvalidMessageError |
| 11 | +from mcp.server.minimcp.managers.context_manager import ScopeT |
| 12 | +from mcp.server.minimcp.minimcp import MiniMCP |
| 13 | +from mcp.server.minimcp.types import MESSAGE_ENCODING, Message, NoMessage |
| 14 | +from mcp.server.minimcp.utils import json_rpc |
| 15 | + |
| 16 | +logger = logging.getLogger(__name__) |
| 17 | + |
| 18 | + |
| 19 | +class StdioTransport(Generic[ScopeT]): |
| 20 | + """stdio transport implementation per MCP specification. |
| 21 | + https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#stdio |
| 22 | +
|
| 23 | + - The server reads JSON-RPC messages from its standard input (stdin) |
| 24 | + - The server sends messages to its standard output (stdout) |
| 25 | + - Messages are individual JSON-RPC requests, notifications, or responses |
| 26 | + - Messages are delimited by newlines and MUST NOT contain embedded newlines |
| 27 | + - The server MUST NOT write anything to stdout that is not a valid MCP message |
| 28 | +
|
| 29 | + **IMPORTANT - Logging Configuration:** |
| 30 | + Applications MUST configure logging to write to stderr (not stdout) to avoid interfering |
| 31 | + with the stdio transport. The specification states: "The server MAY write UTF-8 strings to |
| 32 | + its standard error (stderr) for logging purposes." |
| 33 | +
|
| 34 | + Example logging configuration: |
| 35 | + logging.basicConfig( |
| 36 | + level=logging.DEBUG, |
| 37 | + handlers=[logging.StreamHandler(sys.stderr)] |
| 38 | + ) |
| 39 | +
|
| 40 | + Implementation details: |
| 41 | + - The anyio.wrap_file implementation naturally applies backpressure |
| 42 | + - Concurrent message handling via task groups |
| 43 | + - Concurrency management is enforced by MiniMCP |
| 44 | + - Exceptions are formatted as standard MCP errors, and shouldn't cause the transport to terminate |
| 45 | + """ |
| 46 | + |
| 47 | + minimcp: MiniMCP[ScopeT] |
| 48 | + |
| 49 | + stdin: anyio.AsyncFile[str] |
| 50 | + stdout: anyio.AsyncFile[str] |
| 51 | + |
| 52 | + def __init__( |
| 53 | + self, |
| 54 | + minimcp: MiniMCP[ScopeT], |
| 55 | + stdin: anyio.AsyncFile[str] | None = None, |
| 56 | + stdout: anyio.AsyncFile[str] | None = None, |
| 57 | + ) -> None: |
| 58 | + """ |
| 59 | + Args: |
| 60 | + minimcp: The MiniMCP instance to use. |
| 61 | + stdin: Optional stdin stream to use. |
| 62 | + stdout: Optional stdout stream to use. |
| 63 | + """ |
| 64 | + self.minimcp = minimcp |
| 65 | + |
| 66 | + self.stdin = stdin or anyio.wrap_file(TextIOWrapper(sys.stdin.buffer, encoding=MESSAGE_ENCODING)) |
| 67 | + self.stdout = stdout or anyio.wrap_file( |
| 68 | + TextIOWrapper(sys.stdout.buffer, encoding=MESSAGE_ENCODING, line_buffering=True) |
| 69 | + ) |
| 70 | + |
| 71 | + async def write_msg(self, response_msg: Message) -> None: |
| 72 | + """Write a message to stdout. |
| 73 | +
|
| 74 | + Per the MCP stdio transport specification, messages MUST NOT contain embedded newlines. |
| 75 | + This function validates that constraint before writing. |
| 76 | +
|
| 77 | + Args: |
| 78 | + response_msg: The message to write to stdout. |
| 79 | +
|
| 80 | + Raises: |
| 81 | + ValueError: If the message contains embedded newlines (violates stdio spec). |
| 82 | + """ |
| 83 | + if "\n" in response_msg or "\r" in response_msg: |
| 84 | + raise ValueError("Messages MUST NOT contain embedded newlines") |
| 85 | + |
| 86 | + logger.debug("Writing response message to stdio: %s", response_msg) |
| 87 | + await self.stdout.write(response_msg + "\n") |
| 88 | + |
| 89 | + async def dispatch(self, received_msg: Message) -> None: |
| 90 | + """ |
| 91 | + Dispatch an incoming message to the MiniMCP instance, and write the response to stdout. |
| 92 | + Exceptions are formatted as standard MCP errors, and shouldn't cause the transport to terminate. |
| 93 | +
|
| 94 | + Args: |
| 95 | + received_msg: The message to dispatch to the MiniMCP instance |
| 96 | + """ |
| 97 | + |
| 98 | + response: Message | NoMessage | None = None |
| 99 | + |
| 100 | + try: |
| 101 | + response = await self.minimcp.handle(received_msg, self.write_msg) |
| 102 | + except InvalidMessageError as e: |
| 103 | + response = e.response |
| 104 | + except Exception as e: |
| 105 | + response, error_message = json_rpc.build_error_message( |
| 106 | + e, |
| 107 | + received_msg, |
| 108 | + types.INTERNAL_ERROR, |
| 109 | + include_stack_trace=True, |
| 110 | + ) |
| 111 | + logger.exception(f"Unexpected error in stdio transport: {error_message}") |
| 112 | + |
| 113 | + if isinstance(response, Message): |
| 114 | + await self.write_msg(response) |
| 115 | + |
| 116 | + async def run(self) -> None: |
| 117 | + """ |
| 118 | + Start the stdio transport. |
| 119 | + This will read messages from stdin and dispatch them to the MiniMCP instance, and write |
| 120 | + the response to stdout. The transport must run until the stdin is closed. |
| 121 | + """ |
| 122 | + |
| 123 | + try: |
| 124 | + logger.debug("Starting stdio transport") |
| 125 | + async with anyio.create_task_group() as tg: |
| 126 | + async for line in self.stdin: |
| 127 | + _line = line.strip() |
| 128 | + if _line: |
| 129 | + tg.start_soon(self.dispatch, _line) |
| 130 | + except anyio.ClosedResourceError: |
| 131 | + # Stdin was closed (e.g., during shutdown) |
| 132 | + # Use checkpoint to allow cancellation to be processed |
| 133 | + await anyio.lowlevel.checkpoint() |
| 134 | + finally: |
| 135 | + logger.debug("Stdio transport stopped") |
0 commit comments