Add a client extension API (#3034)
This commit is contained in:
Max
GitHub
parent
7322ca56f4
commit
4df609119f
37 changed files with 3410 additions and 180 deletions
|
|
@ -2,9 +2,10 @@
|
|||
|
||||
An **extension** is an opt-in bundle of MCP behaviour behind one identifier.
|
||||
|
||||
It can contribute tools, resources, and new request methods, and it can wrap `tools/call`.
|
||||
The server advertises it under `capabilities.extensions`, the client opts in the same way,
|
||||
and nothing changes for anyone who didn't ask for it. That is the contract ([SEP-2133](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2133)), and
|
||||
On a server it can contribute tools, resources, and new request methods, and it can wrap
|
||||
`tools/call`. On a client it can claim extra `tools/call` result shapes and observe vendor
|
||||
notifications. Each side advertises under its own `capabilities.extensions`, and nothing
|
||||
changes for anyone who didn't ask for it. That is the contract ([SEP-2133](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2133)), and
|
||||
it has one golden rule: **extensions are off by default**.
|
||||
|
||||
## Using an extension
|
||||
|
|
@ -79,7 +80,7 @@ And `main()` is the proof, an in-memory client straight against `mcp`:
|
|||
An extension can register **new request methods**: its own verbs, served next to the
|
||||
spec's:
|
||||
|
||||
```python title="server.py" hl_lines="15-21 30 39-47"
|
||||
```python title="server.py" hl_lines="16-22 31 40-48"
|
||||
--8<-- "docs_src/extensions/tutorial004.py"
|
||||
```
|
||||
|
||||
|
|
@ -108,19 +109,19 @@ runtime:
|
|||
|
||||
The same file's `main()` is the whole client story, both halves of it:
|
||||
|
||||
```python title="server.py" hl_lines="53-57"
|
||||
```python title="server.py" hl_lines="54-58"
|
||||
--8<-- "docs_src/extensions/tutorial004.py"
|
||||
```
|
||||
|
||||
* `Client(..., extensions={EXTENSION_ID: {}})` declares the extension. That map
|
||||
becomes `ClientCapabilities.extensions`: on a 2026-07-28 connection it travels in
|
||||
the per-request `_meta` envelope, so the server sees it on **every** request; on
|
||||
a legacy connection it rides the `initialize` handshake. Server code doesn't care
|
||||
which: `require_client_extension(ctx, ...)` and
|
||||
* `Client(..., extensions=[advertise(EXTENSION_ID)])` declares the extension. The
|
||||
declarations become `ClientCapabilities.extensions`: on a 2026-07-28 connection
|
||||
the map travels in the per-request `_meta` envelope, so the server sees it on
|
||||
**every** request; on a legacy connection it rides the `initialize` handshake.
|
||||
Server code doesn't care which: `require_client_extension(ctx, ...)` and
|
||||
`ctx.session.check_client_capability(...)` read the right source on both paths.
|
||||
* Vendor methods drop one layer to `client.session.send_request(...)`; `Client`
|
||||
only grows first-class methods for spec verbs. The `cast` is there because
|
||||
`send_request` is typed against the spec's closed request union.
|
||||
only grows first-class methods for spec verbs. `send_request` accepts any
|
||||
`Request` subclass, so the vendor request passes as-is.
|
||||
|
||||
### Intercepting `tools/call`
|
||||
|
||||
|
|
@ -144,15 +145,104 @@ or veto a tool call:
|
|||
The hook wraps `tools/call` and nothing else. For every-message concerns, use
|
||||
[Middleware](middleware.md). That is what it is for.
|
||||
|
||||
## Using a client extension
|
||||
|
||||
A **client extension** is the same contract from the consuming side: a bundle of
|
||||
client-side behaviour behind one identifier. Pass instances to
|
||||
`Client(extensions=[...])` and call tools normally:
|
||||
|
||||
```python title="client.py" hl_lines="67-69"
|
||||
--8<-- "docs_src/extensions/tutorial006.py"
|
||||
```
|
||||
|
||||
`call_tool("buy", ...)` returns a plain `CallToolResult`, like every other call. What
|
||||
the extension changed: the server may now answer `buy` with a `receipt` **result
|
||||
shape** instead of a final result, and `Receipts` finishes it (here by redeeming the
|
||||
receipt with a follow-up call) before `call_tool` returns. Nothing about the call
|
||||
site moves.
|
||||
|
||||
Drop the extension and none of this exists: the server's gate refuses a client
|
||||
that did not declare it (error -32021), and a claimed shape from a server that
|
||||
skips the gate fails validation, exactly as the spec requires for an
|
||||
unrecognized `resultType`. Off by default, on both ends of the wire.
|
||||
|
||||
To advertise an identifier with **no** client-side behaviour (the server gates on
|
||||
the capability, the client does nothing, as in the search client above), use
|
||||
`advertise()`:
|
||||
|
||||
```python
|
||||
from mcp.client import advertise
|
||||
|
||||
client = Client(mcp, extensions=[advertise("com.example/search")])
|
||||
```
|
||||
|
||||
## Writing a client extension
|
||||
|
||||
Subclass `ClientExtension` and override only what you need. Three contribution
|
||||
kinds, each with a default: `settings()`, `claims()`, and `notifications()`.
|
||||
|
||||
```python title="client.py" hl_lines="18-19 44-45 47-48"
|
||||
--8<-- "docs_src/extensions/tutorial006.py"
|
||||
```
|
||||
|
||||
* The identifier follows the same grammar as the server's, validated when the class
|
||||
is defined.
|
||||
* `claims()` returns `ResultClaim`s: a wire tag, the model that parses it, and the
|
||||
resolver that finishes it. The model must pin the tag with
|
||||
`result_type: Literal["receipt"]` and must not subclass the verb's core result
|
||||
types; both are enforced when the claim is constructed. Vendor fields like
|
||||
`receipt_token` ride the wire as-is: a substituted shape reaches the client
|
||||
verbatim.
|
||||
* The resolver receives the parsed model and a `ClaimContext`; `ctx.session` is the
|
||||
same public handle as `client.session`, so follow-ups are ordinary session calls.
|
||||
It returns the verb's normal `CallToolResult`.
|
||||
* `settings()` is the value advertised at `ClientCapabilities.extensions[identifier]`,
|
||||
read once at `Client` construction.
|
||||
|
||||
`notifications()` declares vendor server notifications to observe:
|
||||
|
||||
```python
|
||||
def notifications(self) -> Sequence[NotificationBinding[Any]]:
|
||||
return [NotificationBinding(method="notifications/receipts", params_type=ReceiptEvent, handler=self.on_receipt)]
|
||||
```
|
||||
|
||||
The handler receives validated params one at a time, in dispatch order. It observes; it cannot veto
|
||||
or reply.
|
||||
|
||||
Two quiet rules. Claims are active on 2026-07-28 connections only, and the capability
|
||||
ad follows them: on a legacy connection the claims dissolve and the identifier drops
|
||||
out of the ad with them, so the client never advertises an extension whose shapes it
|
||||
would reject. And when you want the claimed shape yourself instead of the resolver,
|
||||
call `client.session.call_tool(..., allow_claimed=True)`; without that flag, a
|
||||
claimed shape reaching a session-tier caller raises `UnexpectedClaimedResult`.
|
||||
|
||||
### Extension verbs
|
||||
|
||||
An extension's own request methods need no client-side registration. A vendor request
|
||||
type subclasses `mcp_types.Request` and goes through `client.session.send_request`,
|
||||
as in [Serving your own methods](#serving-your-own-methods). One addition: when a
|
||||
params key must ride the `Mcp-Name` header (extension specs such as tasks require
|
||||
this for their verbs), the request type declares `name_param`:
|
||||
|
||||
```python title="client.py" hl_lines="23-26 47-48"
|
||||
--8<-- "docs_src/extensions/tutorial007.py"
|
||||
```
|
||||
|
||||
The session mirrors `params["jobId"]` into `Mcp-Name` on every send path, and a
|
||||
missing value fails loudly rather than silently omitting a required header.
|
||||
|
||||
## What an extension cannot do
|
||||
|
||||
The contribution surface is **closed** on purpose: settings, tools, resources,
|
||||
methods, one `tools/call` interceptor. An extension cannot:
|
||||
The contribution surface is **closed** on purpose. On the server: settings, tools,
|
||||
resources, methods, one `tools/call` interceptor. On the client: settings, result
|
||||
claims, notification bindings. An extension cannot:
|
||||
|
||||
* **Reach into the server.** It declares data; it holds no server reference.
|
||||
* **Replace core behaviour.** Spec methods are rejected at construction, and
|
||||
`initialize` is reserved by the runner outright.
|
||||
* **Register late.** After `MCPServer(...)` returns, the extension set is what it is.
|
||||
* **Reach into the host.** It declares data; it holds no server or client reference.
|
||||
* **Replace core behaviour.** Spec methods and core result tags are rejected at
|
||||
construction (`initialize` is reserved by the runner outright); a notification
|
||||
binding shadowed by core vocabulary goes quiet with a warning instead.
|
||||
* **Register late.** After `MCPServer(...)` or `Client(...)` returns, the extension
|
||||
set is what it is.
|
||||
|
||||
If you are fighting these walls, you are not writing an extension. You are writing
|
||||
a fork. The walls are the feature: a user reading `extensions=[Apps(), Stamps()]`
|
||||
|
|
|
|||
|
|
@ -469,11 +469,42 @@ extension handler can call `mcp.server.mcpserver.require_client_extension(ctx, i
|
|||
to reject a request with the `-32021` (missing required client capability) error
|
||||
when the client did not declare the extension.
|
||||
|
||||
Clients advertise extension support with the new `Client(extensions=...)` /
|
||||
`ClientSession(extensions=...)` argument, mirrored into `ClientCapabilities.extensions`.
|
||||
The extensions capability map is negotiated over `server/discover` (modern path);
|
||||
a legacy `initialize` handshake does not carry it. Extensions are off by default
|
||||
and never alter behaviour unless registered.
|
||||
On the client, `Client(extensions=...)` takes a sequence of
|
||||
`mcp.client.ClientExtension` instances. A client extension contributes its
|
||||
capability ad (mirrored into `ClientCapabilities.extensions`), its result
|
||||
claims (extra `tools/call` result shapes that `Client.call_tool` resolves
|
||||
transparently through the claim's resolver), and its notification bindings
|
||||
(handlers for vendor server notifications). The capability map rides
|
||||
`server/discover` and every modern request's `_meta` envelope; a legacy
|
||||
`initialize` handshake carries only the claim-less identifiers, since claimed
|
||||
result shapes cannot be delivered on a legacy wire. Extensions are off by
|
||||
default and never alter behaviour unless registered. (The low-level
|
||||
`ClientSession(extensions=...)` keeps the raw identifier-to-settings dict.)
|
||||
|
||||
Changed in the v2 pre-releases: earlier alphas took
|
||||
`Client(extensions={identifier: settings})`, an advertisement-only dict.
|
||||
Extensions now contribute behaviour (claims and notification handlers), not
|
||||
just an ad, so the argument is a sequence of declaration objects. An ad-only
|
||||
entry becomes an `advertise()` call:
|
||||
|
||||
**Before (v2 alphas):**
|
||||
|
||||
```python
|
||||
client = Client(server, extensions={"com.example/ui": {"mimeTypes": [...]}})
|
||||
```
|
||||
|
||||
**After:**
|
||||
|
||||
```python
|
||||
from mcp.client import advertise
|
||||
|
||||
client = Client(server, extensions=[advertise("com.example/ui", {"mimeTypes": [...]})])
|
||||
```
|
||||
|
||||
`advertise()` is only for identifiers with no client-side behaviour. For a
|
||||
behavioural extension (e.g. tasks, once its extension ships), construct that
|
||||
extension's object instead; advertising an identifier you do not implement
|
||||
asserts wire support you don't have.
|
||||
|
||||
### `McpError` renamed to `MCPError`
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from mcp import Client
|
||||
from mcp.client import advertise
|
||||
from mcp.server.apps import APP_MIME_TYPE, EXTENSION_ID, Apps, client_supports_apps
|
||||
from mcp.server.mcpserver import MCPServer
|
||||
from mcp.server.mcpserver.context import Context
|
||||
|
|
@ -32,7 +33,7 @@ mcp = MCPServer("clock", extensions=[apps])
|
|||
|
||||
|
||||
async def main() -> None:
|
||||
async with Client(mcp, extensions={EXTENSION_ID: {"mimeTypes": [APP_MIME_TYPE]}}) as client:
|
||||
async with Client(mcp, extensions=[advertise(EXTENSION_ID, {"mimeTypes": [APP_MIME_TYPE]})]) as client:
|
||||
result = await client.call_tool("get_time", {})
|
||||
print(result.content)
|
||||
# [TextContent(text='2026-06-26T12:00:00Z')]
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
from collections.abc import Sequence
|
||||
from typing import Any, Literal, cast
|
||||
from typing import Any, Literal
|
||||
|
||||
import mcp_types as types
|
||||
from pydantic import Field
|
||||
|
||||
from mcp import Client
|
||||
from mcp.client import advertise
|
||||
from mcp.server.context import ServerRequestContext
|
||||
from mcp.server.extension import Extension, MethodBinding
|
||||
from mcp.server.mcpserver import MCPServer, require_client_extension
|
||||
|
|
@ -51,8 +52,8 @@ mcp = MCPServer("catalog", extensions=[Search()])
|
|||
|
||||
|
||||
async def main() -> None:
|
||||
async with Client(mcp, extensions={EXTENSION_ID: {}}) as client:
|
||||
async with Client(mcp, extensions=[advertise(EXTENSION_ID)]) as client:
|
||||
request = SearchRequest(params=SearchParams(query="mcp", limit=3))
|
||||
result = await client.session.send_request(cast("types.ClientRequest", request), SearchResult)
|
||||
result = await client.session.send_request(request, SearchResult)
|
||||
print(result.items)
|
||||
# ['mcp-0', 'mcp-1', 'mcp-2']
|
||||
|
|
|
|||
70
docs_src/extensions/tutorial006.py
Normal file
70
docs_src/extensions/tutorial006.py
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
from collections.abc import Sequence
|
||||
from typing import Any, Literal
|
||||
|
||||
import mcp_types as types
|
||||
|
||||
from mcp import Client
|
||||
from mcp.client import ClaimContext, ClientExtension, ResultClaim
|
||||
from mcp.server.context import CallNext, HandlerResult, ServerRequestContext
|
||||
from mcp.server.extension import Extension
|
||||
from mcp.server.mcpserver import MCPServer, require_client_extension
|
||||
|
||||
EXTENSION_ID = "com.example/receipts"
|
||||
|
||||
|
||||
class ReceiptResult(types.Result):
|
||||
"""The claimed result shape; `result_type` pins the wire tag."""
|
||||
|
||||
result_type: Literal["receipt"] = "receipt"
|
||||
receipt_token: str
|
||||
|
||||
|
||||
class ReceiptIssuer(Extension):
|
||||
"""Server half: answers `buy` with a receipt instead of a final result."""
|
||||
|
||||
identifier = EXTENSION_ID
|
||||
|
||||
async def intercept_tool_call(
|
||||
self,
|
||||
params: types.CallToolRequestParams,
|
||||
ctx: ServerRequestContext[Any, Any],
|
||||
call_next: CallNext,
|
||||
) -> HandlerResult:
|
||||
if params.name != "buy":
|
||||
return await call_next(ctx)
|
||||
require_client_extension(ctx, EXTENSION_ID)
|
||||
return {"resultType": "receipt", "receiptToken": "r-117"}
|
||||
|
||||
|
||||
class Receipts(ClientExtension):
|
||||
"""Client half: claims the `receipt` shape and supplies the code that finishes it."""
|
||||
|
||||
identifier = EXTENSION_ID
|
||||
|
||||
def claims(self) -> Sequence[ResultClaim[Any]]:
|
||||
return [ResultClaim(result_type="receipt", model=ReceiptResult, resolve=self._redeem)]
|
||||
|
||||
async def _redeem(self, claimed: ReceiptResult, ctx: ClaimContext) -> types.CallToolResult:
|
||||
return await ctx.session.call_tool("redeem", {"token": claimed.receipt_token})
|
||||
|
||||
|
||||
mcp = MCPServer("shop", extensions=[ReceiptIssuer()])
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def buy(item: str) -> types.CallToolResult:
|
||||
"""Buy an item."""
|
||||
raise NotImplementedError # ReceiptIssuer answers `buy` before the tool runs
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def redeem(token: str) -> str:
|
||||
"""Exchange a receipt token for the goods."""
|
||||
return f"goods for {token}"
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
async with Client(mcp, extensions=[Receipts()]) as client:
|
||||
result = await client.call_tool("buy", {"item": "lamp"})
|
||||
print(result.content)
|
||||
# [TextContent(text='goods for r-117')]
|
||||
50
docs_src/extensions/tutorial007.py
Normal file
50
docs_src/extensions/tutorial007.py
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
from collections.abc import Sequence
|
||||
from typing import Any, Literal
|
||||
|
||||
import mcp_types as types
|
||||
|
||||
from mcp import Client
|
||||
from mcp.client import advertise
|
||||
from mcp.server.context import ServerRequestContext
|
||||
from mcp.server.extension import Extension, MethodBinding
|
||||
from mcp.server.mcpserver import MCPServer
|
||||
|
||||
EXTENSION_ID = "com.example/jobs"
|
||||
|
||||
|
||||
class JobParams(types.RequestParams):
|
||||
job_id: str
|
||||
|
||||
|
||||
class JobStatus(types.Result):
|
||||
status: str
|
||||
|
||||
|
||||
class JobStatusRequest(types.Request[JobParams, Literal["com.example/jobs.status"]]):
|
||||
method: Literal["com.example/jobs.status"] = "com.example/jobs.status"
|
||||
params: JobParams
|
||||
name_param = "jobId" # params["jobId"] rides the Mcp-Name header
|
||||
|
||||
|
||||
async def job_status(ctx: ServerRequestContext[Any, Any], params: JobParams) -> JobStatus:
|
||||
return JobStatus(status=f"{params.job_id} is running")
|
||||
|
||||
|
||||
class Jobs(Extension):
|
||||
"""An extension whose verb names its subject, so the header can route on it."""
|
||||
|
||||
identifier = EXTENSION_ID
|
||||
|
||||
def methods(self) -> Sequence[MethodBinding]:
|
||||
return [MethodBinding("com.example/jobs.status", JobParams, job_status)]
|
||||
|
||||
|
||||
mcp = MCPServer("worker", extensions=[Jobs()])
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
async with Client(mcp, extensions=[advertise(EXTENSION_ID)]) as client:
|
||||
request = JobStatusRequest(params=JobParams(job_id="job-7"))
|
||||
result = await client.session.send_request(request, JobStatus)
|
||||
print(result.status)
|
||||
# job-7 is running
|
||||
|
|
@ -26,7 +26,7 @@ uv run python -m stories.apps.client --http
|
|||
`text/html;profile=mcp-app`.
|
||||
- `server.py` `client_supports_apps(ctx)` — SEP-2133 graceful degradation: a
|
||||
client that did not negotiate Apps gets a text-only result.
|
||||
- `client.py` `Client(target, extensions={...})` — the client advertises Apps
|
||||
- `client.py` `Client(target, extensions=[advertise(...)])` — the client advertises Apps
|
||||
support so the server returns the UI-enabled result, then reads the tool's
|
||||
`_meta.ui.resourceUri` and fetches that resource.
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
from mcp_types import TextContent, TextResourceContents
|
||||
|
||||
from mcp.client import Client
|
||||
from mcp.client import Client, advertise
|
||||
from mcp.server.apps import APP_MIME_TYPE, EXTENSION_ID
|
||||
from stories._harness import Target, run_client
|
||||
|
||||
|
|
@ -10,7 +10,9 @@ from stories._harness import Target, run_client
|
|||
async def main(target: Target, *, mode: str = "auto") -> None:
|
||||
# Advertise MCP Apps support so the server returns the UI-enabled result; a
|
||||
# client that omits this gets the text-only fallback (graceful degradation).
|
||||
async with Client(target, mode=mode, extensions={EXTENSION_ID: {"mimeTypes": [APP_MIME_TYPE]}}) as client:
|
||||
async with Client(
|
||||
target, mode=mode, extensions=[advertise(EXTENSION_ID, {"mimeTypes": [APP_MIME_TYPE]})]
|
||||
) as client:
|
||||
# The extensions capability map rides `server/discover` (modern only). On a
|
||||
# legacy connection (today's stdio) it is absent, so assert it only when present.
|
||||
if client.server_capabilities.extensions is not None:
|
||||
|
|
|
|||
|
|
@ -28,18 +28,16 @@ uv run python -m stories.custom_methods.client --http
|
|||
method string is the wire `method`; use a vendor prefix so it can never
|
||||
collide with a future spec method.
|
||||
- `client.py` `client.session.send_request(...)` — `Client` only exposes spec
|
||||
verbs, so vendor methods go through the underlying `ClientSession`. The
|
||||
`cast("types.ClientRequest", ...)` is needed because `send_request`'s
|
||||
`request` parameter is currently typed as the closed spec union; widening it
|
||||
(or adding `Client.send_request`) is tracked for beta.
|
||||
verbs, so vendor methods go through the underlying `ClientSession`.
|
||||
`send_request` accepts any `types.Request` subclass.
|
||||
|
||||
## Caveats
|
||||
|
||||
- The TypeScript SDK's equivalent example also shows a custom server→client
|
||||
**notification** (`acme/searchProgress`). The Python client currently drops
|
||||
any notification whose method is not in the spec registry
|
||||
(`ClientSession._on_notify` → `KeyError` → silent drop), and there is no
|
||||
`set_notification_handler` analogue. That half is omitted here.
|
||||
**notification** (`acme/searchProgress`). The Python client can observe
|
||||
vendor notifications via `NotificationBinding` (see
|
||||
`docs/advanced/extensions.md`). That half is omitted here because the
|
||||
lowlevel server has no surface for emitting vendor notifications yet.
|
||||
|
||||
## Spec
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
"""Send a vendor-prefixed request via the `client.session` escape hatch."""
|
||||
|
||||
from typing import Literal, cast
|
||||
from typing import Literal
|
||||
|
||||
import mcp_types as types
|
||||
|
||||
|
|
@ -26,12 +26,10 @@ async def main(target: Target, *, mode: str = "auto") -> None:
|
|||
async with Client(target, mode=mode) as client:
|
||||
# `Client` only exposes spec-defined verbs, so vendor methods have to drop one
|
||||
# layer to `client.session` today — there is no `Client`-level API for them
|
||||
# yet, and whether `.session` stays public is undecided. `send_request` is
|
||||
# typed against the closed `ClientRequest` union, hence the cast; at runtime
|
||||
# the body only calls `.model_dump()` and the unknown method skips the
|
||||
# per-spec result-validation registry.
|
||||
# yet, and whether `.session` stays public is undecided. `send_request`
|
||||
# accepts any `Request` subclass.
|
||||
request = SearchRequest(params=SearchParams(query="mcp", limit=3))
|
||||
result = await client.session.send_request(cast("types.ClientRequest", request), SearchResult)
|
||||
result = await client.session.send_request(request, SearchResult)
|
||||
assert result.items == ["mcp-0", "mcp-1", "mcp-2"], result
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ uv run python -m stories.extensions.client --http
|
|||
rejects clients that did not declare the extension with `-32021` (missing
|
||||
required client capability) and a machine-readable `requiredCapabilities`
|
||||
payload.
|
||||
- `client.py` `Client(target, extensions={EXTENSION_ID: {}})` — the client-side
|
||||
- `client.py` `Client(target, extensions=[advertise(EXTENSION_ID)])` — the client-side
|
||||
half of the negotiation; on 2026-07-28 it travels in the per-request `_meta`
|
||||
envelope.
|
||||
- `client.py` `client.session.send_request(...)` — vendor methods have no
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
"""Discover an extension's capability entry, call its tool, then send its vendor method."""
|
||||
|
||||
from typing import Literal, cast
|
||||
from typing import Literal
|
||||
|
||||
import mcp_types as types
|
||||
from mcp_types import TextContent
|
||||
|
||||
from mcp.client import Client
|
||||
from mcp.client import Client, advertise
|
||||
from stories._harness import Target, run_client
|
||||
|
||||
EXTENSION_ID = "com.example/catalog"
|
||||
|
|
@ -28,7 +28,7 @@ class SearchResult(types.Result):
|
|||
async def main(target: Target, *, mode: str = "auto") -> None:
|
||||
# Declare the extension client-side so the server's `require_client_extension`
|
||||
# gate on `com.example/search` passes.
|
||||
async with Client(target, mode=mode, extensions={EXTENSION_ID: {}}) as client:
|
||||
async with Client(target, mode=mode, extensions=[advertise(EXTENSION_ID)]) as client:
|
||||
# The extensions capability map rides `server/discover` (modern only). On a
|
||||
# legacy connection it is absent, so assert it only when present.
|
||||
if client.server_capabilities.extensions is not None:
|
||||
|
|
@ -43,10 +43,9 @@ async def main(target: Target, *, mode: str = "auto") -> None:
|
|||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "mcp-suggestion", result.content[0].text
|
||||
|
||||
# Vendor methods drop one layer to `client.session` (see custom_methods/);
|
||||
# the cast is needed because `send_request` is typed against the spec union.
|
||||
# Vendor methods drop one layer to `client.session` (see custom_methods/).
|
||||
request = SearchRequest(params=SearchParams(query="mcp", limit=3))
|
||||
found = await client.session.send_request(cast("types.ClientRequest", request), SearchResult)
|
||||
found = await client.session.send_request(request, SearchResult)
|
||||
assert found.items == ["mcp-0", "mcp-1", "mcp-2"], found
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/dr
|
|||
from mcp_types._types import (
|
||||
CLIENT_CAPABILITIES_META_KEY,
|
||||
CLIENT_INFO_META_KEY,
|
||||
CORE_RESULT_TYPES,
|
||||
DEFAULT_NEGOTIATED_VERSION,
|
||||
LOG_LEVEL_META_KEY,
|
||||
PROTOCOL_VERSION_META_KEY,
|
||||
|
|
@ -231,6 +232,7 @@ __all__ = [
|
|||
"CLIENT_CAPABILITIES_META_KEY",
|
||||
"LOG_LEVEL_META_KEY",
|
||||
# Type aliases and variables
|
||||
"CORE_RESULT_TYPES",
|
||||
"ContentBlock",
|
||||
"ElicitRequestedSchema",
|
||||
"ElicitRequestParams",
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ the negotiated version. Per-field docstrings note version availability. The
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Any, Final, Generic, Literal, TypeAlias, TypeVar
|
||||
from typing import Annotated, Any, ClassVar, Final, Generic, Literal, TypeAlias, TypeVar, get_args
|
||||
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
|
|
@ -128,6 +128,12 @@ class Request(MCPModel, Generic[RequestParamsT, MethodT]):
|
|||
method: MethodT
|
||||
params: RequestParamsT
|
||||
|
||||
name_param: ClassVar[str | None] = None
|
||||
"""Wire-params key mirrored into the `Mcp-Name` header on sends; SEP-2663 requires it for tasks/*.
|
||||
|
||||
Subclasses override by bare assignment: re-annotating as `ClassVar` trips pyright's invariance check.
|
||||
"""
|
||||
|
||||
|
||||
class PaginatedRequest(Request[PaginatedRequestParams | None, MethodT], Generic[MethodT]):
|
||||
"""Base class for paginated requests, matching the schema's PaginatedRequest interface."""
|
||||
|
|
@ -144,7 +150,9 @@ class Notification(MCPModel, Generic[NotificationParamsT, MethodT]):
|
|||
params: NotificationParamsT
|
||||
|
||||
|
||||
ResultType = Literal["complete", "input_required"] | str
|
||||
_CoreResultType = Literal["complete", "input_required"]
|
||||
|
||||
ResultType = _CoreResultType | str
|
||||
"""Tags a `Result` so the client knows how to parse it (2026-07-28).
|
||||
|
||||
"complete" means the result is final; "input_required" means it is an
|
||||
|
|
@ -152,6 +160,9 @@ ResultType = Literal["complete", "input_required"] | str
|
|||
Absent `resultType` is equivalent to "complete".
|
||||
"""
|
||||
|
||||
CORE_RESULT_TYPES: Final[frozenset[str]] = frozenset(get_args(_CoreResultType))
|
||||
"""The `resultType` tags owned by the core protocol vocabulary; extension claims may not re-key them."""
|
||||
|
||||
|
||||
class Result(MCPModel):
|
||||
"""Base class for JSON-RPC results.
|
||||
|
|
|
|||
|
|
@ -12,6 +12,14 @@ from mcp.client.caching import (
|
|||
)
|
||||
from mcp.client.client import Client
|
||||
from mcp.client.context import ClientRequestContext
|
||||
from mcp.client.extension import (
|
||||
ClaimContext,
|
||||
ClientExtension,
|
||||
NotificationBinding,
|
||||
ResultClaim,
|
||||
UnexpectedClaimedResult,
|
||||
advertise,
|
||||
)
|
||||
from mcp.client.session import ClientSession
|
||||
|
||||
__all__ = [
|
||||
|
|
@ -19,11 +27,17 @@ __all__ = [
|
|||
"CacheEntry",
|
||||
"CacheKey",
|
||||
"CacheMode",
|
||||
"ClaimContext",
|
||||
"Client",
|
||||
"ClientExtension",
|
||||
"ClientRequestContext",
|
||||
"ClientSession",
|
||||
"InMemoryResponseCacheStore",
|
||||
"InputRequiredRoundsExceededError",
|
||||
"NotificationBinding",
|
||||
"ResponseCacheStore",
|
||||
"ResultClaim",
|
||||
"Transport",
|
||||
"UnexpectedClaimedResult",
|
||||
"advertise",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from __future__ import annotations
|
|||
import hashlib
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import KW_ONLY, dataclass, field
|
||||
from typing import Any, Literal, TypeVar, cast
|
||||
|
|
@ -36,6 +36,7 @@ from mcp_types import (
|
|||
ReadResourceResult,
|
||||
RequestParamsMeta,
|
||||
ResourceTemplateReference,
|
||||
Result,
|
||||
ServerCapabilities,
|
||||
)
|
||||
from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS
|
||||
|
|
@ -46,6 +47,7 @@ from mcp.client._memory import InMemoryTransport
|
|||
from mcp.client._probe import negotiate_auto
|
||||
from mcp.client._transport import Transport
|
||||
from mcp.client.caching import CacheConfig, CacheMode, ClientResponseCache, InMemoryResponseCacheStore
|
||||
from mcp.client.extension import ClaimContext, ClientExtension, NotificationBinding, ResultClaim
|
||||
from mcp.client.session import (
|
||||
ClientRequestContext,
|
||||
ClientSession,
|
||||
|
|
@ -62,6 +64,7 @@ from mcp.server.runner import modern_on_request
|
|||
from mcp.shared.direct_dispatcher import create_direct_dispatcher_pair
|
||||
from mcp.shared.dispatcher import Dispatcher, ProgressFnT
|
||||
from mcp.shared.exceptions import MCPDeprecationWarning, MCPError
|
||||
from mcp.shared.extension import validate_extension_identifier
|
||||
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
|
||||
from mcp.shared.session import RequestResponder
|
||||
|
||||
|
|
@ -188,6 +191,72 @@ async def _no_inbound_client_notifications(_dctx: Any, _method: str, _params: Ma
|
|||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _FoldedExtensions:
|
||||
"""`Client.extensions` instances folded into the shapes `ClientSession` consumes."""
|
||||
|
||||
ad: dict[str, dict[str, Any]] | None
|
||||
claims: dict[str, tuple[ResultClaim[Any], ...]] | None
|
||||
bindings: tuple[NotificationBinding[Any], ...] | None
|
||||
by_model: Mapping[type[Result], ResultClaim[Any]]
|
||||
|
||||
|
||||
def _fold_extensions(extensions: Sequence[ClientExtension] | None) -> _FoldedExtensions:
|
||||
"""Fold extension contributions at construction, naming both owners on duplicate tags or methods."""
|
||||
if isinstance(extensions, Mapping):
|
||||
raise TypeError(
|
||||
"extensions= takes a sequence of ClientExtension instances. The mapping form was "
|
||||
"replaced: use advertise(identifier, settings) for advertise-only entries"
|
||||
)
|
||||
if not extensions:
|
||||
return _FoldedExtensions(ad=None, claims=None, bindings=None, by_model={})
|
||||
ad: dict[str, dict[str, Any]] = {}
|
||||
claims: dict[str, tuple[ResultClaim[Any], ...]] = {}
|
||||
bindings: list[NotificationBinding[Any]] = []
|
||||
by_model: dict[type[Result], ResultClaim[Any]] = {}
|
||||
claim_owners: dict[str, str] = {}
|
||||
binding_owners: dict[str, str] = {}
|
||||
for extension in extensions:
|
||||
identifier = getattr(extension, "identifier", None)
|
||||
if identifier is None:
|
||||
raise ValueError(
|
||||
f"{type(extension).__name__} has no `identifier`; a ClientExtension must set the "
|
||||
"`identifier` class attribute (or assign one in `__init__`) before it can be used"
|
||||
)
|
||||
validate_extension_identifier(identifier, owner=type(extension).__name__)
|
||||
if identifier in ad:
|
||||
raise ValueError(f"extension identifier {identifier!r} is passed more than once")
|
||||
ad[identifier] = extension.settings()
|
||||
extension_claims = tuple(extension.claims())
|
||||
for claim in extension_claims:
|
||||
tag = claim.result_type
|
||||
if tag in claim_owners:
|
||||
owner = claim_owners[tag]
|
||||
both = (
|
||||
f"extension {identifier!r} claims"
|
||||
if owner == identifier
|
||||
else (f"extensions {owner!r} and {identifier!r} both claim")
|
||||
)
|
||||
raise ValueError(f"{both} resultType {tag!r}; a wire tag can have only one resolver")
|
||||
claim_owners[tag] = identifier
|
||||
# Each model pins its result_type Literal to one tag, so this index cannot collide.
|
||||
by_model[claim.model] = claim
|
||||
if extension_claims:
|
||||
claims[identifier] = extension_claims
|
||||
for binding in extension.notifications():
|
||||
if binding.method in binding_owners:
|
||||
owner = binding_owners[binding.method]
|
||||
both = (
|
||||
f"extension {identifier!r} binds"
|
||||
if owner == identifier
|
||||
else (f"extensions {owner!r} and {identifier!r} both bind")
|
||||
)
|
||||
raise ValueError(f"{both} notification method {binding.method!r}; a method can have only one observer")
|
||||
binding_owners[binding.method] = identifier
|
||||
bindings.append(binding)
|
||||
return _FoldedExtensions(ad=ad, claims=claims or None, bindings=tuple(bindings) or None, by_model=by_model)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Client:
|
||||
"""A high-level MCP client for connecting to MCP servers.
|
||||
|
|
@ -268,9 +337,12 @@ class Client:
|
|||
`read_resource` give up. Use `client.session.<method>(..., allow_input_required=True)`
|
||||
to drive the loop manually instead."""
|
||||
|
||||
extensions: dict[str, dict[str, Any]] | None = None
|
||||
"""SEP-2133 extension support to advertise under `ClientCapabilities.extensions`
|
||||
(identifier -> settings), e.g. `{"io.modelcontextprotocol/ui": {"mimeTypes": [...]}}`."""
|
||||
extensions: Sequence[ClientExtension] | None = None
|
||||
"""Opt-in client extensions (SEP-2133).
|
||||
|
||||
Each instance contributes its capability ad, its result claims (resolved
|
||||
transparently by `call_tool`), and its notification bindings. For an
|
||||
ad-only entry use `mcp.client.advertise(identifier, settings)`."""
|
||||
|
||||
cache: CacheConfig | Literal[False] | None = None
|
||||
"""Client-side response caching for the SEP-2549 cacheable methods (2026-07-28).
|
||||
|
|
@ -286,6 +358,7 @@ class Client:
|
|||
_exit_stack: AsyncExitStack | None = field(init=False, default=None)
|
||||
_connect: _Connector = field(init=False, repr=False, compare=False)
|
||||
_response_cache: ClientResponseCache | None = field(init=False, default=None, repr=False, compare=False)
|
||||
_folded_extensions: _FoldedExtensions = field(init=False, repr=False, compare=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.mode not in ("legacy", "auto") and self.mode not in MODERN_PROTOCOL_VERSIONS:
|
||||
|
|
@ -298,6 +371,8 @@ class Client:
|
|||
f"mode must be 'legacy', 'auto', or one of {list(MODERN_PROTOCOL_VERSIONS)}; got {self.mode!r}{hint}"
|
||||
)
|
||||
|
||||
self._folded_extensions = _fold_extensions(self.extensions)
|
||||
|
||||
srv = self.server
|
||||
if isinstance(srv, MCPServer):
|
||||
srv = srv._lowlevel_server # pyright: ignore[reportPrivateUsage]
|
||||
|
|
@ -348,7 +423,9 @@ class Client:
|
|||
message_handler=message_handler,
|
||||
client_info=self.client_info,
|
||||
elicitation_callback=self.elicitation_callback,
|
||||
extensions=self.extensions,
|
||||
extensions=self._folded_extensions.ad,
|
||||
result_claims=self._folded_extensions.claims,
|
||||
notification_bindings=self._folded_extensions.bindings,
|
||||
)
|
||||
|
||||
async def __aenter__(self) -> Client:
|
||||
|
|
@ -613,6 +690,11 @@ class Client:
|
|||
state is still subject to the server's TTL, request binding, and key
|
||||
lifetime; a server on the default process-local key rejects it after a restart.
|
||||
|
||||
Result shapes claimed by this client's `extensions` are finished by the
|
||||
owning claim's resolver, whose `CallToolResult` is returned; resolver
|
||||
exceptions propagate as-is. To receive the claimed shape yourself, use
|
||||
`client.session.call_tool(..., allow_claimed=True)`.
|
||||
|
||||
Args:
|
||||
name: The name of the tool to call.
|
||||
arguments: Arguments to pass to the tool.
|
||||
|
|
@ -631,7 +713,7 @@ class Client:
|
|||
MCPError: A callback returned `ErrorData` for an embedded input request.
|
||||
"""
|
||||
|
||||
async def retry(r: InputResponses | None, s: str | None) -> CallToolResult | InputRequiredResult:
|
||||
async def retry(r: InputResponses | None, s: str | None) -> CallToolResult | InputRequiredResult | Result:
|
||||
return await self.session.call_tool(
|
||||
name,
|
||||
arguments,
|
||||
|
|
@ -641,9 +723,23 @@ class Client:
|
|||
request_state=s,
|
||||
meta=meta,
|
||||
allow_input_required=True,
|
||||
# Input rounds resolve before a claimed result, so a claim may end any round.
|
||||
allow_claimed=True,
|
||||
)
|
||||
|
||||
return await self._drive_input_required(await retry(input_responses, request_state), retry)
|
||||
result = await self._drive_input_required(await retry(input_responses, request_state), retry)
|
||||
if isinstance(result, CallToolResult):
|
||||
return result
|
||||
# Only claimed shapes reach this point, so the lookup is total.
|
||||
claim = self._folded_extensions.by_model[type(result)]
|
||||
final = await claim.resolve(
|
||||
result,
|
||||
ClaimContext(session=self.session, tool_name=name, read_timeout_seconds=read_timeout_seconds),
|
||||
)
|
||||
if not final.is_error:
|
||||
# Match the direct path: revalidate the output schema, but never for isError results.
|
||||
await self.session.validate_tool_result(name, final)
|
||||
return final
|
||||
|
||||
async def list_prompts(
|
||||
self,
|
||||
|
|
@ -717,7 +813,7 @@ class Client:
|
|||
|
||||
async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
|
||||
ctx = ClientRequestContext(session=session, request_id=key, meta=req.params.meta if req.params else None)
|
||||
return await session._dispatch_input_request(ctx, req) # pyright: ignore[reportPrivateUsage]
|
||||
return await session.dispatch_input_request(ctx, req)
|
||||
|
||||
return await run_input_required_driver(
|
||||
first, dispatch=dispatch, retry=retry, max_rounds=self.input_required_max_rounds
|
||||
|
|
|
|||
196
src/mcp/client/extension.py
Normal file
196
src/mcp/client/extension.py
Normal file
|
|
@ -0,0 +1,196 @@
|
|||
"""Opt-in extension interface for MCP clients.
|
||||
|
||||
Subclass `ClientExtension`, set `identifier`, override the hooks you need, and
|
||||
pass instances to `Client(extensions=[...])`. For an identifier-only
|
||||
capability ad, use `advertise()`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeVar, get_args
|
||||
|
||||
from mcp_types import CORE_RESULT_TYPES, CallToolResult, InputRequiredResult, Result
|
||||
from mcp_types.version import MODERN_PROTOCOL_VERSIONS
|
||||
from pydantic import AliasChoices, AliasPath, BaseModel
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.shared.extension import validate_extension_identifier
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mcp.client.session import ClientSession
|
||||
|
||||
__all__ = [
|
||||
"ClaimContext",
|
||||
"ClientExtension",
|
||||
"NotificationBinding",
|
||||
"ResultClaim",
|
||||
"UnexpectedClaimedResult",
|
||||
"advertise",
|
||||
]
|
||||
|
||||
_CLAIM_METHODS: Final[frozenset[str]] = frozenset({"tools/call"})
|
||||
"""The closed set of verbs a claim may attach to; widen together with the `method` Literal."""
|
||||
|
||||
_RESERVED_WIRE_ALIASES: Final[frozenset[str]] = frozenset({"requestState", "inputRequests"})
|
||||
"""Typed optional fields of the core result surface that pre-validates every inbound result."""
|
||||
|
||||
|
||||
def _wire_keys(name: str, field: FieldInfo) -> frozenset[str]:
|
||||
"""Every top-level wire key this field can read from or write to."""
|
||||
keys = {field.alias or name}
|
||||
if field.serialization_alias:
|
||||
keys.add(field.serialization_alias)
|
||||
validation_alias = field.validation_alias
|
||||
choices = validation_alias.choices if isinstance(validation_alias, AliasChoices) else [validation_alias]
|
||||
for choice in choices:
|
||||
if isinstance(choice, AliasPath):
|
||||
choice = choice.path[0]
|
||||
if isinstance(choice, str):
|
||||
keys.add(choice)
|
||||
return frozenset(keys)
|
||||
|
||||
|
||||
ClaimedT = TypeVar("ClaimedT", bound=Result)
|
||||
NotifyParamsT = TypeVar("NotifyParamsT", bound=BaseModel)
|
||||
|
||||
|
||||
@dataclass(frozen=True, kw_only=True)
|
||||
class ClaimContext:
|
||||
"""Host-injected context for one `ResultClaim.resolve` call."""
|
||||
|
||||
session: ClientSession
|
||||
tool_name: str
|
||||
read_timeout_seconds: float | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, kw_only=True)
|
||||
class ResultClaim(Generic[ClaimedT]):
|
||||
"""One extra result shape on one spec verb, keyed by the wire `resultType`.
|
||||
|
||||
Active only while the declaring extension is constructed into the client and
|
||||
the negotiated protocol version admits it. `resolve` finishes a claimed
|
||||
result, may send follow-ups through `ctx.session`, and must return the
|
||||
verb's ordinary result. All field constraints are enforced at construction.
|
||||
"""
|
||||
|
||||
result_type: str
|
||||
model: type[ClaimedT]
|
||||
resolve: Callable[[ClaimedT, ClaimContext], Awaitable[CallToolResult]]
|
||||
method: Literal["tools/call"] = "tools/call"
|
||||
protocol_versions: frozenset[str] | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.method not in _CLAIM_METHODS:
|
||||
raise ValueError(f"claims attach to {sorted(_CLAIM_METHODS)} only; got method {self.method!r}")
|
||||
if self.result_type in CORE_RESULT_TYPES:
|
||||
raise ValueError(f"resultType {self.result_type!r} is core protocol vocabulary")
|
||||
if Result not in self.model.__mro__: # runtime guard; the ClaimedT bound only constrains checked callers
|
||||
raise ValueError(f"{self.model.__name__} must subclass mcp_types.Result")
|
||||
if issubclass(self.model, CallToolResult | InputRequiredResult):
|
||||
raise ValueError("claim models must not subclass core result types")
|
||||
for name, model_field in self.model.model_fields.items():
|
||||
for clash in sorted(_wire_keys(name, model_field) & _RESERVED_WIRE_ALIASES):
|
||||
raise ValueError(
|
||||
f"{self.model.__name__}.{name} aliases {clash!r}, a typed field of the core "
|
||||
"result surface; a colliding value would fail core validation before the "
|
||||
"claim adapter runs"
|
||||
)
|
||||
field = self.model.model_fields.get("result_type")
|
||||
if field is None or get_args(field.annotation) != (self.result_type,):
|
||||
raise ValueError(f"{self.model.__name__}.result_type must be Literal[{self.result_type!r}]")
|
||||
if self.protocol_versions is not None and not self.protocol_versions:
|
||||
raise ValueError("empty protocol_versions could never activate; use None for all")
|
||||
if self.protocol_versions is not None and not self.protocol_versions.issubset(MODERN_PROTOCOL_VERSIONS):
|
||||
unrecognized = sorted(self.protocol_versions.difference(MODERN_PROTOCOL_VERSIONS))
|
||||
raise ValueError(
|
||||
f"protocol_versions {unrecognized} are not modern protocol revisions; claimed shapes "
|
||||
"cannot be delivered on a legacy wire (None means every modern version)"
|
||||
)
|
||||
|
||||
|
||||
class UnexpectedClaimedResult(RuntimeError):
|
||||
"""A claimed (extension) result arrived on a `call_tool` that did not opt in.
|
||||
|
||||
The parsed value is carried as `result`; the server may already hold state it
|
||||
references. Opt in via `Client(extensions=[...])` or `allow_claimed=True`.
|
||||
"""
|
||||
|
||||
def __init__(self, result: Result) -> None:
|
||||
super().__init__(
|
||||
f"Server returned a claimed result ({type(result).__name__}); pass the owning extension to "
|
||||
"Client(extensions=[...]) for transparent resolution, or call with allow_claimed=True "
|
||||
"and handle the shape. The carried result may reference server-side state needing cleanup."
|
||||
)
|
||||
self.result = result
|
||||
|
||||
|
||||
@dataclass(frozen=True, kw_only=True)
|
||||
class NotificationBinding(Generic[NotifyParamsT]):
|
||||
"""Deliver server notifications for `method` (the bare wire name) to `handler`.
|
||||
|
||||
Observation-only: validated params arrive one at a time per binding, in
|
||||
dispatch order, through a bounded queue that drops the oldest with a warning
|
||||
on overflow. Stream transports dispatch each notification independently, so
|
||||
near-simultaneous notifications may be dispatched out of wire order. Methods
|
||||
the negotiated version's core tables handle are never delivered to bindings.
|
||||
"""
|
||||
|
||||
method: str
|
||||
params_type: type[NotifyParamsT]
|
||||
handler: Callable[[NotifyParamsT], Awaitable[None]]
|
||||
|
||||
|
||||
class ClientExtension:
|
||||
"""Base class for an opt-in client extension; override only what you need.
|
||||
|
||||
The surface is declarative, fixed at construction, and never receives the client.
|
||||
"""
|
||||
|
||||
#: Reverse-DNS extension identifier, advertised under `ClientCapabilities.extensions`.
|
||||
identifier: str
|
||||
|
||||
def __init_subclass__(cls, **kwargs: Any) -> None:
|
||||
super().__init_subclass__(**kwargs)
|
||||
# Per-instance identifiers (assigned in __init__) are validated at consumption instead.
|
||||
if (identifier := cls.__dict__.get("identifier")) is not None:
|
||||
validate_extension_identifier(identifier, owner=cls.__name__)
|
||||
|
||||
def settings(self) -> dict[str, Any]:
|
||||
"""Per-extension settings advertised at `ClientCapabilities.extensions[identifier]`.
|
||||
|
||||
Read once at `Client` construction. A claim-bearing extension is
|
||||
advertised only at protocol versions where at least one of its claims
|
||||
is active.
|
||||
"""
|
||||
return {}
|
||||
|
||||
def claims(self) -> Sequence[ResultClaim[Any]]:
|
||||
"""Extra result shapes this extension claims, with their resolvers."""
|
||||
return ()
|
||||
|
||||
def notifications(self) -> Sequence[NotificationBinding[Any]]:
|
||||
"""Server notifications this extension observes."""
|
||||
return ()
|
||||
|
||||
|
||||
class _AdvertiseOnly(ClientExtension):
|
||||
"""Ad-only extension returned by `advertise()`."""
|
||||
|
||||
def __init__(self, identifier: str, settings: dict[str, Any]) -> None:
|
||||
self.identifier = identifier
|
||||
self._settings = settings
|
||||
|
||||
def settings(self) -> dict[str, Any]:
|
||||
return self._settings
|
||||
|
||||
|
||||
def advertise(identifier: str, settings: dict[str, Any] | None = None) -> ClientExtension:
|
||||
"""Advertise an extension identifier (with optional settings) and nothing else.
|
||||
|
||||
Advertising an extension you do not implement asserts wire support you do
|
||||
not have; for behavioral extensions construct the real extension instead.
|
||||
"""
|
||||
validate_extension_identifier(identifier, owner="advertise")
|
||||
return _AdvertiseOnly(identifier, {} if settings is None else settings)
|
||||
|
|
@ -1,15 +1,18 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from operator import or_
|
||||
from types import TracebackType
|
||||
from typing import Any, Literal, Protocol, cast, overload
|
||||
from typing import Annotated, Any, Final, Literal, Protocol, cast, overload
|
||||
|
||||
import anyio
|
||||
import anyio.abc
|
||||
import anyio.lowlevel
|
||||
import mcp_types as types
|
||||
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
||||
from mcp_types import (
|
||||
CLIENT_CAPABILITIES_META_KEY,
|
||||
CLIENT_INFO_META_KEY,
|
||||
|
|
@ -27,10 +30,11 @@ from mcp_types.version import (
|
|||
LATEST_MODERN_VERSION,
|
||||
MODERN_PROTOCOL_VERSIONS,
|
||||
)
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, Discriminator, Tag, TypeAdapter, ValidationError
|
||||
from typing_extensions import Self, TypeVar, deprecated
|
||||
|
||||
from mcp.client._transport import ReadStream, WriteStream
|
||||
from mcp.client.extension import NotificationBinding, ResultClaim, UnexpectedClaimedResult
|
||||
from mcp.shared._compat import resync_tracer
|
||||
from mcp.shared.dispatcher import CallOptions, DispatchContext, Dispatcher, ProgressFnT
|
||||
from mcp.shared.exceptions import MCPDeprecationWarning, MCPError
|
||||
|
|
@ -51,6 +55,7 @@ from mcp.shared.transport_context import TransportContext
|
|||
|
||||
DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
|
||||
DISCOVER_TIMEOUT_SECONDS = 10.0
|
||||
_NOTIFICATION_QUEUE_SIZE: Final = 256
|
||||
|
||||
logger = logging.getLogger("client")
|
||||
|
||||
|
|
@ -189,7 +194,8 @@ async def _default_logging_callback(
|
|||
|
||||
ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(types.ClientResult | types.ErrorData)
|
||||
|
||||
_CallToolResultAdapter: TypeAdapter[types.CallToolResult | types.InputRequiredResult] = TypeAdapter(
|
||||
# Typed against the wide parse union so adopt-built claim adapters share this attribute type.
|
||||
_CallToolResultAdapter: TypeAdapter[types.CallToolResult | types.InputRequiredResult | types.Result] = TypeAdapter(
|
||||
types.CallToolResult | types.InputRequiredResult
|
||||
)
|
||||
_GetPromptResultAdapter: TypeAdapter[types.GetPromptResult | types.InputRequiredResult] = TypeAdapter(
|
||||
|
|
@ -200,6 +206,89 @@ _ReadResourceResultAdapter: TypeAdapter[types.ReadResourceResult | types.InputRe
|
|||
)
|
||||
|
||||
|
||||
def _claim_active(claim: ResultClaim[Any], version: str) -> bool:
|
||||
"""A claim is active at modern versions only, narrowed by its optional version subset."""
|
||||
return version in MODERN_PROTOCOL_VERSIONS and (
|
||||
claim.protocol_versions is None or version in claim.protocol_versions
|
||||
)
|
||||
|
||||
|
||||
def _active_claims_at(
|
||||
claims_by_extension: Mapping[str, tuple[ResultClaim[Any], ...]], version: str
|
||||
) -> dict[str, ResultClaim[Any]]:
|
||||
"""Claims active at `version`, keyed by wire tag; empty at any legacy version."""
|
||||
return {
|
||||
claim.result_type: claim
|
||||
for claims in claims_by_extension.values()
|
||||
for claim in claims
|
||||
if _claim_active(claim, version)
|
||||
}
|
||||
|
||||
|
||||
def _build_call_tool_adapter(
|
||||
active: Mapping[str, ResultClaim[Any]],
|
||||
) -> TypeAdapter[types.CallToolResult | types.InputRequiredResult | types.Result]:
|
||||
"""Build a discriminated tools/call adapter: a core arm plus one arm per active claim."""
|
||||
if not active:
|
||||
return _CallToolResultAdapter
|
||||
tags = frozenset(active)
|
||||
core_arm = "core"
|
||||
while core_arm in tags: # the routing sentinel must never collide with a claimed tag
|
||||
core_arm += "-"
|
||||
|
||||
def _route(value: Any) -> str:
|
||||
# pydantic hands the discriminator either the raw dict or an already-built model.
|
||||
# Unknown or non-string tags route to the core arm and fail core validation there.
|
||||
if isinstance(value, dict):
|
||||
tag = cast("dict[str, Any]", value).get("resultType")
|
||||
else:
|
||||
tag = getattr(value, "result_type", None)
|
||||
return tag if isinstance(tag, str) and tag in tags else core_arm
|
||||
|
||||
arms: list[Any] = [Annotated[types.CallToolResult | types.InputRequiredResult, Tag(core_arm)]]
|
||||
arms += [Annotated[claim.model, Tag(tag)] for tag, claim in active.items()]
|
||||
# reduce(or_) rather than Union star-unpack, which needs py3.11+.
|
||||
return TypeAdapter(Annotated[reduce(or_, arms), Discriminator(_route)])
|
||||
|
||||
|
||||
def _index_claims(
|
||||
result_claims: Mapping[str, Sequence[ResultClaim[Any]]] | None,
|
||||
extensions: dict[str, dict[str, Any]] | None,
|
||||
) -> dict[str, tuple[ResultClaim[Any], ...]]:
|
||||
"""Validate and copy the claims-by-extension mapping."""
|
||||
indexed: dict[str, tuple[ResultClaim[Any], ...]] = {}
|
||||
seen: set[str] = set()
|
||||
for identifier, claims in (result_claims or {}).items():
|
||||
if extensions is None or identifier not in extensions:
|
||||
raise ValueError(
|
||||
f"result_claims key {identifier!r} has no extensions entry; a claim is only "
|
||||
"advertised through its extension's capability ad"
|
||||
)
|
||||
if not claims:
|
||||
raise ValueError(
|
||||
f"result_claims[{identifier!r}] is empty and would drop the extension from "
|
||||
"the capability ad at every version. Omit the key instead"
|
||||
)
|
||||
for claim in claims:
|
||||
if claim.result_type in seen:
|
||||
raise ValueError(f"duplicate result claim for resultType {claim.result_type!r}")
|
||||
seen.add(claim.result_type)
|
||||
indexed[identifier] = tuple(claims)
|
||||
return indexed
|
||||
|
||||
|
||||
def _index_bindings(
|
||||
notification_bindings: Sequence[NotificationBinding[Any]] | None,
|
||||
) -> dict[str, NotificationBinding[Any]]:
|
||||
"""Index bindings by wire method, rejecting duplicates."""
|
||||
indexed: dict[str, NotificationBinding[Any]] = {}
|
||||
for binding in notification_bindings or ():
|
||||
if binding.method in indexed:
|
||||
raise ValueError(f"duplicate notification binding for method {binding.method!r}")
|
||||
indexed[binding.method] = binding
|
||||
return indexed
|
||||
|
||||
|
||||
def _input_required_unexpected(method: str) -> RuntimeError:
|
||||
return RuntimeError(
|
||||
"Server returned InputRequiredResult; pass allow_input_required=True to receive it "
|
||||
|
|
@ -216,6 +305,9 @@ class ClientSession:
|
|||
correlation; this class owns the typed MCP layer and the constructor
|
||||
callbacks. Transport `Exception` items reach `message_handler` only when
|
||||
the session builds its own dispatcher from a stream pair.
|
||||
|
||||
Extension `result_claims` fold into tools/call parsing at `adopt()`;
|
||||
`notification_bindings` observe vendor notifications via bounded FIFOs.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
|
@ -232,13 +324,22 @@ class ClientSession:
|
|||
*,
|
||||
sampling_capabilities: types.SamplingCapability | None = None,
|
||||
extensions: dict[str, dict[str, Any]] | None = None,
|
||||
result_claims: Mapping[str, Sequence[ResultClaim[Any]]] | None = None,
|
||||
notification_bindings: Sequence[NotificationBinding[Any]] | None = None,
|
||||
dispatcher: Dispatcher[Any] | None = None,
|
||||
) -> None:
|
||||
self._session_read_timeout_seconds = read_timeout_seconds
|
||||
self._client_info = client_info or DEFAULT_CLIENT_INFO
|
||||
self._sampling_callback = sampling_callback or _default_sampling_callback
|
||||
self._sampling_capabilities = sampling_capabilities
|
||||
self._extensions = extensions
|
||||
self._extensions = dict(extensions) if extensions is not None else None
|
||||
self._result_claims = _index_claims(result_claims, extensions)
|
||||
self._notification_bindings = _index_bindings(notification_bindings)
|
||||
self._active_claims: dict[str, ResultClaim[Any]] = {}
|
||||
self._call_tool_adapter = _CallToolResultAdapter
|
||||
self._binding_queues: dict[
|
||||
str, tuple[MemoryObjectSendStream[BaseModel], MemoryObjectReceiveStream[BaseModel]]
|
||||
] = {}
|
||||
self._elicitation_callback = elicitation_callback or _default_elicitation_callback
|
||||
self._list_roots_callback = list_roots_callback or _default_list_roots_callback
|
||||
self._logging_callback = logging_callback or _default_logging_callback
|
||||
|
|
@ -274,7 +375,14 @@ class ClientSession:
|
|||
self._task_group = anyio.create_task_group()
|
||||
await self._task_group.__aenter__()
|
||||
try:
|
||||
# Queues must exist before the dispatcher starts: _on_notify enqueues into this dict.
|
||||
for binding in self._notification_bindings.values():
|
||||
send, receive = anyio.create_memory_object_stream[BaseModel](_NOTIFICATION_QUEUE_SIZE)
|
||||
self._binding_queues[binding.method] = (send, receive)
|
||||
await self._task_group.start(self._dispatcher.run, self._on_request, self._on_notify)
|
||||
for binding in self._notification_bindings.values():
|
||||
_, receive = self._binding_queues[binding.method]
|
||||
self._task_group.start_soon(self._deliver_bound_notifications, binding, receive)
|
||||
except BaseException:
|
||||
# Unwind the entered task group before propagating: a cancellation
|
||||
# landing here (e.g. `move_on_after` around connect) would abandon
|
||||
|
|
@ -285,7 +393,10 @@ class ClientSession:
|
|||
# Shield the group's own scope (a new one would break LIFO exit)
|
||||
# so a pending outer cancellation cannot re-fire inside __aexit__.
|
||||
task_group.cancel_scope.shield = True
|
||||
await task_group.__aexit__(None, None, None)
|
||||
try:
|
||||
await task_group.__aexit__(None, None, None)
|
||||
finally:
|
||||
self._close_binding_queues()
|
||||
raise
|
||||
return self
|
||||
|
||||
|
|
@ -295,16 +406,38 @@ class ClientSession:
|
|||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> bool | None:
|
||||
# Exit must not block: cancel the dispatcher and in-flight callbacks.
|
||||
# Exit must not block: cancel the dispatcher, binding consumers, and in-flight callbacks.
|
||||
assert self._task_group is not None
|
||||
self._task_group.cancel_scope.cancel()
|
||||
result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
|
||||
try:
|
||||
result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
|
||||
finally:
|
||||
self._close_binding_queues()
|
||||
await resync_tracer()
|
||||
return result
|
||||
|
||||
def _close_binding_queues(self) -> None:
|
||||
# Unclosed memory object streams warn at garbage collection; close is idempotent.
|
||||
for send, receive in self._binding_queues.values():
|
||||
send.close()
|
||||
receive.close()
|
||||
self._binding_queues.clear()
|
||||
|
||||
async def _deliver_bound_notifications(
|
||||
self, binding: NotificationBinding[Any], receive: MemoryObjectReceiveStream[BaseModel]
|
||||
) -> None:
|
||||
"""Consume one binding's FIFO, decoupled from the dispatcher so handlers can do session I/O."""
|
||||
while True:
|
||||
params = await receive.receive()
|
||||
try:
|
||||
await binding.handler(params)
|
||||
except Exception:
|
||||
# A raising handler costs only that delivery, as in _on_notify.
|
||||
logger.exception("notification binding handler for %r raised", binding.method)
|
||||
|
||||
async def send_request(
|
||||
self,
|
||||
request: types.ClientRequest,
|
||||
request: types.ClientRequest | types.Request[Any, Any],
|
||||
result_type: type[ReceiveResultT] | TypeAdapter[ReceiveResultT],
|
||||
request_read_timeout_seconds: float | None = None,
|
||||
metadata: ClientMessageMetadata | None = None,
|
||||
|
|
@ -318,11 +451,20 @@ class ClientSession:
|
|||
Raises:
|
||||
MCPError: Error response, read timeout, or connection closed.
|
||||
RuntimeError: Called before entering the context manager.
|
||||
ValueError: The request declares `name_param` but its params carry no string name.
|
||||
"""
|
||||
data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
method: str = data["method"]
|
||||
opts: CallOptions = {}
|
||||
self._stamp(data, opts)
|
||||
# The stamp runs first, so its NAME_BEARING_METHODS rows win; a missing name fails loud.
|
||||
headers = opts.setdefault("headers", {})
|
||||
if (key := type(request).name_param) is not None and MCP_NAME_HEADER not in headers:
|
||||
params_data: dict[str, Any] = data.get("params") or {}
|
||||
name = params_data.get(key)
|
||||
if not isinstance(name, str):
|
||||
raise ValueError(f"{method} requires params[{key!r}] for Mcp-Name")
|
||||
headers[MCP_NAME_HEADER] = encode_header_value(name)
|
||||
timeout = (
|
||||
request_read_timeout_seconds
|
||||
if request_read_timeout_seconds is not None
|
||||
|
|
@ -360,7 +502,21 @@ class ClientSession:
|
|||
self._stamp(data, opts)
|
||||
await self._dispatcher.notify(data["method"], data.get("params"), opts)
|
||||
|
||||
def _build_capabilities(self) -> types.ClientCapabilities:
|
||||
def _build_capabilities(self, version: str) -> types.ClientCapabilities:
|
||||
"""Build the capability ad for a wire speaking `version`.
|
||||
|
||||
Claim-bearing identifiers whose claims are all inactive at `version` drop, so
|
||||
the client never advertises result shapes it would reject; claim-less
|
||||
identifiers always advertise.
|
||||
"""
|
||||
extensions = self._extensions
|
||||
if extensions is not None and self._result_claims:
|
||||
extensions = {
|
||||
identifier: settings
|
||||
for identifier, settings in extensions.items()
|
||||
if identifier not in self._result_claims
|
||||
or any(_claim_active(claim, version) for claim in self._result_claims[identifier])
|
||||
} or None
|
||||
sampling = (
|
||||
(self._sampling_capabilities or types.SamplingCapability())
|
||||
if self._sampling_callback is not _default_sampling_callback
|
||||
|
|
@ -380,7 +536,7 @@ class ClientSession:
|
|||
else None
|
||||
)
|
||||
return types.ClientCapabilities(
|
||||
sampling=sampling, elicitation=elicitation, experimental=None, extensions=self._extensions, roots=roots
|
||||
sampling=sampling, elicitation=elicitation, experimental=None, extensions=extensions, roots=roots
|
||||
)
|
||||
|
||||
async def initialize(self) -> types.InitializeResult:
|
||||
|
|
@ -390,7 +546,8 @@ class ClientSession:
|
|||
types.InitializeRequest(
|
||||
params=types.InitializeRequestParams(
|
||||
protocol_version=LATEST_HANDSHAKE_VERSION,
|
||||
capabilities=self._build_capabilities(),
|
||||
# The handshake negotiates only legacy versions, where no claim is active.
|
||||
capabilities=self._build_capabilities(LATEST_HANDSHAKE_VERSION),
|
||||
client_info=self._client_info,
|
||||
),
|
||||
),
|
||||
|
|
@ -424,17 +581,30 @@ class ClientSession:
|
|||
f"No mutually supported modern protocol version "
|
||||
f"(server: {result.supported_versions}, client: {list(MODERN_PROTOCOL_VERSIONS)})"
|
||||
)
|
||||
version = mutual[-1]
|
||||
client_info = self._client_info.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
capabilities = self._build_capabilities().model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
self._stamp = _make_modern_stamp(mutual[-1], client_info, capabilities, self._resolve_param_headers)
|
||||
capabilities = self._build_capabilities(version).model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
self._stamp = _make_modern_stamp(version, client_info, capabilities, self._resolve_param_headers)
|
||||
self._discover_result = result
|
||||
self._initialize_result = None
|
||||
self._negotiated_version = mutual[-1]
|
||||
else:
|
||||
self._stamp = _make_handshake_stamp(result.protocol_version)
|
||||
version = result.protocol_version
|
||||
self._stamp = _make_handshake_stamp(version)
|
||||
self._initialize_result = result
|
||||
self._discover_result = None
|
||||
self._negotiated_version = result.protocol_version
|
||||
self._negotiated_version = version
|
||||
# Both arms reach here, so re-adoption resets cleanly; legacy versions activate no claims.
|
||||
# Core-vocabulary tags are unconstructible (ResultClaim.__post_init__), so no exclusion needed.
|
||||
self._active_claims = _active_claims_at(self._result_claims, version)
|
||||
self._call_tool_adapter = _build_call_tool_adapter(self._active_claims)
|
||||
for method in self._notification_bindings:
|
||||
# Bindings are consulted only for methods core does not know, so this one can never fire.
|
||||
if (method, version) in _methods.SERVER_NOTIFICATIONS:
|
||||
logger.warning(
|
||||
"notification binding for %r will never fire at %s: the core protocol defines this method",
|
||||
method,
|
||||
version,
|
||||
)
|
||||
|
||||
async def send_discover(self, version: str) -> dict[str, Any]:
|
||||
"""Send a single ``server/discover`` at ``version`` and return the raw result dict.
|
||||
|
|
@ -450,7 +620,7 @@ class ClientSession:
|
|||
synthesized into a JSON-RPC error by the transport).
|
||||
"""
|
||||
client_info = self._client_info.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
capabilities = self._build_capabilities().model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
capabilities = self._build_capabilities(version).model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
request = types.DiscoverRequest(
|
||||
params=types.RequestParams(
|
||||
_meta={
|
||||
|
|
@ -704,6 +874,7 @@ class ClientSession:
|
|||
request_state: str | None = None,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
allow_input_required: Literal[False] = False,
|
||||
allow_claimed: Literal[False] = False,
|
||||
) -> types.CallToolResult: ...
|
||||
|
||||
@overload
|
||||
|
|
@ -718,8 +889,39 @@ class ClientSession:
|
|||
request_state: str | None = None,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
allow_input_required: bool,
|
||||
allow_claimed: Literal[False] = False,
|
||||
) -> types.CallToolResult | types.InputRequiredResult: ...
|
||||
|
||||
@overload
|
||||
async def call_tool(
|
||||
self,
|
||||
name: str,
|
||||
arguments: dict[str, Any] | None = None,
|
||||
read_timeout_seconds: float | None = None,
|
||||
progress_callback: ProgressFnT | None = None,
|
||||
*,
|
||||
input_responses: types.InputResponses | None = None,
|
||||
request_state: str | None = None,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
allow_input_required: Literal[False] = False,
|
||||
allow_claimed: bool,
|
||||
) -> types.CallToolResult | types.Result: ...
|
||||
|
||||
@overload
|
||||
async def call_tool(
|
||||
self,
|
||||
name: str,
|
||||
arguments: dict[str, Any] | None = None,
|
||||
read_timeout_seconds: float | None = None,
|
||||
progress_callback: ProgressFnT | None = None,
|
||||
*,
|
||||
input_responses: types.InputResponses | None = None,
|
||||
request_state: str | None = None,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
allow_input_required: bool,
|
||||
allow_claimed: bool,
|
||||
) -> types.CallToolResult | types.InputRequiredResult | types.Result: ...
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
name: str,
|
||||
|
|
@ -731,7 +933,8 @@ class ClientSession:
|
|||
request_state: str | None = None,
|
||||
meta: RequestParamsMeta | None = None,
|
||||
allow_input_required: bool = False,
|
||||
) -> types.CallToolResult | types.InputRequiredResult:
|
||||
allow_claimed: bool = False,
|
||||
) -> types.CallToolResult | types.InputRequiredResult | types.Result:
|
||||
"""Send a tools/call request with optional progress callback support.
|
||||
|
||||
On a modern (2026-07-28) connection, arguments annotated with `x-mcp-header`
|
||||
|
|
@ -745,10 +948,13 @@ class ClientSession:
|
|||
allow_input_required: When ``False`` (default), an `InputRequiredResult`
|
||||
from the server raises `RuntimeError`; when ``True``, it is returned
|
||||
so the caller can resolve the requests and retry.
|
||||
allow_claimed: When `False` (default), a claimed extension result raises
|
||||
`UnexpectedClaimedResult`; when `True`, the parsed claim model is returned.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the server returns an `InputRequiredResult` and
|
||||
``allow_input_required`` is ``False``.
|
||||
UnexpectedClaimedResult: Claimed result with `allow_claimed` False; carries the parsed value.
|
||||
"""
|
||||
result = await self.send_request(
|
||||
types.CallToolRequest(
|
||||
|
|
@ -760,16 +966,19 @@ class ClientSession:
|
|||
_meta=meta,
|
||||
),
|
||||
),
|
||||
_CallToolResultAdapter,
|
||||
self._call_tool_adapter,
|
||||
request_read_timeout_seconds=read_timeout_seconds,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
if isinstance(result, types.CallToolResult) and not result.is_error:
|
||||
await self._validate_tool_result(name, result)
|
||||
await self.validate_tool_result(name, result)
|
||||
|
||||
# The input_required arm stays first; a claimed shape is terminal for the multi-round-trip driver.
|
||||
if isinstance(result, types.InputRequiredResult) and not allow_input_required:
|
||||
raise _input_required_unexpected("call_tool")
|
||||
if not isinstance(result, types.CallToolResult | types.InputRequiredResult) and not allow_claimed:
|
||||
raise UnexpectedClaimedResult(result)
|
||||
return result
|
||||
|
||||
def _resolve_param_headers(self, name: str, arguments: Mapping[str, Any]) -> dict[str, str]:
|
||||
|
|
@ -779,8 +988,12 @@ class ClientSession:
|
|||
return {}
|
||||
return mcp_param_headers(header_map, arguments)
|
||||
|
||||
async def _validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
|
||||
"""Validate the structured content of a tool result against its output schema."""
|
||||
async def validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
|
||||
"""Revalidate a `CallToolResult` against the tool's declared output schema.
|
||||
|
||||
Raises:
|
||||
RuntimeError: Structured content is missing or does not conform to the schema.
|
||||
"""
|
||||
if name not in self._tool_output_schemas:
|
||||
# refresh output schema cache
|
||||
await self.list_tools()
|
||||
|
|
@ -970,7 +1183,7 @@ class ClientSession:
|
|||
ctx = ClientRequestContext(
|
||||
session=self, request_id=dctx.request_id, meta=request.params.meta if request.params else None
|
||||
)
|
||||
response = await self._dispatch_input_request(ctx, request)
|
||||
response = await self.dispatch_input_request(ctx, request)
|
||||
client_response = ClientResponse.validate_python(response)
|
||||
if isinstance(client_response, types.ErrorData):
|
||||
raise MCPError.from_error_data(client_response)
|
||||
|
|
@ -982,16 +1195,18 @@ class ClientSession:
|
|||
raise MCPError(code=INTERNAL_ERROR, message="Client callback returned an invalid result") from None
|
||||
return dumped
|
||||
|
||||
async def _dispatch_input_request(
|
||||
self, ctx: ClientRequestContext, req: types.InputRequest
|
||||
async def dispatch_input_request(
|
||||
self, ctx: ClientRequestContext, request: types.InputRequest
|
||||
) -> types.InputResponse | types.ErrorData:
|
||||
"""Route a server-initiated input request to the matching constructor callback.
|
||||
"""Route an input request through the client's callback table.
|
||||
|
||||
Shared by the legacy server→client RPC path (`_on_request`) and the
|
||||
2026-07-28 multi-round-trip driver, which dispatches the embedded
|
||||
`InputRequiredResult.input_requests` through the same callbacks.
|
||||
|
||||
Returns the callback's `InputResponse`, or `ErrorData` when the callback declines.
|
||||
"""
|
||||
match req:
|
||||
match request:
|
||||
case types.CreateMessageRequest(params=p):
|
||||
return await self._sampling_callback(ctx, p)
|
||||
case types.ElicitRequest(params=p):
|
||||
|
|
@ -1008,7 +1223,26 @@ class ClientSession:
|
|||
try:
|
||||
notification = cast(types.ServerNotification, _methods.parse_server_notification(method, version, params))
|
||||
except KeyError:
|
||||
logger.debug("dropped %r: not defined at %s", method, version)
|
||||
# Only methods unknown to the negotiated version's core tables reach the bindings.
|
||||
binding = self._notification_bindings.get(method)
|
||||
if binding is None:
|
||||
logger.debug("dropped %r: not defined at %s", method, version)
|
||||
return
|
||||
try:
|
||||
bound_params = binding.params_type.model_validate(params or {})
|
||||
except ValidationError:
|
||||
logger.warning("Failed to validate notification: %s", method, exc_info=True)
|
||||
return
|
||||
send, receive = self._binding_queues[method]
|
||||
try:
|
||||
# Must not await: DirectDispatcher calls _on_notify inline; blocking deadlocks in-process servers.
|
||||
send.send_nowait(bound_params)
|
||||
except anyio.WouldBlock:
|
||||
# Evict the oldest event; no checkpoint since the failed send,
|
||||
# so the buffer is still full and the retry cannot block.
|
||||
receive.receive_nowait()
|
||||
logger.warning("notification queue for %r is full; dropped the oldest event", method)
|
||||
send.send_nowait(bound_params)
|
||||
return
|
||||
except ValidationError:
|
||||
logger.warning("Failed to validate notification: %s", method, exc_info=True)
|
||||
|
|
|
|||
|
|
@ -19,7 +19,6 @@ extensions remain importable without constructing an `MCPServer`.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
|
@ -30,31 +29,14 @@ from pydantic import BaseModel
|
|||
|
||||
from mcp.server.context import CallNext, HandlerResult, ServerMiddleware, ServerRequestContext
|
||||
|
||||
# Re-exported from `mcp.shared.extension` (shared with the client surface) for existing importers.
|
||||
from mcp.shared.extension import validate_extension_identifier as validate_extension_identifier
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mcp.server.mcpserver.resources import Resource
|
||||
|
||||
RequestHandler = Callable[[ServerRequestContext[Any, Any], Any], Awaitable[HandlerResult]]
|
||||
|
||||
# Extension identifiers follow the `_meta` key grammar with a mandatory prefix
|
||||
# (SEP-2133 / basic/index.mdx): dot-separated labels, each starting with a
|
||||
# letter and ending with a letter or digit (hyphens interior), then `/`, then a
|
||||
# name that starts and ends alphanumeric (`.`/`_`/`-` interior).
|
||||
_LABEL = r"[A-Za-z](?:[A-Za-z0-9-]*[A-Za-z0-9])?"
|
||||
_NAME = r"[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?"
|
||||
_IDENTIFIER_RE = re.compile(rf"{_LABEL}(?:\.{_LABEL})*/{_NAME}")
|
||||
|
||||
|
||||
def validate_extension_identifier(identifier: Any, *, owner: str) -> None:
|
||||
"""Raise `TypeError` unless `identifier` is a `vendor-prefix/name` string.
|
||||
|
||||
SEP-2133 requires extension identifiers to carry a reverse-DNS prefix.
|
||||
"""
|
||||
if not isinstance(identifier, str) or not _IDENTIFIER_RE.fullmatch(identifier):
|
||||
raise TypeError(
|
||||
f"{owner}.identifier must be a `vendor-prefix/name` string "
|
||||
f"(reverse-DNS prefix required), got {identifier!r}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ToolBinding:
|
||||
|
|
|
|||
28
src/mcp/shared/extension.py
Normal file
28
src/mcp/shared/extension.py
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
"""Extension-identifier grammar shared by the server and client extension surfaces."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
__all__ = ["validate_extension_identifier"]
|
||||
|
||||
# Extension identifiers follow the `_meta` key grammar with a mandatory prefix
|
||||
# (SEP-2133 / basic/index.mdx): dot-separated labels, each starting with a
|
||||
# letter and ending with a letter or digit (hyphens interior), then `/`, then a
|
||||
# name that starts and ends alphanumeric (`.`/`_`/`-` interior).
|
||||
_LABEL = r"[A-Za-z](?:[A-Za-z0-9-]*[A-Za-z0-9])?"
|
||||
_NAME = r"[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?"
|
||||
_IDENTIFIER_RE = re.compile(rf"{_LABEL}(?:\.{_LABEL})*/{_NAME}")
|
||||
|
||||
|
||||
def validate_extension_identifier(identifier: Any, *, owner: str) -> None:
|
||||
"""Raise `TypeError` unless `identifier` is a `vendor-prefix/name` string.
|
||||
|
||||
SEP-2133 requires extension identifiers to carry a reverse-DNS prefix.
|
||||
"""
|
||||
if not isinstance(identifier, str) or not _IDENTIFIER_RE.fullmatch(identifier):
|
||||
raise TypeError(
|
||||
f"{owner}.identifier must be a `vendor-prefix/name` string "
|
||||
f"(reverse-DNS prefix required), got {identifier!r}"
|
||||
)
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue