mcp/services/diagnostic_reporter.py

import hashlib
import hmac
import json
import logging
import queue
import threading
import time
import urllib.request
from urllib.parse import urlsplit
import uuid


logger = logging.getLogger(__name__)


def diagnostic_reporter_from_config(config, **overrides):
    if getattr(config, 'diagnosis_enabled', False) is not True:
        return NullDiagnosticReporter()
    url = str(getattr(config, 'diagnosis_url', '') or '').strip()
    key_id = str(getattr(config, 'diagnosis_key_id', '') or '').strip()
    secret = str(getattr(config, 'diagnosis_secret', '') or '')
    allow_insecure_http = getattr(
        config, 'diagnosis_allow_insecure_http', False
    ) is True
    parsed_url = urlsplit(url)
    if (
        (
            parsed_url.scheme != 'https'
            and not (parsed_url.scheme == 'http' and allow_insecure_http)
        )
        or not parsed_url.netloc
        or not key_id
        or len(secret) < 32
    ):
        logger.warning('MCP diagnostic reporter configuration is invalid')
        return NullDiagnosticReporter()
    options = {
        'url': url,
        'key_id': key_id,
        'secret': secret,
        'queue_size': getattr(config, 'diagnosis_queue_size', 1000),
        'batch_size': getattr(config, 'diagnosis_batch_size', 100),
        'timeout_seconds': getattr(config, 'diagnosis_timeout_seconds', 0.5),
        'initial_backoff': getattr(
            config, 'diagnosis_initial_backoff_seconds', 0.25
        ),
        'max_backoff': getattr(
            config, 'diagnosis_max_backoff_seconds', 5.0
        ),
    }
    options.update(overrides)
    reporter = DiagnosticReporter(**options)
    if reporter.start():
        return reporter
    reporter.close()
    return NullDiagnosticReporter()


def _http_transport(url, body, headers, timeout):
    request = urllib.request.Request(url, data=body, headers=headers, method='POST')
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return response.getcode(), response.read()


class NullDiagnosticReporter:
    drop_count = 0
    pending_count = 0

    def start(self):
        return True

    def report(self, _event):
        return True

    def close(self):
        return None


class DiagnosticReporter:
    def __init__(
        self,
        url,
        key_id,
        secret,
        queue_size=1000,
        batch_size=100,
        timeout_seconds=0.5,
        initial_backoff=0.25,
        max_backoff=5.0,
        transport=None,
        clock=None,
        nonce_factory=None,
        sleeper=None,
        thread_factory=None,
    ):
        self.url = str(url)
        self.key_id = str(key_id)
        self.secret = str(secret)
        self.batch_size = max(1, min(100, int(batch_size)))
        self.timeout_seconds = float(timeout_seconds)
        self.initial_backoff = max(0.0, float(initial_backoff))
        self.max_backoff = max(self.initial_backoff, float(max_backoff))
        self.transport = transport or _http_transport
        self.clock = clock or time.time
        self.nonce_factory = nonce_factory or (lambda: uuid.uuid4().hex)
        self.sleeper = sleeper or time.sleep
        self.thread_factory = thread_factory or threading.Thread
        self._queue = queue.Queue(maxsize=max(1, int(queue_size)))
        self._pending = []
        self._backoff = self.initial_backoff
        self._stop = threading.Event()
        self._thread = None
        self.drop_count = 0

    @property
    def pending_count(self):
        return len(self._pending)

    def report(self, event):
        try:
            self._queue.put_nowait(event)
            return True
        except queue.Full:
            self.drop_count += 1
            return False

    def start(self):
        if self._thread is not None:
            return True
        try:
            thread = self.thread_factory(target=self._run, daemon=True)
            thread.start()
            self._thread = thread
            return True
        except Exception:
            logger.warning('MCP diagnostic reporter thread unavailable')
            return False

    def close(self):
        self._stop.set()
        thread = self._thread
        if thread is not None and thread.is_alive():
            thread.join(timeout=self.timeout_seconds)

    def process_once(self):
        if not self._pending:
            self._pending = self._take_batch()
        if not self._pending:
            return True

        body = json.dumps(
            {'events': self._pending},
            ensure_ascii=False,
            separators=(',', ':'),
        ).encode('utf-8')
        headers = self._signed_headers(body)
        try:
            status, response_body = self.transport(
                self.url,
                body,
                headers,
                self.timeout_seconds,
            )
            payload = json.loads(response_body.decode('utf-8'))
            if status < 200 or status >= 300 or payload.get('code') != 'MCP_DIAG_INGEST_0000':
                raise ValueError('collector rejected batch')
        except Exception:
            delay = self._backoff
            self._backoff = min(
                self.max_backoff,
                max(self.initial_backoff, self._backoff * 2),
            )
            self.sleeper(delay)
            return False

        self._pending = []
        self._backoff = self.initial_backoff
        return True

    def _take_batch(self):
        events = []
        while len(events) < self.batch_size:
            try:
                events.append(self._queue.get_nowait())
            except queue.Empty:
                break
        return events

    def _signed_headers(self, body):
        timestamp = str(int(self.clock()))
        nonce = str(self.nonce_factory())
        path = urlsplit(self.url).path or '/'
        canonical = 'POST\n{0}\n{1}\n{2}\n{3}'.format(
            path,
            timestamp,
            nonce,
            hashlib.sha256(body).hexdigest(),
        )
        signature = hmac.new(
            self.secret.encode('utf-8'),
            canonical.encode('utf-8'),
            hashlib.sha256,
        ).hexdigest()
        return {
            'Content-Type': 'application/json',
            'X-MCP-Source': 'gateway',
            'X-MCP-Timestamp': timestamp,
            'X-MCP-Nonce': nonce,
            'X-MCP-Key-Id': self.key_id,
            'X-MCP-Signature': signature,
        }

    def _run(self):
        while not self._stop.is_set():
            if not self.process_once():
                continue
            self._stop.wait(0.1)