"""Middleware MCP server implementation for MCP Chain."""
from typing import Any, Dict
from .types import DictMCPServer, MetadataTransformer, RequestResponseTransformer
class MiddlewareMCPServer:
"""An MCP Server that works with dict-based transformers and can chain other dict-based servers."""
def __init__(
self,
downstream_server: DictMCPServer,
metadata_transformer: MetadataTransformer = None,
request_transformer: RequestResponseTransformer = None,
):
self._downstream = downstream_server
self._metadata_transformer = metadata_transformer or (
lambda next_server, metadata_dict: next_server.get_metadata()
)
self._request_transformer = request_transformer or (
lambda next_server, request_dict: next_server.handle_request(request_dict)
)
def get_metadata(self) -> Dict[str, Any]:
"""Returns server metadata as dict."""
if self._downstream is None:
raise ValueError("No downstream server configured")
# Call metadata transformer with empty dict (metadata transformers don't use the metadata_dict parameter)
return self._metadata_transformer(self._downstream, {})
def handle_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
"""Handles dict request and returns dict response."""
if self._downstream is None:
raise ValueError("No downstream server configured")
# Call request transformer with the request dict
return self._request_transformer(self._downstream, request)
def then(self, *args):
"""Chain another transformer or server."""
if self._downstream is None:
raise ValueError("No downstream server configured")
# Check if downstream has a then method
if hasattr(self._downstream, "then"):
# Delegate to child's then method
child_result = self._downstream.then(*args)
# Create a new MiddlewareMCPServer that's a copy of ourselves
# but with the child_result as the new downstream
return MiddlewareMCPServer(
downstream_server=child_result,
metadata_transformer=self._metadata_transformer,
request_transformer=self._request_transformer,
)
else:
# Downstream doesn't have `then` method - this should be an error
# since we can't chain further on a terminal server
raise ValueError(
"Cannot chain on downstream server: it has no `then` method"
)
