mcp/services/diagnostic_event.py

import hashlib
import re
import uuid
from datetime import datetime, timezone


_STAGE_CONTEXT_FIELDS = {
    'request_ingress': frozenset(('jsonrpc_method', 'http_status', 'transport')),
    'protocol_validation': frozenset(('jsonrpc_method', 'jsonrpc_code', 'transport')),
    'gateway_session': frozenset(('transport',)),
    'rate_limit': frozenset(('limit_type', 'transport')),
    'backend_call': frozenset(('http_status', 'transport')),
    'response_safety': frozenset(('transport',)),
    'response_write': frozenset(('http_status', 'client_disconnected', 'transport')),
}
_STATUSES = frozenset(('started', 'succeeded', 'failed', 'skipped'))
_CODE_PATTERN = re.compile(r'^[A-Z][A-Z0-9_]*$')


def _utc_timestamp():
    return datetime.now(timezone.utc).isoformat(timespec='milliseconds').replace(
        '+00:00', 'Z'
    )


def _optional_positive_integer(event, name, value):
    if value is None:
        return
    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError('{0} must be a positive integer'.format(name))
    event[name] = value


def _optional_code(event, name, value, max_length):
    if value is None or value == '':
        return
    if (
        not isinstance(value, str)
        or len(value) > max_length
        or _CODE_PATTERN.fullmatch(value) is None
    ):
        raise ValueError('{0} is invalid'.format(name))
    event[name] = value


def _validated_context(stage, context):
    context = {} if context is None else context
    if not isinstance(context, dict):
        raise ValueError('context must be an object')
    unknown = set(context) - _STAGE_CONTEXT_FIELDS[stage]
    if unknown:
        raise ValueError('context contains unsupported fields')

    result = {}
    for key, value in context.items():
        if key == 'http_status':
            if isinstance(value, bool) or not isinstance(value, int) or not 100 <= value <= 599:
                raise ValueError('http_status is invalid')
        elif key == 'jsonrpc_code':
            if isinstance(value, bool) or not isinstance(value, int):
                raise ValueError('jsonrpc_code is invalid')
        elif key == 'client_disconnected':
            if not isinstance(value, bool):
                raise ValueError('client_disconnected is invalid')
        elif not isinstance(value, str) or not value or len(value) > 64:
            raise ValueError('{0} is invalid'.format(key))
        result[key] = value
    return result


def build_diagnostic_event(
    request_id,
    stage,
    status,
    event_code,
    occurred_at=None,
    session_credential='',
    company_id=None,
    admin_id=None,
    tool_code=None,
    response_code=None,
    cost_ms=None,
    summary_code=None,
    context=None,
):
    if not isinstance(request_id, str) or not re.fullmatch(
        r'rq_[A-Za-z0-9._:-]{1,97}', request_id
    ):
        raise ValueError('request_id is invalid')
    if stage not in _STAGE_CONTEXT_FIELDS:
        raise ValueError('stage is invalid')
    if status not in _STATUSES:
        raise ValueError('status is invalid')
    if (
        not isinstance(event_code, str)
        or len(event_code) > 64
        or _CODE_PATTERN.fullmatch(event_code) is None
    ):
        raise ValueError('event_code is invalid')

    event = {
        'schema_version': 1,
        'event_id': 'evt_gateway_{0}'.format(uuid.uuid4().hex),
        'request_id': request_id,
        'source': 'gateway',
        'stage': stage,
        'status': status,
        'event_code': event_code,
        'occurred_at': occurred_at or _utc_timestamp(),
    }
    _optional_positive_integer(event, 'company_id', company_id)
    _optional_positive_integer(event, 'admin_id', admin_id)
    if tool_code is not None and tool_code != '':
        if not isinstance(tool_code, str) or re.fullmatch(r'[a-z][a-z0-9_]{0,63}', tool_code) is None:
            raise ValueError('tool_code is invalid')
        event['tool_code'] = tool_code
    if session_credential:
        if not isinstance(session_credential, str):
            raise ValueError('session_credential must be text')
        event['session_hash'] = hashlib.sha256(
            session_credential.encode('utf-8')
        ).hexdigest()[:12]
    _optional_code(event, 'response_code', response_code, 32)
    _optional_code(event, 'summary_code', summary_code, 64)
    if cost_ms is not None:
        if (
            isinstance(cost_ms, bool)
            or not isinstance(cost_ms, int)
            or not 0 <= cost_ms <= 3600000
        ):
            raise ValueError('cost_ms is invalid')
        event['cost_ms'] = cost_ms
    validated_context = _validated_context(stage, context)
    if validated_context:
        event['context'] = validated_context
    return event


class RequestDiagnosticEmitter:
    MAX_EVENTS = 20

    def __init__(
        self,
        reporter,
        request_id,
        defer_until_identity=False,
        **event_defaults
    ):
        self.reporter = reporter
        self.request_id = request_id
        self.event_defaults = event_defaults
        self.defer_until_identity = bool(defer_until_identity)
        self._deferred = []
        self.emitted_count = 0
        self.dropped_count = 0

    def set_defaults(self, **event_defaults):
        self.event_defaults.update(event_defaults)
        company_id = self.event_defaults.get('company_id')
        if (
            self.defer_until_identity
            and isinstance(company_id, int)
            and not isinstance(company_id, bool)
            and company_id > 0
        ):
            self.flush()

    def emit(self, **fields):
        if self.emitted_count >= self.MAX_EVENTS:
            self.dropped_count += 1
            return False
        event_fields = dict(self.event_defaults)
        event_fields.update(fields)
        event_fields['request_id'] = self.request_id
        self.emitted_count += 1
        if (
            self.defer_until_identity
            and not event_fields.get('company_id')
        ):
            self._deferred.append(event_fields)
            return True
        return self._deliver(event_fields)

    def flush(self):
        deferred = self._deferred
        self._deferred = []
        accepted = True
        for fields in deferred:
            enriched = dict(fields)
            for name in ('company_id', 'admin_id', 'session_credential'):
                if name in self.event_defaults:
                    enriched[name] = self.event_defaults[name]
            if not self._deliver(enriched):
                accepted = False
        return accepted

    def _deliver(self, event_fields):
        try:
            event = build_diagnostic_event(**event_fields)
            accepted = self.reporter.report(event)
        except Exception:
            self.dropped_count += 1
            return False
        if not accepted:
            self.dropped_count += 1
        return accepted