Quellcode durchsuchen

企业agent配置兼容

jackson vor 14 Stunden
Ursprung
Commit
333eacaf29
3 geänderte Dateien mit 384 neuen und 5 gelöschten Zeilen
  1. 3 1
      README.md
  2. 116 4
      public_server.py
  3. 265 0
      tests/test_public_server_stream.py

Datei-Diff unterdrückt, da er zu groß ist
+ 3 - 1
README.md


+ 116 - 4
public_server.py

@@ -27,6 +27,24 @@ UPLOAD_BODY_LIMIT = 102 * 1024 * 1024
 UPLOAD_LIMIT_NAME = 'upload-headhaul-document'
 UPLOAD_DISCARD_CHUNK = 64 * 1024
 UPLOAD_DISCARD_IDLE_SECONDS = 0.5
+# Exact path only. Prefix matching would swallow /mcp and the upload route.
+STREAM_PATH = '/mcp/stream'
+STREAM_KNOWN_PROTOCOL_VERSIONS = ('2025-03-26', '2025-06-18')
+STREAM_DEFAULT_PROTOCOL_VERSION = '2025-03-26'
+
+
+def stream_protocol_version(params):
+    """Return the protocol version advertised on /mcp/stream.
+
+    /mcp keeps its fixed version. This helper only rewrites the copy written
+    for the cloud Streamable HTTP path.
+    """
+    if not isinstance(params, dict):
+        return STREAM_DEFAULT_PROTOCOL_VERSION
+    client_version = str(params.get('protocolVersion') or '').strip()
+    if client_version in STREAM_KNOWN_PROTOCOL_VERSIONS:
+        return client_version
+    return STREAM_DEFAULT_PROTOCOL_VERSION
 
 
 def discard_request_body(rfile, connection, length, chunk=UPLOAD_DISCARD_CHUNK,
@@ -748,9 +766,18 @@ def create_http_handler(gateway_app, rate_limiter=None, reporter=None):
             if self.path == '/health':
                 self._write_json({'ok': True})
                 return
+            if self.path == STREAM_PATH:
+                self._write_empty(405, 'POST')
+                return
             self.send_response(404)
             self.end_headers()
 
+        def do_DELETE(self):
+            if self.path == STREAM_PATH:
+                self._write_empty(405, 'POST')
+                return
+            self.send_error(501, "Unsupported method ('DELETE')")
+
         def do_OPTIONS(self):
             if self.path == '/mcp/upload-headhaul-document':
                 self.send_response(204)
@@ -803,6 +830,16 @@ def create_http_handler(gateway_app, rate_limiter=None, reporter=None):
                         rpc_handler.rate_limiter.release(rate_key)
                 return
             length = int(self.headers.get('Content-Length') or '0')
+            if self.path == STREAM_PATH:
+                self._post_stream(
+                    request_headers,
+                    length,
+                    client_ip,
+                    trace_request_id,
+                    diagnostic_emitter,
+                    client_request_id_hash,
+                )
+                return
             if self.path != '/mcp':
                 logger.warning(f"[HTTP] 404: ip={client_ip}, path={self.path}")
                 if length > 0:
@@ -879,6 +916,81 @@ def create_http_handler(gateway_app, rate_limiter=None, reporter=None):
                 diagnostic_emitter=diagnostic_emitter,
             )
 
+        def _write_empty(self, status, allow=None):
+            self.send_response(status)
+            if allow:
+                self.send_header('Allow', allow)
+            self.send_header('Content-Length', '0')
+            self.end_headers()
+
+        def _post_stream(
+            self,
+            request_headers,
+            length,
+            client_ip,
+            trace_request_id,
+            diagnostic_emitter,
+            client_request_id_hash,
+        ):
+            body = ''
+            if length > 0:
+                body = self.rfile.read(length).decode('utf-8-sig')
+            try:
+                message = json.loads(body)
+            except json.JSONDecodeError:
+                self._write_json(
+                    McpProtocolHandler._error_response(
+                        None,
+                        -32700,
+                        'Parse error',
+                        trace_request_id,
+                    ),
+                    status=400,
+                )
+                return
+            if not isinstance(message, dict):
+                self._write_json(
+                    McpProtocolHandler._error_response(
+                        None,
+                        -32600,
+                        'Invalid Request',
+                        trace_request_id,
+                    ),
+                    status=400,
+                )
+                return
+            if 'id' not in message:
+                self._write_empty(202)
+                return
+            method = str(message.get('method') or '').strip()
+            if method == 'ping':
+                self._write_json(
+                    McpProtocolHandler._success_response(message.get('id'), {}),
+                )
+                return
+            logger.info(
+                'MCP public request',
+                extra={
+                    'request_id': trace_request_id,
+                    'jsonrpc_id': message.get('id'),
+                    'protocol_method': method,
+                    'client_identity': client_ip,
+                    'client_request_id_hash': client_request_id_hash,
+                },
+            )
+            response = rpc_handler.handle_json_rpc(
+                request_headers,
+                message,
+                client_ip,
+                trace_request_id=trace_request_id,
+                diagnostic_emitter=diagnostic_emitter,
+            )
+            if method == 'initialize':
+                response['result']['protocolVersion'] = stream_protocol_version(
+                    message.get('params'),
+                )
+            self._write_json(response, diagnostic_emitter=diagnostic_emitter)
+
         def _write_upload_json(self, status, payload, close=False):
             encoded = json.dumps(payload, ensure_ascii=False).encode('utf-8')
             self.send_response(status)
@@ -892,10 +1004,10 @@ def create_http_handler(gateway_app, rate_limiter=None, reporter=None):
             self.end_headers()
             self.wfile.write(encoded)
 
-        def _write_json(self, payload, diagnostic_emitter=None):
+        def _write_json(self, payload, diagnostic_emitter=None, status=200):
             raw = json.dumps(payload, ensure_ascii=False).encode('utf-8')
             try:
-                self.send_response(200)
+                self.send_response(status)
                 self.send_header('Content-Type', 'application/json; charset=utf-8')
                 self.send_header('Content-Length', str(len(raw)))
                 self.end_headers()
@@ -906,7 +1018,7 @@ def create_http_handler(gateway_app, rate_limiter=None, reporter=None):
                         status='succeeded',
                         event_code='RESPONSE_WRITE_COMPLETED',
                         context={
-                            'http_status': 200,
+                            'http_status': status,
                             'client_disconnected': False,
                             'transport': 'http',
                         },
@@ -919,7 +1031,7 @@ def create_http_handler(gateway_app, rate_limiter=None, reporter=None):
                         status='failed',
                         event_code='CLIENT_DISCONNECTED',
                         context={
-                            'http_status': 200,
+                            'http_status': status,
                             'client_disconnected': True,
                             'transport': 'http',
                         },

+ 265 - 0
tests/test_public_server_stream.py

@@ -0,0 +1,265 @@
+import json
+import unittest
+from http.client import HTTPConnection
+from threading import Thread
+from time import sleep
+from unittest.mock import MagicMock
+
+from constants import DEVICE_INVALID_MESSAGE
+from public_gateway import PublicGatewayApp
+from public_server import serve_public
+from utils.security import generate_gateway_session_id
+
+
+class TestPublicServerStream(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.mock_session_store = MagicMock()
+        cls.mock_api_client = MagicMock()
+        cls.mock_auth_client = MagicMock()
+        cls.gateway_app = PublicGatewayApp(
+            session_store=cls.mock_session_store,
+            api_client=cls.mock_api_client,
+            auth_client=cls.mock_auth_client,
+        )
+        cls.server_thread = Thread(
+            target=serve_public,
+            args=(cls.gateway_app,),
+            kwargs={'host': '127.0.0.1', 'port': 18766, 'enable_rate_limit': False},
+            daemon=True,
+        )
+        cls.server_thread.start()
+        sleep(0.5)
+
+    def setUp(self):
+        self.mock_session_store.reset_mock()
+        self.mock_api_client.reset_mock()
+        self.mock_auth_client.reset_mock()
+        self.mock_session_store.get.return_value = None
+        self.mock_api_client.list_enabled_tools.return_value = {
+            'code': 'MCP_0000',
+            'data': {'tool_codes': ['query_order']},
+        }
+        self.mock_api_client.call_tool.return_value = {
+            'code': '0',
+            'msg': 'success',
+            'data': {'order_no': 'ABC123'},
+        }
+
+    def _exchange(self, method, path, body=None, headers=None):
+        conn = HTTPConnection('127.0.0.1', 18766, timeout=5)
+        try:
+            request_headers = {}
+            if headers:
+                request_headers.update(headers)
+            conn.request(method, path, body, request_headers)
+            response = conn.getresponse()
+            return response.status, response.getheader('Allow'), response.read()
+        finally:
+            conn.close()
+
+    def _post_json(self, path, payload, headers=None):
+        raw = json.dumps(payload).encode('utf-8')
+        status, allow, body = self._exchange('POST', path, raw, headers)
+        parsed = json.loads(body.decode('utf-8')) if body else None
+        return status, allow, parsed
+
+    def test_existing_mcp_initialize_keeps_protocol_version(self):
+        status, _allow, body = self._post_json('/mcp', {
+            'jsonrpc': '2.0',
+            'id': 1,
+            'method': 'initialize',
+            'params': {'protocolVersion': '2025-03-26'},
+        })
+        self.assertEqual(200, status)
+        self.assertEqual('2025-06-18', body['result']['protocolVersion'])
+
+    def test_existing_mcp_notification_stays_method_not_found(self):
+        status, _allow, body = self._post_json('/mcp', {
+            'jsonrpc': '2.0',
+            'method': 'notifications/initialized',
+        })
+        self.assertEqual(200, status)
+        self.assertEqual(-32601, body['error']['code'])
+        self.assertIn('Method not found', body['error']['message'])
+
+    def test_existing_get_mcp_stays_404(self):
+        status, _allow, body = self._exchange('GET', '/mcp')
+        self.assertEqual(404, status)
+        self.assertEqual(b'', body)
+
+    def test_health_stays_200(self):
+        status, _allow, body = self._exchange('GET', '/health')
+        self.assertEqual(200, status)
+        self.assertEqual({'ok': True}, json.loads(body.decode('utf-8')))
+
+    def test_delete_mcp_stays_unimplemented(self):
+        status, _allow, _body = self._exchange('DELETE', '/mcp')
+        self.assertEqual(501, status)
+
+    def test_stream_lookalike_paths_stay_404(self):
+        status, _allow, _body = self._post_json('/mcp/stream-extra', {
+            'jsonrpc': '2.0',
+            'id': 1,
+            'method': 'ping',
+        })
+        self.assertEqual(404, status)
+        status, _allow, _body = self._post_json('/mcp/stream?x=1', {
+            'jsonrpc': '2.0',
+            'id': 1,
+            'method': 'ping',
+        })
+        self.assertEqual(404, status)
+
+    def test_get_stream_returns_405(self):
+        status, allow, body = self._exchange('GET', '/mcp/stream')
+        self.assertEqual(405, status)
+        self.assertEqual('POST', allow)
+        self.assertEqual(b'', body)
+
+    def test_delete_stream_returns_405(self):
+        status, allow, body = self._exchange('DELETE', '/mcp/stream')
+        self.assertEqual(405, status)
+        self.assertEqual('POST', allow)
+        self.assertEqual(b'', body)
+
+    def test_stream_notification_returns_202(self):
+        status, allow, body = self._post_json('/mcp/stream', {
+            'jsonrpc': '2.0',
+            'method': 'notifications/initialized',
+        })
+        self.assertEqual(202, status)
+        self.assertIsNone(allow)
+        self.assertIsNone(body)
+
+    def test_stream_invalid_json_returns_400(self):
+        status, _allow, body = self._exchange(
+            'POST',
+            '/mcp/stream',
+            b'this is not json',
+        )
+        self.assertEqual(400, status)
+        parsed = json.loads(body.decode('utf-8'))
+        self.assertEqual(-32700, parsed['error']['code'])
+        self.assertEqual('Parse error', parsed['error']['message'])
+
+    def test_stream_empty_body_returns_parse_error(self):
+        status, _allow, body = self._exchange('POST', '/mcp/stream', b'')
+        self.assertEqual(400, status)
+        parsed = json.loads(body.decode('utf-8'))
+        self.assertEqual(-32700, parsed['error']['code'])
+
+    def test_stream_non_object_returns_invalid_request(self):
+        status, _allow, body = self._exchange('POST', '/mcp/stream', b'[1]')
+        self.assertEqual(400, status)
+        parsed = json.loads(body.decode('utf-8'))
+        self.assertEqual(-32600, parsed['error']['code'])
+        self.assertEqual('Invalid Request', parsed['error']['message'])
+
+    def test_stream_ping(self):
+        status, _allow, body = self._post_json('/mcp/stream', {
+            'jsonrpc': '2.0',
+            'id': 7,
+            'method': 'ping',
+        })
+        self.assertEqual(200, status)
+        self.assertEqual({'jsonrpc': '2.0', 'id': 7, 'result': {}}, body)
+
+    def test_stream_initialize_echoes_known_versions(self):
+        for version in ('2025-03-26', '2025-06-18'):
+            status, _allow, body = self._post_json('/mcp/stream', {
+                'jsonrpc': '2.0',
+                'id': 1,
+                'method': 'initialize',
+                'params': {'protocolVersion': version},
+            })
+            self.assertEqual(200, status)
+            self.assertEqual(version, body['result']['protocolVersion'])
+            self.assertEqual('fms-mcp-gateway', body['result']['serverInfo']['name'])
+            self.assertEqual(False, body['result']['capabilities']['tools']['listChanged'])
+
+    def test_stream_initialize_unknown_version_falls_back(self):
+        status, _allow, body = self._post_json('/mcp/stream', {
+            'jsonrpc': '2.0',
+            'id': 1,
+            'method': 'initialize',
+            'params': {'protocolVersion': '2099-01-01'},
+        })
+        self.assertEqual(200, status)
+        self.assertEqual('2025-03-26', body['result']['protocolVersion'])
+
+    def test_stream_initialize_without_params_falls_back(self):
+        status, _allow, body = self._post_json('/mcp/stream', {
+            'jsonrpc': '2.0',
+            'id': 1,
+            'method': 'initialize',
+        })
+        self.assertEqual(200, status)
+        self.assertEqual('2025-03-26', body['result']['protocolVersion'])
+
+    def test_stream_initialize_with_empty_params_falls_back(self):
+        status, _allow, body = self._post_json('/mcp/stream', {
+            'jsonrpc': '2.0',
+            'id': 1,
+            'method': 'initialize',
+            'params': {},
+        })
+        self.assertEqual(200, status)
+        self.assertEqual('2025-03-26', body['result']['protocolVersion'])
+
+    def test_stream_initialize_with_non_object_params_falls_back(self):
+        status, _allow, body = self._post_json('/mcp/stream', {
+            'jsonrpc': '2.0',
+            'id': 1,
+            'method': 'initialize',
+            'params': ['2025-06-18'],
+        })
+        self.assertEqual(200, status)
+        self.assertEqual('2025-03-26', body['result']['protocolVersion'])
+
+    def test_stream_request_without_method_uses_existing_error(self):
+        status, _allow, body = self._post_json('/mcp/stream', {
+            'jsonrpc': '2.0',
+            'id': 3,
+        })
+        self.assertEqual(200, status)
+        self.assertEqual(-32601, body['error']['code'])
+
+    def test_stream_tools_list_without_session_matches_existing_error(self):
+        status, _allow, body = self._post_json('/mcp/stream', {
+            'jsonrpc': '2.0',
+            'id': 1,
+            'method': 'tools/list',
+        })
+        self.assertEqual(200, status)
+        self.assertEqual(-32001, body['error']['code'])
+        self.assertIn(DEVICE_INVALID_MESSAGE, body['error']['message'])
+
+    def test_stream_tools_call_uses_existing_session(self):
+        session_id = generate_gateway_session_id()
+        self.mock_session_store.get.return_value = {
+            'mcp_token': 'MT_valid_token',
+            'admin_id': 456,
+            'company_id': 200,
+        }
+        status, _allow, body = self._post_json(
+            '/mcp/stream',
+            {
+                'jsonrpc': '2.0',
+                'id': 1,
+                'method': 'tools/call',
+                'params': {
+                    'name': 'query_order',
+                    'arguments': {'order_no': 'ABC123'},
+                },
+            },
+            {'X-Gateway-Session': session_id},
+        )
+        self.assertEqual(200, status)
+        self.assertIn('result', body)
+        self.assertFalse(body['result']['isError'])
+        self.mock_api_client.call_tool.assert_called_once()
+
+
+if __name__ == '__main__':
+    unittest.main()