fix: update MCP request handling to include source ID and type

Author: liuruibin

apps/application/chat_pipeline/step/chat_step/impl/base_chat_step.py

@@ -227,7 +227,7 @@ def reset_message_list(message_list: List[BaseMessage], answer_text):
         return result
 
     def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_ids, tool_ids,
-                            application_ids, mcp_output_enable, chat_model, message_list):
+                            application_ids, mcp_output_enable, chat_model, message_list, agent_id):
 
         mcp_servers_config = {}
 
@@ -258,7 +258,7 @@ def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_ids, tool_ids,
                     tool_init_params = json.loads(rsa_long_decrypt(tool.init_params))
                 else:
                     params = {}
-                tool_config = executor.get_tool_mcp_config(tool.code, params, tool.name, tool.desc)
+                tool_config = executor.get_tool_mcp_config(tool, params)
 
                 mcp_servers_config[str(tool.id)] = tool_config
 
@@ -290,7 +290,10 @@ def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_ids, tool_ids,
                 mcp_servers_config[app.name] = app_config
 
         if len(mcp_servers_config) > 0:
-            return mcp_response_generator(chat_model, message_list, json.dumps(mcp_servers_config), mcp_output_enable, tool_init_params)
+            source_id = agent_id
+            source_type = 'APPLICATION'
+            return mcp_response_generator(chat_model, message_list, json.dumps(mcp_servers_config), mcp_output_enable,
+                                          tool_init_params, source_id, source_type)
 
         return None
 
@@ -304,7 +307,9 @@ def get_stream_result(self, message_list: List[BaseMessage],
                           mcp_source="referencing",
                           tool_ids=None,
                           application_ids=None,
-                          mcp_output_enable=True):
+                          mcp_output_enable=True,
+                          agent_id=None
+                          ):
         if paragraph_list is None:
             paragraph_list = []
         directly_return_chunk_list = [AIMessageChunk(content=paragraph.content)
@@ -324,7 +329,7 @@ def get_stream_result(self, message_list: List[BaseMessage],
             mcp_result = self._handle_mcp_request(
                 mcp_source, mcp_servers, mcp_tool_ids, tool_ids,
                 application_ids, mcp_output_enable, chat_model,
-                message_list,
+                message_list, agent_id
             )
             if mcp_result:
                 return mcp_result, True
@@ -351,7 +356,7 @@ def execute_stream(self, message_list: List[BaseMessage],
                                                          no_references_setting, problem_text, mcp_tool_ids,
                                                          mcp_servers, mcp_source, tool_ids,
                                                          application_ids,
-                                                         mcp_output_enable)
+                                                         mcp_output_enable, manage.context.get('application_id'))
         chat_record_id = self.context.get('step_args', {}).get('chat_record_id') if self.context.get('step_args',
                                                                                                      {}).get(
             'chat_record_id') else uuid.uuid7()
@@ -375,7 +380,8 @@ def get_block_result(self, message_list: List[BaseMessage],
                          mcp_source="referencing",
                          tool_ids=None,
                          application_ids=None,
-                         mcp_output_enable=True
+                         mcp_output_enable=True,
+                         application_id=None
                          ):
         if paragraph_list is None:
             paragraph_list = []
@@ -395,7 +401,7 @@ def get_block_result(self, message_list: List[BaseMessage],
             mcp_result = self._handle_mcp_request(
                 mcp_source, mcp_servers, mcp_tool_ids, tool_ids,
                 application_ids, mcp_output_enable,
-                chat_model, message_list,
+                chat_model, message_list, application_id
             )
             if mcp_result:
                 return mcp_result, True
@@ -429,7 +435,7 @@ def execute_block(self, message_list: List[BaseMessage],
                                                             no_references_setting, problem_text,
                                                             mcp_tool_ids, mcp_servers, mcp_source,
                                                             tool_ids, application_ids,
-                                                            mcp_output_enable)
+                                                            mcp_output_enable, manage.context.get('application_id'))
             if is_ai_chat:
                 request_token = chat_model.get_num_tokens_from_messages(message_list)
                 response_token = chat_model.get_num_tokens(chat_result.content)

apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py

@@ -243,7 +243,7 @@ def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_id, mcp_tool_ids
                     tool_init_params = json.loads(rsa_long_decrypt(tool.init_params))
                 else:
                     params = {}
-                tool_config = executor.get_tool_mcp_config(tool.code, params, tool.name, tool.desc)
+                tool_config = executor.get_tool_mcp_config(tool, params)
 
                 mcp_servers_config[str(tool.id)] = tool_config
 
@@ -275,7 +275,17 @@ def _handle_mcp_request(self, mcp_source, mcp_servers, mcp_tool_id, mcp_tool_ids
                 mcp_servers_config[app.name] = app_config
 
         if len(mcp_servers_config) > 0:
-            r = mcp_response_generator(chat_model, message_list, json.dumps(mcp_servers_config), mcp_output_enable, tool_init_params)
+            # 安全获取 application
+            application_id = None
+            if (self.workflow_manage and
+                    self.workflow_manage.work_flow_post_handler and
+                    self.workflow_manage.work_flow_post_handler.chat_info):
+                application_id = self.workflow_manage.work_flow_post_handler.chat_info.application.id
+            knowledge_id = self.workflow_params.get('knowledge_id')
+            source_id = application_id or knowledge_id
+            source_type = 'APPLICATION' if application_id else 'KNOWLEDGE'
+            r = mcp_response_generator(chat_model, message_list, json.dumps(mcp_servers_config), mcp_output_enable,
+                                       tool_init_params, source_id, source_type)
             return NodeResult(
                 {'result': r, 'chat_model': chat_model, 'message_list': message_list,
                  'history_message': [{'content': message.content, 'role': message.type} for message in

apps/application/flow/tools.py

@@ -14,6 +14,9 @@
 from functools import reduce
 from typing import Iterator
 
+import uuid_utils.compat as uuid
+from asgiref.sync import sync_to_async
+from django.db.models import QuerySet
 from django.http import StreamingHttpResponse
 from langchain_core.messages import BaseMessageChunk, BaseMessage, ToolMessage, AIMessageChunk
 from langchain_mcp_adapters.client import MultiServerMCPClient
@@ -22,7 +25,9 @@
 from application.flow.i_step_node import WorkFlowPostHandler
 from common.result import result
 from common.utils.logger import maxkb_logger
+from knowledge.models.knowledge_action import State
 from maxkb.const import CONFIG
+from tools.models import ToolRecord, Tool
 
 
 class Reasoning:
@@ -316,7 +321,8 @@ def _extract_tool_id(raw_id):
     return tool_id or raw_id
 
 
-async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_enable=True, tool_init_params={}):
+async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_enable=True, tool_init_params={},
+                              source_id=None, source_type=None):
     try:
         client = MultiServerMCPClient(json.loads(mcp_servers))
         tools = await client.get_tools()
@@ -393,10 +399,14 @@ async def _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_
 
                 if tool_id in tool_calls_info:
                     tool_info = tool_calls_info[tool_id]
+                    tool_result = json.loads(chunk[0].content)
+                    tool_id = tool_result.pop('tool_id')
+                    if tool_id:
+                        await save_tool_record(tool_id, tool_info, tool_result, source_id, source_type)
                     content = generate_tool_message_complete(
                         tool_info['name'],
                         tool_info['input'],
-                        chunk[0].content
+                        json.dumps(tool_result),
                     )
                     chunk[0].content = content
                 else:
@@ -421,14 +431,30 @@ def get_real_error(exc):
         raise RuntimeError(error_msg) from None
 
 
-def mcp_response_generator(chat_model, message_list, mcp_servers, mcp_output_enable=True, tool_init_params={}):
+async def save_tool_record(tool_id, tool_info, tool_result, source_id, source_type):
+    tool = await sync_to_async(lambda: QuerySet(Tool).filter(id=tool_id).first())()
+    tool_record = ToolRecord(
+        id=uuid.uuid7(),
+        workspace_id=tool.workspace_id,
+        tool_id=tool_id,
+        source_type=source_type,
+        source_id=source_id,
+        meta={'input': tool_info['input'], 'output': tool_result},
+        state=State.SUCCESS
+    )
+    await sync_to_async(tool_record.save)()
+
+
+def mcp_response_generator(chat_model, message_list, mcp_servers, mcp_output_enable=True, tool_init_params={},
+                           source_id=None, source_type=None):
     """使用全局事件循环，不创建新实例"""
     result_queue = queue.Queue()
     loop = get_global_loop()  # 使用共享循环
 
     async def _run():
         try:
-            async_gen = _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_enable, tool_init_params)
+            async_gen = _yield_mcp_response(chat_model, message_list, mcp_servers, mcp_output_enable, tool_init_params,
+                                            source_id, source_type)
             async for chunk in async_gen:
                 result_queue.put(('data', chunk))
         except Exception as e:

apps/common/utils/tool_code.py

@@ -137,8 +137,8 @@ def exec_code(self, code_str, keywords, function_name=None):
             return result.get('data')
         raise Exception(result.get('msg') + (f'\n{subprocess_result.stderr}' if subprocess_result.stderr else ''))
 
-    def _generate_mcp_server_code(self, _code, params, name=None, description=None):
-        # 解析代码，提取导入语句和函数定义
+    def _generate_mcp_server_code(self, _code, params, name=None, description=None, tool_id=None):
+        # 解析代码,提取导入语句和函数定义
         try:
             tree = ast.parse(_code)
         except SyntaxError:
@@ -155,7 +155,7 @@ def _generate_mcp_server_code(self, _code, params, name=None, description=None):
                     continue
                 # 修改函数参数以包含 params 中的默认值
                 arg_names = [arg.arg for arg in node.args.args]
-                # 为参数添加默认值，确保参数顺序正确
+                # 为参数添加默认值,确保参数顺序正确
                 defaults = []
                 num_defaults = 0
                 # 从后往前检查哪些参数有默认值
@@ -178,26 +178,72 @@ def _generate_mcp_server_code(self, _code, params, name=None, description=None):
                             else:
                                 defaults.append(ast.Constant(value=str(default_value)))
                         else:
-                            # 如果某个参数没有默认值，需要添加 None 占位
+                            # 如果某个参数没有默认值,需要添加 None 占位
                             defaults.append(ast.Constant(value=None))
                     node.args.defaults = defaults
+
+                # 修改返回类型注解为 Result
+                node.returns = ast.Name(id='Result', ctx=ast.Load())
+
+                # 修改 return 语句为 return Result(result=..., tool_id=...)
+                class ReturnTransformer(ast.NodeTransformer):
+                    def __init__(self, func_name):
+                        self.func_name = func_name
+
+                    def visit_Return(self, node):
+                        if node.value is None:
+                            # return 语句没有返回值
+                            new_return = ast.Return(
+                                value=ast.Call(
+                                    func=ast.Name(id='Result', ctx=ast.Load()),
+                                    args=[],
+                                    keywords=[
+                                        ast.keyword(arg='result', value=ast.Constant(value=None)),
+                                        ast.keyword(arg='tool_id', value=ast.Constant(value=tool_id))
+                                    ]
+                                )
+                            )
+                        else:
+                            # return 语句有返回值
+                            new_return = ast.Return(
+                                value=ast.Call(
+                                    func=ast.Name(id='Result', ctx=ast.Load()),
+                                    args=[],
+                                    keywords=[
+                                        ast.keyword(arg='result', value=node.value),
+                                        ast.keyword(arg='tool_id', value=ast.Constant(value=tool_id))
+                                    ]
+                                )
+                            )
+                        return ast.copy_location(new_return, node)
+
+                transformer = ReturnTransformer(node.name)
+                node = transformer.visit(node)
+                ast.fix_missing_locations(node)
+
                 func_code = ast.unparse(node)
-                # 有些模型不支持name是中文，例如: deepseek, 其他模型未知
+                # 有些模型不支持name是中文,例如: deepseek, 其他模型未知
                 escaped_desc = (name + ' ' + description).replace('\n', ' ').replace("'", " ")
                 functions.append(f"@mcp.tool(description='{escaped_desc}')\n{func_code}\n")
             else:
                 other_code.append(ast.unparse(node))
+
         # 构建完整的 MCP 服务器代码
         code_parts = ["from mcp.server.fastmcp import FastMCP"]
         code_parts.extend(imports)
+        code_parts.append(f"\nfrom pydantic import BaseModel")
+        code_parts.append(f"\nfrom typing import Any")
+        code_parts.append(f"\nclass Result(BaseModel):")
+        code_parts.append(f"\n\tresult: Any")
+        code_parts.append(f"\n\ttool_id: str\n")
         code_parts.append(f"\nmcp = FastMCP(\"{uuid.uuid7()}\")\n")
         code_parts.extend(other_code)
         code_parts.extend(functions)
         code_parts.append("\nmcp.run(transport=\"stdio\")\n")
         return "\n".join(code_parts)
 
-    def generate_mcp_server_code(self, code_str, params, name, description):
-        code = self._generate_mcp_server_code(code_str, params, name, description)
+    def generate_mcp_server_code(self, code_str, params, name, description, tool_id):
+        code = self._generate_mcp_server_code(code_str, params, name, description, tool_id)
         set_run_user = f'os.setgid({pwd.getpwnam(_run_user).pw_gid});os.setuid({pwd.getpwnam(_run_user).pw_uid});' if _enable_sandbox else ''
         return f"""
 import os, sys, logging
@@ -212,8 +258,8 @@ def generate_mcp_server_code(self, code_str, params, name, description):
 exec({dedent(code)!a})
 """
 
-    def get_tool_mcp_config(self, code, params, name, description):
-        _code = self.generate_mcp_server_code(code, params, name, description)
+    def get_tool_mcp_config(self, tool, params):
+        _code = self.generate_mcp_server_code(tool.code, params, tool.name, tool.desc, str(tool.id))
         maxkb_logger.debug(f"Python code of mcp tool: {_code}")
         compressed_and_base64_encoded_code_str = base64.b64encode(gzip.compress(_code.encode())).decode()
         tool_config = {