feat: Add message middleware support for session message transformation

[jerome3o-anthropic]

Summary

Add a middleware pattern that allows transforming JSON-RPC messages before sending and after receiving. This provides a clean way to extend protocol messages (e.g., adding custom capabilities to initialize requests) without needing to subclass or override session methods.

Motivation

When clients need to add custom fields to MCP messages (like protocol extensions or custom capabilities), they currently have to override the initialize() method and duplicate its internal logic. This is error-prone and creates maintenance burden.

A middleware approach provides a clean escape hatch that:

• Doesn't require subclassing
• Works with any message type (requests, responses, notifications)
• Supports both sync and async transformations
• Can be composed (multiple middleware functions)

Example Usage

from mcp import ClientSession, JSONRPCMessage, JSONRPCRequest

def add_ui_extension(message: JSONRPCMessage) -> JSONRPCMessage:
    """Add UI extension capabilities to initialize request."""
    if isinstance(message.root, JSONRPCRequest):
        data = message.root.model_dump(by_alias=True, mode="json", exclude_none=True)
        if data.get("method") == "initialize":
            caps = data.setdefault("params", {}).setdefault("capabilities", {})
            caps["extensions"] = {
                "io.modelcontextprotocol/ui": {"mimeTypes": ["text/html"]}
            }
            return JSONRPCMessage(JSONRPCRequest(**data))
    return message

session = ClientSession(
    read_stream,
    write_stream,
    send_middleware=[add_ui_extension],
)

Implementation

• Added MessageMiddleware type alias: Callable[[JSONRPCMessage], JSONRPCMessage | Awaitable[JSONRPCMessage]]
• Added send_middleware and receive_middleware parameters to BaseSession, ClientSession, and ServerSession
• Middleware is applied in send_request(), send_notification(), _send_response() for outgoing messages
• Middleware is applied in _receive_loop() for incoming messages
• Both sync and async middleware functions are supported

Testing

Added tests verifying:

• Sync middleware can transform outgoing messages
• Async middleware works correctly

[maxisbey]

move to top

[maxisbey]

should inspect the function first instead of doing it on every single message

[jerome3o-anthropic]

Addressed the review comments:

1. Moved import inspect to top of file
2. Now pre-computing whether each middleware is async using inspect.iscoroutinefunction() when middleware is registered, instead of checking isawaitable() on every message
3. Added test for receive_middleware to fix coverage

Thanks for the feedback!