mcp/app.py

import argparse
import json
import sys
import uuid

from config import GatewayConfig
from constants import DEVICE_INVALID_MESSAGE
from mcp_protocol import McpProtocolHandler
from public_gateway import PublicGatewayApp
from public_server import serve_public
from services.api_client import ApiClient
from services.auth_client import AuthClient
from services.gateway_session_store import GatewaySessionStore
from services.scoped_api_client import ScopedApiClient
from services.token_store import FileTokenStore, RedisSocketClient, RedisTokenStore
from tools.list_order_filter_options import ListOrderFilterOptionsTool
from tools.export_pending_outbound_orders import ExportPendingOutboundOrdersTool
from tools.export_out_of_province_port_data import (
    ExportOutOfProvincePortDataTool,
)
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_track import QueryTrackTool


def parse_int_list(value):
    if not value:
        return []
    return [int(item.strip()) for item in value.split(',') if item.strip()]


def parse_string_list(value):
    result = []
    for item in (value or '').split(','):
        item = item.strip()
        if item and item not in result:
            result.append(item)
    return result


class GatewayApp:
    def __init__(self, auth_client=None, api_client=None, token_store=None):
        self.auth_client = auth_client
        self.api_client = api_client
        self.token_store = token_store
        self._tools = {
            'query_order': QueryOrderTool(api_client=api_client),
            'query_track': QueryTrackTool(api_client=api_client),
            'query_order_exact': QueryOrderExactTool(api_client=api_client),
            'query_customs_declaration_files':
                QueryCustomsDeclarationFilesTool(api_client=api_client),
            'list_order_filter_options': ListOrderFilterOptionsTool(
                api_client=api_client
            ),
            'export_pending_outbound_orders': ExportPendingOutboundOrdersTool(
                api_client=api_client
            ),
            'export_out_of_province_port_data':
                ExportOutOfProvincePortDataTool(api_client=api_client),
            'list_pending_outbound_export_filter_options':
                ListPendingOutboundExportFilterOptionsTool(api_client=api_client),
        }

    @classmethod
    def from_config(cls, config, redis_client=None):
        if config.token_store_type == 'redis':
            redis = redis_client or RedisSocketClient(
                host=config.redis_host,
                port=config.redis_port,
                db=config.redis_db,
                password=config.redis_password,
                timeout=config.timeout_seconds,
            )
            token_store = RedisTokenStore(
                redis,
                prefix=config.redis_prefix,
                session_key=config.session_key,
                refresh_skew_seconds=config.refresh_skew_seconds,
            )
        elif config.token_store_type == 'file':
            token_store = FileTokenStore(
                config.token_store_path,
                refresh_skew_seconds=config.refresh_skew_seconds,
            )
        else:
            raise ValueError('unsupported token store type: {0}'.format(config.token_store_type))

        auth_client = AuthClient(
            base_url=config.auth_base_url,
            client_type=config.client_type,
            token_store=token_store,
            timeout=config.timeout_seconds,
            session_key=config.session_key,
        )
        api_client = ApiClient(
            base_url=config.tools_base_url,
            token_store=token_store,
            timeout=config.timeout_seconds,
        )
        return cls(auth_client=auth_client, api_client=api_client, token_store=token_store)

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

    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, request_id=''):
        if self.api_client is None or not hasattr(self.api_client, 'list_enabled_tools'):
            raise RuntimeError('enabled tool client unavailable')
        return self._enabled_tool_names(
            self.api_client.list_enabled_tools(request_id=request_id)
        )

    def list_tools(self, request_id=''):
        request_id = self.build_request_id(request_id)
        enabled = self._load_enabled_tool_names(request_id)
        return [
            tool.metadata()
            for name, tool in self._tools.items()
            if name in enabled
        ]

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

    def ensure_session(self):
        if self.token_store is None:
            return
        session = self.token_store.get()
        if not session or not session.get('token'):
            raise RuntimeError(DEVICE_INVALID_MESSAGE)
        if self.token_store.is_expiring():
            if self.auth_client is None:
                raise RuntimeError('mcp token expiring but auth client missing')
            self.auth_client.refresh(session['token'])

    def call_tool(self, name, arguments=None, request_id=''):
        if name not in self._tools:
            raise KeyError('tool not registered: {0}'.format(name))
        tool = self._tools[name]
        if getattr(tool, 'requires_session', True):
            self.ensure_session()
        request_id = self.build_request_id(request_id)
        if name not in self._load_enabled_tool_names(request_id):
            raise RuntimeError('tool disabled: {0}'.format(name))
        arguments = arguments or {}
        return tool.call(request_id=request_id, **arguments)

    def create_protocol_handler(self):
        return McpProtocolHandler(self)

    def run_cli(self, argv=None, stdin=None, stdout=None):
        stdin = stdin or sys.stdin
        stdout = stdout or sys.stdout
        parser = argparse.ArgumentParser(prog='mcp-gateway')
        subparsers = parser.add_subparsers(dest='command', required=True)

        subparsers.add_parser('list-tools')
        subparsers.add_parser('serve-stdio')

        public_parser = subparsers.add_parser('serve-public')
        public_parser.add_argument('--host', default='0.0.0.0')
        public_parser.add_argument('--port', type=int, default=8765)

        call_parser = subparsers.add_parser('call')
        call_parser.add_argument('--tool', required=True)
        call_parser.add_argument('--keyword', default='')
        call_parser.add_argument('--order-id', type=int, default=0)
        call_parser.add_argument('--order-number', default='')
        call_parser.add_argument('--order-numbers', default='')
        call_parser.add_argument('--tracking-number', default='')
        call_parser.add_argument('--tracking-numbers', default='')
        call_parser.add_argument('--reference-number', default='')
        call_parser.add_argument('--reference-numbers', default='')
        call_parser.add_argument('--outbound-number', default='')
        call_parser.add_argument('--outbound-numbers', default='')
        call_parser.add_argument('--container-code', default='')
        call_parser.add_argument('--container-codes', default='')
        call_parser.add_argument('--so-number', default='')
        call_parser.add_argument('--so-numbers', default='')
        call_parser.add_argument('--shipment-id', default='')
        call_parser.add_argument('--receiver-country', default='')
        call_parser.add_argument('--product-ids', default='')
        call_parser.add_argument('--customer-ids', default='')
        call_parser.add_argument('--sales-id', type=int, default=0)
        call_parser.add_argument('--warehouse-ids', default='')
        call_parser.add_argument('--department-id', type=int, default=0)
        call_parser.add_argument('--inbound-date-start', default='')
        call_parser.add_argument('--inbound-date-end', default='')
        call_parser.add_argument('--outbound-date-start', default='')
        call_parser.add_argument('--outbound-date-end', default='')
        call_parser.add_argument('--filter-type', default='')
        call_parser.add_argument('--page', type=int, default=1)
        call_parser.add_argument('--limit', type=int, default=20)
        call_parser.add_argument('--request-id', default='')

        args = parser.parse_args(argv or [])

        if args.command == 'list-tools':
            payload = self.list_tools()
        elif args.command == 'serve-stdio':
            return self.create_protocol_handler().run_stdio(stdin=stdin, stdout=stdout)
        elif args.command == 'serve-public':
            config = GatewayConfig.from_env()
            redis = RedisSocketClient(
                host=config.redis_host,
                port=config.redis_port,
                db=config.redis_db,
                password=config.redis_password,
                timeout=config.timeout_seconds,
            )
            session_store = GatewaySessionStore(
                redis,
                prefix=config.redis_prefix,
                ttl_seconds=config.gateway_session_ttl_seconds,
            )
            public_app = PublicGatewayApp(
                session_store=session_store,
                api_client=ScopedApiClient(config.tools_base_url, timeout=config.timeout_seconds),
            )
            return serve_public(
                public_app,
                host=args.host,
                port=args.port,
                enable_rate_limit=config.rate_limit_enabled,
                rate_limit_max_requests=config.rate_limit_max_requests,
                rate_limit_window_seconds=config.rate_limit_window_seconds,
            )
        elif args.command == 'call':
            tool_args = {
                'page': args.page,
                'limit': args.limit,
            }

            if args.tool == 'query_order':
                if not args.keyword:
                    raise ValueError('--keyword is required for query_order')
                tool_args['keyword'] = args.keyword
            elif args.tool == 'query_track':
                if args.order_id > 0:
                    tool_args['order_id'] = args.order_id
                if args.order_number:
                    tool_args['order_number'] = args.order_number
                if args.tracking_number:
                    tool_args['tracking_number'] = args.tracking_number
                if args.order_id <= 0 and not args.order_number and not args.tracking_number:
                    raise ValueError('--order-id, --order-number or --tracking-number is required for query_track')
            elif args.tool == 'query_order_exact':
                exact_strings = {
                    'order_number': args.order_number,
                    'reference_number': args.reference_number,
                    'tracking_number': args.tracking_number,
                    'outbound_number': args.outbound_number,
                    'container_code': args.container_code,
                    'so_number': args.so_number,
                    'shipment_id': args.shipment_id,
                    'receiver_country': args.receiver_country,
                    'inbound_date_start': args.inbound_date_start,
                    'inbound_date_end': args.inbound_date_end,
                    'outbound_date_start': args.outbound_date_start,
                    'outbound_date_end': args.outbound_date_end,
                }
                for field, value in exact_strings.items():
                    if value:
                        tool_args[field] = value
                exact_number_lists = {
                    'order_numbers': args.order_numbers,
                    'reference_numbers': args.reference_numbers,
                    'tracking_numbers': args.tracking_numbers,
                    'outbound_numbers': args.outbound_numbers,
                    'container_codes': args.container_codes,
                    'so_numbers': args.so_numbers,
                }
                for field, value in exact_number_lists.items():
                    if value:
                        tool_args[field] = parse_string_list(value)
                exact_lists = {
                    'product_ids': args.product_ids,
                    'customer_ids': args.customer_ids,
                    'warehouse_ids': args.warehouse_ids,
                }
                for field, value in exact_lists.items():
                    if value:
                        tool_args[field] = parse_int_list(value)
                if args.sales_id > 0:
                    tool_args['sales_id'] = args.sales_id
                if args.department_id > 0:
                    tool_args['department_id'] = args.department_id
            elif args.tool == 'query_customs_declaration_files':
                if args.outbound_numbers:
                    tool_args['outbound_numbers'] = parse_string_list(
                        args.outbound_numbers
                    )
                if args.order_numbers:
                    tool_args['order_numbers'] = parse_string_list(
                        args.order_numbers
                    )
            elif args.tool == 'list_order_filter_options':
                if not args.filter_type:
                    raise ValueError(
                        '--filter-type is required for '
                        'list_order_filter_options'
                    )
                tool_args['filter_type'] = args.filter_type
                tool_args['keyword'] = args.keyword
            else:
                if args.keyword:
                    tool_args['keyword'] = args.keyword

            payload = self.call_tool(
                args.tool,
                tool_args,
                request_id=args.request_id,
            )
        else:
            raise RuntimeError('unsupported command')

        stdout.write(json.dumps(payload, ensure_ascii=False))
        return 0


def main(argv=None):
    config = GatewayConfig.from_env()
    app = GatewayApp.from_config(config)
    return app.run_cli(argv=argv)


if __name__ == '__main__':
    raise SystemExit(main(sys.argv[1:]))