forked from modelcontextprotocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmemory.py
More file actions
33 lines (24 loc) · 1.44 KB
/
memory.py
File metadata and controls
33 lines (24 loc) · 1.44 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
"""In-memory transports"""
from __future__ import annotations
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from mcp.shared._compat import resync_tracer
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
from mcp.shared.message import SessionMessage
MessageStream = tuple[ContextReceiveStream[SessionMessage | Exception], ContextSendStream[SessionMessage | Exception]]
@asynccontextmanager
async def create_client_server_memory_streams() -> AsyncGenerator[tuple[MessageStream, MessageStream], None]:
"""Creates a pair of bidirectional memory streams for client-server communication.
Yields:
A tuple of (client_streams, server_streams) where each is a tuple of
(read_stream, write_stream)
"""
# Create streams for both directions
server_to_client_send, server_to_client_receive = create_context_streams[SessionMessage | Exception](1)
client_to_server_send, client_to_server_receive = create_context_streams[SessionMessage | Exception](1)
client_streams = (server_to_client_receive, client_to_server_send)
server_streams = (client_to_server_receive, server_to_client_send)
async with server_to_client_receive, client_to_server_send, client_to_server_receive, server_to_client_send:
yield client_streams, server_streams
# Heals caller-driven cancels; closing memory streams never suspends.
await resync_tracer()