feat: add MCP transport validation to ToolExecutor

Author: liuruibin

apps/common/utils/tool_code.py

@@ -1,12 +1,13 @@
 # coding=utf-8
 import ast
-import os
 import json
+import os
 import subprocess
 import sys
 from textwrap import dedent
 
 import uuid_utils.compat as uuid
+from django.utils.translation import gettext_lazy as _
 
 from maxkb.const import BASE_DIR, CONFIG
 from maxkb.const import PROJECT_DIR
@@ -210,6 +211,12 @@ def validate_banned_keywords(self, code_str):
         if matched:
             raise Exception(f"keyword '{matched}' is banned in the tool.")
 
+    def validate_mcp_transport(self, code_str):
+        servers = json.loads(code_str)
+        for server, config in servers.items():
+            if config.get('transport') not in ['sse', 'streamable_http']:
+                raise Exception(_('Only support transport=sse or transport=streamable_http'))
+
     @staticmethod
     def _exec(_code):
         return subprocess.run([python_directory, '-c', _code], text=True, capture_output=True)

apps/tools/serializers/tool.py

@@ -356,6 +356,7 @@ def insert(self, instance, with_valid=True):
                 ToolCreateRequest(data=instance).is_valid(raise_exception=True)
                 # 校验代码是否包括禁止的关键字
                 ToolExecutor().validate_banned_keywords(instance.get('code', ''))
+                ToolExecutor().validate_mcp_transport(instance.get('code', ''))
 
             tool_id = uuid.uuid7()
             Tool(
@@ -391,6 +392,8 @@ def test_connection(self):
             self.is_valid(raise_exception=True)
             # 校验代码是否包括禁止的关键字
             ToolExecutor().validate_banned_keywords(self.data.get('code', ''))
+            ToolExecutor().validate_mcp_transport(self.data.get('code', ''))
+
             # 校验mcp json
             validate_mcp_config(json.loads(self.data.get('code')))
             return True
@@ -484,7 +487,7 @@ def edit(self, instance, with_valid=True):
                 ToolEditRequest(data=instance).is_valid(raise_exception=True)
                 # 校验代码是否包括禁止的关键字
                 ToolExecutor().validate_banned_keywords(instance.get('code', ''))
-
+                ToolExecutor().validate_mcp_transport(instance.get('code', ''))
 
             if not QuerySet(Tool).filter(id=self.data.get('id')).exists():
                 raise serializers.ValidationError(_('Tool not found'))