mcp/public_server.py

import json
import logging
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from constants import DEVICE_INVALID_MESSAGE
from mcp_protocol import McpProtocolHandler
from services.request_context import RequestContextParser
from utils.rate_limiter import SimpleRateLimiter


logger = logging.getLogger(__name__)

def extract_client_ip(headers, client_address):
    if client_address and len(client_address) > 0:
        return str(client_address[0])
    return ''


class PublicMcpHttpHandler:
    def __init__(self, gateway_app, context_parser=None, rate_limiter=None):
        self.gateway_app = gateway_app
        self.context_parser = context_parser or RequestContextParser()
        self.rate_limiter = rate_limiter

    def handle_json_rpc(self, headers, message, client_ip=''):
        request_id = message.get('id') if isinstance(message, dict) else None
        method = str((message or {}).get('method') or '').strip()

        # Rate limiting by IP
        if self.rate_limiter and client_ip:
            if not self.rate_limiter.is_allowed(client_ip):
                logger.warning(f"[RATE_LIMIT] IP rate limit exceeded: ip={client_ip}, method={method}")
                return McpProtocolHandler._error_response(request_id, -32000, 'Rate limit exceeded. Please try again later.')

        try:
            if method == 'initialize':
                return McpProtocolHandler._success_response(request_id, {
                    'protocolVersion': McpProtocolHandler.protocol_version,
                    'capabilities': {'tools': {'listChanged': False}},
                    'serverInfo': {
                        'name': McpProtocolHandler.server_name,
                        'version': McpProtocolHandler.server_version,
                    },
                })
            if method == 'tools/list':
                tools = []
                for tool in self.gateway_app.list_tools():
                    normalized = dict(tool)
                    if 'input_schema' in normalized:
                        normalized['inputSchema'] = normalized.pop('input_schema')
                    tools.append(normalized)
                return McpProtocolHandler._success_response(request_id, {'tools': tools})
            if method == 'tools/call':
                context = self.context_parser.parse(headers or {})
                if not context.has_session():
                    raise RuntimeError(DEVICE_INVALID_MESSAGE)
                params = message.get('params') or {}
                result = self.gateway_app.call_tool(
                    context.gateway_session_id,
                    params.get('name'),
                    params.get('arguments') or {},
                    request_id='rq_http_{0}_{1}'.format(request_id, uuid.uuid4().hex[:16]),
                )
                structured_content = result.get('data') or {}
                return McpProtocolHandler._success_response(request_id, {
                    'content': [{'type': 'text', 'text': McpProtocolHandler._render_text(structured_content)}],
                    'structuredContent': structured_content,
                    'isError': False,
                })
            return McpProtocolHandler._error_response(request_id, -32601, 'Method not found: {0}'.format(method))
        except Exception as exc:
            if method == 'tools/call':
                return McpProtocolHandler._success_response(request_id, {
                    'content': [{'type': 'text', 'text': str(exc)}],
                    'isError': True,
                })
            return McpProtocolHandler._error_response(request_id, -32000, str(exc))


def create_http_handler(gateway_app, rate_limiter=None):
    rpc_handler = PublicMcpHttpHandler(gateway_app, rate_limiter=rate_limiter)

    class Handler(BaseHTTPRequestHandler):
        def log_message(self, format, *args):
            # Override to use Python logging instead of stderr
            logger.info(f"{self.address_string()} - {format % args}")

        def do_GET(self):
            if self.path == '/health':
                self._write_json({'ok': True})
                return
            self.send_response(404)
            self.end_headers()

        def do_POST(self):
            client_ip = extract_client_ip(dict(self.headers.items()), self.client_address)
            if self.path != '/mcp':
                logger.warning(f"[HTTP] 404: ip={client_ip}, path={self.path}")
                self.send_response(404)
                self.end_headers()
                return
            length = int(self.headers.get('Content-Length') or '0')
            body = self.rfile.read(length).decode('utf-8-sig')
            message = json.loads(body)
            method = message.get('method', '')
            request_id = message.get('id', '')
            logger.info(f"[HTTP] request: ip={client_ip}, method={method}, id={request_id}")
            response = rpc_handler.handle_json_rpc(dict(self.headers.items()), message, client_ip)
            self._write_json(response)

        def _write_json(self, payload):
            raw = json.dumps(payload, ensure_ascii=False).encode('utf-8')
            self.send_response(200)
            self.send_header('Content-Type', 'application/json; charset=utf-8')
            self.send_header('Content-Length', str(len(raw)))
            self.end_headers()
            self.wfile.write(raw)

    return Handler


def serve_public(gateway_app, host='0.0.0.0', port=8765, enable_rate_limit=True):
    # Configure logging
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )

    # Configure rate limiting (default: 60 requests per minute per IP)
    rate_limiter = None
    if enable_rate_limit:
        rate_limiter = SimpleRateLimiter(max_requests=60, window_seconds=60)
        logger.info("Rate limiting enabled: 60 requests/minute per IP")

    logger.info(f"Starting public MCP Gateway on {host}:{port}")
    server = ThreadingHTTPServer((host, int(port)), create_http_handler(gateway_app, rate_limiter))
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        logger.info("Shutting down public MCP Gateway")
        server.shutdown()