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.
This commit is contained in:
David Soria Parra
2024-11-06 22:50:37 +00:00
parent 14addfb872
commit 5497da0afd
6 changed files with 150 additions and 29 deletions

View File

@@ -33,10 +33,13 @@ from .types import (
Notification,
PingRequest,
ProgressNotification,
PromptsCapability,
ReadResourceRequest,
ReadResourceResult,
Resource,
ResourcesCapability,
ResourceUpdatedNotification,
RootsCapability,
SamplingMessage,
ServerCapabilities,
ServerNotification,
@@ -46,6 +49,7 @@ from .types import (
StopReason,
SubscribeRequest,
Tool,
ToolsCapability,
UnsubscribeRequest,
)
from .types import (
@@ -82,10 +86,13 @@ __all__ = [
"Notification",
"PingRequest",
"ProgressNotification",
"PromptsCapability",
"ReadResourceRequest",
"ReadResourceResult",
"ResourcesCapability",
"ResourceUpdatedNotification",
"Resource",
"RootsCapability",
"SamplingMessage",
"SamplingRole",
"ServerCapabilities",
@@ -98,6 +105,7 @@ __all__ = [
"StopReason",
"SubscribeRequest",
"Tool",
"ToolsCapability",
"UnsubscribeRequest",
"stdio_client",
"stdio_server",

View File

@@ -26,6 +26,7 @@ from mcp_python.types import (
PromptReference,
ReadResourceResult,
ResourceReference,
RootsCapability,
ServerNotification,
ServerRequest,
)
@@ -69,12 +70,12 @@ class ClientSession(
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"),
),

View File

@@ -12,9 +12,19 @@ from mcp_python.types import JSONRPCMessage
# 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():

View File

@@ -28,22 +28,26 @@ from mcp_python.types import (
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,6 +58,18 @@ request_ctx: contextvars.ContextVar[RequestContext] = contextvars.ContextVar(
)
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
@@ -61,9 +77,14 @@ class Server:
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 @@ class Server:
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."""
prompts_capability = None
resources_capability = None
tools_capability = None
logging_capability = None
def get_capability(req_type: type) -> dict[str, Any] | None:
return {} if req_type in self.request_handlers else 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

View File

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