Skip to content

Commit 5497da0

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 14addfb commit 5497da0

File tree

6 files changed

+151
-30
lines changed

6 files changed

+151
-30
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/client/session.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
PromptReference,
2727
ReadResourceResult,
2828
ResourceReference,
29+
RootsCapability,
2930
ServerNotification,
3031
ServerRequest,
3132
)
@@ -69,12 +70,12 @@ async def initialize(self) -> InitializeResult:
6970
capabilities=ClientCapabilities(
7071
sampling=None,
7172
experimental=None,
72-
roots={
73+
roots=RootsCapability(
7374
# TODO: Should this be based on whether we
7475
# _will_ send notifications, or only whether
7576
# they're supported?
76-
"listChanged": True
77-
},
77+
listChanged=True
78+
),
7879
),
7980
clientInfo=Implementation(name="mcp_python", version="0.1.0"),
8081
),

mcp_python/client/stdio.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,19 @@
1212

1313
# Environment variables to inherit by default
1414
DEFAULT_INHERITED_ENV_VARS = (
15-
["APPDATA", "HOMEDRIVE", "HOMEPATH", "LOCALAPPDATA", "PATH",
16-
"PROCESSOR_ARCHITECTURE", "SYSTEMDRIVE", "SYSTEMROOT", "TEMP",
17-
"USERNAME", "USERPROFILE"]
15+
[
16+
"APPDATA",
17+
"HOMEDRIVE",
18+
"HOMEPATH",
19+
"LOCALAPPDATA",
20+
"PATH",
21+
"PROCESSOR_ARCHITECTURE",
22+
"SYSTEMDRIVE",
23+
"SYSTEMROOT",
24+
"TEMP",
25+
"USERNAME",
26+
"USERPROFILE",
27+
]
1828
if sys.platform == "win32"
1929
else ["HOME", "LOGNAME", "PATH", "SHELL", "TERM", "USER"]
2030
)
@@ -74,7 +84,7 @@ async def stdio_client(server: StdioServerParameters):
7484
process = await anyio.open_process(
7585
[server.command, *server.args],
7686
env=server.env if server.env is not None else get_default_environment(),
77-
stderr=sys.stderr
87+
stderr=sys.stderr,
7888
)
7989

8090
async def stdout_reader():

mcp_python/server/__init__.py

Lines changed: 62 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,22 +28,26 @@
2828
ListResourcesResult,
2929
ListToolsRequest,
3030
ListToolsResult,
31+
LoggingCapability,
3132
LoggingLevel,
3233
PingRequest,
3334
ProgressNotification,
3435
Prompt,
3536
PromptMessage,
3637
PromptReference,
38+
PromptsCapability,
3739
ReadResourceRequest,
3840
ReadResourceResult,
3941
Resource,
4042
ResourceReference,
43+
ResourcesCapability,
4144
ServerCapabilities,
4245
ServerResult,
4346
SetLevelRequest,
4447
SubscribeRequest,
4548
TextContent,
4649
Tool,
50+
ToolsCapability,
4751
UnsubscribeRequest,
4852
)
4953

@@ -54,16 +58,33 @@
5458
)
5559

5660

61+
class NotificationOptions:
62+
def __init__(
63+
self,
64+
prompts_changed: bool = False,
65+
resources_changed: bool = False,
66+
tools_changed: bool = False,
67+
):
68+
self.prompts_changed = prompts_changed
69+
self.resources_changed = resources_changed
70+
self.tools_changed = tools_changed
71+
72+
5773
class Server:
5874
def __init__(self, name: str):
5975
self.name = name
6076
self.request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]] = {
6177
PingRequest: _ping_handler,
6278
}
6379
self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {}
80+
self.notification_options = NotificationOptions()
6481
logger.debug(f"Initializing server '{name}'")
6582

66-
def create_initialization_options(self) -> types.InitializationOptions:
83+
def create_initialization_options(
84+
self,
85+
notification_options: NotificationOptions | None = None,
86+
experimental_capabilities: dict[str, dict[str, Any]] | None = None,
87+
) -> types.InitializationOptions:
6788
"""Create initialization options from this server instance."""
6889

6990
def pkg_version(package: str) -> str:
@@ -81,20 +102,51 @@ def pkg_version(package: str) -> str:
81102
return types.InitializationOptions(
82103
server_name=self.name,
83104
server_version=pkg_version("mcp_python"),
84-
capabilities=self.get_capabilities(),
105+
capabilities=self.get_capabilities(
106+
notification_options or NotificationOptions(),
107+
experimental_capabilities or {},
108+
),
85109
)
86110

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

93144
return ServerCapabilities(
94-
prompts=get_capability(ListPromptsRequest),
95-
resources=get_capability(ListResourcesRequest),
96-
tools=get_capability(ListToolsRequest),
97-
logging=get_capability(SetLevelRequest),
145+
prompts=prompts_capability,
146+
resources=resources_capability,
147+
tools=tools_capability,
148+
logging=logging_capability,
149+
experimental=experimental_capabilities,
98150
)
99151

100152
@property

mcp_python/types.py

Lines changed: 52 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -184,30 +184,76 @@ 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+
195+
class SamplingCapability(BaseModel):
196+
"""Capability for logging operations."""
197+
198+
model_config = ConfigDict(extra="allow")
199+
200+
187201
class ClientCapabilities(BaseModel):
188202
"""Capabilities a client may support."""
189203

190204
experimental: dict[str, dict[str, Any]] | None = None
191205
"""Experimental, non-standard capabilities that the client supports."""
192-
sampling: dict[str, Any] | None = None
206+
sampling: SamplingCapability | None = None
193207
"""Present if the client supports sampling from an LLM."""
194-
roots: dict[str, Any] | None = None
208+
roots: RootsCapability | None = None
195209
"""Present if the client supports listing roots."""
196210
model_config = ConfigDict(extra="allow")
197211

198212

213+
class PromptsCapability(BaseModel):
214+
"""Capability for prompts operations."""
215+
216+
listChanged: bool | None = None
217+
"""Whether this server supports notifications for changes to the prompt list."""
218+
model_config = ConfigDict(extra="allow")
219+
220+
221+
class ResourcesCapability(BaseModel):
222+
"""Capability for resources operations."""
223+
224+
subscribe: bool | None = None
225+
"""Whether this server supports subscribing to resource updates."""
226+
listChanged: bool | None = None
227+
"""Whether this server supports notifications for changes to the resource list."""
228+
model_config = ConfigDict(extra="allow")
229+
230+
231+
class ToolsCapability(BaseModel):
232+
"""Capability for tools operations."""
233+
234+
listChanged: bool | None = None
235+
"""Whether this server supports notifications for changes to the tool list."""
236+
model_config = ConfigDict(extra="allow")
237+
238+
239+
class LoggingCapability(BaseModel):
240+
"""Capability for logging operations."""
241+
242+
model_config = ConfigDict(extra="allow")
243+
244+
199245
class ServerCapabilities(BaseModel):
200246
"""Capabilities that a server may support."""
201247

202248
experimental: dict[str, dict[str, Any]] | None = None
203249
"""Experimental, non-standard capabilities that the server supports."""
204-
logging: dict[str, Any] | None = None
250+
logging: LoggingCapability | None = None
205251
"""Present if the server supports sending log messages to the client."""
206-
prompts: dict[str, Any] | None = None
252+
prompts: PromptsCapability | None = None
207253
"""Present if the server offers any prompt templates."""
208-
resources: dict[str, Any] | None = None
254+
resources: ResourcesCapability | None = None
209255
"""Present if the server offers any resources to read."""
210-
tools: dict[str, Any] | None = None
256+
tools: ToolsCapability | None = None
211257
"""Present if the server offers any tools to call."""
212258
model_config = ConfigDict(extra="allow")
213259

tests/server/test_session.py

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,15 @@
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,
12+
PromptsCapability,
13+
ResourcesCapability,
1214
ServerCapabilities,
1315
)
1416

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