feat: add structured capability types

Replace generic capability dictionaries with structured types for prompts,
resources, tools, and roots. This improves type safety and makes capability
features like listChanged and subscribe more explicit in the protocol.

Author: dsp-ant

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",

mcp_python/server/__init__.py

@@ -27,20 +27,24 @@
     ListResourcesResult,
     ListToolsRequest,
     ListToolsResult,
+    LoggingCapability,
     LoggingLevel,
     PingRequest,
     ProgressNotification,
     Prompt,
     PromptReference,
+    PromptsCapability,
     ReadResourceRequest,
     ReadResourceResult,
     Resource,
     ResourceReference,
+    ResourcesCapability,
     ServerCapabilities,
     ServerResult,
     SetLevelRequest,
     SubscribeRequest,
     Tool,
+    ToolsCapability,
     UnsubscribeRequest,
 )
 
@@ -51,16 +55,33 @@
 )
 
 
+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
         self.request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]] = {
             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:
@@ -78,20 +99,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
@@ -169,9 +221,7 @@ def decorator(func: Callable[[], Awaitable[list[Resource]]]):
 
             async def handler(_: Any):
                 resources = await func()
-                return ServerResult(
-                    ListResourcesResult(resources=resources)
-                )
+                return ServerResult(ListResourcesResult(resources=resources))
 
             self.request_handlers[ListResourcesRequest] = handler
             return func
@@ -216,7 +266,6 @@ async def handler(req: ReadResourceRequest):
 
         return decorator
 
-
     def set_logging_level(self):
         from mcp_python.types import EmptyResult
 

mcp_python/shared/memory.py

@@ -15,14 +15,14 @@
 
 MessageStream = tuple[
     MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-    MemoryObjectSendStream[JSONRPCMessage]
+    MemoryObjectSendStream[JSONRPCMessage],
 ]
 
+
 @asynccontextmanager
-async def create_client_server_memory_streams() -> AsyncGenerator[
-    tuple[MessageStream, MessageStream],
-    None
-]:
+async def create_client_server_memory_streams() -> (
+    AsyncGenerator[tuple[MessageStream, MessageStream], None]
+):
     """
     Creates a pair of bidirectional memory streams for client-server communication.
 

mcp_python/shared/session.py

@@ -154,7 +154,8 @@ async def send_request(
 
         try:
             with anyio.fail_after(
-                None if self._read_timeout_seconds is None
+                None
+                if self._read_timeout_seconds is None
                 else self._read_timeout_seconds.total_seconds()
             ):
                 response_or_error = await response_stream_reader.receive()
@@ -168,7 +169,6 @@ async def send_request(
                         f"{self._read_timeout_seconds} seconds."
                     ),
                 )
-
             )
 
         if isinstance(response_or_error, JSONRPCError):

mcp_python/types.py

@@ -184,28 +184,70 @@ 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 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
     """Present if the client supports sampling from an LLM."""
+    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."""
+
+    pass
+
+
 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")
 

tests/conftest.py

@@ -11,6 +11,7 @@
     capabilities=ServerCapabilities(),
 )
 
+
 @pytest.fixture
 def mcp_server() -> Server:
     server = Server(name="test_server")
@@ -21,7 +22,7 @@ async def handle_list_resources():
             Resource(
                 uri=AnyUrl("memory://test"),
                 name="Test Resource",
-                description="A test resource"
+                description="A test resource",
             )
         ]
 

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,15 +86,15 @@ 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
     @server.list_resources()
     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)