From acb89b4bcc5221419b0d761aa0db6c97caa75bb5 Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Wed, 30 Oct 2024 20:54:33 -0400 Subject: [PATCH 01/12] Add in-memory transport --- mcp_python/client/session.py | 3 +- mcp_python/server/__init__.py | 24 +++++++++--- mcp_python/server/memory.py | 71 +++++++++++++++++++++++++++++++++++ mcp_python/shared/session.py | 11 +++++- mcp_python/types.py | 4 ++ tests/conftest.py | 30 +++++++++++++++ tests/test_memory.py | 37 ++++++++++++++++++ uv.lock | 2 +- 8 files changed, 174 insertions(+), 8 deletions(-) create mode 100644 mcp_python/server/memory.py create mode 100644 tests/conftest.py create mode 100644 tests/test_memory.py diff --git a/mcp_python/client/session.py b/mcp_python/client/session.py index 266e741c9..eba4923dd 100644 --- a/mcp_python/client/session.py +++ b/mcp_python/client/session.py @@ -36,8 +36,9 @@ def __init__( self, read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception], write_stream: MemoryObjectSendStream[JSONRPCMessage], + read_timeout_seconds: int | float | None = None, ) -> None: - super().__init__(read_stream, write_stream, ServerRequest, ServerNotification) + super().__init__(read_stream, write_stream, ServerRequest, ServerNotification, read_timeout_seconds=read_timeout_seconds) async def initialize(self) -> InitializeResult: from mcp_python.types import ( diff --git a/mcp_python/server/__init__.py b/mcp_python/server/__init__.py index 71b373d63..a386e46bc 100644 --- a/mcp_python/server/__init__.py +++ b/mcp_python/server/__init__.py @@ -63,9 +63,13 @@ def pkg_version(package: str) -> str: try: from importlib.metadata import version - return version(package) + v = version(package) + if v is not None: + return v except Exception: - return "unknown" + pass + + return "unknown" return types.InitializationOptions( server_name=self.name, @@ -330,6 +334,11 @@ async def run( read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception], write_stream: MemoryObjectSendStream[JSONRPCMessage], initialization_options: types.InitializationOptions, + # When True, exceptions are returned as messages to the client. + # When False, exceptions are raised, which will cause the server to shut down + # but also make tracing exceptions much easier during testing and when using + # in-process servers. + raise_exceptions: bool = False, ): with warnings.catch_warnings(record=True) as w: async with ServerSession( @@ -349,6 +358,7 @@ async def run( f"Dispatching request of type {type(req).__name__}" ) + token = None try: # Set our global state that can be retrieved via # app.get_request_context() @@ -360,12 +370,16 @@ async def run( ) ) response = await handler(req) - # Reset the global state after we are done - request_ctx.reset(token) except Exception as err: + if raise_exceptions: + raise err response = ErrorData( code=0, message=str(err), data=None ) + finally: + # Reset the global state after we are done + if token is not None: + request_ctx.reset(token) await message.respond(response) else: @@ -373,7 +387,7 @@ async def run( ErrorData( code=METHOD_NOT_FOUND, message="Method not found", - ) + ), ) logger.debug("Response sent") diff --git a/mcp_python/server/memory.py b/mcp_python/server/memory.py new file mode 100644 index 000000000..2384c535d --- /dev/null +++ b/mcp_python/server/memory.py @@ -0,0 +1,71 @@ +""" +In-memory transports +""" + +from contextlib import asynccontextmanager +from typing import AsyncGenerator, Tuple + +import anyio +from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream + +from mcp_python.client.session import ClientSession +from mcp_python.server import Server +from mcp_python.server.session import ServerSession +from mcp_python.types import ErrorData, JSONRPCMessage + +@asynccontextmanager +async def create_client_server_memory_streams() -> AsyncGenerator[ Tuple[ + Tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]], + Tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]] +], None]: + """ + Creates a pair of bidirectional memory streams for client-server communication. + + Returns: + 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 = anyio.create_memory_object_stream[JSONRPCMessage | Exception](1) + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[JSONRPCMessage | Exception](1) + + # Return streams grouped by client/server + 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 + + +@asynccontextmanager +async def create_connected_server_and_client_session(server: Server) -> AsyncGenerator[ClientSession, None]: + """Creates a ServerSession that is connected to the `server`.""" + async with create_client_server_memory_streams() as (client_streams, server_streams): + # Unpack the streams + client_read, client_write = client_streams + server_read, server_write = server_streams + print("stream-1") + + # Create a cancel scope for the server task + async with anyio.create_task_group() as tg: + + tg.start_soon( + server.run, + server_read, + server_write, + server.create_initialization_options() + ) + + print("stream2") + + try: + # Client session could be created here using client_read and client_write + # This would allow testing the server with a client in the same process + async with ClientSession( + read_stream=client_read, write_stream=client_write + ) as client_session: + await client_session.initialize() + yield client_session + finally: + tg.cancel_scope.cancel() diff --git a/mcp_python/shared/session.py b/mcp_python/shared/session.py index 3bc66fcd0..4e8ff005b 100644 --- a/mcp_python/shared/session.py +++ b/mcp_python/shared/session.py @@ -87,6 +87,8 @@ def __init__( write_stream: MemoryObjectSendStream[JSONRPCMessage], receive_request_type: type[ReceiveRequestT], receive_notification_type: type[ReceiveNotificationT], + # If none, reading will never time out + read_timeout_seconds: int | float | None = None, ) -> None: self._read_stream = read_stream self._write_stream = write_stream @@ -94,6 +96,7 @@ def __init__( self._request_id = 0 self._receive_request_type = receive_request_type self._receive_notification_type = receive_notification_type + self._read_timeout_seconds = read_timeout_seconds self._incoming_message_stream_writer, self._incoming_message_stream_reader = ( anyio.create_memory_object_stream[ @@ -147,7 +150,13 @@ async def send_request( await self._write_stream.send(JSONRPCMessage(jsonrpc_request)) - response_or_error = await response_stream_reader.receive() + try: + with anyio.fail_after(self._read_timeout_seconds): + response_or_error = await response_stream_reader.receive() + except TimeoutError: + # TODO: make sure this response comes back correctly to the client + raise McpError(ErrorData(code=408, message=f"Timed out while waiting for response to {request.__class__.__name__}. Waited {READ_TIMEOUT_SECONDS} seconds.")) + if isinstance(response_or_error, JSONRPCError): raise McpError(response_or_error.error) else: diff --git a/mcp_python/types.py b/mcp_python/types.py index 012122ed7..858bf0dd0 100644 --- a/mcp_python/types.py +++ b/mcp_python/types.py @@ -1,6 +1,7 @@ from typing import Any, Generic, Literal, TypeVar from pydantic import BaseModel, ConfigDict, RootModel +from pydantic.fields import Field from pydantic.networks import AnyUrl """ @@ -141,16 +142,19 @@ class ErrorData(BaseModel): code: int """The error type that occurred.""" + message: str """ A short description of the error. The message SHOULD be limited to a concise single sentence. """ + data: Any | None = None """ Additional information about the error. The value of this member is defined by the sender (e.g. detailed error information, nested errors etc.). """ + model_config = ConfigDict(extra="allow") diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..b29b5e1cd --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,30 @@ +import pytest +from pydantic import HttpUrl + +from mcp_python.server import Server +from mcp_python.server.types import InitializationOptions +from mcp_python.types import Resource, ServerCapabilities + + +TEST_INITIALIZATION_OPTIONS = InitializationOptions( + server_name="my_mcp_server", + server_version="0.1.0", + capabilities=ServerCapabilities(), +) + +@pytest.fixture +def mcp_server() -> Server: + server = Server(name="test_server") + + # Add a simple resource for testing + @server.list_resources() + async def handle_list_resources(): + return [ + Resource( + uri=HttpUrl("memory://test"), + name="Test Resource", + description="A test resource" + ) + ] + + return server diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 000000000..3e563b519 --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,37 @@ +from typing_extensions import AsyncGenerator +import anyio +from pydantic import HttpUrl +from pydantic_core import Url +import pytest + +from mcp_python.client.session import ClientSession +from mcp_python.server import Server +from mcp_python.server.memory import create_client_server_memory_streams, create_connected_server_and_client_session +from mcp_python.server.session import ServerSession +from mcp_python.server.types import InitializationOptions +from mcp_python.types import ( + ClientNotification, + EmptyResult, + InitializedNotification, + JSONRPCMessage, + Resource, + ServerCapabilities, +) + + +@pytest.fixture +async def client_connected_to_server(mcp_server: Server) -> AsyncGenerator[ClientSession, None]: + print('11111') + async with create_connected_server_and_client_session(mcp_server) as client_session: + print('2222k') + yield client_session + print('33') + + +@pytest.mark.anyio +async def test_memory_server_and_client_connection(client_connected_to_server: ClientSession): + """Shows how a client and server can communicate over memory streams.""" + response = await client_connected_to_server.send_ping() + print('foo') + assert isinstance(response, EmptyResult) + print('bar') diff --git a/uv.lock b/uv.lock index e085cced4..81c40d0c1 100644 --- a/uv.lock +++ b/uv.lock @@ -163,7 +163,7 @@ wheels = [ [[package]] name = "mcp-python" -version = "0.3.0.dev0" +version = "0.4.0.dev0" source = { editable = "." } dependencies = [ { name = "anyio" }, From 85e9e19357cc49f1654752c5e18a22845bf56453 Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Mon, 4 Nov 2024 15:31:57 -0500 Subject: [PATCH 02/12] handle ping requests --- mcp_python/server/__init__.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/mcp_python/server/__init__.py b/mcp_python/server/__init__.py index a386e46bc..2e4fc026b 100644 --- a/mcp_python/server/__init__.py +++ b/mcp_python/server/__init__.py @@ -18,6 +18,7 @@ ClientNotification, ClientRequest, CompleteRequest, + EmptyResult, ErrorData, JSONRPCMessage, ListPromptsRequest, @@ -27,6 +28,7 @@ ListToolsRequest, ListToolsResult, LoggingLevel, + PingRequest, ProgressNotification, Prompt, PromptReference, @@ -52,7 +54,9 @@ class Server: def __init__(self, name: str): self.name = name - self.request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]] = {} + self.request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]] = { + PingRequest: _ping_handler, + } self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {} logger.info(f"Initializing server '{name}'") @@ -413,3 +417,7 @@ async def run( logger.info( f"Warning: {warning.category.__name__}: {warning.message}" ) + + +async def _ping_handler(request: PingRequest) -> ServerResult: + return ServerResult(EmptyResult()) From b544f938b009ce7c0ef1e7f3ded6f794af241421 Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Tue, 5 Nov 2024 09:52:57 -0500 Subject: [PATCH 03/12] info => debug --- mcp_python/server/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mcp_python/server/__init__.py b/mcp_python/server/__init__.py index 2e4fc026b..618d75f4d 100644 --- a/mcp_python/server/__init__.py +++ b/mcp_python/server/__init__.py @@ -58,7 +58,7 @@ def __init__(self, name: str): PingRequest: _ping_handler, } self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {} - logger.info(f"Initializing server '{name}'") + logger.debug(f"Initializing server '{name}'") def create_initialization_options(self) -> types.InitializationOptions: """Create initialization options from this server instance.""" From 4f1bca6869922075f737cc082942529b6f3293a9 Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Tue, 5 Nov 2024 09:55:22 -0500 Subject: [PATCH 04/12] ruff --- mcp_python/server/memory.py | 4 ++-- mcp_python/types.py | 1 - tests/conftest.py | 1 - tests/test_memory.py | 16 ++++------------ 4 files changed, 6 insertions(+), 16 deletions(-) diff --git a/mcp_python/server/memory.py b/mcp_python/server/memory.py index 2384c535d..859fd1337 100644 --- a/mcp_python/server/memory.py +++ b/mcp_python/server/memory.py @@ -10,8 +10,8 @@ from mcp_python.client.session import ClientSession from mcp_python.server import Server -from mcp_python.server.session import ServerSession -from mcp_python.types import ErrorData, JSONRPCMessage +from mcp_python.types import JSONRPCMessage + @asynccontextmanager async def create_client_server_memory_streams() -> AsyncGenerator[ Tuple[ diff --git a/mcp_python/types.py b/mcp_python/types.py index 858bf0dd0..b3ab4dd98 100644 --- a/mcp_python/types.py +++ b/mcp_python/types.py @@ -1,7 +1,6 @@ from typing import Any, Generic, Literal, TypeVar from pydantic import BaseModel, ConfigDict, RootModel -from pydantic.fields import Field from pydantic.networks import AnyUrl """ diff --git a/tests/conftest.py b/tests/conftest.py index b29b5e1cd..e96e25281 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,7 +5,6 @@ from mcp_python.server.types import InitializationOptions from mcp_python.types import Resource, ServerCapabilities - TEST_INITIALIZATION_OPTIONS = InitializationOptions( server_name="my_mcp_server", server_version="0.1.0", diff --git a/tests/test_memory.py b/tests/test_memory.py index 3e563b519..410cafee6 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -1,21 +1,13 @@ -from typing_extensions import AsyncGenerator -import anyio -from pydantic import HttpUrl -from pydantic_core import Url import pytest +from typing_extensions import AsyncGenerator from mcp_python.client.session import ClientSession from mcp_python.server import Server -from mcp_python.server.memory import create_client_server_memory_streams, create_connected_server_and_client_session -from mcp_python.server.session import ServerSession -from mcp_python.server.types import InitializationOptions +from mcp_python.server.memory import ( + create_connected_server_and_client_session, +) from mcp_python.types import ( - ClientNotification, EmptyResult, - InitializedNotification, - JSONRPCMessage, - Resource, - ServerCapabilities, ) From 488a424ae3bffd3669bccc80bcee52adb9b76943 Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Tue, 5 Nov 2024 09:59:16 -0500 Subject: [PATCH 05/12] ruff E501 --- mcp_python/client/session.py | 8 +++++- mcp_python/server/memory.py | 52 ++++++++++++++++++++++++++---------- mcp_python/shared/session.py | 12 ++++++++- tests/test_memory.py | 18 ++++++++----- 4 files changed, 67 insertions(+), 23 deletions(-) diff --git a/mcp_python/client/session.py b/mcp_python/client/session.py index eba4923dd..c8acf33fd 100644 --- a/mcp_python/client/session.py +++ b/mcp_python/client/session.py @@ -38,7 +38,13 @@ def __init__( write_stream: MemoryObjectSendStream[JSONRPCMessage], read_timeout_seconds: int | float | None = None, ) -> None: - super().__init__(read_stream, write_stream, ServerRequest, ServerNotification, read_timeout_seconds=read_timeout_seconds) + super().__init__( + read_stream, + write_stream, + ServerRequest, + ServerNotification, + read_timeout_seconds=read_timeout_seconds, + ) async def initialize(self) -> InitializeResult: from mcp_python.types import ( diff --git a/mcp_python/server/memory.py b/mcp_python/server/memory.py index 859fd1337..582462139 100644 --- a/mcp_python/server/memory.py +++ b/mcp_python/server/memory.py @@ -14,10 +14,21 @@ @asynccontextmanager -async def create_client_server_memory_streams() -> AsyncGenerator[ Tuple[ - Tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]], - Tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]] -], None]: +async def create_client_server_memory_streams() -> ( + AsyncGenerator[ + Tuple[ + Tuple[ + MemoryObjectReceiveStream[JSONRPCMessage | Exception], + MemoryObjectSendStream[JSONRPCMessage], + ], + Tuple[ + MemoryObjectReceiveStream[JSONRPCMessage | Exception], + MemoryObjectSendStream[JSONRPCMessage], + ], + ], + None, + ] +): """ Creates a pair of bidirectional memory streams for client-server communication. @@ -26,22 +37,35 @@ async def create_client_server_memory_streams() -> AsyncGenerator[ Tuple[ (read_stream, write_stream) """ # Create streams for both directions - server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[JSONRPCMessage | Exception](1) - client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[JSONRPCMessage | Exception](1) + server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[ + JSONRPCMessage | Exception + ](1) + client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[ + JSONRPCMessage | Exception + ](1) # Return streams grouped by client/server 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): + async with ( + server_to_client_receive, + client_to_server_send, + client_to_server_receive, + server_to_client_send, + ): yield client_streams, server_streams @asynccontextmanager -async def create_connected_server_and_client_session(server: Server) -> AsyncGenerator[ClientSession, None]: +async def create_connected_server_and_client_session( + server: Server, +) -> AsyncGenerator[ClientSession, None]: """Creates a ServerSession that is connected to the `server`.""" - async with create_client_server_memory_streams() as (client_streams, server_streams): + async with create_client_server_memory_streams() as ( + client_streams, + server_streams, + ): # Unpack the streams client_read, client_write = client_streams server_read, server_write = server_streams @@ -49,19 +73,19 @@ async def create_connected_server_and_client_session(server: Server) -> AsyncGen # Create a cancel scope for the server task async with anyio.create_task_group() as tg: - tg.start_soon( server.run, server_read, server_write, - server.create_initialization_options() + server.create_initialization_options(), ) print("stream2") try: - # Client session could be created here using client_read and client_write - # This would allow testing the server with a client in the same process + # Client session could be created here using client_read and + # client_write This would allow testing the server with a client + # in the same process async with ClientSession( read_stream=client_read, write_stream=client_write ) as client_session: diff --git a/mcp_python/shared/session.py b/mcp_python/shared/session.py index 4e8ff005b..e4c0b2d0c 100644 --- a/mcp_python/shared/session.py +++ b/mcp_python/shared/session.py @@ -155,7 +155,17 @@ async def send_request( response_or_error = await response_stream_reader.receive() except TimeoutError: # TODO: make sure this response comes back correctly to the client - raise McpError(ErrorData(code=408, message=f"Timed out while waiting for response to {request.__class__.__name__}. Waited {READ_TIMEOUT_SECONDS} seconds.")) + raise McpError( + ErrorData( + code=408, + message=( + f"Timed out while waiting for response to " + f"{request.__class__.__name__}. Waited " + f"{self._read_timeout_seconds} seconds." + ), + ) + + ) if isinstance(response_or_error, JSONRPCError): raise McpError(response_or_error.error) diff --git a/tests/test_memory.py b/tests/test_memory.py index 410cafee6..058104735 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -12,18 +12,22 @@ @pytest.fixture -async def client_connected_to_server(mcp_server: Server) -> AsyncGenerator[ClientSession, None]: - print('11111') +async def client_connected_to_server( + mcp_server: Server, +) -> AsyncGenerator[ClientSession, None]: + print("11111") async with create_connected_server_and_client_session(mcp_server) as client_session: - print('2222k') + print("2222k") yield client_session - print('33') + print("33") @pytest.mark.anyio -async def test_memory_server_and_client_connection(client_connected_to_server: ClientSession): +async def test_memory_server_and_client_connection( + client_connected_to_server: ClientSession, +): """Shows how a client and server can communicate over memory streams.""" response = await client_connected_to_server.send_ping() - print('foo') + print("foo") assert isinstance(response, EmptyResult) - print('bar') + print("bar") From 593efcbf781640b5d505ca4cf497843d7628db65 Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Tue, 5 Nov 2024 10:02:04 -0500 Subject: [PATCH 06/12] fix type --- tests/conftest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index e96e25281..77cf78594 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,5 @@ +from pydantic import AnyUrl import pytest -from pydantic import HttpUrl from mcp_python.server import Server from mcp_python.server.types import InitializationOptions @@ -20,7 +20,7 @@ def mcp_server() -> Server: async def handle_list_resources(): return [ Resource( - uri=HttpUrl("memory://test"), + uri=AnyUrl("memory://test"), name="Test Resource", description="A test resource" ) From 21d0ded5743272ce7334de38be86b0fdb56a16ea Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Tue, 5 Nov 2024 10:05:59 -0500 Subject: [PATCH 07/12] ruff --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 77cf78594..849b2dca1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,5 @@ -from pydantic import AnyUrl import pytest +from pydantic import AnyUrl from mcp_python.server import Server from mcp_python.server.types import InitializationOptions From ff86ee7952aaf19ebcd618cdd35df220c33f8854 Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Tue, 5 Nov 2024 10:16:25 -0500 Subject: [PATCH 08/12] self review --- mcp_python/server/__init__.py | 2 +- mcp_python/server/memory.py | 30 +++++++++--------------------- mcp_python/shared/session.py | 1 - tests/conftest.py | 1 - tests/test_memory.py | 5 ----- 5 files changed, 10 insertions(+), 29 deletions(-) diff --git a/mcp_python/server/__init__.py b/mcp_python/server/__init__.py index 618d75f4d..ff8900fc7 100644 --- a/mcp_python/server/__init__.py +++ b/mcp_python/server/__init__.py @@ -391,7 +391,7 @@ async def run( ErrorData( code=METHOD_NOT_FOUND, message="Method not found", - ), + ) ) logger.debug("Response sent") diff --git a/mcp_python/server/memory.py b/mcp_python/server/memory.py index 582462139..8a36a8ef5 100644 --- a/mcp_python/server/memory.py +++ b/mcp_python/server/memory.py @@ -3,7 +3,7 @@ """ from contextlib import asynccontextmanager -from typing import AsyncGenerator, Tuple +from typing import AsyncGenerator import anyio from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream @@ -12,23 +12,16 @@ from mcp_python.server import Server from mcp_python.types import JSONRPCMessage +MessageStream = tuple[ + MemoryObjectReceiveStream[JSONRPCMessage | Exception], + MemoryObjectSendStream[JSONRPCMessage] +] @asynccontextmanager -async def create_client_server_memory_streams() -> ( - AsyncGenerator[ - Tuple[ - Tuple[ - MemoryObjectReceiveStream[JSONRPCMessage | Exception], - MemoryObjectSendStream[JSONRPCMessage], - ], - Tuple[ - MemoryObjectReceiveStream[JSONRPCMessage | Exception], - MemoryObjectSendStream[JSONRPCMessage], - ], - ], - None, - ] -): +async def create_client_server_memory_streams() -> AsyncGenerator[ + tuple[MessageStream, MessageStream], + None +]: """ Creates a pair of bidirectional memory streams for client-server communication. @@ -44,7 +37,6 @@ async def create_client_server_memory_streams() -> ( JSONRPCMessage | Exception ](1) - # Return streams grouped by client/server client_streams = (server_to_client_receive, client_to_server_send) server_streams = (client_to_server_receive, server_to_client_send) @@ -66,10 +58,8 @@ async def create_connected_server_and_client_session( client_streams, server_streams, ): - # Unpack the streams client_read, client_write = client_streams server_read, server_write = server_streams - print("stream-1") # Create a cancel scope for the server task async with anyio.create_task_group() as tg: @@ -80,8 +70,6 @@ async def create_connected_server_and_client_session( server.create_initialization_options(), ) - print("stream2") - try: # Client session could be created here using client_read and # client_write This would allow testing the server with a client diff --git a/mcp_python/shared/session.py b/mcp_python/shared/session.py index e4c0b2d0c..b9ba6da57 100644 --- a/mcp_python/shared/session.py +++ b/mcp_python/shared/session.py @@ -154,7 +154,6 @@ async def send_request( with anyio.fail_after(self._read_timeout_seconds): response_or_error = await response_stream_reader.receive() except TimeoutError: - # TODO: make sure this response comes back correctly to the client raise McpError( ErrorData( code=408, diff --git a/tests/conftest.py b/tests/conftest.py index 849b2dca1..37ff5a4ec 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,7 +15,6 @@ def mcp_server() -> Server: server = Server(name="test_server") - # Add a simple resource for testing @server.list_resources() async def handle_list_resources(): return [ diff --git a/tests/test_memory.py b/tests/test_memory.py index 058104735..dcb069c71 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -15,11 +15,8 @@ async def client_connected_to_server( mcp_server: Server, ) -> AsyncGenerator[ClientSession, None]: - print("11111") async with create_connected_server_and_client_session(mcp_server) as client_session: - print("2222k") yield client_session - print("33") @pytest.mark.anyio @@ -28,6 +25,4 @@ async def test_memory_server_and_client_connection( ): """Shows how a client and server can communicate over memory streams.""" response = await client_connected_to_server.send_ping() - print("foo") assert isinstance(response, EmptyResult) - print("bar") From bef37ce6f7a08ded4a0378df2be8aad50f85c6ef Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Tue, 5 Nov 2024 14:46:10 -0500 Subject: [PATCH 09/12] add options to client server connection --- mcp_python/server/memory.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/mcp_python/server/memory.py b/mcp_python/server/memory.py index 8a36a8ef5..6885d0d5c 100644 --- a/mcp_python/server/memory.py +++ b/mcp_python/server/memory.py @@ -52,6 +52,8 @@ async def create_client_server_memory_streams() -> AsyncGenerator[ @asynccontextmanager async def create_connected_server_and_client_session( server: Server, + read_timeout_seconds: int | float | None = None, + raise_exceptions: bool = False, ) -> AsyncGenerator[ClientSession, None]: """Creates a ServerSession that is connected to the `server`.""" async with create_client_server_memory_streams() as ( @@ -64,10 +66,12 @@ async def create_connected_server_and_client_session( # Create a cancel scope for the server task async with anyio.create_task_group() as tg: tg.start_soon( - server.run, - server_read, - server_write, - server.create_initialization_options(), + lambda: server.run( + server_read, + server_write, + server.create_initialization_options(), + raise_exceptions=raise_exceptions, + ) ) try: @@ -75,7 +79,8 @@ async def create_connected_server_and_client_session( # client_write This would allow testing the server with a client # in the same process async with ClientSession( - read_stream=client_read, write_stream=client_write + read_stream=client_read, write_stream=client_write, + read_timeout_seconds=read_timeout_seconds, ) as client_session: await client_session.initialize() yield client_session From 0ae03fa8db8dfe55ca0e085b9ef181bbbe5bda40 Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Tue, 5 Nov 2024 14:57:46 -0500 Subject: [PATCH 10/12] cleanup from review comments --- mcp_python/client/session.py | 4 +++- mcp_python/server/memory.py | 8 +++----- mcp_python/shared/session.py | 11 ++++++++--- pyproject.toml | 2 +- uv.lock | 2 +- 5 files changed, 16 insertions(+), 11 deletions(-) diff --git a/mcp_python/client/session.py b/mcp_python/client/session.py index c8acf33fd..82ec0b2bd 100644 --- a/mcp_python/client/session.py +++ b/mcp_python/client/session.py @@ -1,3 +1,5 @@ +from datetime import timedelta + from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from pydantic import AnyUrl @@ -36,7 +38,7 @@ def __init__( self, read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception], write_stream: MemoryObjectSendStream[JSONRPCMessage], - read_timeout_seconds: int | float | None = None, + read_timeout_seconds: timedelta | None = None, ) -> None: super().__init__( read_stream, diff --git a/mcp_python/server/memory.py b/mcp_python/server/memory.py index 6885d0d5c..7a69764f1 100644 --- a/mcp_python/server/memory.py +++ b/mcp_python/server/memory.py @@ -3,6 +3,7 @@ """ from contextlib import asynccontextmanager +from datetime import timedelta from typing import AsyncGenerator import anyio @@ -52,10 +53,10 @@ async def create_client_server_memory_streams() -> AsyncGenerator[ @asynccontextmanager async def create_connected_server_and_client_session( server: Server, - read_timeout_seconds: int | float | None = None, + read_timeout_seconds: timedelta | None = None, raise_exceptions: bool = False, ) -> AsyncGenerator[ClientSession, None]: - """Creates a ServerSession that is connected to the `server`.""" + """Creates a ClientSession that is connected to a running MCP server.""" async with create_client_server_memory_streams() as ( client_streams, server_streams, @@ -75,9 +76,6 @@ async def create_connected_server_and_client_session( ) try: - # Client session could be created here using client_read and - # client_write This would allow testing the server with a client - # in the same process async with ClientSession( read_stream=client_read, write_stream=client_write, read_timeout_seconds=read_timeout_seconds, diff --git a/mcp_python/shared/session.py b/mcp_python/shared/session.py index b9ba6da57..0b87d429e 100644 --- a/mcp_python/shared/session.py +++ b/mcp_python/shared/session.py @@ -1,5 +1,7 @@ from contextlib import AbstractAsyncContextManager +from datetime import timedelta from typing import Generic, TypeVar +import httpx import anyio import anyio.lowlevel @@ -88,7 +90,7 @@ def __init__( receive_request_type: type[ReceiveRequestT], receive_notification_type: type[ReceiveNotificationT], # If none, reading will never time out - read_timeout_seconds: int | float | None = None, + read_timeout_seconds: timedelta | None = None, ) -> None: self._read_stream = read_stream self._write_stream = write_stream @@ -151,12 +153,15 @@ async def send_request( await self._write_stream.send(JSONRPCMessage(jsonrpc_request)) try: - with anyio.fail_after(self._read_timeout_seconds): + with anyio.fail_after( + None if self._read_timeout_seconds is None + else self._read_timeout_seconds.total_seconds() + ): response_or_error = await response_stream_reader.receive() except TimeoutError: raise McpError( ErrorData( - code=408, + code=httpx.codes.REQUEST_TIMEOUT, message=( f"Timed out while waiting for response to " f"{request.__class__.__name__}. Waited " diff --git a/pyproject.toml b/pyproject.toml index bebdd827b..fa4ab28e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "mcp-python" -version = "0.4.0.dev" +version = "0.3.1.dev" description = "Model Context Protocol implementation for Python" readme = "README.md" requires-python = ">=3.10" diff --git a/uv.lock b/uv.lock index 81c40d0c1..ed545091e 100644 --- a/uv.lock +++ b/uv.lock @@ -163,7 +163,7 @@ wheels = [ [[package]] name = "mcp-python" -version = "0.4.0.dev0" +version = "0.3.1.dev0" source = { editable = "." } dependencies = [ { name = "anyio" }, From 5dd79b2b12add590d62b4bc238ed7c2170a369df Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Tue, 5 Nov 2024 14:59:16 -0500 Subject: [PATCH 11/12] move memory.py from server/ to shared/ --- mcp_python/{server => shared}/memory.py | 0 mcp_python/shared/session.py | 2 +- tests/{ => shared}/test_memory.py | 2 +- 3 files changed, 2 insertions(+), 2 deletions(-) rename mcp_python/{server => shared}/memory.py (100%) rename tests/{ => shared}/test_memory.py (95%) diff --git a/mcp_python/server/memory.py b/mcp_python/shared/memory.py similarity index 100% rename from mcp_python/server/memory.py rename to mcp_python/shared/memory.py diff --git a/mcp_python/shared/session.py b/mcp_python/shared/session.py index 0b87d429e..f063a33bd 100644 --- a/mcp_python/shared/session.py +++ b/mcp_python/shared/session.py @@ -1,10 +1,10 @@ from contextlib import AbstractAsyncContextManager from datetime import timedelta from typing import Generic, TypeVar -import httpx import anyio import anyio.lowlevel +import httpx from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from pydantic import BaseModel diff --git a/tests/test_memory.py b/tests/shared/test_memory.py similarity index 95% rename from tests/test_memory.py rename to tests/shared/test_memory.py index dcb069c71..60a90de3d 100644 --- a/tests/test_memory.py +++ b/tests/shared/test_memory.py @@ -3,7 +3,7 @@ from mcp_python.client.session import ClientSession from mcp_python.server import Server -from mcp_python.server.memory import ( +from mcp_python.shared.memory import ( create_connected_server_and_client_session, ) from mcp_python.types import ( From a64d1ee3b9138154632173227759690927b5b3f3 Mon Sep 17 00:00:00 2001 From: Nick Merrill Date: Tue, 5 Nov 2024 18:18:54 -0500 Subject: [PATCH 12/12] bump minor version and relock --- mcp_python/shared/memory.py | 3 ++- pyproject.toml | 2 +- uv.lock | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/mcp_python/shared/memory.py b/mcp_python/shared/memory.py index 7a69764f1..a2917499a 100644 --- a/mcp_python/shared/memory.py +++ b/mcp_python/shared/memory.py @@ -77,7 +77,8 @@ async def create_connected_server_and_client_session( try: async with ClientSession( - read_stream=client_read, write_stream=client_write, + read_stream=client_read, + write_stream=client_write, read_timeout_seconds=read_timeout_seconds, ) as client_session: await client_session.initialize() diff --git a/pyproject.toml b/pyproject.toml index fa4ab28e6..96f24bbf9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "mcp-python" -version = "0.3.1.dev" +version = "0.5.0dev" description = "Model Context Protocol implementation for Python" readme = "README.md" requires-python = ">=3.10" diff --git a/uv.lock b/uv.lock index ed545091e..f0fac56f8 100644 --- a/uv.lock +++ b/uv.lock @@ -163,7 +163,7 @@ wheels = [ [[package]] name = "mcp-python" -version = "0.3.1.dev0" +version = "0.5.0.dev0" source = { editable = "." } dependencies = [ { name = "anyio" },