diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 0f3e313e..27ca74d8 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -3,9 +3,9 @@ from datetime import timedelta from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from pydantic import AnyUrl +import mcp.types as types from mcp.shared.session import BaseSession from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS -import mcp.types as types class ClientSession( diff --git a/src/mcp/client/sse.py b/src/mcp/client/sse.py index c79f48a7..8a90e588 100644 --- a/src/mcp/client/sse.py +++ b/src/mcp/client/sse.py @@ -84,10 +84,8 @@ async def sse_client( case "message": try: - message = ( - types.JSONRPCMessage.model_validate_json( - sse.data - ) + message = types.JSONRPCMessage.model_validate_json( # noqa: E501 + sse.data ) logger.debug( f"Received server message: {message}" diff --git a/src/mcp/server/__init__.py b/src/mcp/server/__init__.py index abfc40da..3ce8b7bf 100644 --- a/src/mcp/server/__init__.py +++ b/src/mcp/server/__init__.py @@ -7,12 +7,12 @@ from typing import Any, Sequence from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from pydantic import AnyUrl +import mcp.types as types from mcp.server.models import InitializationOptions from mcp.server.session import ServerSession from mcp.server.stdio import stdio_server as stdio_server from mcp.shared.context import RequestContext from mcp.shared.session import RequestResponder -import mcp.types as types logger = logging.getLogger(__name__) @@ -36,7 +36,9 @@ class NotificationOptions: class Server: def __init__(self, name: str): self.name = name - self.request_handlers: dict[type, Callable[..., Awaitable[types.ServerResult]]] = { + self.request_handlers: dict[ + type, Callable[..., Awaitable[types.ServerResult]] + ] = { types.PingRequest: _ping_handler, } self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {} @@ -153,7 +155,9 @@ class Server: async def handler(_: Any): resources = await func() - return types.ServerResult(types.ListResourcesResult(resources=resources)) + return types.ServerResult( + types.ListResourcesResult(resources=resources) + ) self.request_handlers[types.ListResourcesRequest] = handler return func @@ -249,7 +253,11 @@ class Server: def decorator( func: Callable[ ..., - Awaitable[Sequence[types.TextContent | types.ImageContent | types.EmbeddedResource]], + Awaitable[ + Sequence[ + types.TextContent | types.ImageContent | types.EmbeddedResource + ] + ], ], ): logger.debug("Registering handler for CallToolRequest") @@ -261,7 +269,9 @@ class Server: for result in results: match result: case str() as text: - content.append(types.TextContent(type="text", text=text)) + content.append( + types.TextContent(type="text", text=text) + ) case types.ImageContent() as img: content.append( types.ImageContent( @@ -277,7 +287,9 @@ class Server: ) ) - return types.ServerResult(types.CallToolResult(content=content, isError=False)) + return types.ServerResult( + types.CallToolResult(content=content, isError=False) + ) except Exception as e: return types.ServerResult( types.CallToolResult( @@ -312,7 +324,10 @@ class Server: def decorator( func: Callable[ - [types.PromptReference | types.ResourceReference, types.CompletionArgument], + [ + types.PromptReference | types.ResourceReference, + types.CompletionArgument, + ], Awaitable[types.Completion | None], ], ): diff --git a/src/mcp/server/__main__.py b/src/mcp/server/__main__.py index 417380d6..1970eca7 100644 --- a/src/mcp/server/__main__.py +++ b/src/mcp/server/__main__.py @@ -4,9 +4,9 @@ import sys import anyio +from mcp.server.models import InitializationOptions from mcp.server.session import ServerSession from mcp.server.stdio import stdio_server -from mcp.server.models import InitializationOptions from mcp.types import ServerCapabilities if not sys.warnoptions: diff --git a/src/mcp/server/models.py b/src/mcp/server/models.py index 89205978..377ed517 100644 --- a/src/mcp/server/models.py +++ b/src/mcp/server/models.py @@ -3,9 +3,6 @@ This module provides simpler types to use with the server for managing prompts and tools. """ -from dataclasses import dataclass -from typing import Literal - from pydantic import BaseModel from mcp.types import ( diff --git a/src/mcp/server/session.py b/src/mcp/server/session.py index 97b70bda..dd663078 100644 --- a/src/mcp/server/session.py +++ b/src/mcp/server/session.py @@ -6,12 +6,12 @@ import anyio.lowlevel from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from pydantic import AnyUrl +import mcp.types as types from mcp.server.models import InitializationOptions from mcp.shared.session import ( BaseSession, RequestResponder, ) -import mcp.types as types class InitializationState(Enum): @@ -37,7 +37,9 @@ class ServerSession( write_stream: MemoryObjectSendStream[types.JSONRPCMessage], init_options: InitializationOptions, ) -> None: - super().__init__(read_stream, write_stream, types.ClientRequest, types.ClientNotification) + super().__init__( + read_stream, write_stream, types.ClientRequest, types.ClientNotification + ) self._initialization_state = InitializationState.NotInitialized self._init_options = init_options @@ -65,7 +67,9 @@ class ServerSession( "Received request before initialization was complete" ) - async def _received_notification(self, notification: types.ClientNotification) -> None: + async def _received_notification( + self, notification: types.ClientNotification + ) -> None: # Need this to avoid ASYNC910 await anyio.lowlevel.checkpoint() match notification.root: diff --git a/src/mcp/server/sse.py b/src/mcp/server/sse.py index 4074fdbf..f6e90bff 100644 --- a/src/mcp/server/sse.py +++ b/src/mcp/server/sse.py @@ -30,7 +30,9 @@ class SseServerTransport: """ _endpoint: str - _read_stream_writers: dict[UUID, MemoryObjectSendStream[types.JSONRPCMessage | Exception]] + _read_stream_writers: dict[ + UUID, MemoryObjectSendStream[types.JSONRPCMessage | Exception] + ] def __init__(self, endpoint: str) -> None: """ diff --git a/tests/server/test_session.py b/tests/server/test_session.py index 728d2766..a78ca90e 100644 --- a/tests/server/test_session.py +++ b/tests/server/test_session.py @@ -3,8 +3,8 @@ import pytest from mcp.client.session import ClientSession from mcp.server import NotificationOptions, Server -from mcp.server.session import ServerSession from mcp.server.models import InitializationOptions +from mcp.server.session import ServerSession from mcp.types import ( ClientNotification, InitializedNotification,