diff --git a/README.md b/README.md index 0ca039ae3..26fb9b62c 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,7 @@ - [Writing MCP Clients](#writing-mcp-clients) - [MCP Primitives](#mcp-primitives) - [Server Capabilities](#server-capabilities) + - [Custom blocks](#custom-blocks) - [Documentation](#documentation) - [Contributing](#contributing) - [License](#license) @@ -643,6 +644,38 @@ MCP servers declare capabilities during initialization: | `logging` | - | Server logging configuration | | `completion`| - | Argument completion suggestions | +### Custom blocks + +The SDK ships with a set of built-in blocks, but you can define additional +kinds just as easily. Implement a subclass of `Block` and register it with the +registry helper: + +```python +from mcp.blocks.base import Block +from mcp.blocks.registry import register_block, get_block_class + + +@register_block("sticker") +class Badge(Block): + """A lightweight block that references a badge graphic.""" + + def __init__(self, url: str, label: str | None = None): + self.url = url + self.label = label + + +# Later on ────────────────────────────────────────────────────────── +payload = {"kind": "badge", "url": "gold.png", "label": "premium"} + +cls = get_block_class(payload["kind"]) # Resolve to *Badge* +block = cls(**{k: v for k, v in payload.items() if k != "kind"}) +``` + +> **Note** +> Re-using an existing *kind* name replaces the previous class and raises a +> `RuntimeWarning`. Choose globally unique identifiers unless you intend to +> override a built-in block. + ## Documentation - [Model Context Protocol documentation](https://modelcontextprotocol.io) diff --git a/src/mcp/blocks/__init__.py b/src/mcp/blocks/__init__.py new file mode 100644 index 000000000..a3b660802 --- /dev/null +++ b/src/mcp/blocks/__init__.py @@ -0,0 +1,23 @@ +"""Block classes for the MCP server. + +This module contains block classes for the MCP server. +""" + +# Import all built-in block classes here + +# Import registry—public helpers ensure the module is loaded for side-effects. +from .registry import ( + UnknownBlockKindError, + get_block_class, + is_block_kind_registered, + list_block_kinds, + register_block, +) + +__all__ = [ + "register_block", + "get_block_class", + "list_block_kinds", + "is_block_kind_registered", + "UnknownBlockKindError", +] diff --git a/src/mcp/blocks/base.py b/src/mcp/blocks/base.py new file mode 100644 index 000000000..eca37e792 --- /dev/null +++ b/src/mcp/blocks/base.py @@ -0,0 +1,9 @@ +class Block: + """Base class for all Blocks. + + This is currently an empty marker class but may be expanded in the future to + include common functionality or interfaces required by all block + implementations. + """ + + pass diff --git a/src/mcp/blocks/registry.py b/src/mcp/blocks/registry.py new file mode 100644 index 000000000..bfdd8d8a4 --- /dev/null +++ b/src/mcp/blocks/registry.py @@ -0,0 +1,140 @@ +from __future__ import annotations as _annotations + +import warnings +from typing import TypeVar + +from mcp.blocks.base import Block + +"""Block registry utilities for the MCP Python SDK. + +This module provides a minimal—but extensible—mechanism for third-party +packages to register custom *Block* implementations with the core SDK. The +registry exposes two public helper functions—:func:`register_block` and +:func:`get_block_class`—which allow downstream libraries to add new block +"kinds" and retrieve them at runtime, respectively. + +The design is intentionally lightweight: + +* The registry is just an in-memory ``dict`` mapping *kind* strings to Python + classes. +* Registration is performed via a decorator for ergonomic usage on the block + class definition:: + + from mcp.blocks.registry import register_block + + @register_block("my-cool-block") + class MyBlock(Block): + ... + +* The registry is populated for built-in block types when + ``mcp.blocks.__init__`` is imported. Third-party packages can import + ``mcp.blocks.registry`` at import-time (or call :func:`register_block` during + plugin initialization) to extend the mapping. + +* Thread-safety: Mutating a ``dict`` is atomic in CPython, so casual concurrent + registration *should* be safe. However, if your application registers block + types from multiple threads, you may still wish to provide an external lock + to coordinate access during import time. + +""" + +__all__ = [ + "register_block", + "get_block_class", + "list_block_kinds", + "is_block_kind_registered", + "UnknownBlockKindError", +] + +_BlockT = TypeVar("_BlockT", bound=Block) + +# NOTE: keep registry private — public API is via helper functions. +_BLOCK_REGISTRY: dict[str, type[Block]] = {} + + +def register_block(kind: str): # noqa: D401 + """Return a decorator that registers *cls* under *kind* and yields it back. + + The primary call-site is as a class decorator. The function also supports + direct invocation for dynamic registration:: + + MyBlock = create_block_cls() + register_block("my-block")(MyBlock) + + If *kind* is already present, the previous entry will be silently + overwritten—mirroring Python's module import semantics. Duplicate kinds + are thus the caller's responsibility. + """ + + def _inner(cls: type[_BlockT]) -> type[_BlockT]: + if kind in _BLOCK_REGISTRY: + warnings.warn( + f"Block kind {kind!r} is already registered and will be " + "overwritten.", + RuntimeWarning, + stacklevel=2, + ) + _BLOCK_REGISTRY[kind] = cls + # Intentionally do *not* mutate the class object beyond registration to + # keep the hook minimal and avoid leaking extra attributes into user + # classes. Downstream packages can attach helpers if they need them. + return cls + + return _inner + + +def get_block_class(kind: str) -> type[Block]: + """Return the class registered for *kind*. + + Raises + ------ + KeyError + If *kind* has not been registered (either built-in or via + :func:`register_block`). + """ + + try: + return _BLOCK_REGISTRY[kind] + except KeyError as exc: + # Re-raise as a more specific exception while preserving backward + # compatibility with ``except KeyError`` clauses. + raise UnknownBlockKindError(kind) from exc + + +# === Public utility helpers ================================================== + + +def list_block_kinds() -> list[str]: # noqa: D401 + """Return a *copy* of all currently registered block *kind* strings. + + The returned list is a snapshot—mutating it will **not** affect the global + registry. The order of kinds is implementation-defined and should not be + relied upon. + """ + + return list(_BLOCK_REGISTRY.keys()) + + +def is_block_kind_registered(kind: str) -> bool: # noqa: D401 + """Return ``True`` if *kind* is currently registered. + + This is equivalent to ``kind in list_block_kinds()`` but avoids the + intermediate list allocation. + """ + + return kind in _BLOCK_REGISTRY + + +# === Exceptions ============================================================== + + +class UnknownBlockKindError(KeyError): + """Raised when :func:`get_block_class` is called with an unregistered kind. + + Subclasses :class:`KeyError` for backward-compatibility so that existing + `except KeyError:` handlers continue to work while allowing callers to catch + this more specific error. + """ + + def __init__(self, kind: str): + super().__init__(kind) diff --git a/tests/test_registry.py b/tests/test_registry.py new file mode 100644 index 000000000..dbd525d2c --- /dev/null +++ b/tests/test_registry.py @@ -0,0 +1,48 @@ +from __future__ import annotations as _annotations + +import pytest + +from mcp.blocks.base import Block +from mcp.blocks.registry import get_block_class, register_block + + +@register_block("foo") +class DummyBlock(Block): ... + + +def test_register_and_fetch(): + assert get_block_class("foo") is DummyBlock + + +def test_register_overwrite(): + with pytest.warns(RuntimeWarning): + + @register_block("foo") + class DummyBlock2(Block): ... + + assert get_block_class("foo") is DummyBlock2 + + +@pytest.mark.parametrize("kind", ["unknown", "bar"]) +def test_get_block_class_missing(kind: str): + with pytest.raises(KeyError): + _ = get_block_class(kind) + + +def test_overwrite_warning(recwarn): + @register_block("bar") + class DummyA(Block): ... + + # second registration should raise RuntimeWarning + @register_block("bar") + class DummyB(Block): ... + + w = recwarn.pop(RuntimeWarning) + assert "overwritten" in str(w.message) + + +def test_custom_exception(): + from mcp.blocks.registry import UnknownBlockKindError + + with pytest.raises(UnknownBlockKindError): + _ = get_block_class("does-not-exist")