Merge pull request #30 from modelcontextprotocol/davidsp/capabilities

feat: add structured capability types

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

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():

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,16 +58,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:
@@ -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

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

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)