diff --git a/apps/application/workflow/nodes/ai_chat_node/tools/application.py b/apps/application/workflow/nodes/ai_chat_node/tools/application.py index c3091a13293..353984990e6 100644 --- a/apps/application/workflow/nodes/ai_chat_node/tools/application.py +++ b/apps/application/workflow/nodes/ai_chat_node/tools/application.py @@ -26,7 +26,7 @@ def get_application_args(): """ 应用(Agent)工具对模型暴露的入参:固定单个必填 message。 - 与 chat.mcp.tools.MCPToolHandler.list_tools 的 inputSchema 保持一致, + 与 chat.mcp.application.MCPToolHandler.list_tools 的 inputSchema 保持一致, 这样从远端 MCP 代理切换为进程内直调时,模型侧契约不变。 """ return build_schema( diff --git a/apps/chat/mcp/tools.py b/apps/chat/mcp/application.py similarity index 96% rename from apps/chat/mcp/tools.py rename to apps/chat/mcp/application.py index 36d9cc13c15..03136bc05f1 100644 --- a/apps/chat/mcp/tools.py +++ b/apps/chat/mcp/application.py @@ -1,3 +1,5 @@ +"""Application MCP tools: application API key (v2) and chat user API key (v3) handlers.""" + import base64 import binascii import json @@ -10,6 +12,7 @@ from django.db.models import QuerySet from django.utils import timezone +from chat.mcp.common import initialize_result, json_result from chat.serializers.chat import ChatSerializers from common.auth.constants.group_constants import Group from knowledge.models import File, FileSourceType @@ -32,8 +35,6 @@ "video_list": ("MP4", "AVI", "MKV", "MOV", "FLV", "WMV"), } -MCP_PROTOCOL_VERSION = "2025-06-18" - AGENT_TOOL_PREFIX = "agent_" UPLOAD_FILE_TOOL = "upload_file" @@ -167,12 +168,8 @@ def decode_form_data(form_data): return {} return form_data - def initialize(self): - return { - "protocolVersion": MCP_PROTOCOL_VERSION, - "serverInfo": {"name": "maxkb-mcp", "version": "1.0.0"}, - "capabilities": {"tools": {}}, - } + def initialize(self, params=None): + return initialize_result(params, "maxkb-mcp") def build_description(self): """ @@ -229,12 +226,8 @@ def __init__(self, chat_user, auth, application_id=None): # 单应用 MCP 地址只暴露该应用 self.application_id = application_id - def initialize(self): - return { - "protocolVersion": MCP_PROTOCOL_VERSION, - "serverInfo": {"name": "maxkb-mcp", "version": "1.0.0"}, - "capabilities": {"tools": {}}, - } + def initialize(self, params=None): + return initialize_result(params, "maxkb-mcp") def get_application_ids(self): prefix = f"{Group.CHAT_USER}:r:" @@ -360,10 +353,7 @@ def list_agents(self): {"application_id": str(application.id), "name": application.name, "desc": application.desc or ""} for application in self.get_applications() ] - return { - "content": [{"type": "text", "text": json.dumps(agents, ensure_ascii=False)}], - "structuredContent": {"agents": agents}, - } + return json_result(agents, {"agents": agents}) def check_chat_id(self, chat_id, application_id=None): """ @@ -449,10 +439,7 @@ def upload_file(self, args): "field": get_chat_file_field(file_name), "chat_id": chat_id, } - return { - "content": [{"type": "text", "text": json.dumps(result, ensure_ascii=False)}], - "structuredContent": result, - } + return json_result(result) def call_tool(self, params): name = str(params.get("name") or "") diff --git a/apps/chat/mcp/common.py b/apps/chat/mcp/common.py new file mode 100644 index 00000000000..24c20c538f7 --- /dev/null +++ b/apps/chat/mcp/common.py @@ -0,0 +1,29 @@ +"""MCP protocol helpers shared by the application and knowledge MCP servers.""" + +import json + +PROTOCOL_VERSIONS = ("2025-03-26", "2025-06-18", "2025-11-25") + +DEFAULT_PROTOCOL_VERSION = "2025-06-18" + + +def initialize_result(params, server_name): + # 客户端请求的版本受支持时原样返回,否则返回默认版本,由客户端决定是否继续 + version = (params or {}).get("protocolVersion") + return { + "protocolVersion": version if version in PROTOCOL_VERSIONS else DEFAULT_PROTOCOL_VERSION, + "serverInfo": {"name": server_name, "version": "1.0.0"}, + "capabilities": {"tools": {}}, + } + + +def json_result(data, structured_content=None): + """工具结果:文本内容为 data 的 JSON,structuredContent 默认与 data 相同""" + return { + "content": [{"type": "text", "text": json.dumps(data, ensure_ascii=False)}], + "structuredContent": data if structured_content is None else structured_content, + } + + +def tool_error(text): + return {"isError": True, "content": [{"type": "text", "text": text}]} diff --git a/apps/chat/mcp/knowledge.py b/apps/chat/mcp/knowledge.py index 223b40a1f40..40dee8f5555 100644 --- a/apps/chat/mcp/knowledge.py +++ b/apps/chat/mcp/knowledge.py @@ -1,13 +1,58 @@ """Knowledge equivalent of the application's MCPToolHandler.""" import json +import uuid +from django.db.models import QuerySet from rest_framework.exceptions import ValidationError +from chat.mcp.common import initialize_result, json_result, tool_error +from knowledge.models import Knowledge from knowledge.services.external_retrieval import retrieve -from knowledge.services.retrieval_access import RetrievalError +from knowledge.services.retrieval_access import RetrievalError, authorized_ids -PROTOCOL_VERSIONS = ("2025-03-26", "2025-06-18", "2025-11-25") +# 检索参数,单知识库 MCP 与全局知识库 MCP 的 search_knowledge 共用 +RETRIEVAL_INPUT_PROPERTIES = { + "query_text": {"type": "string", "minLength": 1, "maxLength": 8000}, + "top_number": {"type": "integer", "minimum": 1, "maximum": 50, "default": 5}, + "similarity": {"type": "number", "minimum": 0, "maximum": 1, "default": 0}, + "search_mode": { + "type": "string", + "enum": ["embedding", "keywords", "blend"], + "default": "embedding", + }, +} + +# 全局知识库 MCP 地址:先列出可检索的知识库,再按 knowledge_id 检索 +LIST_KNOWLEDGE_TOOL = "list_knowledge" + +SEARCH_KNOWLEDGE_TOOL = "search_knowledge" + + +def initialize_knowledge(params): + if ( + not isinstance(params.get("protocolVersion"), str) + or not isinstance(params.get("capabilities"), dict) + or not isinstance(params.get("clientInfo"), dict) + ): + raise ValidationError("Invalid initialization parameters.") + return initialize_result(params, "maxkb-knowledge-mcp") + + +def retrieve_result(knowledge_id, identity, arguments): + """ + 执行检索并组装工具结果。参数不合法、开关关闭或授权被撤销等以工具错误返回并附原因, + 调用方模型可据此修正后重试,而不是让 MCP 会话因 HTTP 错误中断。 + """ + try: + output = retrieve(knowledge_id, identity, arguments) + except RetrievalError as e: + return tool_error(e.message) + except ValidationError as e: + return tool_error(json.dumps(e.detail, ensure_ascii=False, default=str)) + except Exception: + return tool_error("Knowledge retrieval failed.") + return json_result(output) class KnowledgeMCPToolHandler: @@ -16,39 +61,23 @@ def __init__(self, knowledge, identity): self.tool_name = f"knowledge_{knowledge.id}" def initialize(self, params): - version = params.get("protocolVersion") - if ( - not isinstance(version, str) - or not isinstance(params.get("capabilities"), dict) - or not isinstance(params.get("clientInfo"), dict) - ): - raise ValidationError("Invalid initialization parameters.") - return { - "protocolVersion": version if version in PROTOCOL_VERSIONS else "2025-06-18", - "serverInfo": {"name": "maxkb-knowledge-mcp", "version": "1.0.0"}, - "capabilities": {"tools": {}}, - } + return initialize_knowledge(params) def list_tools(self): + # 带上知识库描述,调用方模型才能判断什么问题该检索该知识库 + description = f"检索知识库:{self.knowledge.name}" + if self.knowledge.desc: + description = f"{description}\n{self.knowledge.desc}" return { "tools": [ { "name": self.tool_name, - "description": f"检索知识库:{self.knowledge.name}", + "description": description, "inputSchema": { "type": "object", "additionalProperties": False, "required": ["query_text"], - "properties": { - "query_text": {"type": "string", "minLength": 1, "maxLength": 8000}, - "top_number": {"type": "integer", "minimum": 1, "maximum": 50, "default": 5}, - "similarity": {"type": "number", "minimum": 0, "maximum": 1, "default": 0}, - "search_mode": { - "type": "string", - "enum": ["embedding", "keywords", "blend"], - "default": "embedding", - }, - }, + "properties": RETRIEVAL_INPUT_PROPERTIES, }, } ] @@ -57,10 +86,83 @@ def list_tools(self): def call_tool(self, params): if params.get("name") != self.tool_name: raise ValidationError("Unknown tool.") + return retrieve_result(self.knowledge.id, self.identity, params.get("arguments", {})) + + +class ChatUserKnowledgeMCPToolHandler: + """ + 全局知识库 MCP:以对话用户 API Key 的身份列出并检索其可访问的知识库,知识库再多工具数量也固定。 + 可访问与单知识库 MCP 地址的判定一致:外部检索服务已开启,且未开启身份认证或当前对话用户已获授权。 + """ + + def __init__(self, identity): + self.identity = identity + + def initialize(self, params): + return initialize_knowledge(params) + + def list_tools(self): + return { + "tools": [ + { + "name": LIST_KNOWLEDGE_TOOL, + "description": "List the knowledge bases you can search, with their id, name and description. " + f"Pick the one that fits the question and search it with {SEARCH_KNOWLEDGE_TOOL}.", + "inputSchema": {"type": "object", "properties": {}}, + }, + { + "name": SEARCH_KNOWLEDGE_TOOL, + "description": f"Search a knowledge base from {LIST_KNOWLEDGE_TOOL} by knowledge_id. " + "Returns the matching paragraphs with similarity scores and citations.", + "inputSchema": { + "type": "object", + "additionalProperties": False, + "required": ["knowledge_id", "query_text"], + "properties": { + "knowledge_id": { + "type": "string", + "description": f"Id of the knowledge base to search, from {LIST_KNOWLEDGE_TOOL}.", + }, + **RETRIEVAL_INPUT_PROPERTIES, + }, + }, + }, + ] + } + + def get_knowledge(self): + knowledge = list(QuerySet(Knowledge).filter(external_service__enabled=True).order_by("name")) + restricted = [str(item.id) for item in knowledge if item.external_service.get("authentication", False)] + authorized = authorized_ids(self.identity.user_id, restricted) if restricted else set() + return [ + item + for item in knowledge + if not item.external_service.get("authentication", False) or str(item.id) in authorized + ] + + def list_knowledge(self): + knowledge = [ + {"knowledge_id": str(item.id), "name": item.name, "desc": item.desc or ""} for item in self.get_knowledge() + ] + return json_result(knowledge, {"knowledge": knowledge}) + + def search_knowledge(self, arguments): + arguments = dict(arguments) + knowledge_id = arguments.pop("knowledge_id", None) try: - output = retrieve(self.knowledge.id, self.identity, params.get("arguments", {})) - except (RetrievalError, ValidationError): - raise - except Exception: - return {"isError": True, "content": [{"type": "text", "text": "Knowledge retrieval failed."}]} - return {"content": [{"type": "text", "text": json.dumps(output, ensure_ascii=False)}]} + knowledge_id = str(uuid.UUID(str(knowledge_id))) + except ValueError: + return tool_error(f"Unknown knowledge_id: {knowledge_id}") + # retrieve 内部会再次校验开关与授权,并在检索前后复核 API Key + return retrieve_result(knowledge_id, self.identity, arguments) + + def call_tool(self, params): + name = params.get("name") + arguments = params.get("arguments") or {} + if name == LIST_KNOWLEDGE_TOOL: + return self.list_knowledge() + if name == SEARCH_KNOWLEDGE_TOOL: + if not isinstance(arguments, dict): + raise ValidationError("Invalid arguments.") + return self.search_knowledge(arguments) + raise ValidationError("Unknown tool.") diff --git a/apps/chat/tests/test_mcp.py b/apps/chat/tests/test_mcp.py index 79161aca96d..4b8bb7e77ee 100644 --- a/apps/chat/tests/test_mcp.py +++ b/apps/chat/tests/test_mcp.py @@ -6,7 +6,7 @@ from django.urls import resolve from application.models import ChatSourceChoices, ChatUserType -from chat.mcp.tools import ChatUserMCPToolHandler, InvalidParamsError, chat_application, get_chat_file_field +from chat.mcp.application import ChatUserMCPToolHandler, InvalidParamsError, chat_application, get_chat_file_field from chat.views.v3.mcp import mcp_view from common.auth.struct.auth import Auth from common.exception.app_exception import AppAuthenticationFailed @@ -28,7 +28,7 @@ class ChatApplicationTest(SimpleTestCase): @patch("common.init.init_template.run") - @patch("chat.mcp.tools.ChatSerializers") + @patch("chat.mcp.application.ChatSerializers") def test_chats_as_given_chat_user_in_a_new_chat(self, chat_serializers, _init): chat_serializers.return_value.chat.return_value = SimpleNamespace( content=json.dumps({"data": {"content": "answertool"}}).encode() @@ -50,7 +50,7 @@ def test_chats_as_given_chat_user_in_a_new_chat(self, chat_serializers, _init): self.assertEqual(result, ("answer", chat_info["chat_id"])) @patch("common.init.init_template.run") - @patch("chat.mcp.tools.ChatSerializers") + @patch("chat.mcp.application.ChatSerializers") def test_continues_given_chat(self, chat_serializers, _init): chat_serializers.return_value.chat.return_value = SimpleNamespace(content=b'{"data": {"content": "a"}}') @@ -88,7 +88,7 @@ def test_application_scoped_handler_only_exposes_that_application(self): 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") + @patch("chat.mcp.application.QuerySet") 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_ID, name="Agent", desc="desc") @@ -119,7 +119,7 @@ def test_list_agents_returns_accessible_applications(self, get_applications): ) self.assertEqual(json.loads(result["content"][0]["text"]), result["structuredContent"]["agents"]) - @patch("chat.mcp.tools.chat_application", return_value=("answer", "chat-id")) + @patch("chat.mcp.application.chat_application", return_value=("answer", "chat-id")) @patch.object(ChatUserMCPToolHandler, "get_applications") 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 @@ -131,7 +131,7 @@ def test_call_agent_tool_chats_as_chat_user_and_returns_chat_id(self, get_applic 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("chat.mcp.application.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 @@ -155,7 +155,7 @@ def test_call_tool_rejects_unknown_tools_and_applications(self, get_applications 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") + @patch("chat.mcp.application.QuerySet") def test_check_chat_id(self, query_set): chat_id = "01a17d06-3bd0-7661-91dc-9d990282d863" query_set.return_value.filter.return_value.first.return_value = None @@ -177,7 +177,7 @@ def test_check_chat_id(self, query_set): with self.assertRaises(InvalidParamsError): self.handler.check_chat_id("not-a-uuid") - @patch("chat.mcp.tools.QuerySet") + @patch("chat.mcp.application.QuerySet") 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"), @@ -200,7 +200,7 @@ def test_chat_files_keep_caller_fields_and_infer_chat_id(self, query_set): 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") + @patch("chat.mcp.application.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")] diff --git a/apps/chat/tests/test_mcp_common.py b/apps/chat/tests/test_mcp_common.py new file mode 100644 index 00000000000..82c5d3e6787 --- /dev/null +++ b/apps/chat/tests/test_mcp_common.py @@ -0,0 +1,19 @@ +from django.test import SimpleTestCase + +from chat.mcp.common import DEFAULT_PROTOCOL_VERSION, initialize_result, json_result, tool_error + + +class MCPCommonTest(SimpleTestCase): + def test_initialize_negotiates_protocol_version(self): + self.assertEqual(initialize_result({"protocolVersion": "2025-11-25"}, "s")["protocolVersion"], "2025-11-25") + for params in ({"protocolVersion": "1999-01-01"}, {}, None): + self.assertEqual(initialize_result(params, "s")["protocolVersion"], DEFAULT_PROTOCOL_VERSION) + self.assertEqual(initialize_result({}, "maxkb-mcp")["serverInfo"]["name"], "maxkb-mcp") + + def test_json_result_and_tool_error(self): + self.assertEqual( + json_result(["中文"], {"items": ["中文"]}), + {"content": [{"type": "text", "text": '["中文"]'}], "structuredContent": {"items": ["中文"]}}, + ) + self.assertEqual(json_result({"a": 1})["structuredContent"], {"a": 1}) + self.assertEqual(tool_error("bad"), {"isError": True, "content": [{"type": "text", "text": "bad"}]}) diff --git a/apps/chat/urls.py b/apps/chat/urls.py index e0ad7485542..67be68c0154 100644 --- a/apps/chat/urls.py +++ b/apps/chat/urls.py @@ -2,7 +2,7 @@ from application.views import ChatRecordDetailView, ChatRecordLinkView from chat.views import v2 as v2_views, v3 as v3_views -from chat.views.v3.knowledge import knowledge_mcp_view, retrieve_view +from chat.views.v3.knowledge import knowledge_global_mcp_view, knowledge_mcp_view, retrieve_view app_name = "chat" # @formatter:off @@ -31,7 +31,9 @@ ])), # ---- 全局(非 application 作用域)---- path('embed', v3_views.ChatEmbedView.as_view()), + # 这块是application的MCP 因为knowledge和application时分开控制的所以mcp也是分开的所以就按照接口分开 path('mcp', v3_views.mcp_view), + path('knowledge_mcp', knowledge_global_mcp_view), path('auth/anonymous', v3_views.AnonymousAuthentication.as_view()), path('auth/login', v3_views.LocalLoginView.as_view()), path('auth/logout', v3_views.Logout.as_view(), name='v3_logout'), diff --git a/apps/chat/views/v2/mcp.py b/apps/chat/views/v2/mcp.py index f4733a8c050..8a9bf84b048 100644 --- a/apps/chat/views/v2/mcp.py +++ b/apps/chat/views/v2/mcp.py @@ -3,7 +3,7 @@ from django.http import HttpResponse, JsonResponse from django.views.decorators.csrf import csrf_exempt -from chat.mcp.tools import MCPToolHandler +from chat.mcp.application import MCPToolHandler @csrf_exempt @@ -27,7 +27,7 @@ def mcp_view(request): # 路由方法 if method == "initialize": - result = handler.initialize() + result = handler.initialize(params) elif method == "tools/list": result = handler.list_tools() diff --git a/apps/chat/views/v3/knowledge.py b/apps/chat/views/v3/knowledge.py index 34cdade3f4a..e8a1271d5de 100644 --- a/apps/chat/views/v3/knowledge.py +++ b/apps/chat/views/v3/knowledge.py @@ -8,9 +8,15 @@ from django.views.decorators.http import require_POST from rest_framework.exceptions import ValidationError -from chat.mcp.knowledge import KnowledgeMCPToolHandler, PROTOCOL_VERSIONS +from chat.mcp.common import PROTOCOL_VERSIONS +from chat.mcp.knowledge import ChatUserKnowledgeMCPToolHandler, KnowledgeMCPToolHandler from knowledge.services.external_retrieval import retrieve -from knowledge.services.retrieval_access import RetrievalError, authenticate_key, authorize_external +from knowledge.services.retrieval_access import ( + RetrievalError, + authenticate_key, + authorize_external, + check_knowledge_api_enabled, +) def response_headers(response): @@ -21,14 +27,22 @@ def response_headers(response): return response -def read_request(request, knowledge_id): +def check_origin(request): origin = request.headers.get("Origin") if origin: expected = urlsplit(request.build_absolute_uri("/")) if origin != f"{expected.scheme}://{expected.netloc}": raise RetrievalError("invalid_origin", "Origin is not allowed.", 403) + + +def read_request(request, knowledge_id): + check_origin(request) identity = authenticate_key(request.headers.get("Authorization")) knowledge = authorize_external(knowledge_id, identity) + return knowledge, identity, read_json(request) + + +def read_json(request): if request.content_type != "application/json": raise RetrievalError("invalid_content_type", "Use application/json.", 415) if int(request.META.get("CONTENT_LENGTH") or 0) > 65536: @@ -36,7 +50,7 @@ def read_request(request, knowledge_id): body = request.body if len(body) > 65536: raise RetrievalError("request_too_large", "Request body is too large.", 413) - return knowledge, identity, json.loads(body) + return json.loads(body) def error_response(error): @@ -68,9 +82,36 @@ def rpc_response(request_id, result=None, code=None, message=None, status=200): @csrf_exempt @require_POST def knowledge_mcp_view(request, knowledge_id): + def read(): + knowledge, identity, data = read_request(request, knowledge_id) + return KnowledgeMCPToolHandler(knowledge, identity), data + + return mcp_response(request, read) + + +@csrf_exempt +@require_POST +def knowledge_global_mcp_view(request): + """ + 全局知识库 MCP:不绑定知识库,必须带对话用户 API Key,列出并检索该用户可访问的知识库 + """ + + def read(): + check_origin(request) + # list_knowledge 不经过 authorize_external,入口处先检查开关 + check_knowledge_api_enabled() + identity = authenticate_key(request.headers.get("Authorization")) + if not identity.user_id: + raise RetrievalError("authentication_required", "A chat user API key is required.", 401) + return ChatUserKnowledgeMCPToolHandler(identity), read_json(request) + + return mcp_response(request, read) + + +def mcp_response(request, read): request_id = None try: - knowledge, identity, data = read_request(request, knowledge_id) + handler, data = read() accept = request.headers.get("Accept", "") if not all(value in accept for value in ("application/json", "text/event-stream")): return rpc_response( @@ -92,7 +133,6 @@ def knowledge_mcp_view(request, knowledge_id): ) if "id" not in data: return response_headers(HttpResponse(status=202)) - handler = KnowledgeMCPToolHandler(knowledge, identity) method = data["method"] if method == "initialize": output = handler.initialize(params) diff --git a/apps/chat/views/v3/mcp.py b/apps/chat/views/v3/mcp.py index 971b3abd3b9..aed1e4dd825 100644 --- a/apps/chat/views/v3/mcp.py +++ b/apps/chat/views/v3/mcp.py @@ -7,7 +7,7 @@ from django.views.decorators.csrf import csrf_exempt from django.views.decorators.http import require_POST -from chat.mcp.tools import ChatUserMCPToolHandler, InvalidParamsError +from chat.mcp.application import ChatUserMCPToolHandler, InvalidParamsError from common.auth.handle.impl.chat_user_key import ChatUserKey from common.exception.app_exception import AppAuthenticationFailed @@ -121,7 +121,7 @@ def mcp_view(request, application_id=None): # 路由方法 if method == "initialize": - result = handler.initialize() + result = handler.initialize(params) elif method == "ping": result = {} diff --git a/apps/common/auth/handle/impl/chat_user_key.py b/apps/common/auth/handle/impl/chat_user_key.py index d64ea71b68e..0eb67c0c7ab 100644 --- a/apps/common/auth/handle/impl/chat_user_key.py +++ b/apps/common/auth/handle/impl/chat_user_key.py @@ -15,6 +15,7 @@ from common.auth.handle.impl.chat_user_token import get_auth from common.auth.struct.auth import Principal from common.exception.app_exception import AppAuthenticationFailed +from portal.models import get_portal_setting from system_manage.models import ChatUser, ChatUserApiKey @@ -28,6 +29,9 @@ def handle(self, request, token: str, get_token_details): chat_user = QuerySet(ChatUser).filter(id=chat_user_api_key.user_id).first() if chat_user is None or not chat_user.is_active: raise AppAuthenticationFailed(500, _("Secret key is invalid")) + # 门户「智能体后端 API 访问」开关控制对话用户 API Key 的全部对话接口(含应用 MCP),知识库检索由独立开关控制 + if not get_portal_setting("enable_api"): + raise AppAuthenticationFailed(500, _("API service is disabled")) auth = get_auth(chat_user.source, chat_user.id, None) return Principal(str(chat_user.id), ChatUserType.CHAT_USER, profile=chat_user), auth diff --git a/apps/common/database_model_manage/handle/impl/default_base_model_handle.py b/apps/common/database_model_manage/handle/impl/default_base_model_handle.py index 1c02b1dd891..74158c8ff50 100644 --- a/apps/common/database_model_manage/handle/impl/default_base_model_handle.py +++ b/apps/common/database_model_manage/handle/impl/default_base_model_handle.py @@ -1,14 +1,56 @@ # coding=utf-8 """ - @project: MaxKB - @Author:虎虎 - @file: default_base_model_handle.py - @date:2025/4/15 11:20 - @desc: +@project: MaxKB +@Author:虎虎 +@file: default_base_model_handle.py +@date:2025/4/15 11:20 +@desc: """ + +from typing import List + +from django.db.models import QuerySet + from common.database_model_manage.handle.base_handle import IBaseModelHandle +from system_manage.models import ( + ResourceChatUserAuthorize, + ResourceChatUserGroupAuthorize, + ResourceType, + UserGroupRelation, +) + + +def get_knowledge_list_of_authorized(chat_user_id, knowledge_id_list: List): + """ + 从 knowledge_id_list 中筛出对话用户被授权(直接授权或所在用户组授权)的知识库 id + """ + knowledge_id_list = [str(knowledge_id) for knowledge_id in knowledge_id_list] + group_id_list = QuerySet(UserGroupRelation).filter(user_id=chat_user_id).values_list("group_id", flat=True) + # 通过对话用户组授权的知识库id + auth_knowledge_id_by_group_list = [ + str(resource_id) + for resource_id in QuerySet(ResourceChatUserGroupAuthorize) + .filter( + user_group_id__in=group_id_list, + resource_type=ResourceType.KNOWLEDGE, + resource_id__in=knowledge_id_list, + is_auth=True, + ) + .values_list("resource_id", flat=True) + ] + # 通过对话用户直接授权的知识库id + auth_knowledge_id_by_user_list = [ + str(resource_id) + for resource_id in QuerySet(ResourceChatUserAuthorize) + .filter( + user_id=chat_user_id, resource_type=ResourceType.KNOWLEDGE, resource_id__in=knowledge_id_list, is_auth=True + ) + .values_list("resource_id", flat=True) + ] + authorized = {*auth_knowledge_id_by_group_list, *auth_knowledge_id_by_user_list} + return [knowledge_id for knowledge_id in knowledge_id_list if knowledge_id in authorized] class DefaultBaseModelHandle(IBaseModelHandle): def get_model_dict(self): - return {} + return {"get_knowledge_list_of_authorized": get_knowledge_list_of_authorized} diff --git a/apps/common/tests/test_chat_user_key.py b/apps/common/tests/test_chat_user_key.py new file mode 100644 index 00000000000..5bba2f403bf --- /dev/null +++ b/apps/common/tests/test_chat_user_key.py @@ -0,0 +1,47 @@ +from types import SimpleNamespace +from unittest.mock import patch + +from django.test import SimpleTestCase + +from common.auth.handle.impl import chat_user_key +from common.auth.handle.impl.chat_user_key import ChatUserKey +from common.exception.app_exception import AppAuthenticationFailed +from portal.models import get_portal_setting + +USER_ID = "40000000-0000-0000-0000-000000000001" + + +class ChatUserKeyPortalSwitchTest(SimpleTestCase): + def setUp(self): + api_key = SimpleNamespace(is_active=True, user_id=USER_ID) + chat_user = SimpleNamespace(id=USER_ID, is_active=True, source="LOCAL") + query_set = patch.object(chat_user_key, "QuerySet").start() + query_set.side_effect = lambda model: SimpleNamespace( + filter=lambda **kwargs: SimpleNamespace( + first=lambda: api_key if model is chat_user_key.ChatUserApiKey else chat_user + ) + ) + self.get_auth = patch.object(chat_user_key, "get_auth", return_value="auth").start() + self.portal = patch.object(chat_user_key, "get_portal_setting", return_value=True).start() + self.addCleanup(patch.stopall) + + def test_api_enabled_authenticates(self): + principal, auth = ChatUserKey().handle(None, "chat-key", None) + self.portal.assert_called_once_with("enable_api") + self.assertEqual((principal.profile.id, auth), (USER_ID, "auth")) + + def test_api_disabled_rejects_key(self): + self.portal.return_value = False + with self.assertRaises(AppAuthenticationFailed): + ChatUserKey().handle(None, "chat-key", None) + self.get_auth.assert_not_called() + + +class PortalSettingTest(SimpleTestCase): + @patch("portal.models.portal.Portal.objects") + def test_uses_field_default_without_portal(self, objects): + objects.values_list.return_value.first.return_value = None + self.assertTrue(get_portal_setting("enable_api")) + objects.values_list.return_value.first.return_value = False + self.assertFalse(get_portal_setting("enable_knowledge_base_api")) + objects.values_list.assert_called_with("enable_knowledge_base_api", flat=True) diff --git a/apps/common/tests/test_database_model_handle.py b/apps/common/tests/test_database_model_handle.py new file mode 100644 index 00000000000..9753a022660 --- /dev/null +++ b/apps/common/tests/test_database_model_handle.py @@ -0,0 +1,52 @@ +from unittest.mock import MagicMock, patch + +from django.conf import settings +from django.test import SimpleTestCase + +from common.database_model_manage.database_model_manage import DatabaseModelManage +from common.database_model_manage.handle.impl import default_base_model_handle as handle +from system_manage.models import ResourceChatUserAuthorize, ResourceChatUserGroupAuthorize, ResourceType + +KB_GROUP = "01a186b2-bed3-7cd0-9ba3-669bd89e59b1" +KB_USER = "01a186b2-bed3-7cd0-9ba3-669bd89e59b2" +KB_BOTH = "01a186b2-bed3-7cd0-9ba3-669bd89e59b3" +KB_NONE = "01a186b2-bed3-7cd0-9ba3-669bd89e59b4" + + +class DefaultModelHandleTest(SimpleTestCase): + def test_community_settings_register_knowledge_authorization(self): + self.assertIn( + "common.database_model_manage.handle.impl.default_base_model_handle.DefaultBaseModelHandle", + settings.MODEL_HANDLES, + ) + with patch.object(DatabaseModelManage, "model_dict", {}): + DatabaseModelManage.init() + self.assertIs( + DatabaseModelManage.get_model("get_knowledge_list_of_authorized"), + handle.get_knowledge_list_of_authorized, + ) + + def test_knowledge_authorized_by_user_or_group(self): + query_sets = {} + + def query_set(model): + query_sets[model] = MagicMock() + values = { + ResourceChatUserGroupAuthorize: [KB_GROUP, KB_BOTH], + ResourceChatUserAuthorize: [KB_USER, KB_BOTH], + }.get(model, ["group-1"]) + query_sets[model].filter.return_value.values_list.return_value = values + return query_sets[model] + + with patch.object(handle, "QuerySet", side_effect=query_set): + result = handle.get_knowledge_list_of_authorized("user-1", [KB_GROUP, KB_USER, KB_BOTH, KB_NONE]) + + self.assertEqual(result, [KB_GROUP, KB_USER, KB_BOTH]) + group_filter = query_sets[ResourceChatUserGroupAuthorize].filter.call_args.kwargs + self.assertEqual(group_filter["resource_type"], ResourceType.KNOWLEDGE) + self.assertTrue(group_filter["is_auth"]) + user_filter = query_sets[ResourceChatUserAuthorize].filter.call_args.kwargs + self.assertEqual( + (user_filter["user_id"], user_filter["resource_type"], user_filter["is_auth"]), + ("user-1", ResourceType.KNOWLEDGE, True), + ) diff --git a/apps/knowledge/services/retrieval_access.py b/apps/knowledge/services/retrieval_access.py index 1954a02d317..0f8003df9fa 100644 --- a/apps/knowledge/services/retrieval_access.py +++ b/apps/knowledge/services/retrieval_access.py @@ -6,6 +6,7 @@ from common.database_model_manage.database_model_manage import DatabaseModelManage from knowledge.models import Knowledge +from portal.models import get_portal_setting from system_manage.models import ChatUser, ChatUserApiKey @@ -50,7 +51,14 @@ def authorized_ids(user_id, ids): return set(map(str, handler(user_id, ids))) if handler else set() +def check_knowledge_api_enabled(): + """门户「知识库后端 API 检索」开关,控制对外的知识库检索与 MCP 接口,不影响工作流内部检索""" + if not get_portal_setting("enable_knowledge_base_api"): + raise RetrievalError("api_disabled", "Knowledge retrieval API is disabled.", 403) + + def authorize_external(knowledge_id, identity): + check_knowledge_api_enabled() knowledge = QuerySet(Knowledge).filter(id=knowledge_id).first() if knowledge is None or knowledge.external_service.get("enabled") is not True: raise RetrievalError("service_unavailable", "Retrieval service is unavailable.", 404) diff --git a/apps/knowledge/tests/retrieval/test_external_retrieval.py b/apps/knowledge/tests/retrieval/test_external_retrieval.py index ba0b9f2e10f..fd0b3a44cee 100644 --- a/apps/knowledge/tests/retrieval/test_external_retrieval.py +++ b/apps/knowledge/tests/retrieval/test_external_retrieval.py @@ -11,6 +11,7 @@ from django.urls import resolve from rest_framework.exceptions import ValidationError +from chat.mcp import knowledge as knowledge_mcp from chat.views.v3 import knowledge as views from knowledge.models import DocumentResourceType, Knowledge, KnowledgeType, SourceType from knowledge.serializers.external_retrieval import ( @@ -31,7 +32,7 @@ def knowledge(**settings): return SimpleNamespace( - id=KID, name="manual", workspace_id="default", type=KnowledgeType.BASE, external_service=settings + id=KID, name="manual", desc="", workspace_id="default", type=KnowledgeType.BASE, external_service=settings ) @@ -155,11 +156,14 @@ def test_routes_resolve_to_knowledge_and_preserve_application_mcp(self): resolve(f"/admin/api/workspace/default/knowledge/{KID}/external_service").kwargs["workspace_id"], "default" ) self.assertEqual(resolve("/chat/api/v3/mcp").func.__module__, "chat.views.v3.mcp") + self.assertIs(resolve("/chat/api/v3/knowledge_mcp").func, views.knowledge_global_mcp_view) class RetrievalAccessTests(SimpleTestCase): def setUp(self): self.key = SimpleNamespace(id=KEY, user_id=UID, is_active=True, user=SimpleNamespace(is_active=True)) + self.portal = patch.object(access, "get_portal_setting", return_value=True).start() + self.addCleanup(patch.stopall) @patch.object(access, "QuerySet") def test_existing_api_key_uses_original_secret_without_writes(self, queryset): @@ -209,6 +213,15 @@ def test_external_switch_and_authorization_matrix(self, queryset, authorized): access.authorize_external(KID, access.RetrievalIdentity(UID)) self.assertEqual(error.exception.status, 403) + @patch.object(access, "QuerySet") + def test_portal_switch_disables_external_retrieval(self, queryset): + self.portal.return_value = False + with self.assertRaises(access.RetrievalError) as error: + access.authorize_external(KID, access.RetrievalIdentity(UID, KEY)) + self.assertEqual((error.exception.code, error.exception.status), ("api_disabled", 403)) + self.portal.assert_called_once_with("enable_knowledge_base_api") + queryset.assert_not_called() + @patch.object(access.DatabaseModelManage, "get_model", return_value=None) @patch.object(access, "QuerySet") def test_missing_handler_or_inactive_user_denies(self, queryset, handler): @@ -458,3 +471,118 @@ async def dispatch(request): self.assertEqual(tools.tools[0].name, f"knowledge_{KID}") result = await session.call_tool(tools.tools[0].name, {"query_text": "hello"}) self.assertEqual(json.loads(result.content[0].text), self.output) + + +class KnowledgeMCPToolResultTests(SimpleTestCase): + def call(self, side_effect): + handler = knowledge_mcp.KnowledgeMCPToolHandler(knowledge(enabled=True), access.RetrievalIdentity()) + with patch.object(knowledge_mcp, "retrieve", side_effect=side_effect): + return handler.call_tool({"name": f"knowledge_{KID}", "arguments": {"query": "x"}}) + + def test_invalid_arguments_and_revoked_access_are_tool_errors_with_reason(self): + result = self.call(ValidationError({"query_text": ["This field is required."]})) + self.assertTrue(result["isError"]) + self.assertIn("query_text", result["content"][0]["text"]) + result = self.call(access.RetrievalError("access_denied", "Knowledge access denied.", 403)) + self.assertEqual(result, {"isError": True, "content": [{"type": "text", "text": "Knowledge access denied."}]}) + + def test_description_includes_knowledge_desc(self): + item = SimpleNamespace(**{**vars(knowledge(enabled=True)), "desc": "产品手册"}) + handler = knowledge_mcp.KnowledgeMCPToolHandler(item, access.RetrievalIdentity()) + self.assertEqual(handler.list_tools()["tools"][0]["description"], "检索知识库:manual\n产品手册") + + +class GlobalKnowledgeMCPTests(SimpleTestCase): + def setUp(self): + self.factory = RequestFactory() + self.identity = access.RetrievalIdentity(user_id=UID, api_key_id=KEY) + self.handler = knowledge_mcp.ChatUserKnowledgeMCPToolHandler(self.identity) + self.portal = patch.object(access, "get_portal_setting", return_value=True).start() + self.addCleanup(patch.stopall) + + def rpc(self, data, **headers): + request = self.factory.post( + "/knowledge_mcp", + json.dumps(data), + content_type="application/json", + HTTP_ACCEPT="application/json, text/event-stream", + **headers, + ) + return views.knowledge_global_mcp_view(request) + + def test_requires_chat_user_api_key(self): + with patch.object(views, "authenticate_key", return_value=access.RetrievalIdentity()): + response = self.rpc({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}) + self.assertEqual(response.status_code, 401) + self.assertEqual(response["WWW-Authenticate"], "Bearer") + + def test_lists_tools_with_api_key(self): + with patch.object(views, "authenticate_key", return_value=self.identity) as authenticate: + response = self.rpc({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, HTTP_AUTHORIZATION="Bearer k") + authenticate.assert_called_once_with("Bearer k") + tools = json.loads(response.content)["result"]["tools"] + self.assertEqual([tool["name"] for tool in tools], ["list_knowledge", "search_knowledge"]) + self.assertEqual(tools[1]["inputSchema"]["required"], ["knowledge_id", "query_text"]) + + def test_rejects_foreign_origin(self): + response = self.rpc({"jsonrpc": "2.0", "id": 1, "method": "ping"}, HTTP_ORIGIN="https://foreign.example") + self.assertEqual(response.status_code, 403) + + def test_portal_switch_disables_global_mcp(self): + self.portal.return_value = False + with patch.object(views, "authenticate_key", return_value=self.identity) as authenticate: + response = self.rpc({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, HTTP_AUTHORIZATION="Bearer k") + self.assertEqual(response.status_code, 403) + self.assertEqual(json.loads(response.content)["error"]["code"], "api_disabled") + authenticate.assert_not_called() + + def test_list_knowledge_includes_public_and_authorized(self): + items = [ + SimpleNamespace(id="k-public", name="a", desc=None, external_service={"enabled": True}), + SimpleNamespace( + id="k-granted", name="b", desc="d", external_service={"enabled": True, "authentication": True} + ), + SimpleNamespace( + id="k-denied", name="c", desc="", external_service={"enabled": True, "authentication": True} + ), + ] + with ( + patch.object(knowledge_mcp, "QuerySet") as query_set, + patch.object(knowledge_mcp, "authorized_ids", return_value={"k-granted"}) as authorized, + ): + query_set.return_value.filter.return_value.order_by.return_value = items + result = self.handler.call_tool({"name": "list_knowledge"}) + + query_set.return_value.filter.assert_called_once_with(external_service__enabled=True) + authorized.assert_called_once_with(UID, ["k-granted", "k-denied"]) + self.assertEqual( + result["structuredContent"]["knowledge"], + [ + {"knowledge_id": "k-public", "name": "a", "desc": ""}, + {"knowledge_id": "k-granted", "name": "b", "desc": "d"}, + ], + ) + + def test_search_knowledge_retrieves_with_full_identity(self): + output = {"knowledge_id": KID, "hits": []} + with patch.object(knowledge_mcp, "retrieve", return_value=output) as retrieve: + result = self.handler.call_tool( + {"name": "search_knowledge", "arguments": {"knowledge_id": KID, "query_text": "hello"}} + ) + retrieve.assert_called_once_with(KID, self.identity, {"query_text": "hello"}) + self.assertEqual(result["structuredContent"], output) + + def test_search_knowledge_errors_are_tool_errors(self): + result = self.handler.call_tool({"name": "search_knowledge", "arguments": {"knowledge_id": "bad"}}) + self.assertTrue(result["isError"]) + with patch.object( + knowledge_mcp, + "retrieve", + side_effect=access.RetrievalError("access_denied", "Knowledge access denied.", 403), + ): + result = self.handler.call_tool( + {"name": "search_knowledge", "arguments": {"knowledge_id": KID, "query_text": "x"}} + ) + self.assertEqual(result["content"][0]["text"], "Knowledge access denied.") + with self.assertRaises(ValidationError): + self.handler.call_tool({"name": "agent_x"}) diff --git a/apps/locales/en_US/LC_MESSAGES/django.po b/apps/locales/en_US/LC_MESSAGES/django.po index a33e7edf1a9..569285fb221 100644 --- a/apps/locales/en_US/LC_MESSAGES/django.po +++ b/apps/locales/en_US/LC_MESSAGES/django.po @@ -8991,6 +8991,9 @@ msgstr "" msgid "Secret key is expired" msgstr "" +msgid "API service is disabled" +msgstr "" + msgid "Online Usage" msgstr "" diff --git a/apps/locales/zh_CN/LC_MESSAGES/django.po b/apps/locales/zh_CN/LC_MESSAGES/django.po index bb79b3a1674..b57336df890 100644 --- a/apps/locales/zh_CN/LC_MESSAGES/django.po +++ b/apps/locales/zh_CN/LC_MESSAGES/django.po @@ -9116,6 +9116,9 @@ msgstr "密钥无效" msgid "Secret key is expired" msgstr "密钥已过期" +msgid "API service is disabled" +msgstr "API 服务未开启" + msgid "Online Usage" msgstr "线上使用" diff --git a/apps/locales/zh_Hant/LC_MESSAGES/django.po b/apps/locales/zh_Hant/LC_MESSAGES/django.po index bfe0860ea76..7d85ac57cd5 100644 --- a/apps/locales/zh_Hant/LC_MESSAGES/django.po +++ b/apps/locales/zh_Hant/LC_MESSAGES/django.po @@ -9116,6 +9116,9 @@ msgstr "密鑰無效" msgid "Secret key is expired" msgstr "密鑰已過期" +msgid "API service is disabled" +msgstr "API 服務未開啟" + msgid "Online Usage" msgstr "線上使用" diff --git a/apps/maxkb/settings/__init__.py b/apps/maxkb/settings/__init__.py index 547d2f849a1..0ded3363c58 100644 --- a/apps/maxkb/settings/__init__.py +++ b/apps/maxkb/settings/__init__.py @@ -1,11 +1,12 @@ # coding=utf-8 """ - @project: MaxKB - @Author:虎虎 - @file: __init__.py.py - @date:2025/4/11 16:39 - @desc: +@project: MaxKB +@Author:虎虎 +@file: __init__.py.py +@date:2025/4/11 16:39 +@desc: """ + import warnings warnings.filterwarnings("ignore", message="pkg_resources is deprecated as an API") @@ -14,4 +15,5 @@ from .logging import * from .auth import * from .lib import * -from .mem import * \ No newline at end of file +from .model import * +from .mem import * diff --git a/apps/maxkb/settings/model.py b/apps/maxkb/settings/model.py new file mode 100644 index 00000000000..94697fb1f19 --- /dev/null +++ b/apps/maxkb/settings/model.py @@ -0,0 +1,11 @@ +# coding=utf-8 +""" +@project: MaxKB +@Author:虎虎 +@file: model.py +@date:2026/11/2 16:00 +@desc: DatabaseModelManage 的模型/函数注册;企业版以同路径文件覆盖并注册自己的实现 +""" + +DEFAULT_MODEL = "common.database_model_manage.handle.impl.default_base_model_handle.DefaultBaseModelHandle" +MODEL_HANDLES = [DEFAULT_MODEL] diff --git a/apps/portal/models/portal.py b/apps/portal/models/portal.py index dac64f7dd47..6d38ec7c99a 100644 --- a/apps/portal/models/portal.py +++ b/apps/portal/models/portal.py @@ -39,3 +39,11 @@ class Portal(AppModelMixin): class Meta: db_table = "portal" + + +def get_portal_setting(field): + """ + 读取门户开关;门户未初始化时取字段默认值 + """ + value = Portal.objects.values_list(field, flat=True).first() + return Portal._meta.get_field(field).default if value is None else value