Skip to content

Commit b255dd7

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 b255dd7

File tree

6 files changed

+121
-26
lines changed

6 files changed

+121
-26
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: 58 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,31 @@
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, notification_options: NotificationOptions | None = None
82+
) -> types.InitializationOptions:
6483
"""Create initialization options from this server instance."""
6584

6685
def pkg_version(package: str) -> str:
@@ -78,20 +97,48 @@ def pkg_version(package: str) -> str:
7897
return types.InitializationOptions(
7998
server_name=self.name,
8099
server_version=pkg_version("mcp_python"),
81-
capabilities=self.get_capabilities(),
100+
capabilities=self.get_capabilities(
101+
notification_options or NotificationOptions()
102+
),
82103
)
83104

84-
def get_capabilities(self) -> ServerCapabilities:
105+
def get_capabilities(
106+
self, notification_options: NotificationOptions
107+
) -> ServerCapabilities:
85108
"""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
109+
prompts_capability = None
110+
resources_capability = None
111+
tools_capability = None
112+
logging_capability = None
113+
114+
# Set prompt capabilities if handler exists
115+
if ListPromptsRequest in self.request_handlers:
116+
prompts_capability = PromptsCapability(
117+
listChanged=notification_options.prompts_changed
118+
)
119+
120+
# Set resource capabilities if handler exists
121+
if ListResourcesRequest in self.request_handlers:
122+
resources_capability = ResourcesCapability(
123+
subscribe=False,
124+
listChanged=notification_options.resources_changed
125+
)
126+
127+
# Set tool capabilities if handler exists
128+
if ListToolsRequest in self.request_handlers:
129+
tools_capability = ToolsCapability(
130+
listChanged=notification_options.tools_changed
131+
)
132+
133+
# Set logging capabilities if handler exists
134+
if SetLevelRequest in self.request_handlers:
135+
logging_capability = LoggingCapability()
89136

90137
return ServerCapabilities(
91-
prompts=get_capability(ListPromptsRequest),
92-
resources=get_capability(ListResourcesRequest),
93-
tools=get_capability(ListToolsRequest),
94-
logging=get_capability(SetLevelRequest),
138+
prompts=prompts_capability,
139+
resources=resources_capability,
140+
tools=tools_capability,
141+
logging=logging_capability
95142
)
96143

97144
@property
@@ -169,9 +216,7 @@ def decorator(func: Callable[[], Awaitable[list[Resource]]]):
169216

170217
async def handler(_: Any):
171218
resources = await func()
172-
return ServerResult(
173-
ListResourcesResult(resources=resources)
174-
)
219+
return ServerResult(ListResourcesResult(resources=resources))
175220

176221
self.request_handlers[ListResourcesRequest] = handler
177222
return func
@@ -216,7 +261,6 @@ async def handler(req: ReadResourceRequest):
216261

217262
return decorator
218263

219-
220264
def set_logging_level(self):
221265
from mcp_python.types import EmptyResult
222266

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

0 commit comments

Comments
 (0)