mcp/mcp_protocol.py

import json
import logging
import sys
import uuid

from services.output_presenter import OutputPresenter
from services.diagnostic_event import RequestDiagnosticEmitter
from services.diagnostic_reporter import NullDiagnosticReporter


logger = logging.getLogger(__name__)


class McpProtocolHandler:
    protocol_version = '2025-06-18'
    server_name = 'fms-mcp-gateway'
    server_version = '0.1.0'
    output_presenter = OutputPresenter()

    def __init__(self, gateway_app, reporter=None):
        self.gateway_app = gateway_app
        self.reporter = reporter or NullDiagnosticReporter()
        self.initialized = False

    def handle_message(self, message):
        if not isinstance(message, dict):
            return self._error_response(
                None,
                -32600,
                'Invalid Request',
                self._build_trace_request_id(),
            )
        if 'id' in message:
            return self.handle_request(message)
        method = str(message.get('method') or '').strip()
        if method == 'notifications/initialized':
            self.initialized = True
            return None
        return None

    def handle_request(self, request):
        request_id = request.get('id')
        method = str(request.get('method') or '').strip()
        tool_name = ''
        trace_request_id = self._build_trace_request_id()
        emitter = RequestDiagnosticEmitter(self.reporter, trace_request_id)
        emitter.emit(
            stage='request_ingress',
            status='started',
            event_code='REQUEST_RECEIVED',
            context={
                'jsonrpc_method': method or 'unknown',
                'transport': 'stdio',
            },
        )
        backend_started = False
        try:
            if method == 'initialize':
                emitter.emit(
                    stage='protocol_validation',
                    status='succeeded',
                    event_code='PROTOCOL_VALIDATION_COMPLETED',
                    context={
                        'jsonrpc_method': method,
                        'transport': 'stdio',
                    },
                )
                self.initialized = True
                return self._success_response(
                    request_id,
                    {
                        'protocolVersion': self.protocol_version,
                        'capabilities': {
                            'tools': {
                                'listChanged': False,
                            }
                        },
                        'serverInfo': {
                            'name': self.server_name,
                            'version': self.server_version,
                        },
                    },
                )
            if method == 'tools/list':
                emitter.emit(
                    stage='protocol_validation',
                    status='succeeded',
                    event_code='PROTOCOL_VALIDATION_COMPLETED',
                    context={
                        'jsonrpc_method': method,
                        'transport': 'stdio',
                    },
                )
                tools = [
                    self._normalize_tool(tool)
                    for tool in self.gateway_app.list_tools(
                        request_id=trace_request_id,
                    )
                ]
                return self._success_response(request_id, {'tools': tools})
            if method == 'tools/call':
                params = request.get('params') or {}
                if not isinstance(params, dict):
                    raise ValueError('tool parameters must be an object')
                raw_tool_name = params.get('name')
                if not isinstance(raw_tool_name, str) or not raw_tool_name.strip():
                    emitter.emit(
                        stage='protocol_validation',
                        status='failed',
                        event_code='PARAM_VALIDATION_FAILED',
                        context={
                            'jsonrpc_method': method,
                            'jsonrpc_code': -32602,
                            'transport': 'stdio',
                        },
                    )
                    return self._error_response(
                        request_id,
                        -32602,
                        'Invalid params',
                        trace_request_id,
                    )
                tool_name = raw_tool_name.strip()
                arguments = params.get('arguments') or {}
                emitter.emit(
                    stage='protocol_validation',
                    status='succeeded',
                    event_code='PROTOCOL_VALIDATION_COMPLETED',
                    tool_code=tool_name,
                    context={
                        'jsonrpc_method': method,
                        'transport': 'stdio',
                    },
                )
                backend_started = True
                emitter.emit(
                    stage='backend_call',
                    status='started',
                    event_code='BACKEND_CALL_STARTED',
                    tool_code=tool_name,
                    context={'transport': 'stdio'},
                )
                tool_result = self.gateway_app.call_tool(
                    tool_name,
                    arguments,
                    request_id=trace_request_id,
                )
                backend_started = False
                emitter.emit(
                    stage='backend_call',
                    status='succeeded',
                    event_code='BACKEND_CALL_COMPLETED',
                    tool_code=tool_name,
                    response_code=(
                        tool_result.get('code')
                        if isinstance(tool_result, dict)
                        else None
                    ),
                    context={'transport': 'stdio'},
                )
                try:
                    response = self._tool_call_response(
                        request_id,
                        tool_name,
                        tool_result,
                    )
                except Exception:
                    emitter.emit(
                        stage='response_safety',
                        status='failed',
                        event_code='RESPONSE_SAFETY_REJECTED',
                        tool_code=tool_name,
                        context={'transport': 'stdio'},
                    )
                    raise
                emitter.emit(
                    stage='response_safety',
                    status='succeeded',
                    event_code='RESPONSE_SAFETY_COMPLETED',
                    tool_code=tool_name,
                    context={'transport': 'stdio'},
                )
                return response
            return self._error_response(
                request_id,
                -32601,
                'Method not found: {0}'.format(method or '<empty>'),
                trace_request_id,
            )
        except Exception as exc:
            if method == 'tools/call':
                diagnostic_reason = (
                    'PARAM_VALIDATION_FAILED'
                    if isinstance(exc, ValueError)
                    else 'UNEXPECTED_EXCEPTION'
                )
                logger.error(
                    'MCP stdio 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__,
                    },
                )
                if backend_started:
                    emitter.emit(
                        stage='backend_call',
                        status='failed',
                        event_code=diagnostic_reason,
                        tool_code=tool_name or None,
                        response_code='MCP_9001',
                        context={'transport': 'stdio'},
                    )
                return self._tool_exception_response(
                    request_id,
                    tool_name,
                    exc,
                    trace_request_id,
                )
            logger.error(
                "MCP stdio 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 self._error_response(
                request_id,
                -32000,
                'Gateway request failed. Please try again later.',
                trace_request_id,
            )

    def run_stdio(self, stdin=None, stdout=None):
        stdin = stdin or sys.stdin
        stdout = stdout or sys.stdout
        for raw_line in stdin:
            line = str(raw_line).strip()
            if not line:
                continue
            try:
                message = json.loads(line)
            except ValueError:
                response = self._error_response(
                    None,
                    -32700,
                    'Parse error',
                    self._build_trace_request_id(),
                )
            else:
                response = self.handle_message(message)
            if response is None:
                continue
            stdout.write(json.dumps(response, ensure_ascii=True) + '\n')
            if hasattr(stdout, 'flush'):
                stdout.flush()
        return 0

    def _normalize_tool(self, tool):
        normalized = dict(tool)
        if 'input_schema' in normalized:
            normalized['inputSchema'] = normalized.pop('input_schema')
        return normalized

    @classmethod
    def _tool_call_response(cls, request_id, tool_name, tool_result):
        if tool_name == 'query_order':
            return cls._legacy_tool_call_response(request_id, tool_result)

        presented = cls.output_presenter.present(tool_name, tool_result)
        return cls._presented_response(request_id, presented)

    @classmethod
    def _tool_exception_response(
        cls,
        request_id,
        tool_name,
        exception,
        trace_request_id='',
    ):
        if tool_name == 'query_order':
            return cls._success_response(
                request_id,
                {
                    'content': [{
                        'type': 'text',
                        'text': str(exception),
                    }],
                    'isError': True,
                },
            )

        presented = cls.output_presenter.present_exception(
            tool_name,
            exception,
        )
        if trace_request_id:
            presented['meta'] = {'request_id': trace_request_id}
        return cls._presented_response(request_id, presented)

    @classmethod
    def _presented_response(cls, request_id, presented):
        result = {
            'content': [{
                'type': 'text',
                'text': presented['text'],
            }],
            'structuredContent': presented['structured_content'],
            'isError': presented['is_error'],
        }
        if presented['meta']:
            result['_meta'] = presented['meta']
        return cls._success_response(request_id, result)

    @classmethod
    def _legacy_tool_call_response(cls, request_id, tool_result):
        if not isinstance(tool_result, dict):
            raise RuntimeError('invalid tool response')

        raw_code = tool_result.get('code')
        code = str(raw_code).strip() if raw_code is not None else ''
        message = str(tool_result.get('msg') or '').strip()
        data = tool_result.get('data')
        meta = tool_result.get('meta')

        if code in ('MCP_0000', '0'):
            if isinstance(data, dict):
                structured_content = dict(data)
            elif data:
                structured_content = {'data': data}
            else:
                structured_content = {}
            if isinstance(meta, dict) and meta:
                structured_content['meta'] = dict(meta)

            return cls._success_response(request_id, {
                'content': [{
                    'type': 'text',
                    'text': cls._render_text(structured_content),
                }],
                'structuredContent': structured_content,
                'isError': False,
            })

        error_content = {
            'code': code or 'MCP_9001',
            'msg': message or 'tool call failed',
        }
        if data not in (None, [], {}):
            error_content['data'] = data
        if isinstance(meta, dict) and meta:
            error_content['meta'] = dict(meta)

        return cls._success_response(request_id, {
            'content': [{
                'type': 'text',
                'text': '{0}: {1}'.format(
                    error_content['code'],
                    error_content['msg'],
                ),
            }],
            'structuredContent': error_content,
            'isError': True,
        })

    @staticmethod
    def _render_text(structured_content):
        if not structured_content:
            return 'ok'
        columns = structured_content.get('columns') if isinstance(structured_content, dict) else None
        records = structured_content.get('records') if isinstance(structured_content, dict) else None
        if isinstance(columns, list) and isinstance(records, list):
            return McpProtocolHandler._render_table_like_text(structured_content, columns, records)
        if isinstance(records, list) and records:
            # Check if this is a track-like structure (no columns, but has records)
            first_record = records[0] if records else {}
            if 'status' in first_record and 'content' in first_record and 'time' in first_record:
                return McpProtocolHandler._render_track_like_text(structured_content, records)
        return json.dumps(structured_content, ensure_ascii=False)

    @staticmethod
    def _render_table_like_text(structured_content, columns, records):
        lines = []
        summary = str(structured_content.get('summary') or '').strip()
        if summary:
            lines.append(summary)
        lines.append('表头共 {0} 列：'.format(len(columns)))
        for index, column in enumerate(columns, start=1):
            key = str(column.get('key') or '').strip()
            name = str(column.get('name') or key).strip()
            lines.append('{0}. {1} ({2})'.format(index, name, key))

        if not records:
            tips = structured_content.get('tips') or []
            if tips:
                lines.append('提示：{0}'.format('；'.join(str(tip) for tip in tips)))
            return '\n'.join(lines)

        for record_index, record in enumerate(records, start=1):
            lines.append('订单 {0}：'.format(record_index))
            for column in columns:
                key = str(column.get('key') or '').strip()
                name = str(column.get('name') or key).strip()
                value = record.get(key, '') if isinstance(record, dict) else ''
                if value is None:
                    value = ''
                lines.append('- {0}: {1}'.format(name, value))

        tips = structured_content.get('tips') or []
        if tips:
            lines.append('提示：{0}'.format('；'.join(str(tip) for tip in tips)))
        return '\n'.join(lines)

    @staticmethod
    def _render_track_like_text(structured_content, records):
        lines = []
        summary = str(structured_content.get('summary') or '').strip()
        if summary:
            lines.append(summary)

        if not records:
            tips = structured_content.get('tips') or []
            if tips:
                lines.append('提示：{0}'.format('；'.join(str(tip) for tip in tips)))
            return '\n'.join(lines)

        for record_index, record in enumerate(records, start=1):
            status = str(record.get('status') or '').strip()
            content = str(record.get('content') or '').strip()
            location = str(record.get('location') or '').strip()
            time = str(record.get('time') or '').strip()
            tracking_number = str(record.get('tracking_number') or '').strip()
            shipment_id = str(record.get('shipment_id') or '').strip()
            sub_track = int(record.get('sub_track') or 0)

            lines.append('轨迹 {0}：'.format(record_index))
            lines.append('- 状态: {0}'.format(status or '无'))
            lines.append('- 内容: {0}'.format(content or '无'))
            if location:
                lines.append('- 地点: {0}'.format(location))
            lines.append('- 时间: {0}'.format(time or '无'))
            if tracking_number:
                lines.append('- 快递单号: {0}'.format(tracking_number))
            if shipment_id:
                lines.append('- Shipment ID: {0}'.format(shipment_id))
            if sub_track:
                lines.append('- 子单轨迹: 是')

        tips = structured_content.get('tips') or []
        if tips:
            lines.append('提示：{0}'.format('；'.join(str(tip) for tip in tips)))
        return '\n'.join(lines)

    @staticmethod
    def _success_response(request_id, result):
        return {
            'jsonrpc': '2.0',
            'id': request_id,
            'result': result,
        }

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

    @staticmethod
    def _error_response(request_id, code, message, trace_request_id=''):
        error = {
            'code': code,
            'message': message,
        }
        if trace_request_id:
            error['data'] = {'request_id': trace_request_id}
        return {
            'jsonrpc': '2.0',
            'id': request_id,
            'error': error,
        }