Skip to content

feat: add structured capability types #30

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 7, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions mcp_python/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,13 @@
Notification,
PingRequest,
ProgressNotification,
PromptsCapability,
ReadResourceRequest,
ReadResourceResult,
Resource,
ResourcesCapability,
ResourceUpdatedNotification,
RootsCapability,
SamplingMessage,
ServerCapabilities,
ServerNotification,
Expand All @@ -46,6 +49,7 @@
StopReason,
SubscribeRequest,
Tool,
ToolsCapability,
UnsubscribeRequest,
)
from .types import (
Expand Down Expand Up @@ -82,10 +86,13 @@
"Notification",
"PingRequest",
"ProgressNotification",
"PromptsCapability",
"ReadResourceRequest",
"ReadResourceResult",
"ResourcesCapability",
"ResourceUpdatedNotification",
"Resource",
"RootsCapability",
"SamplingMessage",
"SamplingRole",
"ServerCapabilities",
Expand All @@ -98,6 +105,7 @@
"StopReason",
"SubscribeRequest",
"Tool",
"ToolsCapability",
"UnsubscribeRequest",
"stdio_client",
"stdio_server",
Expand Down
7 changes: 4 additions & 3 deletions mcp_python/client/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
PromptReference,
ReadResourceResult,
ResourceReference,
RootsCapability,
ServerNotification,
ServerRequest,
)
Expand Down Expand Up @@ -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"),
),
Expand Down
18 changes: 14 additions & 4 deletions mcp_python/client/stdio.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
)
Expand Down Expand Up @@ -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():
Expand Down
72 changes: 62 additions & 10 deletions mcp_python/server/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand All @@ -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:
Expand All @@ -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
Expand Down
58 changes: 52 additions & 6 deletions mcp_python/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down
18 changes: 11 additions & 7 deletions tests/server/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down Expand Up @@ -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

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