mcp/services/api_client.py

import json
import uuid
import urllib.request


class JsonTransport:
    def post_json(self, url, payload, headers, timeout):
        request = urllib.request.Request(
            url=url,
            data=json.dumps(payload).encode('utf-8'),
            headers=dict({'Content-Type': 'application/json'}, **headers),
            method='POST',
        )
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode('utf-8-sig'))

    def post_multipart(self, url, fields, files, headers, timeout):
        boundary = '----McpFormBoundary{0}'.format(uuid.uuid4().hex)
        body = bytearray()
        for name, value in (fields or {}).items():
            body.extend(('--{0}\r\n'.format(boundary)).encode('utf-8'))
            body.extend(
                'Content-Disposition: form-data; name="{0}"\r\n\r\n'.format(
                    name
                ).encode('utf-8')
            )
            body.extend(str(value).encode('utf-8'))
            body.extend(b'\r\n')
        for name, spec in (files or {}).items():
            filename, file_bytes, content_type = spec
            safe_name = str(filename or 'file').replace('"', '')
            body.extend(('--{0}\r\n'.format(boundary)).encode('utf-8'))
            body.extend(
                (
                    'Content-Disposition: form-data; name="{0}"; '
                    'filename="{1}"\r\n'
                ).format(name, safe_name).encode('utf-8')
            )
            body.extend(
                'Content-Type: {0}\r\n\r\n'.format(
                    content_type or 'application/octet-stream'
                ).encode('utf-8')
            )
            body.extend(file_bytes if isinstance(file_bytes, (bytes, bytearray)) else b'')
            body.extend(b'\r\n')
        body.extend(('--{0}--\r\n'.format(boundary)).encode('utf-8'))
        request_headers = dict(headers or {})
        request_headers['Content-Type'] = (
            'multipart/form-data; boundary={0}'.format(boundary)
        )
        request = urllib.request.Request(
            url=url,
            data=bytes(body),
            headers=request_headers,
            method='POST',
        )
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode('utf-8-sig'))


class ApiClient:
    def __init__(self, base_url, token_store, transport=None, timeout=10):
        self.base_url = (base_url or '').rstrip('/')
        self.token_store = token_store
        self.transport = transport or JsonTransport()
        self.timeout = int(timeout)

    def call_tool(self, tool_code, route_path, payload, request_id):
        token = self.token_store.require_token()
        url = self.base_url + '/' + route_path.lstrip('/')
        headers = {
            'Authorization': 'Bearer {0}'.format(token),
            'X-MCP-Tool-Code': tool_code,
            'X-Request-Id': request_id,
        }
        return self.transport.post_json(url, payload, headers, self.timeout)

    def list_enabled_tools(self, request_id=''):
        token = self.token_store.require_token()
        url = self.base_url + '/mcp/tools/listEnabledTools'
        headers = {
            'Authorization': 'Bearer {0}'.format(token),
            'X-Request-Id': str(request_id or '').strip(),
        }
        return self.transport.post_json(url, {}, headers, self.timeout)

    def upload_headhaul_document(
        self,
        upload_token,
        filename,
        file_bytes,
        content_type='',
        request_id='',
    ):
        token = self.token_store.require_token()
        url = self.base_url + '/mcp/tools/uploadHeadhaulDocument'
        headers = {
            'Authorization': 'Bearer {0}'.format(token),
            'X-MCP-Tool-Code': 'prepare_headhaul_document_upload',
            'X-Request-Id': str(request_id or '').strip(),
        }
        return self.transport.post_multipart(
            url,
            {'upload_token': upload_token},
            {
                'file': (
                    filename,
                    file_bytes,
                    content_type or 'application/octet-stream',
                ),
            },
            headers,
            max(self.timeout, 120),
        )