Coverage for mcpgateway/transports/stdio_transport.py: 100%

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

2"""stdio Transport Implementation.

3

4 Copyright 2025

5 SPDX-License-Identifier: Apache-2.0

6 Authors: Mihai Criveti

7

8 This module implements stdio transport for MCP, handling

9 communication over standard input/output streams.

10"""

11

12# Standard

13 import asyncio

14 import json

15 import logging

16 import sys

17 from typing import Any, AsyncGenerator, Dict, Optional

18

19# First-Party

20 from mcpgateway.transports.base import Transport

21

22 logger = logging.getLogger(__name__)

23

24

25 class 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