|
15 | 15 | from mcp.shared.message import SessionMessage |
16 | 16 |
|
17 | 17 |
|
18 | | -@asynccontextmanager |
19 | | -async def create_in_memory_transport( |
20 | | - server: Server[Any] | FastMCP, |
21 | | - *, |
22 | | - raise_exceptions: bool = False, |
23 | | -) -> AsyncGenerator[ |
24 | | - tuple[ |
25 | | - MemoryObjectReceiveStream[SessionMessage | Exception], |
26 | | - MemoryObjectSendStream[SessionMessage], |
27 | | - ], |
28 | | - None, |
29 | | -]: |
| 18 | +class InMemoryTransport: |
30 | 19 | """ |
31 | | - Create an in-memory transport connected to a server. |
| 20 | + In-memory transport for testing MCP servers without network overhead. |
32 | 21 |
|
33 | | - This starts the server in a background task and returns streams |
34 | | - for client-side communication. The server is automatically stopped |
35 | | - when the context manager exits. |
36 | | -
|
37 | | - Args: |
38 | | - server: The MCP server to connect to (Server or FastMCP instance) |
39 | | - raise_exceptions: Whether to raise exceptions from the server |
40 | | -
|
41 | | - Yields: |
42 | | - A tuple of (read_stream, write_stream) for communicating with the server |
| 22 | + This transport starts the server in a background task and provides |
| 23 | + streams for client-side communication. The server is automatically |
| 24 | + stopped when the context manager exits. |
43 | 25 |
|
44 | 26 | Example: |
45 | 27 | server = FastMCP("test") |
| 28 | + transport = InMemoryTransport(server) |
46 | 29 |
|
47 | | - async with create_in_memory_transport(server) as (read, write): |
48 | | - async with ClientSession(read, write) as session: |
| 30 | + async with transport.connect() as (read_stream, write_stream): |
| 31 | + async with ClientSession(read_stream, write_stream) as session: |
49 | 32 | await session.initialize() |
50 | 33 | # Use the session... |
| 34 | +
|
| 35 | + Or more commonly, use with Client: |
| 36 | + async with Client.from_server(server) as client: |
| 37 | + result = await client.call_tool("my_tool", {...}) |
51 | 38 | """ |
52 | | - # Unwrap FastMCP to get underlying Server |
53 | | - actual_server: Server[Any] |
54 | | - if isinstance(server, FastMCP): |
55 | | - actual_server = server._mcp_server # type: ignore[reportPrivateUsage] |
56 | | - else: |
57 | | - actual_server = server |
58 | | - |
59 | | - async with create_client_server_memory_streams() as (client_streams, server_streams): |
60 | | - client_read, client_write = client_streams |
61 | | - server_read, server_write = server_streams |
62 | | - |
63 | | - async with anyio.create_task_group() as tg: |
64 | | - # Start server in background |
65 | | - tg.start_soon( |
66 | | - lambda: actual_server.run( |
67 | | - server_read, |
68 | | - server_write, |
69 | | - actual_server.create_initialization_options(), |
70 | | - raise_exceptions=raise_exceptions, |
| 39 | + |
| 40 | + def __init__( |
| 41 | + self, |
| 42 | + server: Server[Any] | FastMCP, |
| 43 | + *, |
| 44 | + raise_exceptions: bool = False, |
| 45 | + ) -> None: |
| 46 | + """ |
| 47 | + Initialize the in-memory transport. |
| 48 | +
|
| 49 | + Args: |
| 50 | + server: The MCP server to connect to (Server or FastMCP instance) |
| 51 | + raise_exceptions: Whether to raise exceptions from the server |
| 52 | + """ |
| 53 | + self._server = server |
| 54 | + self._raise_exceptions = raise_exceptions |
| 55 | + |
| 56 | + @asynccontextmanager |
| 57 | + async def connect( |
| 58 | + self, |
| 59 | + ) -> AsyncGenerator[ |
| 60 | + tuple[ |
| 61 | + MemoryObjectReceiveStream[SessionMessage | Exception], |
| 62 | + MemoryObjectSendStream[SessionMessage], |
| 63 | + ], |
| 64 | + None, |
| 65 | + ]: |
| 66 | + """ |
| 67 | + Connect to the server and return streams for communication. |
| 68 | +
|
| 69 | + Yields: |
| 70 | + A tuple of (read_stream, write_stream) for bidirectional communication |
| 71 | + """ |
| 72 | + # Unwrap FastMCP to get underlying Server |
| 73 | + actual_server: Server[Any] |
| 74 | + if isinstance(self._server, FastMCP): |
| 75 | + actual_server = self._server._mcp_server # type: ignore[reportPrivateUsage] |
| 76 | + else: |
| 77 | + actual_server = self._server |
| 78 | + |
| 79 | + async with create_client_server_memory_streams() as (client_streams, server_streams): |
| 80 | + client_read, client_write = client_streams |
| 81 | + server_read, server_write = server_streams |
| 82 | + |
| 83 | + async with anyio.create_task_group() as tg: |
| 84 | + # Start server in background |
| 85 | + tg.start_soon( |
| 86 | + lambda: actual_server.run( |
| 87 | + server_read, |
| 88 | + server_write, |
| 89 | + actual_server.create_initialization_options(), |
| 90 | + raise_exceptions=self._raise_exceptions, |
| 91 | + ) |
71 | 92 | ) |
72 | | - ) |
73 | 93 |
|
74 | | - try: |
75 | | - yield client_read, client_write |
76 | | - finally: |
77 | | - tg.cancel_scope.cancel() |
| 94 | + try: |
| 95 | + yield client_read, client_write |
| 96 | + finally: |
| 97 | + tg.cancel_scope.cancel() |
0 commit comments