diff --git a/mcp_python/__init__.py b/mcp_python/__init__.py index 78a62be65..0d3c372ce 100644 --- a/mcp_python/__init__.py +++ b/mcp_python/__init__.py @@ -33,10 +33,13 @@ Notification, PingRequest, ProgressNotification, + PromptsCapability, ReadResourceRequest, ReadResourceResult, Resource, + ResourcesCapability, ResourceUpdatedNotification, + RootsCapability, SamplingMessage, ServerCapabilities, ServerNotification, @@ -46,6 +49,7 @@ StopReason, SubscribeRequest, Tool, + ToolsCapability, UnsubscribeRequest, ) from .types import ( @@ -82,10 +86,13 @@ "Notification", "PingRequest", "ProgressNotification", + "PromptsCapability", "ReadResourceRequest", "ReadResourceResult", + "ResourcesCapability", "ResourceUpdatedNotification", "Resource", + "RootsCapability", "SamplingMessage", "SamplingRole", "ServerCapabilities", @@ -98,6 +105,7 @@ "StopReason", "SubscribeRequest", "Tool", + "ToolsCapability", "UnsubscribeRequest", "stdio_client", "stdio_server", diff --git a/mcp_python/client/session.py b/mcp_python/client/session.py index 6c6d01fb3..3f50503e9 100644 --- a/mcp_python/client/session.py +++ b/mcp_python/client/session.py @@ -26,6 +26,7 @@ PromptReference, ReadResourceResult, ResourceReference, + RootsCapability, ServerNotification, ServerRequest, ) @@ -69,12 +70,12 @@ async def initialize(self) -> InitializeResult: capabilities=ClientCapabilities( sampling=None, experimental=None, - roots={ + roots=RootsCapability( # TODO: Should this be based on whether we # _will_ send notifications, or only whether # they're supported? - "listChanged": True - }, + listChanged=True + ), ), clientInfo=Implementation(name="mcp_python", version="0.1.0"), ), diff --git a/mcp_python/client/stdio.py b/mcp_python/client/stdio.py index f69578e83..6fc728f7a 100644 --- a/mcp_python/client/stdio.py +++ b/mcp_python/client/stdio.py @@ -12,9 +12,19 @@ # Environment variables to inherit by default DEFAULT_INHERITED_ENV_VARS = ( - ["APPDATA", "HOMEDRIVE", "HOMEPATH", "LOCALAPPDATA", "PATH", - "PROCESSOR_ARCHITECTURE", "SYSTEMDRIVE", "SYSTEMROOT", "TEMP", - "USERNAME", "USERPROFILE"] + [ + "APPDATA", + "HOMEDRIVE", + "HOMEPATH", + "LOCALAPPDATA", + "PATH", + "PROCESSOR_ARCHITECTURE", + "SYSTEMDRIVE", + "SYSTEMROOT", + "TEMP", + "USERNAME", + "USERPROFILE", + ] if sys.platform == "win32" else ["HOME", "LOGNAME", "PATH", "SHELL", "TERM", "USER"] ) @@ -74,7 +84,7 @@ async def stdio_client(server: StdioServerParameters): process = await anyio.open_process( [server.command, *server.args], env=server.env if server.env is not None else get_default_environment(), - stderr=sys.stderr + stderr=sys.stderr, ) async def stdout_reader(): diff --git a/mcp_python/server/__init__.py b/mcp_python/server/__init__.py index 26d690250..f3183fdb3 100644 --- a/mcp_python/server/__init__.py +++ b/mcp_python/server/__init__.py @@ -28,22 +28,26 @@ ListResourcesResult, ListToolsRequest, ListToolsResult, + LoggingCapability, LoggingLevel, PingRequest, ProgressNotification, Prompt, PromptMessage, PromptReference, + PromptsCapability, ReadResourceRequest, ReadResourceResult, Resource, ResourceReference, + ResourcesCapability, ServerCapabilities, ServerResult, SetLevelRequest, SubscribeRequest, TextContent, Tool, + ToolsCapability, UnsubscribeRequest, ) @@ -54,6 +58,18 @@ ) +class NotificationOptions: + def __init__( + self, + prompts_changed: bool = False, + resources_changed: bool = False, + tools_changed: bool = False, + ): + self.prompts_changed = prompts_changed + self.resources_changed = resources_changed + self.tools_changed = tools_changed + + class Server: def __init__(self, name: str): self.name = name @@ -61,9 +77,14 @@ def __init__(self, name: str): PingRequest: _ping_handler, } self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {} + self.notification_options = NotificationOptions() logger.debug(f"Initializing server '{name}'") - def create_initialization_options(self) -> types.InitializationOptions: + def create_initialization_options( + self, + notification_options: NotificationOptions | None = None, + experimental_capabilities: dict[str, dict[str, Any]] | None = None, + ) -> types.InitializationOptions: """Create initialization options from this server instance.""" def pkg_version(package: str) -> str: @@ -81,20 +102,51 @@ def pkg_version(package: str) -> str: return types.InitializationOptions( server_name=self.name, server_version=pkg_version("mcp_python"), - capabilities=self.get_capabilities(), + capabilities=self.get_capabilities( + notification_options or NotificationOptions(), + experimental_capabilities or {}, + ), ) - def get_capabilities(self) -> ServerCapabilities: + def get_capabilities( + self, + notification_options: NotificationOptions, + experimental_capabilities: dict[str, dict[str, Any]], + ) -> ServerCapabilities: """Convert existing handlers to a ServerCapabilities object.""" - - def get_capability(req_type: type) -> dict[str, Any] | None: - return {} if req_type in self.request_handlers else None + prompts_capability = None + resources_capability = None + tools_capability = None + logging_capability = None + + # Set prompt capabilities if handler exists + if ListPromptsRequest in self.request_handlers: + prompts_capability = PromptsCapability( + listChanged=notification_options.prompts_changed + ) + + # Set resource capabilities if handler exists + if ListResourcesRequest in self.request_handlers: + resources_capability = ResourcesCapability( + subscribe=False, listChanged=notification_options.resources_changed + ) + + # Set tool capabilities if handler exists + if ListToolsRequest in self.request_handlers: + tools_capability = ToolsCapability( + listChanged=notification_options.tools_changed + ) + + # Set logging capabilities if handler exists + if SetLevelRequest in self.request_handlers: + logging_capability = LoggingCapability() return ServerCapabilities( - prompts=get_capability(ListPromptsRequest), - resources=get_capability(ListResourcesRequest), - tools=get_capability(ListToolsRequest), - logging=get_capability(SetLevelRequest), + prompts=prompts_capability, + resources=resources_capability, + tools=tools_capability, + logging=logging_capability, + experimental=experimental_capabilities, ) @property diff --git a/mcp_python/types.py b/mcp_python/types.py index ffcf75e00..a2b897403 100644 --- a/mcp_python/types.py +++ b/mcp_python/types.py @@ -184,30 +184,76 @@ class Implementation(BaseModel): model_config = ConfigDict(extra="allow") +class RootsCapability(BaseModel): + """Capability for root operations.""" + + listChanged: bool | None = None + """Whether the client supports notifications for changes to the roots list.""" + model_config = ConfigDict(extra="allow") + + +class SamplingCapability(BaseModel): + """Capability for logging operations.""" + + model_config = ConfigDict(extra="allow") + + class ClientCapabilities(BaseModel): """Capabilities a client may support.""" experimental: dict[str, dict[str, Any]] | None = None """Experimental, non-standard capabilities that the client supports.""" - sampling: dict[str, Any] | None = None + sampling: SamplingCapability | None = None """Present if the client supports sampling from an LLM.""" - roots: dict[str, Any] | None = None + roots: RootsCapability | None = None """Present if the client supports listing roots.""" model_config = ConfigDict(extra="allow") +class PromptsCapability(BaseModel): + """Capability for prompts operations.""" + + listChanged: bool | None = None + """Whether this server supports notifications for changes to the prompt list.""" + model_config = ConfigDict(extra="allow") + + +class ResourcesCapability(BaseModel): + """Capability for resources operations.""" + + subscribe: bool | None = None + """Whether this server supports subscribing to resource updates.""" + listChanged: bool | None = None + """Whether this server supports notifications for changes to the resource list.""" + model_config = ConfigDict(extra="allow") + + +class ToolsCapability(BaseModel): + """Capability for tools operations.""" + + listChanged: bool | None = None + """Whether this server supports notifications for changes to the tool list.""" + model_config = ConfigDict(extra="allow") + + +class LoggingCapability(BaseModel): + """Capability for logging operations.""" + + model_config = ConfigDict(extra="allow") + + class ServerCapabilities(BaseModel): """Capabilities that a server may support.""" experimental: dict[str, dict[str, Any]] | None = None """Experimental, non-standard capabilities that the server supports.""" - logging: dict[str, Any] | None = None + logging: LoggingCapability | None = None """Present if the server supports sending log messages to the client.""" - prompts: dict[str, Any] | None = None + prompts: PromptsCapability | None = None """Present if the server offers any prompt templates.""" - resources: dict[str, Any] | None = None + resources: ResourcesCapability | None = None """Present if the server offers any resources to read.""" - tools: dict[str, Any] | None = None + tools: ToolsCapability | None = None """Present if the server offers any tools to call.""" model_config = ConfigDict(extra="allow") diff --git a/tests/server/test_session.py b/tests/server/test_session.py index addf0f5f7..4813fbdaf 100644 --- a/tests/server/test_session.py +++ b/tests/server/test_session.py @@ -2,13 +2,15 @@ import pytest from mcp_python.client.session import ClientSession -from mcp_python.server import Server +from mcp_python.server import NotificationOptions, Server from mcp_python.server.session import ServerSession from mcp_python.server.types import InitializationOptions from mcp_python.types import ( ClientNotification, InitializedNotification, JSONRPCMessage, + PromptsCapability, + ResourcesCapability, ServerCapabilities, ) @@ -71,9 +73,11 @@ async def run_server(): @pytest.mark.anyio async def test_server_capabilities(): server = Server("test") + notification_options = NotificationOptions() + experimental_capabilities = {} # Initially no capabilities - caps = server.get_capabilities() + caps = server.get_capabilities(notification_options, experimental_capabilities) assert caps.prompts is None assert caps.resources is None @@ -82,8 +86,8 @@ async def test_server_capabilities(): async def list_prompts(): return [] - caps = server.get_capabilities() - assert caps.prompts == {} + caps = server.get_capabilities(notification_options, experimental_capabilities) + assert caps.prompts == PromptsCapability(listChanged=False) assert caps.resources is None # Add a resources handler @@ -91,6 +95,6 @@ async def list_prompts(): async def list_resources(): return [] - caps = server.get_capabilities() - assert caps.prompts == {} - assert caps.resources == {} + caps = server.get_capabilities(notification_options, experimental_capabilities) + assert caps.prompts == PromptsCapability(listChanged=False) + assert caps.resources == ResourcesCapability(subscribe=False, listChanged=False)