mcp/public_server.py

import hashlib
import json
import logging
import socket
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__)

UPLOAD_CORS_HEADERS = (
    ('Access-Control-Allow-Origin', '*'),
    ('Access-Control-Allow-Methods', 'POST, OPTIONS'),
    ('Access-Control-Allow-Headers', 'Content-Type'),
    ('Access-Control-Max-Age', '600'),
)
UPLOAD_BODY_LIMIT = 102 * 1024 * 1024
UPLOAD_LIMIT_NAME = 'upload-headhaul-document'
UPLOAD_DISCARD_CHUNK = 64 * 1024
UPLOAD_DISCARD_IDLE_SECONDS = 0.5
# Exact path only. Prefix matching would swallow /mcp and the upload route.
STREAM_PATH = '/mcp/stream'
STREAM_KNOWN_PROTOCOL_VERSIONS = ('2025-03-26', '2025-06-18')
STREAM_DEFAULT_PROTOCOL_VERSION = '2025-03-26'


def stream_protocol_version(params):
    """Return the protocol version advertised on /mcp/stream.

    /mcp keeps its fixed version. This helper only rewrites the copy written
    for the cloud Streamable HTTP path.
    """
    if not isinstance(params, dict):
        return STREAM_DEFAULT_PROTOCOL_VERSION
    client_version = str(params.get('protocolVersion') or '').strip()
    if client_version in STREAM_KNOWN_PROTOCOL_VERSIONS:
        return client_version
    return STREAM_DEFAULT_PROTOCOL_VERSION


def discard_request_body(rfile, connection, length, chunk=UPLOAD_DISCARD_CHUNK,
                         idle_timeout=UPLOAD_DISCARD_IDLE_SECONDS):
    if length <= 0 or rfile is None:
        return
    remaining = length
    old_timeout = None
    if connection is not None:
        try:
            old_timeout = connection.gettimeout()
            connection.settimeout(idle_timeout)
        except OSError:
            old_timeout = None
    try:
        while remaining > 0:
            data = rfile.read(min(chunk, remaining))
            if not data:
                break
            remaining -= len(data)
    except (OSError, socket.timeout):
        pass
    finally:
        if connection is not None and old_timeout is not None:
            try:
                connection.settimeout(old_timeout)
            except OSError:
                pass

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

    @staticmethod
    def _upload_error(code, message, trace_request_id):
        return {
            'code': code,
            'msg': message,
            'data': {},
            'meta': {'request_id': trace_request_id},
        }

    def gate_upload(self, headers, content_length, client_ip='', trace_request_id=''):
        if (
            isinstance(content_length, bool)
            or not isinstance(content_length, int)
            or content_length < 0
        ):
            return 400, self._upload_error(
                'MCP_1401',
                '工具参数不正确，请检查后重试',
                trace_request_id,
            ), None
        if content_length > UPLOAD_BODY_LIMIT:
            return 413, self._upload_error(
                'MCP_1401',
                '文件不可大于100M',
                trace_request_id,
            ), None
        context = self.context_parser.parse(headers or {})
        session_id = context.gateway_session_id if context.has_session() else ''
        identity = session_id or ('ip:' + str(client_ip or 'unknown'))
        rate_key = '{0}:{1}'.format(identity, UPLOAD_LIMIT_NAME)
        if self.rate_limiter is not None:
            if not self.rate_limiter.is_allowed(rate_key):
                logger.warning(
                    "MCP public rate limit exceeded",
                    extra={
                        'request_id': trace_request_id,
                        'protocol_method': 'upload',
                        'tool_code': UPLOAD_LIMIT_NAME,
                        'client_identity': identity,
                        'diagnostic_reason': 'RATE_LIMIT_EXCEEDED',
                    },
                )
                return 429, self._upload_error(
                    'MCP_9001',
                    '请求过于频繁，请稍后重试',
                    trace_request_id,
                ), None
            if not self.rate_limiter.try_acquire(rate_key):
                logger.warning(
                    'MCP public concurrency limit exceeded',
                    extra={
                        'request_id': trace_request_id,
                        'protocol_method': 'upload',
                        'tool_code': UPLOAD_LIMIT_NAME,
                        'client_identity': identity,
                        'diagnostic_reason': 'CONCURRENCY_LIMIT_EXCEEDED',
                    },
                )
                return 429, self._upload_error(
                    'MCP_9001',
                    '请求过于频繁，请稍后重试',
                    trace_request_id,
                ), None
        return None, None, rate_key

    def handle_upload(
        self,
        headers,
        body,
        client_ip='',
        trace_request_id='',
    ):
        if len(body or b'') > UPLOAD_BODY_LIMIT:
            return 413, self._upload_error(
                'MCP_1401',
                '文件不可大于100M',
                trace_request_id,
            )
        fields, files = self._parse_multipart(headers, body or b'')
        upload_token = str(fields.get('upload_token') or '').strip()
        file_spec = files.get('file')
        context = self.context_parser.parse(headers or {})
        session_id = context.gateway_session_id if context.has_session() else ''
        if (
            not session_id
            and upload_token
            and hasattr(self.gateway_app, 'session_for_headhaul_token')
        ):
            session_id = str(
                self.gateway_app.session_for_headhaul_token(upload_token) or ''
            ).strip()
        if not session_id:
            return 401, {
                'code': 'MCP_1101',
                'msg': DEVICE_INVALID_MESSAGE,
                'data': {},
                'meta': {'request_id': trace_request_id},
            }
        if not upload_token or not file_spec:
            return 400, {
                'code': 'MCP_1401',
                'msg': '工具参数不正确，请检查后重试',
                'data': {},
                'meta': {'request_id': trace_request_id},
            }
        filename, file_bytes, content_type = file_spec
        try:
            payload = self.gateway_app.upload_headhaul_document(
                session_id,
                upload_token,
                filename,
                file_bytes,
                content_type,
                request_id=trace_request_id,
                client_ip=client_ip,
            )
        except RuntimeError as exc:
            message = str(exc)
            if message == DEVICE_INVALID_MESSAGE:
                return 401, {
                    'code': 'MCP_1101',
                    'msg': DEVICE_INVALID_MESSAGE,
                    'data': {},
                    'meta': {'request_id': trace_request_id},
                }
            if message.lower().startswith('tool disabled'):
                return 403, {
                    'code': 'MCP_1202',
                    'msg': '工具当前不可用',
                    'data': {},
                    'meta': {'request_id': trace_request_id},
                }
            return 500, {
                'code': 'MCP_9001',
                'msg': '系统繁忙，请稍后重试',
                'data': {},
                'meta': {'request_id': trace_request_id},
            }
        if not isinstance(payload, dict):
            return 500, {
                'code': 'MCP_9001',
                'msg': '系统繁忙，请稍后重试',
                'data': {},
                'meta': {'request_id': trace_request_id},
            }
        return 200, payload

    @staticmethod
    def _parse_multipart(headers, body):
        content_type = ''
        for key, value in (headers or {}).items():
            if str(key).lower() == 'content-type':
                content_type = str(value or '')
                break
        if 'multipart/form-data' not in content_type or 'boundary=' not in content_type:
            return {}, {}
        boundary = content_type.split('boundary=', 1)[1].strip().strip('"')
        delimiter = b'--' + boundary.encode('utf-8')
        fields = {}
        files = {}
        for raw_part in body.split(delimiter):
            part = raw_part.strip()
            if not part or part == b'--':
                continue
            header_blob, separator, content = part.partition(b'\r\n\r\n')
            if separator == b'':
                continue
            header_text = header_blob.decode('utf-8', 'replace')
            disposition = ''
            part_type = 'application/octet-stream'
            for line in header_text.split('\r\n'):
                lower = line.lower()
                if lower.startswith('content-disposition:'):
                    disposition = line.split(':', 1)[1].strip()
                if lower.startswith('content-type:'):
                    part_type = line.split(':', 1)[1].strip()
            name = ''
            filename = ''
            for item in disposition.split(';'):
                item = item.strip()
                if item.startswith('name='):
                    name = item.split('=', 1)[1].strip().strip('"')
                elif item.startswith('filename='):
                    filename = item.split('=', 1)[1].strip().strip('"')
            if not name:
                continue
            if filename:
                files[name] = (filename, content, part_type)
            else:
                fields[name] = content.decode('utf-8', 'replace')
        return fields, files

    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': McpProtocolHandler.server_capabilities(),
                    '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 in ('resources/list', 'resources/read'):
                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'},
                    )
                    return McpProtocolHandler._error_response(
                        request_id,
                        -32001,
                        DEVICE_INVALID_MESSAGE,
                        trace_request_id,
                    )
                if method == 'resources/list':
                    resources = []
                    if hasattr(self.gateway_app, 'list_resources'):
                        resources = list(self.gateway_app.list_resources() or [])
                    return McpProtocolHandler._success_response(
                        request_id,
                        {'resources': resources},
                    )
                uri = ''
                if isinstance(message_params, dict):
                    uri = str(message_params.get('uri') or '').strip()
                resource = None
                if uri and hasattr(self.gateway_app, 'read_resource'):
                    resource = self.gateway_app.read_resource(
                        uri,
                        context.gateway_session_id,
                    )
                if not isinstance(resource, dict):
                    return McpProtocolHandler._error_response(
                        request_id,
                        -32602,
                        'Invalid params',
                        trace_request_id,
                    )
                return McpProtocolHandler._success_response(
                    request_id,
                    {'contents': [resource]},
                )
            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,
                )
            if str(exc) == DEVICE_INVALID_MESSAGE:
                parsed = self.context_parser.parse(headers or {})
                emitter.emit(
                    stage='gateway_session',
                    status='failed',
                    event_code='GATEWAY_SESSION_NOT_FOUND',
                    session_credential=(
                        parsed.gateway_session_id if parsed.has_session() else None
                    ),
                    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,
                )
            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
            if self.path == STREAM_PATH:
                self._write_empty(405, 'POST')
                return
            self.send_response(404)
            self.end_headers()

        def do_DELETE(self):
            if self.path == STREAM_PATH:
                self._write_empty(405, 'POST')
                return
            self.send_error(501, "Unsupported method ('DELETE')")

        def do_OPTIONS(self):
            if self.path == '/mcp/upload-headhaul-document':
                self.send_response(204)
                for key, value in UPLOAD_CORS_HEADERS:
                    self.send_header(key, value)
                self.end_headers()
                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
            )
            if self.path == '/mcp/upload-headhaul-document':
                try:
                    length = int(self.headers.get('Content-Length') or '0')
                except (TypeError, ValueError):
                    length = -1
                status, payload, rate_key = rpc_handler.gate_upload(
                    request_headers,
                    length,
                    client_ip=client_ip,
                    trace_request_id=trace_request_id,
                )
                if status is not None:
                    if length > 0:
                        discard_request_body(self.rfile, self.connection, length)
                    self._write_upload_json(status, payload, close=True)
                    return
                try:
                    body = self.rfile.read(length) if length > 0 else b''
                    status, payload = rpc_handler.handle_upload(
                        request_headers,
                        body,
                        client_ip=client_ip,
                        trace_request_id=trace_request_id,
                    )
                    self._write_upload_json(status, payload)
                finally:
                    if rate_key and rpc_handler.rate_limiter is not None:
                        rpc_handler.rate_limiter.release(rate_key)
                return
            length = int(self.headers.get('Content-Length') or '0')
            if self.path == STREAM_PATH:
                self._post_stream(
                    request_headers,
                    length,
                    client_ip,
                    trace_request_id,
                    diagnostic_emitter,
                    client_request_id_hash,
                )
                return
            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_empty(self, status, allow=None):
            self.send_response(status)
            if allow:
                self.send_header('Allow', allow)
            self.send_header('Content-Length', '0')
            self.end_headers()

        def _post_stream(
            self,
            request_headers,
            length,
            client_ip,
            trace_request_id,
            diagnostic_emitter,
            client_request_id_hash,
        ):
            body = ''
            if length > 0:
                body = self.rfile.read(length).decode('utf-8-sig')
            try:
                message = json.loads(body)
            except json.JSONDecodeError:
                self._write_json(
                    McpProtocolHandler._error_response(
                        None,
                        -32700,
                        'Parse error',
                        trace_request_id,
                    ),
                    status=400,
                )
                return
            if not isinstance(message, dict):
                self._write_json(
                    McpProtocolHandler._error_response(
                        None,
                        -32600,
                        'Invalid Request',
                        trace_request_id,
                    ),
                    status=400,
                )
                return
            if 'id' not in message:
                self._write_empty(202)
                return
            method = str(message.get('method') or '').strip()
            if method == 'ping':
                self._write_json(
                    McpProtocolHandler._success_response(message.get('id'), {}),
                )
                return
            logger.info(
                'MCP public request',
                extra={
                    'request_id': trace_request_id,
                    'jsonrpc_id': message.get('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,
            )
            if method == 'initialize':
                response['result']['protocolVersion'] = stream_protocol_version(
                    message.get('params'),
                )
            self._write_json(response, diagnostic_emitter=diagnostic_emitter)

        def _write_upload_json(self, status, payload, close=False):
            encoded = json.dumps(payload, ensure_ascii=False).encode('utf-8')
            self.send_response(status)
            self.send_header('Content-Type', 'application/json; charset=utf-8')
            for key, value in UPLOAD_CORS_HEADERS:
                self.send_header(key, value)
            if close:
                self.send_header('Connection', 'close')
                self.close_connection = True
            self.send_header('Content-Length', str(len(encoded)))
            self.end_headers()
            self.wfile.write(encoded)

        def _write_json(self, payload, diagnostic_emitter=None, status=200):
            raw = json.dumps(payload, ensure_ascii=False).encode('utf-8')
            try:
                self.send_response(status)
                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': status,
                            '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': status,
                            '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 and upload HTTP)")

    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()