mcp/public_server.py

import hashlib
import json
import logging
import threading
import time
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from constants import DEVICE_INVALID_MESSAGE
from mcp_protocol import McpProtocolHandler
from services.diagnostic_event import RequestDiagnosticEmitter
from services.diagnostic_reporter import NullDiagnosticReporter
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,
        reporter=None,
    ):
        self.gateway_app = gateway_app
        self.context_parser = context_parser or RequestContextParser()
        self.rate_limiter = rate_limiter
        self.reporter = reporter or NullDiagnosticReporter()
        # Cache registered tool names at startup for rate-key validation;
        # unknown names fall back to the bare IP bucket, preventing bucket explosion
        self._known_tools = frozenset(gateway_app.registered_tool_names())

    @staticmethod
    def _build_trace_request_id(headers):
        return 'rq_http_{0}'.format(uuid.uuid4().hex[:16])

    @staticmethod
    def _client_request_id(headers):
        incoming = ''
        for key, value in (headers or {}).items():
            if str(key).lower() == 'x-request-id':
                incoming = str(value or '').strip()
                break
        if not incoming:
            return ''
        return hashlib.sha256(incoming.encode('utf-8')).hexdigest()[:16]

    def _check_rate_limit(
        self,
        rate_key,
        method,
        trace_request_id,
        jsonrpc_id=None,
        tool_name='',
        *,
        log_identity=None,
        diagnostic_emitter=None,
    ):
        """Returns an error response if rate limit exceeded, else None.

        rate_key     — the session-based bucket key used for tools/call
        log_identity — optional string shown in warning logs (e.g. client_ip)
        """
        if self.rate_limiter and rate_key and not self.rate_limiter.is_allowed(rate_key):
            if diagnostic_emitter is not None:
                diagnostic_emitter.emit(
                    stage='rate_limit',
                    status='failed',
                    event_code='RATE_LIMIT_EXCEEDED',
                    tool_code=tool_name or None,
                    context={
                        'limit_type': 'request_window',
                        'transport': 'http',
                    },
                )
            logger.warning(
                "MCP public rate limit exceeded",
                extra={
                    'request_id': trace_request_id,
                    'jsonrpc_id': jsonrpc_id,
                    'protocol_method': method,
                    'tool_code': tool_name,
                    'client_identity': log_identity or rate_key,
                    'protocol_code': -32029,
                    'diagnostic_reason': 'RATE_LIMIT_EXCEEDED',
                },
            )
            return McpProtocolHandler._error_response(
                jsonrpc_id,
                -32029,
                'Rate limit exceeded. Please try again later.',
                trace_request_id,
            )
        return None

    def handle_json_rpc(
        self,
        headers,
        message,
        client_ip='',
        trace_request_id='',
        diagnostic_emitter=None,
    ):
        request_id = message.get('id') if isinstance(message, dict) else None
        trace_request_id = str(trace_request_id or '').strip() \
            or self._build_trace_request_id(headers)
        method = str((message or {}).get('method') or '').strip()
        message_params = (message or {}).get('params') or {}
        tool_name = str(message_params.get('name') or '').strip() \
            if isinstance(message_params, dict) else ''
        emitter = diagnostic_emitter or RequestDiagnosticEmitter(
            self.reporter,
            trace_request_id,
            defer_until_identity=True,
        )
        emitter.emit(
            stage='request_ingress',
            status='started',
            event_code='REQUEST_RECEIVED',
            context={'jsonrpc_method': method or 'unknown', 'transport': 'http'},
        )
        protocol_validated = False

        try:
            if method == 'initialize':
                emitter.emit(
                    stage='protocol_validation',
                    status='succeeded',
                    event_code='PROTOCOL_VALIDATION_COMPLETED',
                    context={
                        'jsonrpc_method': method,
                        'transport': 'http',
                    },
                )
                protocol_validated = True
                # initialize is a stateless handshake that only returns server metadata;
                # rate-limiting it would block clients from connecting at all, so we skip it.
                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':
                emitter.emit(
                    stage='protocol_validation',
                    status='succeeded',
                    event_code='PROTOCOL_VALIDATION_COMPLETED',
                    context={
                        'jsonrpc_method': method,
                        'transport': 'http',
                    },
                )
                protocol_validated = True
                context = self.context_parser.parse(headers or {})
                if not context.has_session():
                    emitter.emit(
                        stage='gateway_session',
                        status='failed',
                        event_code='GATEWAY_SESSION_NOT_FOUND',
                        context={'transport': 'http'},
                    )
                    logger.warning(
                        'MCP public device session unavailable',
                        extra={
                            'request_id': trace_request_id,
                            'jsonrpc_id': request_id,
                            'protocol_method': method,
                            'tool_code': '',
                            'protocol_code': -32001,
                            'diagnostic_reason': 'GATEWAY_SESSION_NOT_FOUND',
                        },
                    )
                    return McpProtocolHandler._error_response(
                        request_id,
                        -32001,
                        DEVICE_INVALID_MESSAGE,
                        trace_request_id,
                    )
                tools = []
                for tool in self.gateway_app.list_tools(
                    context.gateway_session_id,
                    request_id=trace_request_id,
                ):
                    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':
                if not isinstance(message_params, dict):
                    raise ValueError('tool parameters must be an object')
                if not tool_name:
                    emitter.emit(
                        stage='protocol_validation',
                        status='failed',
                        event_code='PARAM_VALIDATION_FAILED',
                        context={
                            'jsonrpc_method': method,
                            'jsonrpc_code': -32602,
                            'transport': 'http',
                        },
                    )
                    return McpProtocolHandler._error_response(
                        request_id,
                        -32602,
                        'Invalid params',
                        trace_request_id,
                    )
                emitter.emit(
                    stage='protocol_validation',
                    status='succeeded',
                    event_code='PROTOCOL_VALIDATION_COMPLETED',
                    tool_code=tool_name,
                    context={
                        'jsonrpc_method': method,
                        'transport': 'http',
                    },
                )
                protocol_validated = True
                # Parse context first — tools/call always requires a valid session
                context = self.context_parser.parse(headers or {})
                if not context.has_session():
                    emitter.emit(
                        stage='gateway_session',
                        status='failed',
                        event_code='GATEWAY_SESSION_NOT_FOUND',
                        context={'transport': 'http'},
                    )
                    logger.warning(
                        'MCP public device session unavailable',
                        extra={
                            'request_id': trace_request_id,
                            'jsonrpc_id': request_id,
                            'protocol_method': method,
                            'tool_code': tool_name,
                            'protocol_code': -32001,
                            'diagnostic_reason': 'GATEWAY_SESSION_NOT_FOUND',
                        },
                    )
                    return McpProtocolHandler._error_response(
                        request_id,
                        -32001,
                        DEVICE_INVALID_MESSAGE,
                        trace_request_id,
                    )
                # Session-based rate limiting: each employee gets an independent quota per tool.
                # Unknown tool names fall back to the bare session bucket to prevent key explosion.
                session_id = context.gateway_session_id
                rate_key = '{0}:{1}'.format(session_id, tool_name) if tool_name in self._known_tools else session_id
                blocked = self._check_rate_limit(
                    rate_key,
                    method,
                    trace_request_id,
                    request_id,
                    tool_name,
                    log_identity=client_ip,
                    diagnostic_emitter=emitter,
                )
                if blocked:
                    return blocked
                params = message_params
                acquired = self.rate_limiter is None \
                    or self.rate_limiter.try_acquire(rate_key)
                if not acquired:
                    emitter.emit(
                        stage='rate_limit',
                        status='failed',
                        event_code='CONCURRENCY_LIMIT_EXCEEDED',
                        tool_code=tool_name or None,
                        context={
                            'limit_type': 'concurrency',
                            'transport': 'http',
                        },
                    )
                    logger.warning(
                        'MCP public concurrency limit exceeded',
                        extra={
                            'request_id': trace_request_id,
                            'jsonrpc_id': request_id,
                            'protocol_method': method,
                            'tool_code': tool_name,
                            'client_identity': client_ip,
                            'protocol_code': -32029,
                            'diagnostic_reason': 'CONCURRENCY_LIMIT_EXCEEDED',
                        },
                    )
                    return McpProtocolHandler._error_response(
                        request_id,
                        -32029,
                        'Too many requests in progress. Please try again later.',
                        trace_request_id,
                    )
                emitter.emit(
                    stage='rate_limit',
                    status='succeeded',
                    event_code='RATE_LIMIT_ALLOWED',
                    tool_code=tool_name or None,
                    context={
                        'limit_type': 'tools_call',
                        'transport': 'http',
                    },
                )
                try:
                    result = self.gateway_app.call_tool(
                        session_id,
                        params.get('name'),
                        params.get('arguments') or {},
                        request_id=trace_request_id,
                        client_ip=client_ip,
                        diagnostic_emitter=emitter,
                    )
                finally:
                    if self.rate_limiter is not None:
                        self.rate_limiter.release(rate_key)
                try:
                    response = McpProtocolHandler._tool_call_response(
                        request_id,
                        tool_name,
                        result,
                    )
                except Exception:
                    emitter.emit(
                        stage='response_safety',
                        status='failed',
                        event_code='RESPONSE_SAFETY_REJECTED',
                        context={'transport': 'http'},
                    )
                    raise
                emitter.emit(
                    stage='response_safety',
                    status='succeeded',
                    event_code='RESPONSE_SAFETY_COMPLETED',
                    context={'transport': 'http'},
                )
                return response
            emitter.emit(
                stage='protocol_validation',
                status='failed',
                event_code='METHOD_NOT_FOUND',
                context={
                    'jsonrpc_method': method or 'unknown',
                    'jsonrpc_code': -32601,
                    'transport': 'http',
                },
            )
            return McpProtocolHandler._error_response(
                request_id,
                -32601,
                'Method not found: {0}'.format(method),
                trace_request_id,
            )
        except Exception as exc:
            if not protocol_validated:
                emitter.emit(
                    stage='protocol_validation',
                    status='failed',
                    event_code='PARAM_VALIDATION_FAILED',
                    tool_code=tool_name or None,
                    context={
                        'jsonrpc_method': method or 'unknown',
                        'jsonrpc_code': -32602,
                        'transport': 'http',
                    },
                )
            if method == 'tools/call':
                diagnostic_reason = (
                    'PARAM_VALIDATION_FAILED'
                    if isinstance(exc, ValueError)
                    else 'UNEXPECTED_EXCEPTION'
                )
                logger.error(
                    'MCP public tool request failed',
                    extra={
                        'request_id': trace_request_id,
                        'jsonrpc_id': request_id,
                        'protocol_method': method,
                        'tool_code': tool_name,
                        'response_code': 'MCP_9001',
                        'diagnostic_reason': diagnostic_reason,
                        'exception_class': exc.__class__.__name__,
                    },
                )
                return McpProtocolHandler._tool_exception_response(
                    request_id,
                    tool_name,
                    exc,
                    trace_request_id,
                )
            logger.error(
                "MCP public request failed",
                extra={
                    'request_id': trace_request_id,
                    'jsonrpc_id': request_id,
                    'protocol_method': method,
                    'tool_code': tool_name,
                    'protocol_code': -32000,
                    'diagnostic_reason': 'UNEXPECTED_EXCEPTION',
                    'exception_class': exc.__class__.__name__,
                },
            )
            return McpProtocolHandler._error_response(
                request_id,
                -32000,
                'Gateway request failed. Please try again later.',
                trace_request_id,
            )
        finally:
            emitter.flush()


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

    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):
            request_headers = dict(self.headers.items())
            client_ip = extract_client_ip(request_headers, self.client_address)
            trace_request_id = rpc_handler._build_trace_request_id(request_headers)
            diagnostic_emitter = RequestDiagnosticEmitter(
                rpc_handler.reporter,
                trace_request_id,
                defer_until_identity=True,
            )
            client_request_id_hash = rpc_handler._client_request_id(
                request_headers
            )
            length = int(self.headers.get('Content-Length') or '0')
            if self.path != '/mcp':
                logger.warning(f"[HTTP] 404: ip={client_ip}, path={self.path}")
                if length > 0:
                    self.rfile.read(length)
                self.send_response(404)
                self.end_headers()
                return
            body = self.rfile.read(length).decode('utf-8-sig')
            try:
                message = json.loads(body)
            except json.JSONDecodeError as exc:
                diagnostic_emitter.emit(
                    stage='request_ingress',
                    status='started',
                    event_code='REQUEST_RECEIVED',
                    context={
                        'jsonrpc_method': 'unknown',
                        'transport': 'http',
                    },
                )
                diagnostic_emitter.emit(
                    stage='protocol_validation',
                    status='failed',
                    event_code='PARAM_VALIDATION_FAILED',
                    context={
                        'jsonrpc_code': -32700,
                        'transport': 'http',
                    },
                )
                logger.warning(
                    'MCP public invalid JSON',
                    extra={
                        'request_id': trace_request_id,
                        'jsonrpc_id': None,
                        'protocol_method': '',
                        'tool_code': '',
                        'protocol_code': -32700,
                        'diagnostic_reason': 'PARAM_VALIDATION_FAILED',
                        'exception_class': exc.__class__.__name__,
                        'client_identity': client_ip,
                    },
                )
                self._write_json(
                    McpProtocolHandler._error_response(
                        None,
                        -32700,
                        'Parse error',
                        trace_request_id,
                    ),
                    diagnostic_emitter=diagnostic_emitter,
                )
                return
            method = message.get('method', '')
            request_id = message.get('id') if message.get('id') is not None else ''
            logger.info(
                'MCP public request',
                extra={
                    'request_id': trace_request_id,
                    'jsonrpc_id': request_id,
                    'protocol_method': method,
                    'client_identity': client_ip,
                    'client_request_id_hash': client_request_id_hash,
                },
            )
            response = rpc_handler.handle_json_rpc(
                request_headers,
                message,
                client_ip,
                trace_request_id=trace_request_id,
                diagnostic_emitter=diagnostic_emitter,
            )
            self._write_json(
                response,
                diagnostic_emitter=diagnostic_emitter,
            )

        def _write_json(self, payload, diagnostic_emitter=None):
            raw = json.dumps(payload, ensure_ascii=False).encode('utf-8')
            try:
                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)
                if diagnostic_emitter is not None:
                    diagnostic_emitter.emit(
                        stage='response_write',
                        status='succeeded',
                        event_code='RESPONSE_WRITE_COMPLETED',
                        context={
                            'http_status': 200,
                            'client_disconnected': False,
                            'transport': 'http',
                        },
                    )
            except (BrokenPipeError, ConnectionResetError):
                self.close_connection = True
                if diagnostic_emitter is not None:
                    diagnostic_emitter.emit(
                        stage='response_write',
                        status='failed',
                        event_code='CLIENT_DISCONNECTED',
                        context={
                            'http_status': 200,
                            'client_disconnected': True,
                            'transport': 'http',
                        },
                    )
                logger.info('MCP client disconnected before response was written')
            finally:
                if diagnostic_emitter is not None:
                    diagnostic_emitter.flush()

    return Handler


def serve_public(gateway_app, host='0.0.0.0', port=8765, enable_rate_limit=True,
                 rate_limit_max_requests=60, rate_limit_window_seconds=60,
                 max_in_flight_per_tool=2, reporter=None):
    # 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
    rate_limiter = None
    if enable_rate_limit:
        rate_limiter = SimpleRateLimiter(
            max_requests=rate_limit_max_requests,
            window_seconds=rate_limit_window_seconds,
            max_in_flight=max_in_flight_per_tool,
        )
        logger.info(f"Rate limiting enabled: {rate_limit_max_requests} requests/{rate_limit_window_seconds}s and {max_in_flight_per_tool} in-flight per session and tool (tools/call only)")

    logger.info(f"Starting public MCP Gateway on {host}:{port}")
    server = ThreadingHTTPServer(
        (host, int(port)),
        create_http_handler(gateway_app, rate_limiter, reporter=reporter),
    )

    # Schedule periodic cleanup to prevent unbounded memory growth in the rate limiter
    if rate_limiter is not None:
        def _cleanup_loop():
            while True:
                time.sleep(300)
                # Use the actual window as max_age to avoid deleting entries still within the window
                rate_limiter.cleanup(max_age_seconds=rate_limit_window_seconds)
        t = threading.Thread(target=_cleanup_loop, daemon=True)
        t.start()

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        logger.info("Shutting down public MCP Gateway")
        server.shutdown()