# -*- coding: utf-8 -*-
"""Federation Request Forwarding.
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti
This module implements request forwarding for federated MCP Gateways.
It handles:
- Request routing to appropriate gateways
- Response aggregation
- Error handling and retry logic
- Request/response transformation
"""
# Standard
import asyncio
from datetime import datetime, timezone
import logging
from typing import Any, Dict, List, Optional, Set, Tuple, Union
# Third-Party
import httpx
from sqlalchemy import select
from sqlalchemy.orm import Session
# First-Party
from mcpgateway.config import settings
from mcpgateway.db import Gateway as DbGateway
from mcpgateway.db import Tool as DbTool
from mcpgateway.models import ToolResult
logger = logging.getLogger(__name__)
class ForwardingError(Exception):
"""Base class for forwarding-related errors."""
class ForwardingService:
"""Service for handling request forwarding across gateways.
Handles:
- Request routing
- Response aggregation
- Error handling
- Request transformation
"""
def __init__(self):
"""Initialize forwarding service."""
self._http_client = httpx.AsyncClient(timeout=settings.federation_timeout, verify=not settings.skip_ssl_verify)
# Track active requests
self._active_requests: Dict[str, asyncio.Task] = {}
# Request history for rate limiting
self._request_history: Dict[str, List[datetime]] = {}
# Cache gateway information
self._gateway_tools: Dict[int, Set[str]] = {}
async def start(self) -> None:
"""Start forwarding service."""
logger.info("Request forwarding service started")
async def stop(self) -> None:
"""Stop forwarding service."""
# Cancel active requests
for request_id, task in self._active_requests.items():
logger.info(f"Cancelling request {request_id}")
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
await self._http_client.aclose()
logger.info("Request forwarding service stopped")
async def forward_request(
self,
db: Session,
method: str,
params: Optional[Dict[str, Any]] = None,
target_gateway_id: Optional[int] = None,
) -> Any:
"""Forward a request to gateway(s).
Args:
db: Database session
method: RPC method name
params: Optional method parameters
target_gateway_id: Optional specific gateway ID
Returns:
Forwarded response(s)
Raises:
ForwardingError: If forwarding fails
"""
try:
if target_gateway_id:
# Forward to specific gateway
return await self._forward_to_gateway(db, target_gateway_id, method, params)
# Forward to all relevant gateways
return await self._forward_to_all(db, method, params)
except Exception as e:
raise ForwardingError(f"Forward request failed: {str(e)}")
async def forward_tool_request(self, db: Session, tool_name: str, arguments: Dict[str, Any]) -> ToolResult:
"""Forward a tool invocation request.
Args:
db: Database session
tool_name: Tool to invoke
arguments: Tool arguments
Returns:
Tool result
Raises:
ForwardingError: If forwarding fails
"""
try:
# Find tool
tool = db.execute(select(DbTool).where(DbTool.name == tool_name).where(DbTool.enabled)).scalar_one_or_none()
if not tool:
raise ForwardingError(f"Tool not found: {tool_name}")
if not tool.gateway_id:
raise ForwardingError(f"Tool {tool_name} is not federated")
# Forward to gateway
result = await self._forward_to_gateway(
db,
tool.gateway_id,
"tools/invoke",
{"name": tool_name, "arguments": arguments},
)
# Parse result
return ToolResult(
content=result.get("content", []),
is_error=result.get("is_error", False),
)
except Exception as e:
raise ForwardingError(f"Failed to forward tool request: {str(e)}")
async def forward_resource_request(self, db: Session, uri: str) -> Tuple[Union[str, bytes], str]:
"""Forward a resource read request.
Args:
db: Database session
uri: Resource URI
Returns:
Tuple of (content, mime_type)
Raises:
ForwardingError: If forwarding fails
"""
try:
# Find gateway for resource
gateway = await self._find_resource_gateway(db, uri)
if not gateway:
raise ForwardingError(f"No gateway found for resource: {uri}")
# Forward request
result = await self._forward_to_gateway(db, gateway.id, "resources/read", {"uri": uri})
# Parse result
if "text" in result:
return result["text"], result.get("mime_type", "text/plain")
if "blob" in result:
return result["blob"], result.get("mime_type", "application/octet-stream")
raise ForwardingError("Invalid resource response format")
except Exception as e:
raise ForwardingError(f"Failed to forward resource request: {str(e)}")
async def _forward_to_gateway(
self,
db: Session,
gateway_id: str,
method: str,
params: Optional[Dict[str, Any]] = None,
) -> Any:
"""Forward request to a specific gateway.
Args:
db: Database session
gateway_id: Gateway to forward to
method: RPC method name
params: Optional method parameters
Returns:
Gateway response
Raises:
ForwardingError: If forwarding fails
httpx.TimeoutException: If unable to connect after retries
"""
# Get gateway
gateway = db.get(DbGateway, gateway_id)
if not gateway or not gateway.enabled:
raise ForwardingError(f"Gateway not found: {gateway_id}")
# Check rate limits
if not self._check_rate_limit(gateway.url):
raise ForwardingError("Rate limit exceeded")
try:
# Build request
request = {"jsonrpc": "2.0", "id": 1, "method": method}
if params:
request["params"] = params
# Send request with retries using the persistent client directly
for attempt in range(settings.max_tool_retries):
try:
response = await self._http_client.post(
f"{gateway.url}/rpc",
json=request,
headers=self._get_auth_headers(),
)
response.raise_for_status()
result = response.json()
# Update last seen
gateway.last_seen = datetime.now(timezone.utc)
# Handle response
if "error" in result:
raise ForwardingError(f"Gateway error: {result['error'].get('message')}")
return result.get("result")
except httpx.TimeoutException:
if attempt == settings.max_tool_retries - 1:
raise
await asyncio.sleep(1 * (attempt + 1))
except Exception as e:
raise ForwardingError(f"Failed to forward to {gateway.name}: {str(e)}")
async def _forward_to_all(self, db: Session, method: str, params: Optional[Dict[str, Any]] = None) -> List[Any]:
"""Forward request to all active gateways.
Args:
db: Database session
method: RPC method name
params: Optional method parameters
Returns:
List of responses
Raises:
ForwardingError: If all forwards fail
"""
# Get active gateways
gateways = db.execute(select(DbGateway).where(DbGateway.enabled)).scalars().all()
# Forward to each gateway
results = []
errors = []
for gateway in gateways:
try:
result = await self._forward_to_gateway(db, gateway.id, method, params)
results.append(result)
except Exception as e:
errors.append(str(e))
if not results and errors:
raise ForwardingError(f"All forwards failed: {'; '.join(errors)}")
return results
async def _find_resource_gateway(self, db: Session, uri: str) -> Optional[DbGateway]:
"""Find gateway hosting a resource.
Args:
db: Database session
uri: Resource URI
Returns:
Gateway record or None
"""
# Get active gateways
gateways = db.execute(select(DbGateway).where(DbGateway.enabled)).scalars().all()
# Check each gateway
for gateway in gateways:
try:
resources = await self._forward_to_gateway(db, gateway.id, "resources/list")
for resource in resources:
if resource.get("uri") == uri:
return gateway
except Exception as e:
logger.error(f"Failed to check gateway {gateway.name} for resource {uri}: {str(e)}")
continue
return None
def _check_rate_limit(self, gateway_url: str) -> bool:
"""Check if gateway request is within rate limits.
Args:
gateway_url: Gateway URL
Returns:
True if request allowed
"""
now = datetime.now(timezone.utc)
# Clean old history
self._request_history[gateway_url] = [t for t in self._request_history.get(gateway_url, []) if (now - t).total_seconds() < 60]
# Check limit
if len(self._request_history[gateway_url]) >= settings.tool_rate_limit:
return False
# Record request
self._request_history[gateway_url].append(now)
return True
def _get_auth_headers(self) -> Dict[str, str]:
"""
Get headers for gateway authentication.
Returns:
dict: Authorization header dict
"""
api_key = f"{settings.basic_auth_user}:{settings.basic_auth_password}"
return {"Authorization": f"Basic {api_key}", "X-API-Key": api_key}
