feat(client): Add capability_extensions parameter to ClientSession #1901
Conversation
Add a new `capability_extensions` parameter to `ClientSession.__init__()` that
allows clients to include additional capability fields in the initialize request.
This enables clients to advertise protocol extensions (like `io.modelcontextprotocol/ui`)
without having to override the `initialize()` method.
Example usage:
```python
session = ClientSession(
read_stream,
write_stream,
capability_extensions={
"extensions": {
"io.modelcontextprotocol/ui": {
"mimeTypes": ["text/html;profile=mcp-app"]
}
}
}
)
```
The extensions are merged into `ClientCapabilities` using Pydantic's extra fields
feature (`model_config = {'extra': 'allow'}`).
jerome3o-anthropic
force-pushed
the
jerome/add-capability-extensions-parameter
branch
from
8deee3f to
9830794
Compare
src/mcp/client/session.py
Outdated
| experimental=None, | ||
| roots=roots, | ||
| tasks=self._task_handlers.build_capability(), | ||
| **self._capability_extensions, |
There was a problem hiding this comment.
I don't think it should allow any key here, this will be an issue if in the future ClientCapabilities has a new key but it's used by someone.
| **self._capability_extensions, | |
| extensions=self._capability_extensions, |
If the spec doesn't have this explicit, I think the python SDK should.
There was a problem hiding this comment.
Uh, I don't actually think we should specify extension here, I think this is just something made up for MCP apps (not sure, the SEP for extensions isn't in yet)
But - in principle we should be able to allow any property to be added to any MCP message, that was the intention of making all the objects allow_extra=true etc
There was a problem hiding this comment.
This is what we're trying to build, but idk if "extensions" is the standard way to add things here
There was a problem hiding this comment.
Actually modelcontextprotocol/modelcontextprotocol#1724 looks like it will be specifically under "extensions"
There was a problem hiding this comment.
I'm not sure I agree on it being explicitly called "extensions" in the SDK as SEP-1724 (modelcontextprotocol/modelcontextprotocol#1724) hasn't landed yet, and this kinda feels like adding a convention which should really be in the spec.
Address PR review feedback: Instead of using **capability_extensions to spread arbitrary keys into ClientCapabilities, explicitly use the 'extensions' field. This prevents potential conflicts if new fields are added to ClientCapabilities in the future. Changes: - Add 'extensions' field to ClientCapabilities type definition - Update ClientSession to pass capability_extensions as 'extensions' kwarg - Update test to pass extensions dict directly (without 'extensions' key wrapper)
src/mcp/client/session.py
Outdated
| roots=roots, | ||
| tasks=self._task_handlers.build_capability(), | ||
| **self._capability_extensions, | ||
| extensions=self._capability_extensions or None, |
There was a problem hiding this comment.
it's already dict[str, Any] | None, the or is not needed, is it?
| async def mock_server(): | ||
| nonlocal received_capabilities | ||
|
|
||
| session_message = await client_to_server_receive.receive() | ||
| jsonrpc_request = session_message.message | ||
| assert isinstance(jsonrpc_request.root, JSONRPCRequest) | ||
| request = ClientRequest.model_validate( | ||
| jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True) | ||
| ) | ||
| assert isinstance(request.root, InitializeRequest) | ||
| received_capabilities = request.root.params.capabilities | ||
|
|
||
| result = ServerResult( | ||
| InitializeResult( | ||
| protocol_version=LATEST_PROTOCOL_VERSION, | ||
| capabilities=ServerCapabilities(), | ||
| server_info=Implementation(name="mock-server", version="0.1.0"), | ||
| ) | ||
| ) | ||
|
|
||
| async with server_to_client_send: | ||
| await server_to_client_send.send( | ||
| SessionMessage( | ||
| JSONRPCMessage( | ||
| JSONRPCResponse( | ||
| jsonrpc="2.0", | ||
| id=jsonrpc_request.root.id, | ||
| result=result.model_dump(by_alias=True, mode="json", exclude_none=True), | ||
| ) | ||
| ) | ||
| ) | ||
| ) |
There was a problem hiding this comment.
Is the server supposed to do something when it's aware of the extension?
Summary
Add a new
capability_extensionsparameter toClientSession.__init__()that allows clients to include additional capability fields in the initialize request.Motivation
Currently, if a client wants to advertise protocol extensions (like
io.modelcontextprotocol/uifor UI rendering) in the initialize request, they must override the entireinitialize()method and duplicate its internal logic. This is error-prone and creates maintenance burden when the SDK'sinitialize()implementation changes.This PR provides a clean escape hatch for clients to extend capabilities without subclassing.
Example Usage
Implementation
The extensions are merged into
ClientCapabilitiesusing Pydantic's extra fields feature (model_config = {'extra': 'allow'}onMCPModel).Testing
Added a test that verifies capability extensions are properly included in the initialize request.