|
| 1 | +"""Opt-in extension interface for MCP clients. |
| 2 | +
|
| 3 | +To make an extension: subclass `ClientExtension`, set `identifier`, and |
| 4 | +override whichever of `settings()` / `claims()` / `notifications()` apply. To |
| 5 | +use one: pass instances to `Client(extensions=[...])` — the client folds the |
| 6 | +declarations into its own machinery; the extension never receives the client. |
| 7 | +To advertise an extension identifier with no client-side behaviour, use |
| 8 | +`advertise()`. |
| 9 | +""" |
| 10 | + |
| 11 | +from __future__ import annotations |
| 12 | + |
| 13 | +from collections.abc import Awaitable, Callable, Sequence |
| 14 | +from dataclasses import dataclass |
| 15 | +from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, get_args |
| 16 | + |
| 17 | +from mcp_types import CallToolResult, InputRequiredResult, Result |
| 18 | +from mcp_types.version import MODERN_PROTOCOL_VERSIONS |
| 19 | +from pydantic import BaseModel |
| 20 | + |
| 21 | +from mcp.shared.extension import validate_extension_identifier |
| 22 | + |
| 23 | +if TYPE_CHECKING: |
| 24 | + from mcp.client.session import ClientSession |
| 25 | + |
| 26 | +__all__ = [ |
| 27 | + "ClaimContext", |
| 28 | + "ClientExtension", |
| 29 | + "NotificationBinding", |
| 30 | + "ResultClaim", |
| 31 | + "advertise", |
| 32 | +] |
| 33 | + |
| 34 | +ClaimedT = TypeVar("ClaimedT", bound=Result) |
| 35 | +NotifyParamsT = TypeVar("NotifyParamsT", bound=BaseModel) |
| 36 | + |
| 37 | + |
| 38 | +@dataclass(frozen=True, kw_only=True) |
| 39 | +class ClaimContext: |
| 40 | + """Host-injected context for one `ResultClaim.resolve` call. |
| 41 | +
|
| 42 | + `session` is the sanctioned public low-level handle — the same one users |
| 43 | + already reach via `client.session`; the resolver gets no `Client` and no |
| 44 | + new authority. |
| 45 | + """ |
| 46 | + |
| 47 | + session: ClientSession |
| 48 | + tool_name: str |
| 49 | + read_timeout_seconds: float | None |
| 50 | + |
| 51 | + |
| 52 | +@dataclass(frozen=True, kw_only=True) |
| 53 | +class ResultClaim(Generic[ClaimedT]): |
| 54 | + """One extra result shape on one spec verb, keyed by the wire `resultType`. |
| 55 | +
|
| 56 | + A claim is active only while the declaring extension is constructed in AND |
| 57 | + the negotiated version admits it; otherwise parsing stays byte-identical to |
| 58 | + a claim-less client, so an undeclared shape still fails validation — the |
| 59 | + supported `resultType` set is always core plus declared claims. |
| 60 | +
|
| 61 | + `resolve` finishes a claimed result on the transparent path: it may send |
| 62 | + follow-ups through `ctx.session` and must return the verb's ordinary |
| 63 | + result. It is required — a claim nothing can finish would be useless. A |
| 64 | + package that wants explicit-only handling ships a resolver that raises a |
| 65 | + typed error naming `session.call_tool(allow_claimed=True)`, which is also |
| 66 | + how callers reach the undriven shape per-call. |
| 67 | +
|
| 68 | + `model` must declare `result_type` as a Literal of exactly the claimed tag, |
| 69 | + and must not subclass a core result type — a core subclass would satisfy |
| 70 | + the session's isinstance branches and bypass claim routing. `protocol_versions`, |
| 71 | + when set, restricts the claim to a subset of the modern protocol revisions; |
| 72 | + `None` (the default) means every modern version. The modern floor is |
| 73 | + structural, not a restriction: claimed shapes cannot be delivered on a |
| 74 | + legacy wire. All of this is enforced at construction. |
| 75 | + """ |
| 76 | + |
| 77 | + result_type: str |
| 78 | + model: type[ClaimedT] |
| 79 | + resolve: Callable[[ClaimedT, ClaimContext], Awaitable[CallToolResult]] |
| 80 | + method: Literal["tools/call"] = "tools/call" |
| 81 | + protocol_versions: frozenset[str] | None = None |
| 82 | + |
| 83 | + def __post_init__(self) -> None: |
| 84 | + if self.result_type in ("complete", "input_required"): |
| 85 | + raise ValueError(f"resultType {self.result_type!r} is core protocol vocabulary") |
| 86 | + if issubclass(self.model, CallToolResult | InputRequiredResult): |
| 87 | + raise ValueError("claim models must not subclass core result types") |
| 88 | + field = self.model.model_fields.get("result_type") |
| 89 | + if field is None or get_args(field.annotation) != (self.result_type,): |
| 90 | + raise ValueError(f"{self.model.__name__}.result_type must be Literal[{self.result_type!r}]") |
| 91 | + if self.protocol_versions is not None and not self.protocol_versions: |
| 92 | + raise ValueError("empty protocol_versions could never activate; use None for all") |
| 93 | + if self.protocol_versions is not None and not self.protocol_versions.issubset(MODERN_PROTOCOL_VERSIONS): |
| 94 | + unrecognized = sorted(self.protocol_versions.difference(MODERN_PROTOCOL_VERSIONS)) |
| 95 | + raise ValueError( |
| 96 | + f"protocol_versions {unrecognized} are not modern protocol revisions; claimed shapes " |
| 97 | + "cannot be delivered on a legacy wire (None means every modern version)" |
| 98 | + ) |
| 99 | + |
| 100 | + |
| 101 | +@dataclass(frozen=True, kw_only=True) |
| 102 | +class NotificationBinding(Generic[NotifyParamsT]): |
| 103 | + """Deliver server notifications for `method` to `handler` (today: silently dropped). |
| 104 | +
|
| 105 | + Observation-only: the handler receives validated params, returns None, and |
| 106 | + cannot short-circuit anything. Delivery is per-binding serialized through a |
| 107 | + bounded FIFO — one consumer task per binding, so a handler sees events in |
| 108 | + arrival order and may do session I/O without deadlocking the in-process |
| 109 | + dispatch path; on overflow the oldest event is dropped with a warning |
| 110 | + (observation semantics make the drop acceptable). |
| 111 | +
|
| 112 | + There is deliberately no spec-table check at construction: bindings are |
| 113 | + consulted only for methods the negotiated version's core tables do NOT |
| 114 | + know, so they are additive by construction. If a future core version |
| 115 | + adopts the method, the binding goes quiet — detected and warned once at |
| 116 | + activation, not per delivery — instead of import-erroring every package. |
| 117 | +
|
| 118 | + `method` is the bare wire name (e.g. `notifications/tasks`); `params_type` |
| 119 | + validates the notification params before `handler` runs. |
| 120 | + """ |
| 121 | + |
| 122 | + method: str |
| 123 | + params_type: type[NotifyParamsT] |
| 124 | + handler: Callable[[NotifyParamsT], Awaitable[None]] |
| 125 | + |
| 126 | + |
| 127 | +class ClientExtension: |
| 128 | + """Base class for an opt-in client extension. Override only what you need. |
| 129 | +
|
| 130 | + Mirror of `mcp.server.extension.Extension` in feel: a closed declarative |
| 131 | + surface, fixed at construction, that never receives the client. The |
| 132 | + contribution kinds are the ones a 2026 client actually has — there is |
| 133 | + deliberately no served-request kind (servers do not initiate requests) and |
| 134 | + no open interceptor (the only sanctioned augmentation is extension |
| 135 | + `resultType` values, and a claim already names its owner, so composition |
| 136 | + and ordering questions dissolve by construction). |
| 137 | + """ |
| 138 | + |
| 139 | + #: Reverse-DNS extension identifier, advertised under `ClientCapabilities.extensions`. |
| 140 | + identifier: str |
| 141 | + |
| 142 | + def __init_subclass__(cls, **kwargs: Any) -> None: |
| 143 | + super().__init_subclass__(**kwargs) |
| 144 | + # Validate a class-level `identifier` at definition time. A subclass may |
| 145 | + # instead assign `identifier` in `__init__` (per-instance ids); that case |
| 146 | + # is validated when the extension is consumed, since no class attribute |
| 147 | + # exists to inspect here. |
| 148 | + if (identifier := cls.__dict__.get("identifier")) is not None: |
| 149 | + validate_extension_identifier(identifier, owner=cls.__name__) |
| 150 | + |
| 151 | + def settings(self) -> dict[str, Any]: |
| 152 | + """Per-extension settings advertised at `ClientCapabilities.extensions[identifier]`. |
| 153 | +
|
| 154 | + Read ONCE at `Client` construction — dynamic per-request settings are |
| 155 | + out of scope. An empty dict (the default) advertises the extension with |
| 156 | + no settings. |
| 157 | +
|
| 158 | + A claim-bearing extension's identifier is advertised only at protocol |
| 159 | + versions where at least one of its claims is active: the ad and the |
| 160 | + claims dissolve together, so the client never advertises an extension |
| 161 | + on a request whose claimed result shapes it would reject. Claim-less |
| 162 | + extensions advertise at every version. |
| 163 | + """ |
| 164 | + return {} |
| 165 | + |
| 166 | + def claims(self) -> Sequence[ResultClaim[Any]]: |
| 167 | + """Extra result shapes this extension claims, with their resolvers.""" |
| 168 | + return () |
| 169 | + |
| 170 | + def notifications(self) -> Sequence[NotificationBinding[Any]]: |
| 171 | + """Server notifications this extension observes.""" |
| 172 | + return () |
| 173 | + |
| 174 | + |
| 175 | +class _AdvertiseOnly(ClientExtension): |
| 176 | + """Ad-only extension returned by `advertise()`: an identifier plus captured settings.""" |
| 177 | + |
| 178 | + def __init__(self, identifier: str, settings: dict[str, Any]) -> None: |
| 179 | + self.identifier = identifier |
| 180 | + self._settings = settings |
| 181 | + |
| 182 | + def settings(self) -> dict[str, Any]: |
| 183 | + return self._settings |
| 184 | + |
| 185 | + |
| 186 | +def advertise(identifier: str, settings: dict[str, Any] | None = None) -> ClientExtension: |
| 187 | + """Advertise an extension identifier (with optional settings) and nothing else. |
| 188 | +
|
| 189 | + Returns an extension that contributes only the capability ad: no claims, no |
| 190 | + notification bindings. The identifier is validated eagerly, at this call. |
| 191 | +
|
| 192 | + WARNING: advertising an extension you do not implement asserts wire support |
| 193 | + you don't have — for behavioral extensions (e.g. tasks) construct the real |
| 194 | + extension object instead. |
| 195 | + """ |
| 196 | + validate_extension_identifier(identifier, owner="advertise") |
| 197 | + return _AdvertiseOnly(identifier, {} if settings is None else settings) |
0 commit comments