Skip to content

Commit 004581b

Browse files
committed
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.
1 parent 60e9c7a commit 004581b

File tree

7 files changed

+137
-33
lines changed

7 files changed

+137
-33
lines changed

mcp_python/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,13 @@
3333
Notification,
3434
PingRequest,
3535
ProgressNotification,
36+
PromptsCapability,
3637
ReadResourceRequest,
3738
ReadResourceResult,
3839
Resource,
40+
ResourcesCapability,
3941
ResourceUpdatedNotification,
42+
RootsCapability,
4043
SamplingMessage,
4144
ServerCapabilities,
4245
ServerNotification,
@@ -46,6 +49,7 @@
4649
StopReason,
4750
SubscribeRequest,
4851
Tool,
52+
ToolsCapability,
4953
UnsubscribeRequest,
5054
)
5155
from .types import (
@@ -82,10 +86,13 @@
8286
"Notification",
8387
"PingRequest",
8488
"ProgressNotification",
89+
"PromptsCapability",
8590
"ReadResourceRequest",
8691
"ReadResourceResult",
92+
"ResourcesCapability",
8793
"ResourceUpdatedNotification",
8894
"Resource",
95+
"RootsCapability",
8996
"SamplingMessage",
9097
"SamplingRole",
9198
"ServerCapabilities",
@@ -98,6 +105,7 @@
98105
"StopReason",
99106
"SubscribeRequest",
100107
"Tool",
108+
"ToolsCapability",
101109
"UnsubscribeRequest",
102110
"stdio_client",
103111
"stdio_server",

mcp_python/server/__init__.py

Lines changed: 63 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -27,20 +27,24 @@
2727
ListResourcesResult,
2828
ListToolsRequest,
2929
ListToolsResult,
30+
LoggingCapability,
3031
LoggingLevel,
3132
PingRequest,
3233
ProgressNotification,
3334
Prompt,
3435
PromptReference,
36+
PromptsCapability,
3537
ReadResourceRequest,
3638
ReadResourceResult,
3739
Resource,
3840
ResourceReference,
41+
ResourcesCapability,
3942
ServerCapabilities,
4043
ServerResult,
4144
SetLevelRequest,
4245
SubscribeRequest,
4346
Tool,
47+
ToolsCapability,
4448
UnsubscribeRequest,
4549
)
4650

@@ -51,16 +55,33 @@
5155
)
5256

5357

58+
class NotificationOptions:
59+
def __init__(
60+
self,
61+
prompts_changed: bool = False,
62+
resources_changed: bool = False,
63+
tools_changed: bool = False,
64+
):
65+
self.prompts_changed = prompts_changed
66+
self.resources_changed = resources_changed
67+
self.tools_changed = tools_changed
68+
69+
5470
class Server:
5571
def __init__(self, name: str):
5672
self.name = name
5773
self.request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]] = {
5874
PingRequest: _ping_handler,
5975
}
6076
self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {}
77+
self.notification_options = NotificationOptions()
6178
logger.debug(f"Initializing server '{name}'")
6279

63-
def create_initialization_options(self) -> types.InitializationOptions:
80+
def create_initialization_options(
81+
self,
82+
notification_options: NotificationOptions | None = None,
83+
experimental_capabilities: dict[str, dict[str, Any]] | None = None,
84+
) -> types.InitializationOptions:
6485
"""Create initialization options from this server instance."""
6586

6687
def pkg_version(package: str) -> str:
@@ -78,20 +99,51 @@ def pkg_version(package: str) -> str:
7899
return types.InitializationOptions(
79100
server_name=self.name,
80101
server_version=pkg_version("mcp_python"),
81-
capabilities=self.get_capabilities(),
102+
capabilities=self.get_capabilities(
103+
notification_options or NotificationOptions(),
104+
experimental_capabilities or {},
105+
),
82106
)
83107

84-
def get_capabilities(self) -> ServerCapabilities:
108+
def get_capabilities(
109+
self,
110+
notification_options: NotificationOptions,
111+
experimental_capabilities: dict[str, dict[str, Any]],
112+
) -> ServerCapabilities:
85113
"""Convert existing handlers to a ServerCapabilities object."""
86-
87-
def get_capability(req_type: type) -> dict[str, Any] | None:
88-
return {} if req_type in self.request_handlers else None
114+
prompts_capability = None
115+
resources_capability = None
116+
tools_capability = None
117+
logging_capability = None
118+
119+
# Set prompt capabilities if handler exists
120+
if ListPromptsRequest in self.request_handlers:
121+
prompts_capability = PromptsCapability(
122+
listChanged=notification_options.prompts_changed
123+
)
124+
125+
# Set resource capabilities if handler exists
126+
if ListResourcesRequest in self.request_handlers:
127+
resources_capability = ResourcesCapability(
128+
subscribe=False, listChanged=notification_options.resources_changed
129+
)
130+
131+
# Set tool capabilities if handler exists
132+
if ListToolsRequest in self.request_handlers:
133+
tools_capability = ToolsCapability(
134+
listChanged=notification_options.tools_changed
135+
)
136+
137+
# Set logging capabilities if handler exists
138+
if SetLevelRequest in self.request_handlers:
139+
logging_capability = LoggingCapability()
89140

90141
return ServerCapabilities(
91-
prompts=get_capability(ListPromptsRequest),
92-
resources=get_capability(ListResourcesRequest),
93-
tools=get_capability(ListToolsRequest),
94-
logging=get_capability(SetLevelRequest),
142+
prompts=prompts_capability,
143+
resources=resources_capability,
144+
tools=tools_capability,
145+
logging=logging_capability,
146+
experimental=experimental_capabilities,
95147
)
96148

97149
@property
@@ -169,9 +221,7 @@ def decorator(func: Callable[[], Awaitable[list[Resource]]]):
169221

170222
async def handler(_: Any):
171223
resources = await func()
172-
return ServerResult(
173-
ListResourcesResult(resources=resources)
174-
)
224+
return ServerResult(ListResourcesResult(resources=resources))
175225

176226
self.request_handlers[ListResourcesRequest] = handler
177227
return func
@@ -216,7 +266,6 @@ async def handler(req: ReadResourceRequest):
216266

217267
return decorator
218268

219-
220269
def set_logging_level(self):
221270
from mcp_python.types import EmptyResult
222271

mcp_python/shared/memory.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,14 @@
1515

1616
MessageStream = tuple[
1717
MemoryObjectReceiveStream[JSONRPCMessage | Exception],
18-
MemoryObjectSendStream[JSONRPCMessage]
18+
MemoryObjectSendStream[JSONRPCMessage],
1919
]
2020

21+
2122
@asynccontextmanager
22-
async def create_client_server_memory_streams() -> AsyncGenerator[
23-
tuple[MessageStream, MessageStream],
24-
None
25-
]:
23+
async def create_client_server_memory_streams() -> (
24+
AsyncGenerator[tuple[MessageStream, MessageStream], None]
25+
):
2626
"""
2727
Creates a pair of bidirectional memory streams for client-server communication.
2828

mcp_python/shared/session.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,8 @@ async def send_request(
154154

155155
try:
156156
with anyio.fail_after(
157-
None if self._read_timeout_seconds is None
157+
None
158+
if self._read_timeout_seconds is None
158159
else self._read_timeout_seconds.total_seconds()
159160
):
160161
response_or_error = await response_stream_reader.receive()
@@ -168,7 +169,6 @@ async def send_request(
168169
f"{self._read_timeout_seconds} seconds."
169170
),
170171
)
171-
172172
)
173173

174174
if isinstance(response_or_error, JSONRPCError):

mcp_python/types.py

Lines changed: 46 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -184,28 +184,70 @@ class Implementation(BaseModel):
184184
model_config = ConfigDict(extra="allow")
185185

186186

187+
class RootsCapability(BaseModel):
188+
"""Capability for root operations."""
189+
190+
listChanged: bool | None = None
191+
"""Whether the client supports notifications for changes to the roots list."""
192+
model_config = ConfigDict(extra="allow")
193+
194+
187195
class ClientCapabilities(BaseModel):
188196
"""Capabilities a client may support."""
189197

190198
experimental: dict[str, dict[str, Any]] | None = None
191199
"""Experimental, non-standard capabilities that the client supports."""
192200
sampling: dict[str, Any] | None = None
193201
"""Present if the client supports sampling from an LLM."""
202+
roots: RootsCapability | None = None
203+
"""Present if the client supports listing roots."""
204+
model_config = ConfigDict(extra="allow")
205+
206+
207+
class PromptsCapability(BaseModel):
208+
"""Capability for prompts operations."""
209+
210+
listChanged: bool | None = None
211+
"""Whether this server supports notifications for changes to the prompt list."""
212+
model_config = ConfigDict(extra="allow")
213+
214+
215+
class ResourcesCapability(BaseModel):
216+
"""Capability for resources operations."""
217+
218+
subscribe: bool | None = None
219+
"""Whether this server supports subscribing to resource updates."""
220+
listChanged: bool | None = None
221+
"""Whether this server supports notifications for changes to the resource list."""
222+
model_config = ConfigDict(extra="allow")
223+
224+
225+
class ToolsCapability(BaseModel):
226+
"""Capability for tools operations."""
227+
228+
listChanged: bool | None = None
229+
"""Whether this server supports notifications for changes to the tool list."""
194230
model_config = ConfigDict(extra="allow")
195231

196232

233+
class LoggingCapability(BaseModel):
234+
"""Capability for logging operations."""
235+
236+
pass
237+
238+
197239
class ServerCapabilities(BaseModel):
198240
"""Capabilities that a server may support."""
199241

200242
experimental: dict[str, dict[str, Any]] | None = None
201243
"""Experimental, non-standard capabilities that the server supports."""
202-
logging: dict[str, Any] | None = None
244+
logging: LoggingCapability | None = None
203245
"""Present if the server supports sending log messages to the client."""
204-
prompts: dict[str, Any] | None = None
246+
prompts: PromptsCapability | None = None
205247
"""Present if the server offers any prompt templates."""
206-
resources: dict[str, Any] | None = None
248+
resources: ResourcesCapability | None = None
207249
"""Present if the server offers any resources to read."""
208-
tools: dict[str, Any] | None = None
250+
tools: ToolsCapability | None = None
209251
"""Present if the server offers any tools to call."""
210252
model_config = ConfigDict(extra="allow")
211253

tests/conftest.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
capabilities=ServerCapabilities(),
1212
)
1313

14+
1415
@pytest.fixture
1516
def mcp_server() -> Server:
1617
server = Server(name="test_server")
@@ -21,7 +22,7 @@ async def handle_list_resources():
2122
Resource(
2223
uri=AnyUrl("memory://test"),
2324
name="Test Resource",
24-
description="A test resource"
25+
description="A test resource",
2526
)
2627
]
2728

tests/server/test_session.py

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,16 @@
22
import pytest
33

44
from mcp_python.client.session import ClientSession
5-
from mcp_python.server import Server
5+
from mcp_python.server import NotificationOptions, Server
66
from mcp_python.server.session import ServerSession
77
from mcp_python.server.types import InitializationOptions
88
from mcp_python.types import (
99
ClientNotification,
1010
InitializedNotification,
1111
JSONRPCMessage,
1212
ServerCapabilities,
13+
PromptsCapability,
14+
ResourcesCapability,
1315
)
1416

1517

@@ -71,9 +73,11 @@ async def run_server():
7173
@pytest.mark.anyio
7274
async def test_server_capabilities():
7375
server = Server("test")
76+
notification_options = NotificationOptions()
77+
experimental_capabilities = {}
7478

7579
# Initially no capabilities
76-
caps = server.get_capabilities()
80+
caps = server.get_capabilities(notification_options, experimental_capabilities)
7781
assert caps.prompts is None
7882
assert caps.resources is None
7983

@@ -82,15 +86,15 @@ async def test_server_capabilities():
8286
async def list_prompts():
8387
return []
8488

85-
caps = server.get_capabilities()
86-
assert caps.prompts == {}
89+
caps = server.get_capabilities(notification_options, experimental_capabilities)
90+
assert caps.prompts == PromptsCapability(listChanged=False)
8791
assert caps.resources is None
8892

8993
# Add a resources handler
9094
@server.list_resources()
9195
async def list_resources():
9296
return []
9397

94-
caps = server.get_capabilities()
95-
assert caps.prompts == {}
96-
assert caps.resources == {}
98+
caps = server.get_capabilities(notification_options, experimental_capabilities)
99+
assert caps.prompts == PromptsCapability(listChanged=False)
100+
assert caps.resources == ResourcesCapability(subscribe=False, listChanged=False)

0 commit comments

Comments
 (0)