Coverage for mcpgateway/transports/stdio_transport.py: 71%
55 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-22 12:53 +0100
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-22 12:53 +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"""
12import asyncio
13import json
14import logging
15import sys
16from typing import Any, AsyncGenerator, Dict, Optional
18from mcpgateway.transports.base import Transport
20logger = logging.getLogger(__name__)
23class StdioTransport(Transport):
24 """Transport implementation using stdio streams."""
26 def __init__(self):
27 """Initialize stdio transport."""
28 self._stdin_reader: Optional[asyncio.StreamReader] = None
29 self._stdout_writer: Optional[asyncio.StreamWriter] = None
30 self._connected = False
32 async def connect(self) -> None:
33 """Set up stdio streams."""
34 loop = asyncio.get_running_loop()
36 # Set up stdin reader
37 reader = asyncio.StreamReader()
38 protocol = asyncio.StreamReaderProtocol(reader)
39 await loop.connect_read_pipe(lambda: protocol, sys.stdin)
40 self._stdin_reader = reader
42 # Set up stdout writer
43 transport, protocol = await loop.connect_write_pipe(asyncio.streams.FlowControlMixin, sys.stdout)
44 self._stdout_writer = asyncio.StreamWriter(transport, protocol, reader, loop)
46 self._connected = True
47 logger.info("stdio transport connected")
49 async def disconnect(self) -> None:
50 """Clean up stdio streams."""
51 if self._stdout_writer: 51 ↛ 54line 51 didn't jump to line 54 because the condition on line 51 was always true
52 self._stdout_writer.close()
53 await self._stdout_writer.wait_closed()
54 self._connected = False
55 logger.info("stdio transport disconnected")
57 async def send_message(self, message: Dict[str, Any]) -> None:
58 """Send a message over stdout.
60 Args:
61 message: Message to send
63 Raises:
64 RuntimeError: If transport is not connected
65 Exception: If unable to write to stdio writer
66 """
67 if not self._stdout_writer:
68 raise RuntimeError("Transport not connected")
70 try:
71 data = json.dumps(message)
72 self._stdout_writer.write(f"{data}\n".encode())
73 await self._stdout_writer.drain()
74 except Exception as e:
75 logger.error(f"Failed to send message: {e}")
76 raise
78 async def receive_message(self) -> AsyncGenerator[Dict[str, Any], None]:
79 """Receive messages from stdin.
81 Yields:
82 Received messages
84 Raises:
85 RuntimeError: If transport is not connected
86 """
87 if not self._stdin_reader:
88 raise RuntimeError("Transport not connected")
90 while True:
91 try:
92 # Read line from stdin
93 line = await self._stdin_reader.readline()
94 if not line:
95 break
97 # Parse JSON message
98 message = json.loads(line.decode().strip())
99 yield message
101 except asyncio.CancelledError:
102 break
103 except Exception as e:
104 logger.error(f"Failed to receive message: {e}")
105 continue
107 async def is_connected(self) -> bool:
108 """Check if transport is connected.
110 Returns:
111 True if connected
112 """
113 return self._connected