mirror of
https://github.com/mcp-use/mcp-use.git
synced 2026-09-14 18:06:01 +08:00
Pietro Zullo
a9bc4db3ae
* feat(python): add protocol handlers for logging, completions, subscriptions, and DNS protection Server SDK changes: - logging/setLevel: stores client level, Context.log() filters by RFC 5424 severity - completion/complete: default empty handler, overridable via mcp.completion() - resources/subscribe + unsubscribe: tracks per-session subscriptions - notify_resource_updated(): broadcasts to all subscribed sessions via WeakSet - dns_rebinding_protection: explicit constructor parameter (default: False) - Capabilities patch for subscribe=True Middleware: - Added hooks: on_set_logging_level, on_subscribe_resource, on_unsubscribe_resource, on_complete Tests: - 9 unit tests for log level filtering (RFC 5424) - Updated DNS protection tests for explicit flag Docs: - logging.mdx: MCP protocol logging section with level filtering - resources.mdx: subscription section with example - middleware.mdx: new hooks documented Refs: - https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/logging - https://modelcontextprotocol.io/specification/2025-11-25/server/resources#subscriptions * style: format adapter files for ruff 0.15
71 lines
3.0 KiB
Python
71 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import weakref
|
|
from datetime import datetime
|
|
|
|
import mcp.types as types
|
|
from mcp.server.session import ServerSession
|
|
from mcp.shared.session import RequestResponder
|
|
|
|
from mcp_use.server.middleware import MiddlewareManager, ServerMiddlewareContext
|
|
|
|
|
|
class MiddlewareServerSession(ServerSession):
|
|
"""Extended ServerSession that routes initialize requests through middleware.
|
|
|
|
This class intercepts the MCP initialize handshake at the protocol layer,
|
|
allowing middleware to:
|
|
- Validate client capabilities or metadata during connection
|
|
- Log connection attempts and protocol versions
|
|
- Reject connections early based on initialization parameters
|
|
|
|
It also tracks all active sessions so the server can broadcast notifications
|
|
(e.g., ``notifications/resources/updated``) to all connected clients.
|
|
|
|
The class is injected via monkey-patching of ``mcp.server.session.ServerSession``
|
|
to intercept session creation within the MCP SDK.
|
|
|
|
Note:
|
|
This implementation uses class attributes for middleware injection,
|
|
which means only ONE MCPServer instance is supported per process.
|
|
This is compatible with standard deployment patterns (e.g., uvicorn workers
|
|
run in separate processes), but multiple MCPServer instances in the
|
|
same process will share/override each other's middleware configuration.
|
|
|
|
Attributes:
|
|
_middleware_manager: The middleware manager to process initialize requests.
|
|
Injected by MCPServer during initialization.
|
|
_transport_type: The transport type (e.g., "stdio", "streamable-http").
|
|
Injected by MCPServer during initialization.
|
|
_active_sessions: Weak set of all active sessions for broadcasting notifications.
|
|
"""
|
|
|
|
_middleware_manager: MiddlewareManager | None = None
|
|
_transport_type: str = "unknown"
|
|
_active_sessions: weakref.WeakSet[MiddlewareServerSession] = weakref.WeakSet()
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
MiddlewareServerSession._active_sessions.add(self)
|
|
|
|
async def _received_request(self, responder: RequestResponder[types.ClientRequest, types.ServerResult]) -> None:
|
|
if responder.request.root.method and responder.request.root.method == "initialize":
|
|
if not self._middleware_manager:
|
|
# Fallback to normal behavior if middleware isn't injected yet
|
|
return await ServerSession._received_request(self, responder)
|
|
|
|
ctx = ServerMiddlewareContext(
|
|
message=responder.request.root.params,
|
|
method="initialize",
|
|
timestamp=datetime.now(),
|
|
transport=self._transport_type,
|
|
session_id=getattr(self, "session_id", None),
|
|
)
|
|
|
|
async def call_original(_):
|
|
return await ServerSession._received_request(self, responder)
|
|
|
|
return await self._middleware_manager.process_request(ctx, call_original)
|
|
|
|
return await ServerSession._received_request(self, responder)
|