diff --git a/apps/application/tests/test_ai_chat_tool_output.py b/apps/application/tests/test_ai_chat_tool_output.py index 3c4e1a22bdf..7efa59bb809 100644 --- a/apps/application/tests/test_ai_chat_tool_output.py +++ b/apps/application/tests/test_ai_chat_tool_output.py @@ -403,3 +403,56 @@ async def next_model_input(state): self.assertEqual(cards[-1].status, Status.SUCCESS) self.assertEqual(cards[-1].name, "display-tool") self.assertNotIn("image-payload", cards[-1].content) + + +class AIChatAgentBackendTests(SimpleTestCase): + def _create_deep_agent_kwargs(self, skill_tool_ids): + async def stream(*args, **kwargs): + return + yield + + with ( + TemporaryDirectory(prefix="maxkb-backend-test-") as directory, + patch.object(agent_module, "init_skills", new=AsyncMock()), + patch.object( + agent_module, "MultiServerMCPClient", return_value=SimpleNamespace(get_tools=AsyncMock(return_value=[])) + ), + patch.object(agent_module, "SandboxShellBackend", return_value="shell-backend"), + patch.object( + agent_module, "create_deep_agent", return_value=SimpleNamespace(astream=stream) + ) as create_deep_agent, + ): + agent_module.create_agent( + MagicMock(), + "system", + [], + "{}", + SimpleNamespace(on_next=lambda chunk: None, on_complete=lambda error: None), + chat_id=directory, + skill_tool_ids=skill_tool_ids, + async_runner=SimpleNamespace(run_async=lambda fn: asyncio.run(fn())), + ) + return create_deep_agent.call_args.kwargs + + def test_mcp_or_tools_only_do_not_expose_server_shell(self): + for skill_tool_ids in (None, []): + with self.subTest(skill_tool_ids=skill_tool_ids): + kwargs = self._create_deep_agent_kwargs(skill_tool_ids) + self.assertIsNone(kwargs["backend"]) + self.assertIsNone(kwargs["skills"]) + self.assertIsInstance(kwargs["middleware"][0], agent_module.HideFilesystemToolsMiddleware) + + def test_skills_keep_shell_backend(self): + kwargs = self._create_deep_agent_kwargs(["skill-id"]) + + self.assertEqual(kwargs["backend"], "shell-backend") + self.assertEqual(kwargs["skills"], ["/skills"]) + self.assertEqual(kwargs["middleware"], ()) + + def test_hide_filesystem_tools_keeps_configured_tools(self): + request = MagicMock(tools=[SimpleNamespace(name="ls"), {"name": "execute"}, SimpleNamespace(name="agent_x")]) + + agent_module.HideFilesystemToolsMiddleware().wrap_model_call(request, lambda req: req) + + request.override.assert_called_once() + self.assertEqual([getattr(t, "name", None) for t in request.override.call_args.kwargs["tools"]], ["agent_x"]) diff --git a/apps/application/workflow/nodes/ai_chat_node/agent.py b/apps/application/workflow/nodes/ai_chat_node/agent.py index c4dafbfca48..35b9c554a4c 100644 --- a/apps/application/workflow/nodes/ai_chat_node/agent.py +++ b/apps/application/workflow/nodes/ai_chat_node/agent.py @@ -18,6 +18,7 @@ import langchain_core.messages.ai as _lc_ai_module import uuid_utils.compat as uuid from deepagents import create_deep_agent +from langchain.agents.middleware import AgentMiddleware from langchain_core.utils._merge import merge_lists as _original_merge_lists from langchain_mcp_adapters.client import MultiServerMCPClient from langgraph.checkpoint.memory import MemorySaver @@ -139,6 +140,33 @@ def set_tool_id_name(self, tool_id, name): self.id_name_map[tool_id] = name +# deepagents 内置的文件与 Shell 工具 +_FILESYSTEM_TOOL_NAMES = frozenset({"ls", "read_file", "write_file", "edit_file", "delete", "glob", "grep", "execute"}) + + +class HideFilesystemToolsMiddleware(AgentMiddleware): + """ + 未启用技能时隐藏 deepagents 内置的文件与 Shell 工具，模型只能使用节点配置的工具与 MCP， + 避免绕过 MCP 自己写文件或在服务器上执行命令。 + """ + + @staticmethod + def _filter(request): + tools = [ + tool + for tool in request.tools + if (tool.get("name") if isinstance(tool, dict) else getattr(tool, "name", None)) + not in _FILESYSTEM_TOOL_NAMES + ] + return request.override(tools=tools) if len(tools) != len(request.tools) else request + + def wrap_model_call(self, request, handler): + return handler(self._filter(request)) + + async def awrap_model_call(self, request, handler): + return await handler(self._filter(request)) + + def create_agent( chat_model, system_prompt, @@ -169,10 +197,13 @@ async def _run(): for tool in extra_tools: tools.append(tool) + # 只有技能需要读写技能文件并执行脚本，才挂载本机 Shell 后端与内置文件工具； + # 仅 MCP/工具时使用默认的内存后端并隐藏内置文件工具，模型只能通过配置的工具完成任务 agent = create_deep_agent( model=chat_model, - backend=SandboxShellBackend(root_dir=temp_dir, virtual_mode=True), - skills=["/skills"], + backend=SandboxShellBackend(root_dir=temp_dir, virtual_mode=True) if skill_tool_ids else None, + skills=["/skills"] if skill_tool_ids else None, + middleware=() if skill_tool_ids else (HideFilesystemToolsMiddleware(),), tools=tools, system_prompt=system_prompt, interrupt_on={"write_file": False, "read_file": False, "edit_file": False}, diff --git a/apps/chat/mcp/tools.py b/apps/chat/mcp/tools.py index 29cd7ec3071..36d9cc13c15 100644 --- a/apps/chat/mcp/tools.py +++ b/apps/chat/mcp/tools.py @@ -38,6 +38,11 @@ UPLOAD_FILE_TOOL = "upload_file" +# 全局 MCP 地址（不带应用 id）使用：先列出可访问的智能体，再按 application_id 调用 +LIST_AGENTS_TOOL = "list_agents" + +CALL_AGENT_TOOL = "call_agent" + class InvalidParamsError(Exception): """工具参数不合法，对应 JSON-RPC -32602""" @@ -211,9 +216,11 @@ def call_tool(self, params): class ChatUserMCPToolHandler: """ - 对话用户 API Key（chat-xxx）的 MCP 接口：Key 不绑定应用， - 以对话用户身份暴露其有权限访问的所有已发布应用，每个应用一个工具，另加一个上传文件工具。 - 传入上一次返回的 chat_id 即可续聊；上传的文件挂在 chat_id 下，调用应用时通过 file_ids 携带。 + 对话用户 API Key（chat-xxx）的 MCP 接口，以对话用户身份调用其有权限访问的已发布应用： + - 单应用地址：暴露该应用的 agent_ 工具和上传文件工具； + - 全局地址：暴露 list_agents（列出可访问的应用）、call_agent（按 application_id 调用）和上传文件工具， + 应用再多工具数量也固定。 + 传入上一次返回的 chat_id 即可续聊；上传的文件挂在 chat_id 下，调用应用时通过 image_list 等附件字段携带。 """ def __init__(self, chat_user, auth, application_id=None): @@ -241,40 +248,91 @@ def get_application_ids(self): def get_applications(self): return QuerySet(Application).filter(id__in=self.get_application_ids(), is_publish=True).order_by("name") + def is_accessible_application(self, application_id): + try: + application_id = str(std_uuid.UUID(str(application_id))) + except ValueError: + return False + return self.get_applications().filter(id=application_id).exists() + @staticmethod - def build_agent_tool(application): + def build_agent_input_schema(with_application_id=False): + application_id = ( + { + "application_id": { + "type": "string", + "description": f"Id of the agent to call, from {LIST_AGENTS_TOOL}.", + } + } + if with_application_id + else {} + ) return { - "name": f"{AGENT_TOOL_PREFIX}{str(application.id)}", - "description": f"{application.name} {application.desc or ''}".strip(), - "inputSchema": { - "type": "object", - "properties": { - "message": {"type": "string", "description": "The message to send to the AI."}, - "chat_id": { - "type": "string", - "description": "Conversation id returned by a previous call or by upload_file. " - "Pass it to continue the same conversation with history; omit it to start a new one.", - }, - "file_ids": { + "type": "object", + "properties": { + **application_id, + "message": {"type": "string", "description": "The message to send to the AI."}, + "chat_id": { + "type": "string", + "description": "Conversation id returned by a previous call or by upload_file. " + "Pass it to continue the same conversation with history; omit it to start a new one.", + }, + **{ + field: { "type": "array", - "items": {"type": "string"}, - "description": "Ids of files uploaded with upload_file to attach to this message.", - }, - "form_data": { - "type": "object", - "description": "Values for the application's input parameters.", - }, + "items": { + "type": "object", + "properties": {"file_id": {"type": "string"}}, + "required": ["file_id"], + }, + "description": f"{CHAT_FILE_TYPE_LABELS[field].capitalize()} files uploaded with " + 'upload_file to attach to this message, e.g. [{"file_id": "..."}].', + } + for field in CHAT_FILE_LIST_FIELDS + }, + "form_data": { + "type": "object", + "description": "Values for the application's input parameters.", }, - "required": ["message"], }, + "required": [*application_id, "message"], } + @classmethod + def build_agent_tool(cls, application): + # 应用描述可能为空，补充通用说明，调用方模型才知道把用户的任务交给该智能体 + description = f"{application.name} {application.desc or ''}".strip() + return { + "name": f"{AGENT_TOOL_PREFIX}{str(application.id)}", + "description": f"{description}\nAn AI agent. Send the user's request as message; " + "the agent completes the task and returns its answer.", + "inputSchema": cls.build_agent_input_schema(), + } + + @classmethod + def build_global_tools(cls): + return [ + { + "name": LIST_AGENTS_TOOL, + "description": "List the AI agents you can call, with their id, name and description. " + f"Pick the agent that fits the user's request and call it with {CALL_AGENT_TOOL}.", + "inputSchema": {"type": "object", "properties": {}}, + }, + { + "name": CALL_AGENT_TOOL, + "description": f"Call an AI agent from {LIST_AGENTS_TOOL} by application_id. " + "Send the user's request as message; the agent completes the task and returns its answer.", + "inputSchema": cls.build_agent_input_schema(with_application_id=True), + }, + ] + @staticmethod def build_upload_tool(): return { "name": UPLOAD_FILE_TOOL, - "description": "Upload a file to attach to a conversation. Returns file_id and chat_id; " - "pass both to an agent tool (file_ids, chat_id) so the AI can read the file.", + "description": "Upload a file to attach to a conversation. Returns file_id, field and chat_id; " + 'pass {"file_id": ...} in the returned field (e.g. document_list) together with chat_id ' + "to an agent tool so the AI can read the file.", "inputSchema": { "type": "object", "properties": { @@ -290,11 +348,21 @@ def build_upload_tool(): } def list_tools(self): + agent_tools = ( + [self.build_agent_tool(application) for application in self.get_applications()] + if self.application_id + else self.build_global_tools() + ) + return {"tools": [*agent_tools, self.build_upload_tool()]} + + def list_agents(self): + agents = [ + {"application_id": str(application.id), "name": application.name, "desc": application.desc or ""} + for application in self.get_applications() + ] return { - "tools": [ - *[self.build_agent_tool(application) for application in self.get_applications()], - self.build_upload_tool(), - ] + "content": [{"type": "text", "text": json.dumps(agents, ensure_ascii=False)}], + "structuredContent": {"agents": agents}, } def check_chat_id(self, chat_id, application_id=None): @@ -315,29 +383,42 @@ def check_chat_id(self, chat_id, application_id=None): raise InvalidParamsError(f"Invalid chat_id: {chat_id}") return chat_id - def get_chat_files(self, chat_id, file_ids): + def get_chat_files(self, chat_id, args): """ - 根据 file_ids 查出当前用户上传到该会话的文件，按类型组装成消息附件；未传 chat_id 时取文件所在会话。 + 与 MaxKB 对话接口一致，按 image_list / document_list 等字段接收附件 [{"file_id": ...}]， + 校验为当前用户上传到同一会话的文件后补全 name、url；未传 chat_id 时取文件所在会话。 """ - if not file_ids: + field_file_ids = {} + for field in CHAT_FILE_LIST_FIELDS: + attachments = args.get(field) + if not attachments: + continue + if not isinstance(attachments, list) or not all( + isinstance(item, dict) and item.get("file_id") for item in attachments + ): + raise InvalidParamsError(f'{field} must be an array of {{"file_id": ...}}') + field_file_ids[field] = [str(item["file_id"]) for item in attachments] + if not field_file_ids: return chat_id, {} - if not isinstance(file_ids, list): - raise InvalidParamsError("file_ids must be an array") - files = list( - QuerySet(File).filter( - id__in=[str(file_id) for file_id in file_ids], + file_ids = {file_id for ids in field_file_ids.values() for file_id in ids} + files = { + str(file.id): file + for file in QuerySet(File).filter( + id__in=file_ids, source_type=FileSourceType.CHAT, meta__user_id=str(self.chat_user.id), ) - ) - source_ids = {str(file.source_id) for file in files} - if len(files) != len(set(map(str, file_ids))) or len(source_ids) != 1 or chat_id not in (None, *source_ids): - raise InvalidParamsError("file_ids must be files uploaded to the same chat_id") - chat_files = {} - for file in files: - chat_files.setdefault(get_chat_file_field(file.file_name), []).append( - {"name": file.file_name, "url": f"./oss/file/{file.id}", "file_id": str(file.id)} - ) + } + source_ids = {str(file.source_id) for file in files.values()} + if len(files) != len(file_ids) or len(source_ids) != 1 or chat_id not in (None, *source_ids): + raise InvalidParamsError("Attached files must be uploaded with upload_file to the same chat_id") + chat_files = { + field: [ + {"name": files[file_id].file_name, "url": f"./oss/file/{file_id}", "file_id": file_id} + for file_id in ids + ] + for field, ids in field_file_ids.items() + } return source_ids.pop(), chat_files def upload_file(self, args): @@ -360,7 +441,14 @@ def upload_file(self, args): "source_type": FileSourceType.CHAT, } ).upload(user_id=str(self.chat_user.id)) - result = {"file_id": file_url.split("/")[-1], "name": file_name, "url": file_url, "chat_id": chat_id} + result = { + "file_id": file_url.split("/")[-1], + "name": file_name, + "url": file_url, + # 按后缀给出建议的附件字段，调用方可据此放入对应的 *_list + "field": get_chat_file_field(file_name), + "chat_id": chat_id, + } return { "content": [{"type": "text", "text": json.dumps(result, ensure_ascii=False)}], "structuredContent": result, @@ -371,15 +459,23 @@ def call_tool(self, params): args = params.get("arguments") or {} if name == UPLOAD_FILE_TOOL: return self.upload_file(args) + if not self.application_id and name == LIST_AGENTS_TOOL: + return self.list_agents() + if not self.application_id and name == CALL_AGENT_TOOL: + application_id = str(args.get("application_id") or "") + if not self.is_accessible_application(application_id): + raise InvalidParamsError(f"Unknown application_id: {application_id}") + return self.call_agent(application_id, args) application_id = name[len(AGENT_TOOL_PREFIX) :] if name.startswith(AGENT_TOOL_PREFIX) else "" - if not application_id or not self.get_applications().filter(id=application_id).exists(): + if not self.application_id or not self.is_accessible_application(application_id): raise InvalidParamsError(f"Unknown tool: {name}") + return self.call_agent(application_id, args) + + def call_agent(self, application_id, args): form_data = args.get("form_data") or {} if not isinstance(form_data, dict): raise InvalidParamsError("form_data must be an object") - chat_id, chat_files = self.get_chat_files( - self.check_chat_id(args.get("chat_id"), application_id), args.get("file_ids") - ) + chat_id, chat_files = self.get_chat_files(self.check_chat_id(args.get("chat_id"), application_id), args) # 文件可能先上传到别的应用的会话，这里再按当前应用校验一次 chat_id = self.check_chat_id(chat_id, application_id) answer, chat_id = chat_application( diff --git a/apps/chat/tests/test_mcp.py b/apps/chat/tests/test_mcp.py index e74696f8538..79161aca96d 100644 --- a/apps/chat/tests/test_mcp.py +++ b/apps/chat/tests/test_mcp.py @@ -12,6 +12,20 @@ from common.exception.app_exception import AppAuthenticationFailed +APP_ID = "01a17d06-3bd0-7661-91dc-9d990282d863" + +AGENT_PARAMS = { + "message", + "chat_id", + "form_data", + "image_list", + "document_list", + "audio_list", + "video_list", + "other_list", +} + + class ChatApplicationTest(SimpleTestCase): @patch("common.init.init_template.run") @patch("chat.mcp.tools.ChatSerializers") @@ -71,36 +85,75 @@ def test_application_scoped_handler_only_exposes_that_application(self): ) self.assertEqual(ChatUserMCPToolHandler(SimpleNamespace(id="u"), auth, "app-9").get_application_ids(), []) + def scoped_handler(self): + return ChatUserMCPToolHandler(SimpleNamespace(id="user-id"), Auth(set(), {f"CHAT_USER:r:{APP_ID}": 1}), APP_ID) + @patch("chat.mcp.tools.QuerySet") - def test_list_tools_exposes_one_tool_per_application_and_upload(self, query_set): + def test_scoped_list_tools_exposes_the_agent_and_upload(self, query_set): query_set.return_value.filter.return_value.order_by.return_value = [ - SimpleNamespace(id="app-1", name="Agent", desc="desc") + SimpleNamespace(id=APP_ID, name="Agent", desc="desc") ] + tools = self.scoped_handler().list_tools()["tools"] + + self.assertEqual([tool["name"] for tool in tools], [f"agent_{APP_ID}", "upload_file"]) + self.assertTrue(tools[0]["description"].startswith("Agent desc\nAn AI agent.")) + self.assertEqual(set(tools[0]["inputSchema"]["properties"]), AGENT_PARAMS) + self.assertEqual(tools[0]["inputSchema"]["required"], ["message"]) + + def test_global_list_tools_exposes_list_and_call_agent(self): tools = self.handler.list_tools()["tools"] - self.assertEqual([tool["name"] for tool in tools], ["agent_app-1", "upload_file"]) - self.assertEqual(tools[0]["description"], "Agent desc") - self.assertEqual(set(tools[0]["inputSchema"]["properties"]), {"message", "chat_id", "file_ids", "form_data"}) + self.assertEqual([tool["name"] for tool in tools], ["list_agents", "call_agent", "upload_file"]) + self.assertEqual(set(tools[1]["inputSchema"]["properties"]), {"application_id", *AGENT_PARAMS}) + self.assertEqual(tools[1]["inputSchema"]["required"], ["application_id", "message"]) + + @patch.object(ChatUserMCPToolHandler, "get_applications") + def test_list_agents_returns_accessible_applications(self, get_applications): + get_applications.return_value = [SimpleNamespace(id=APP_ID, name="Agent", desc=None)] + + result = self.handler.call_tool({"name": "list_agents", "arguments": {}}) + + self.assertEqual( + result["structuredContent"], {"agents": [{"application_id": APP_ID, "name": "Agent", "desc": ""}]} + ) + self.assertEqual(json.loads(result["content"][0]["text"]), result["structuredContent"]["agents"]) @patch("chat.mcp.tools.chat_application", return_value=("answer", "chat-id")) @patch.object(ChatUserMCPToolHandler, "get_applications") - def test_call_tool_chats_as_chat_user_and_returns_chat_id(self, get_applications, chat_app): + def test_call_agent_tool_chats_as_chat_user_and_returns_chat_id(self, get_applications, chat_app): get_applications.return_value.filter.return_value.exists.return_value = True - result = self.handler.call_tool({"name": "agent_app-1", "arguments": {"message": "hi"}}) + result = self.scoped_handler().call_tool({"name": f"agent_{APP_ID}", "arguments": {"message": "hi"}}) - get_applications.return_value.filter.assert_called_once_with(id="app-1") - chat_app.assert_called_once_with("app-1", "user-id", ChatUserType.CHAT_USER.value, "hi", {}, {}, None) + get_applications.return_value.filter.assert_called_once_with(id=APP_ID) + chat_app.assert_called_once_with(APP_ID, "user-id", ChatUserType.CHAT_USER.value, "hi", {}, {}, None) self.assertEqual(result["content"][0], {"type": "text", "text": "answer"}) self.assertEqual(result["structuredContent"], {"answer": "answer", "chat_id": "chat-id"}) + @patch("chat.mcp.tools.chat_application", return_value=("answer", "chat-id")) + @patch.object(ChatUserMCPToolHandler, "get_applications") + def test_global_call_agent_uses_application_id(self, get_applications, chat_app): + get_applications.return_value.filter.return_value.exists.return_value = True + + self.handler.call_tool({"name": "call_agent", "arguments": {"application_id": APP_ID, "message": "hi"}}) + + chat_app.assert_called_once_with(APP_ID, "user-id", ChatUserType.CHAT_USER.value, "hi", {}, {}, None) + @patch.object(ChatUserMCPToolHandler, "get_applications") - def test_call_tool_rejects_unpermitted_application(self, get_applications): + def test_call_tool_rejects_unknown_tools_and_applications(self, get_applications): get_applications.return_value.filter.return_value.exists.return_value = False - - with self.assertRaises(InvalidParamsError): - self.handler.call_tool({"name": "agent_app-9", "arguments": {"message": "hi"}}) + cases = [ + (self.handler, "call_agent", {"application_id": APP_ID, "message": "hi"}), + (self.handler, "call_agent", {"application_id": "not-a-uuid", "message": "hi"}), + (self.handler, f"agent_{APP_ID}", {"message": "hi"}), + (self.scoped_handler(), "list_agents", {}), + (self.scoped_handler(), "call_agent", {"application_id": APP_ID, "message": "hi"}), + (self.scoped_handler(), "agent_not-a-uuid", {"message": "hi"}), + ] + for handler, name, arguments in cases: + with self.subTest(name=name, scoped=bool(handler.application_id)), self.assertRaises(InvalidParamsError): + handler.call_tool({"name": name, "arguments": arguments}) @patch("chat.mcp.tools.QuerySet") def test_check_chat_id(self, query_set): @@ -125,31 +178,41 @@ def test_check_chat_id(self, query_set): self.handler.check_chat_id("not-a-uuid") @patch("chat.mcp.tools.QuerySet") - def test_chat_files_are_grouped_and_infer_chat_id(self, query_set): + def test_chat_files_keep_caller_fields_and_infer_chat_id(self, query_set): query_set.return_value.filter.return_value = [ SimpleNamespace(id="f1", file_name="a.png", source_id="chat-1"), SimpleNamespace(id="f2", file_name="b.pdf", source_id="chat-1"), ] - chat_id, chat_files = self.handler.get_chat_files(None, ["f1", "f2"]) + chat_id, chat_files = self.handler.get_chat_files( + None, {"image_list": [{"file_id": "f1"}], "other_list": [{"file_id": "f2", "name": "ignored"}]} + ) self.assertEqual(chat_id, "chat-1") self.assertEqual( chat_files, { "image_list": [{"name": "a.png", "url": "./oss/file/f1", "file_id": "f1"}], - "document_list": [{"name": "b.pdf", "url": "./oss/file/f2", "file_id": "f2"}], + "other_list": [{"name": "b.pdf", "url": "./oss/file/f2", "file_id": "f2"}], }, ) + def test_chat_files_without_attachments(self): + self.assertEqual(self.handler.get_chat_files("chat-1", {"message": "hi"}), ("chat-1", {})) + @patch("chat.mcp.tools.QuerySet") def test_chat_files_must_belong_to_the_chat(self, query_set): query_set.return_value.filter.return_value = [SimpleNamespace(id="f1", file_name="a.png", source_id="chat-1")] with self.assertRaises(InvalidParamsError): - self.handler.get_chat_files("chat-2", ["f1"]) + self.handler.get_chat_files("chat-2", {"image_list": [{"file_id": "f1"}]}) with self.assertRaises(InvalidParamsError): - self.handler.get_chat_files(None, ["f1", "missing"]) + self.handler.get_chat_files(None, {"image_list": [{"file_id": "f1"}, {"file_id": "missing"}]}) + + def test_chat_files_reject_invalid_format(self): + for attachments in (["f1"], {"file_id": "f1"}, [{"name": "a.png"}]): + with self.subTest(attachments=attachments), self.assertRaises(InvalidParamsError): + self.handler.get_chat_files(None, {"image_list": attachments}) @patch("oss.serializers.file.FileSerializer") def test_upload_file_returns_file_and_new_chat_id(self, file_serializer): @@ -162,6 +225,7 @@ def test_upload_file_returns_file_and_new_chat_id(self, file_serializer): self.assertEqual(data["source_id"], result["chat_id"]) file_serializer.return_value.upload.assert_called_once_with(user_id="user-id") self.assertEqual(result["file_id"], "file-1") + self.assertEqual(result["field"], "document_list") def test_upload_file_rejects_invalid_content(self): with self.assertRaises(InvalidParamsError): @@ -219,3 +283,66 @@ def test_application_scoped_route(self): self.assertIs(match.func, mcp_view) self.assertEqual(match.kwargs, {"application_id": "app-1"}) + + def post_sse(self, body): + request = self.factory.post( + "/chat/api/v3/mcp", + data=json.dumps(body), + content_type="application/json", + HTTP_AUTHORIZATION="Bearer chat-key", + HTTP_ACCEPT="application/json, text/event-stream", + ) + return mcp_view(request) + + @staticmethod + def read_sse(response): + return b"".join(response.streaming_content).decode() + + @patch("chat.views.v3.mcp.SSE_HEARTBEAT_SECONDS", 0.01) + @patch("chat.views.v3.mcp.ChatUserMCPToolHandler") + @patch("chat.views.v3.mcp.ChatUserKey.handle", return_value=(SimpleNamespace(profile=None), MagicMock())) + def test_tool_call_streams_progress_then_result(self, _handle, handler_class): + import time + + def slow_call(params): + time.sleep(0.05) + return {"content": [{"type": "text", "text": "done"}]} + + handler_class.return_value.call_tool.side_effect = slow_call + + response = self.post_sse( + { + "jsonrpc": "2.0", + "id": 7, + "method": "tools/call", + "params": {"name": "agent_x", "arguments": {}, "_meta": {"progressToken": "p1"}}, + } + ) + + self.assertEqual(response["Content-Type"], "text/event-stream; charset=utf-8") + events = [json.loads(line[6:]) for line in self.read_sse(response).splitlines() if line.startswith("data: ")] + self.assertEqual(events[0]["method"], "notifications/progress") + self.assertEqual(events[0]["params"]["progressToken"], "p1") + self.assertEqual( + events[-1], {"jsonrpc": "2.0", "id": 7, "result": {"content": [{"type": "text", "text": "done"}]}} + ) + + @patch("chat.views.v3.mcp.SSE_HEARTBEAT_SECONDS", 0.01) + @patch("chat.views.v3.mcp.ChatUserMCPToolHandler") + @patch("chat.views.v3.mcp.ChatUserKey.handle", return_value=(SimpleNamespace(profile=None), MagicMock())) + def test_tool_call_stream_sends_ping_without_progress_token_and_maps_errors(self, _handle, handler_class): + import time + + def failing_call(params): + time.sleep(0.05) + raise InvalidParamsError("Unknown tool: x") + + handler_class.return_value.call_tool.side_effect = failing_call + + body = self.read_sse( + self.post_sse({"jsonrpc": "2.0", "id": 8, "method": "tools/call", "params": {"name": "x"}}) + ) + + self.assertIn(": ping", body) + last = json.loads([line for line in body.splitlines() if line.startswith("data: ")][-1][6:]) + self.assertEqual(last["error"], {"code": -32602, "message": "Unknown tool: x"}) diff --git a/apps/chat/views/v3/mcp.py b/apps/chat/views/v3/mcp.py index 8cc887a2302..971b3abd3b9 100644 --- a/apps/chat/views/v3/mcp.py +++ b/apps/chat/views/v3/mcp.py @@ -1,6 +1,9 @@ import json +import queue +import threading -from django.http import JsonResponse, HttpResponse +from django.db import connection +from django.http import JsonResponse, HttpResponse, StreamingHttpResponse from django.views.decorators.csrf import csrf_exempt from django.views.decorators.http import require_POST @@ -18,6 +21,74 @@ def rpc_response(request_id, result=None, code=None, message=None, status=200): return response +# 智能体执行期间推送进度/心跳的间隔（秒），保证客户端读超时不会在长任务中触发 +SSE_HEARTBEAT_SECONDS = 15 + + +def rpc_error(error): + if isinstance(error, InvalidParamsError): + return -32602, str(error) + return -32603, f"Internal error: {str(error)}" + + +def sse_event(data): + return f"event: message\ndata: {json.dumps(data, ensure_ascii=False)}\n\n" + + +def stream_tool_call(handler, params, request_id): + """ + tools/call 以 SSE 返回：后台线程执行智能体，期间定时推送进度通知（客户端带 progressToken 时）或心跳注释， + 结束后推送 JSON-RPC 结果。长任务不会因为长时间收不到数据而触发客户端 ReadTimeout。 + """ + result_queue = queue.Queue() + + def run(): + try: + result_queue.put(("result", handler.call_tool(params))) + except Exception as e: + result_queue.put(("error", e)) + finally: + connection.close() + + threading.Thread(target=run, daemon=True).start() + progress_token = (params.get("_meta") or {}).get("progressToken") + progress = 0 + while True: + try: + kind, value = result_queue.get(timeout=SSE_HEARTBEAT_SECONDS) + except queue.Empty: + progress += 1 + if progress_token is None: + yield ": ping\n\n" + else: + yield sse_event( + { + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": { + "progressToken": progress_token, + "progress": progress, + "message": "Agent is working", + }, + } + ) + continue + if kind == "result": + yield sse_event({"jsonrpc": "2.0", "id": request_id, "result": value}) + else: + code, message = rpc_error(value) + yield sse_event({"jsonrpc": "2.0", "id": request_id, "error": {"code": code, "message": message}}) + return + + +def sse_response(stream): + response = StreamingHttpResponse(stream, content_type="text/event-stream; charset=utf-8") + response["Cache-Control"] = "no-cache" + # 关闭 Nginx 等反向代理的缓冲，心跳才能及时到达客户端 + response["X-Accel-Buffering"] = "no" + return response + + def authenticate(request, application_id=None): """ MCP 接口使用对话用户 API Key（chat-xxx）认证；带 application_id 时只暴露该应用 @@ -59,6 +130,9 @@ def mcp_view(request, application_id=None): result = handler.list_tools() elif method == "tools/call": + # 智能体执行可能长达数分钟，客户端支持 SSE 时流式返回 + if "text/event-stream" in request.headers.get("Accept", ""): + return sse_response(stream_tool_call(handler, params, request_id)) result = handler.call_tool(params) else: @@ -69,9 +143,8 @@ def mcp_view(request, application_id=None): except AppAuthenticationFailed as e: return rpc_response(request_id, code=-32001, message=str(e.message), status=401) - except InvalidParamsError as e: - return rpc_response(request_id, code=-32602, message=str(e)) except (json.JSONDecodeError, UnicodeDecodeError): return rpc_response(None, code=-32700, message="Parse error", status=400) except Exception as e: - return rpc_response(request_id, code=-32603, message=f"Internal error: {str(e)}") + code, message = rpc_error(e) + return rpc_response(request_id, code=code, message=message)