Coverage for mcpgateway/transports/stdio_transport.py: 100%
55 statements
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-09 11:03 +0100
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-09 11:03 +0100
1# -*- coding: utf-8 -*-
2"""stdio Transport Implementation.
4Copyright 2025
5SPDX-License-Identifier: Apache-2.0
6Authors: Mihai Criveti
8This module implements stdio transport for MCP, handling
9communication over standard input/output streams.
10"""
12# Standard
13import asyncio
14import json
15import logging
16import sys
17from typing import Any, AsyncGenerator, Dict, Optional
19# First-Party
20from mcpgateway.transports.base import Transport
22logger = logging.getLogger(__name__)
25class StdioTransport(Transport):
26 """Transport implementation using stdio streams."""
28 def __init__(self):
29 """Initialize stdio transport."""
30 self._stdin_reader: Optional[asyncio.StreamReader] = None
31 self._stdout_writer: Optional[asyncio.StreamWriter] = None
32 self._connected = False
34 async def connect(self) -> None:
35 """Set up stdio streams."""
36 loop = asyncio.get_running_loop()
38 # Set up stdin reader
39 reader = asyncio.StreamReader()
40 protocol = asyncio.StreamReaderProtocol(reader)
41 await loop.connect_read_pipe(lambda: protocol, sys.stdin)
42 self._stdin_reader = reader
44 # Set up stdout writer
45 transport, protocol = await loop.connect_write_pipe(asyncio.streams.FlowControlMixin, sys.stdout)
46 self._stdout_writer = asyncio.StreamWriter(transport, protocol, reader, loop)
48 self._connected = True
49 logger.info("stdio transport connected")
51 async def disconnect(self) -> None:
52 """Clean up stdio streams."""
53 if self._stdout_writer:
54 self._stdout_writer.close()
55 await self._stdout_writer.wait_closed()
56 self._connected = False
57 logger.info("stdio transport disconnected")
59 async def send_message(self, message: Dict[str, Any]) -> None:
60 """Send a message over stdout.
62 Args:
63 message: Message to send
65 Raises:
66 RuntimeError: If transport is not connected
67 Exception: If unable to write to stdio writer
68 """
69 if not self._stdout_writer:
70 raise RuntimeError("Transport not connected")
72 try:
73 data = json.dumps(message)
74 self._stdout_writer.write(f"{data}\n".encode())
75 await self._stdout_writer.drain()
76 except Exception as e:
77 logger.error(f"Failed to send message: {e}")
78 raise
80 async def receive_message(self) -> AsyncGenerator[Dict[str, Any], None]:
81 """Receive messages from stdin.
83 Yields:
84 Received messages
86 Raises:
87 RuntimeError: If transport is not connected
88 """
89 if not self._stdin_reader:
90 raise RuntimeError("Transport not connected")
92 while True:
93 try:
94 # Read line from stdin
95 line = await self._stdin_reader.readline()
96 if not line:
97 break
99 # Parse JSON message
100 message = json.loads(line.decode().strip())
101 yield message
103 except asyncio.CancelledError:
104 break
105 except Exception as e:
106 logger.error(f"Failed to receive message: {e}")
107 continue
109 async def is_connected(self) -> bool:
110 """Check if transport is connected.
112 Returns:
113 True if connected
114 """
115 return self._connected