mcp/public_gateway.py

import logging
import time
import uuid

from constants import DEVICE_INVALID_MESSAGE
from tools.list_order_filter_options import ListOrderFilterOptionsTool
from tools.list_customer_filter_options import ListCustomerFilterOptionsTool
from tools.list_receivable_cost_filter_options import (
    ListReceivableCostFilterOptionsTool,
)
from tools.list_outbound_filter_options import ListOutboundFilterOptionsTool
from tools.export_pending_outbound_orders import ExportPendingOutboundOrdersTool
from tools.export_out_of_province_port_data import (
    ExportOutOfProvincePortDataTool,
)
from tools.export_receivable_cost_list import ExportReceivableCostListTool
from tools.list_pending_outbound_export_filter_options import (
    ListPendingOutboundExportFilterOptionsTool,
)
from tools.query_order import QueryOrderTool
from tools.query_customs_declaration_files import (
    QueryCustomsDeclarationFilesTool,
)
from tools.query_order_exact import QueryOrderExactTool
from tools.query_order_detail import QueryOrderDetailTool
from tools.query_export_task import QueryExportTaskTool
from tools.query_outbound_detail import QueryOutboundDetailTool
from tools.query_outbound_list import QueryOutboundListTool
from tools.query_track import QueryTrackTool
from tools.query_customer_list import QueryCustomerListTool
from tools.query_customer_payment_followup import QueryCustomerPaymentFollowupTool
from tools.query_customer_unverified_bill_details import QueryCustomerUnverifiedBillDetailsTool
from tools.query_customer_payment_records import QueryCustomerPaymentRecordsTool
from tools.query_order_receivable_cost_details import (
    QueryOrderReceivableCostDetailsTool,
)
from tools.query_receivable_cost_list import QueryReceivableCostListTool
from tools.query_payable_cost_list import QueryPayableCostListTool
from tools.list_payable_cost_filter_options import ListPayableCostFilterOptionsTool
from tools.export_payable_cost_list import ExportPayableCostListTool
from tools.export_pallet_data import ExportPalletDataTool
from tools.query_destination_trailer_list import QueryDestinationTrailerListTool
from tools.list_destination_trailer_filter_options import (
    ListDestinationTrailerFilterOptionsTool,
)
from tools.query_order_abnormal_list import QueryOrderAbnormalListTool
from tools.list_order_abnormal_filter_options import (
    ListOrderAbnormalFilterOptionsTool,
)
from tools.query_receive_volume_list import QueryReceiveVolumeListTool
from tools.list_receive_volume_filter_options import (
    ListReceiveVolumeFilterOptionsTool,
)
from tools.query_container_timeliness_list import QueryContainerTimelinessListTool
from tools.export_container_timeliness_report import (
    ExportContainerTimelinessReportTool,
)
from tools.query_headhaul_document_list import QueryHeadhaulDocumentListTool
from tools.list_headhaul_document_filter_options import (
    ListHeadhaulDocumentFilterOptionsTool,
)
from tools.prepare_headhaul_document_upload import (
    PrepareHeadhaulDocumentUploadTool,
)
from tools.save_headhaul_document import SaveHeadhaulDocumentTool
from tools.delete_headhaul_document import DeleteHeadhaulDocumentTool
from utils.security import hash_gateway_session_id


logger = logging.getLogger(__name__)


class PublicGatewayApp:
    def __init__(self, session_store, api_client, auth_client=None):
        self.session_store = session_store
        self.api_client = api_client
        self._tools = {
            'query_order': QueryOrderTool(api_client=None),
            'query_track': QueryTrackTool(api_client=None),
            'query_order_exact': QueryOrderExactTool(api_client=None),
            'query_order_detail': QueryOrderDetailTool(api_client=None),
            'query_customs_declaration_files':
                QueryCustomsDeclarationFilesTool(api_client=None),
            'query_outbound_list': QueryOutboundListTool(api_client=None),
            'query_outbound_detail': QueryOutboundDetailTool(api_client=None),
            'query_customer_list': QueryCustomerListTool(api_client=None),
            'query_customer_payment_followup': QueryCustomerPaymentFollowupTool(
                api_client=None
            ),
            'query_customer_unverified_bill_details':
                QueryCustomerUnverifiedBillDetailsTool(api_client=None),
            'query_customer_payment_records': QueryCustomerPaymentRecordsTool(
                api_client=None
            ),
            'query_order_receivable_cost_details':
                QueryOrderReceivableCostDetailsTool(api_client=None),
            'query_receivable_cost_list':
                QueryReceivableCostListTool(api_client=None),
            'query_payable_cost_list':
                QueryPayableCostListTool(api_client=None),
            'list_outbound_filter_options': ListOutboundFilterOptionsTool(
                api_client=None
            ),
            'list_order_filter_options': ListOrderFilterOptionsTool(
                api_client=None
            ),
            'list_customer_filter_options': ListCustomerFilterOptionsTool(
                api_client=None
            ),
            'list_receivable_cost_filter_options':
                ListReceivableCostFilterOptionsTool(api_client=None),
            'list_payable_cost_filter_options':
                ListPayableCostFilterOptionsTool(api_client=None),
            'export_pending_outbound_orders': ExportPendingOutboundOrdersTool(
                api_client=None
            ),
            'export_out_of_province_port_data':
                ExportOutOfProvincePortDataTool(api_client=None),
            'export_receivable_cost_list': ExportReceivableCostListTool(
                api_client=None
            ),
            'export_payable_cost_list': ExportPayableCostListTool(
                api_client=None
            ),
            'export_pallet_data': ExportPalletDataTool(
                api_client=None
            ),
            'query_export_task': QueryExportTaskTool(api_client=None),
            'list_pending_outbound_export_filter_options':
                ListPendingOutboundExportFilterOptionsTool(api_client=None),
            'query_destination_trailer_list': QueryDestinationTrailerListTool(
                api_client=None
            ),
            'list_destination_trailer_filter_options':
                ListDestinationTrailerFilterOptionsTool(api_client=None),
            'query_order_abnormal_list': QueryOrderAbnormalListTool(
                api_client=None
            ),
            'list_order_abnormal_filter_options':
                ListOrderAbnormalFilterOptionsTool(api_client=None),
            'query_receive_volume_list': QueryReceiveVolumeListTool(
                api_client=None
            ),
            'list_receive_volume_filter_options':
                ListReceiveVolumeFilterOptionsTool(api_client=None),
            'query_container_timeliness_list': QueryContainerTimelinessListTool(
                api_client=None
            ),
            'export_container_timeliness_report':
                ExportContainerTimelinessReportTool(api_client=None),
            'query_headhaul_document_list': QueryHeadhaulDocumentListTool(
                api_client=None
            ),
            'list_headhaul_document_filter_options':
                ListHeadhaulDocumentFilterOptionsTool(api_client=None),
            'prepare_headhaul_document_upload':
                PrepareHeadhaulDocumentUploadTool(api_client=None),
            'save_headhaul_document': SaveHeadhaulDocumentTool(
                api_client=None
            ),
            'delete_headhaul_document': DeleteHeadhaulDocumentTool(
                api_client=None
            ),
        }
        from services.headhaul_document_app import HeadhaulSlotMemory
        self._headhaul_slots = HeadhaulSlotMemory()

    def registered_tool_names(self):
        return tuple(self._tools.keys())

    def _require_session(self, gateway_session_id, diagnostic_emitter=None):
        session = self.session_store.get(gateway_session_id)
        if not session or not session.get('mcp_token'):
            if diagnostic_emitter is not None:
                diagnostic_emitter.emit(
                    stage='gateway_session',
                    status='failed',
                    event_code='GATEWAY_SESSION_NOT_FOUND',
                    session_credential=gateway_session_id,
                    context={'transport': 'http'},
                )
            raise RuntimeError(DEVICE_INVALID_MESSAGE)
        if diagnostic_emitter is not None:
            diagnostic_emitter.set_defaults(
                session_credential=gateway_session_id,
                admin_id=session.get('admin_id'),
                company_id=session.get('company_id'),
                context={'transport': 'http'},
            )
            diagnostic_emitter.emit(
                stage='gateway_session',
                status='succeeded',
                event_code='GATEWAY_SESSION_RESOLVED',
            )
        return session

    def _enabled_tool_names(self, response):
        if not isinstance(response, dict):
            raise RuntimeError('invalid enabled tool response')
        if response.get('code') != 'MCP_0000':
            raise RuntimeError(response.get('msg') or 'list enabled tools failed')
        data = response.get('data')
        codes = data.get('tool_codes') if isinstance(data, dict) else None
        if not isinstance(codes, list):
            raise RuntimeError('invalid enabled tool response')
        return {
            code.strip().lower()
            for code in codes
            if isinstance(code, str) and code.strip()
        }

    def _load_enabled_tool_names(self, token, request_id=''):
        response = self.api_client.list_enabled_tools(
            token,
            request_id=request_id,
        )
        return self._enabled_tool_names(response)

    def list_tools(self, gateway_session_id, request_id=''):
        session = self._require_session(gateway_session_id)
        if hasattr(self.session_store, 'touch_session'):
            self.session_store.touch_session(gateway_session_id)
        request_id = self.build_request_id(request_id)
        enabled = self._load_enabled_tool_names(
            session['mcp_token'],
            request_id,
        )
        return [
            tool.metadata()
            for name, tool in self._tools.items()
            if name in enabled
        ]

    def list_resources(self):
        return []

    def remember_headhaul_slot(self, gateway_session_id, token):
        self._headhaul_slots.remember(gateway_session_id, token)

    def session_for_headhaul_token(self, token):
        return self._headhaul_slots.session_for_token(token)

    def read_resource(self, uri, gateway_session_id=''):
        return None

    def upload_headhaul_document(
        self,
        gateway_session_id,
        upload_token,
        filename,
        file_bytes,
        content_type='',
        request_id='',
        client_ip='',
    ):
        session = self._require_session(gateway_session_id)
        request_id = self.build_request_id(request_id)
        enabled = self._load_enabled_tool_names(
            session['mcp_token'],
            request_id,
        )
        if 'prepare_headhaul_document_upload' not in enabled:
            raise RuntimeError('tool disabled: prepare_headhaul_document_upload')
        if not hasattr(self.api_client, 'upload_headhaul_document'):
            raise RuntimeError('upload client unavailable')
        return self.api_client.upload_headhaul_document(
            token=session['mcp_token'],
            upload_token=upload_token,
            filename=filename,
            file_bytes=file_bytes,
            content_type=content_type,
            request_id=request_id,
            client_ip=client_ip,
        )

    def build_request_id(self, request_id=''):
        request_id = str(request_id or '').strip()
        return request_id or 'rq_{0}'.format(uuid.uuid4().hex[:16])

    def call_tool(
        self,
        gateway_session_id,
        name,
        arguments=None,
        request_id='',
        client_ip='',
        diagnostic_emitter=None,
    ):
        session = self._require_session(gateway_session_id, diagnostic_emitter)
        request_id = self.build_request_id(request_id)
        if name not in self._tools:
            if diagnostic_emitter is not None:
                diagnostic_emitter.emit(
                    stage='backend_call',
                    status='failed',
                    event_code='TOOL_NOT_REGISTERED',
                    context={'transport': 'http'},
                )
            raise KeyError('tool not registered: {0}'.format(name))

        try:
            enabled_tools = self._load_enabled_tool_names(
                session['mcp_token'],
                request_id,
            )
        except Exception:
            if diagnostic_emitter is not None:
                diagnostic_emitter.emit(
                    stage='backend_call',
                    status='failed',
                    event_code='ENABLED_TOOL_LOOKUP_FAILED',
                    tool_code=name,
                    context={'transport': 'http'},
                )
            raise
        if name not in enabled_tools:
            if diagnostic_emitter is not None:
                diagnostic_emitter.emit(
                    stage='backend_call',
                    status='failed',
                    event_code='TOOL_DISABLED',
                    tool_code=name,
                    context={'transport': 'http'},
                )
            raise RuntimeError('tool disabled: {0}'.format(name))

        tool = self._tools[name]

        if diagnostic_emitter is not None:
            diagnostic_emitter.set_defaults(tool_code=name)

        session_hash = hash_gateway_session_id(gateway_session_id)[:12]
        admin_id = session.get('admin_id')
        company_id = session.get('company_id')
        logger.info(
            'MCP public tool call',
            extra={
                'request_id': request_id,
                'tool_code': name,
                'session_hash': session_hash,
                'admin_id': admin_id,
                'company_id': company_id,
            },
        )

        try:
            started_at = time.monotonic()
            if diagnostic_emitter is not None:
                diagnostic_emitter.emit(
                    stage='backend_call',
                    status='started',
                    event_code='BACKEND_CALL_STARTED',
                )
            result = self.api_client.call_tool(
                token=session['mcp_token'],
                tool_code=tool.name,
                route_path=tool.route_path,
                payload=arguments or {},
                request_id=request_id,
                client_ip=client_ip,
            )
            if hasattr(self.session_store, 'touch_session'):
                self.session_store.touch_session(gateway_session_id)
            if name == 'prepare_headhaul_document_upload' and isinstance(result, dict):
                data = result.get('data')
                token = ''
                if isinstance(data, dict):
                    token = str(data.get('upload_token') or '').strip()
                if token:
                    self.remember_headhaul_slot(gateway_session_id, token)
            logger.info(
                'MCP public tool success',
                extra={
                    'request_id': request_id,
                    'tool_code': name,
                    'session_hash': session_hash,
                    'response_code': result.get('code'),
                },
            )
            if diagnostic_emitter is not None:
                diagnostic_emitter.emit(
                    stage='backend_call',
                    status='succeeded',
                    event_code='BACKEND_CALL_COMPLETED',
                    response_code=(
                        result.get('code') if isinstance(result, dict) else None
                    ),
                    cost_ms=max(0, int((time.monotonic() - started_at) * 1000)),
                )
            return result
        except Exception as e:
            logger.error(
                'MCP public tool failed',
                extra={
                    'request_id': request_id,
                    'tool_code': name,
                    'session_hash': session_hash,
                    'response_code': 'MCP_9001',
                    'diagnostic_reason': 'UNEXPECTED_EXCEPTION',
                    'exception_class': e.__class__.__name__,
                },
            )
            if diagnostic_emitter is not None:
                diagnostic_emitter.emit(
                    stage='backend_call',
                    status='failed',
                    event_code='UNEXPECTED_EXCEPTION',
                    response_code='MCP_9001',
                )
            raise