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

1# -*- coding: utf-8 -*- 

2"""stdio Transport Implementation. 

3 

4Copyright 2025 

5SPDX-License-Identifier: Apache-2.0 

6Authors: Mihai Criveti 

7 

8This module implements stdio transport for MCP, handling 

9communication over standard input/output streams. 

10""" 

11 

12import asyncio 

13import json 

14import logging 

15import sys 

16from typing import Any, AsyncGenerator, Dict, Optional 

17 

18from mcpgateway.transports.base import Transport 

19 

20logger = logging.getLogger(__name__) 

21 

22 

23class StdioTransport(Transport): 

24 """Transport implementation using stdio streams.""" 

25 

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 

31 

32 async def connect(self) -> None: 

33 """Set up stdio streams.""" 

34 loop = asyncio.get_running_loop() 

35 

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 

41 

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) 

45 

46 self._connected = True 

47 logger.info("stdio transport connected") 

48 

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") 

56 

57 async def send_message(self, message: Dict[str, Any]) -> None: 

58 """Send a message over stdout. 

59 

60 Args: 

61 message: Message to send 

62 

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") 

69 

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 

77 

78 async def receive_message(self) -> AsyncGenerator[Dict[str, Any], None]: 

79 """Receive messages from stdin. 

80 

81 Yields: 

82 Received messages 

83 

84 Raises: 

85 RuntimeError: If transport is not connected 

86 """ 

87 if not self._stdin_reader: 

88 raise RuntimeError("Transport not connected") 

89 

90 while True: 

91 try: 

92 # Read line from stdin 

93 line = await self._stdin_reader.readline() 

94 if not line: 

95 break 

96 

97 # Parse JSON message 

98 message = json.loads(line.decode().strip()) 

99 yield message 

100 

101 except asyncio.CancelledError: 

102 break 

103 except Exception as e: 

104 logger.error(f"Failed to receive message: {e}") 

105 continue 

106 

107 async def is_connected(self) -> bool: 

108 """Check if transport is connected. 

109 

110 Returns: 

111 True if connected 

112 """ 

113 return self._connected