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

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 

12# Standard 

13import asyncio 

14import json 

15import logging 

16import sys 

17from typing import Any, AsyncGenerator, Dict, Optional 

18 

19# First-Party 

20from mcpgateway.transports.base import Transport 

21 

22logger = logging.getLogger(__name__) 

23 

24 

25class StdioTransport(Transport): 

26 """Transport implementation using stdio streams.""" 

27 

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 

33 

34 async def connect(self) -> None: 

35 """Set up stdio streams.""" 

36 loop = asyncio.get_running_loop() 

37 

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 

43 

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) 

47 

48 self._connected = True 

49 logger.info("stdio transport connected") 

50 

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

58 

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

60 """Send a message over stdout. 

61 

62 Args: 

63 message: Message to send 

64 

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

71 

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 

79 

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

81 """Receive messages from stdin. 

82 

83 Yields: 

84 Received messages 

85 

86 Raises: 

87 RuntimeError: If transport is not connected 

88 """ 

89 if not self._stdin_reader: 

90 raise RuntimeError("Transport not connected") 

91 

92 while True: 

93 try: 

94 # Read line from stdin 

95 line = await self._stdin_reader.readline() 

96 if not line: 

97 break 

98 

99 # Parse JSON message 

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

101 yield message 

102 

103 except asyncio.CancelledError: 

104 break 

105 except Exception as e: 

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

107 continue 

108 

109 async def is_connected(self) -> bool: 

110 """Check if transport is connected. 

111 

112 Returns: 

113 True if connected 

114 """ 

115 return self._connected