From 716132505cdff7497461e77f90c9003421cf8f75 Mon Sep 17 00:00:00 2001 From: Last-emo-boy Date: Thu, 12 Mar 2026 22:55:16 +0800 Subject: [PATCH 1/5] feat: add MCP client sub-capability bridges MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 补齐 AstrBot 对 MCP client 子能力的接线，并把能力接入现有 bot / dashboard 交互链路，在尽量不影响现有工具流的前提下支持更完整的 MCP 使用场景。 --- astrbot/core/agent/mcp_client.py | 299 +++- .../core/agent/mcp_elicitation_registry.py | 103 ++ astrbot/core/agent/mcp_prompt_bridge.py | 289 ++++ astrbot/core/agent/mcp_resource_bridge.py | 361 +++++ astrbot/core/agent/mcp_stdio_client.py | 162 ++ .../core/agent/mcp_subcapability_bridge.py | 1302 +++++++++++++++++ .../method/agent_sub_stages/internal.py | 9 + .../platform/sources/webchat/webchat_event.py | 11 + astrbot/core/provider/func_tool_manager.py | 82 +- astrbot/dashboard/routes/chat.py | 100 +- astrbot/dashboard/routes/live_chat.py | 13 + astrbot/dashboard/routes/open_api.py | 13 + astrbot/dashboard/routes/tools.py | 10 +- dashboard/src/components/chat/Chat.vue | 13 + dashboard/src/components/chat/MessageList.vue | 14 + .../src/components/chat/StandaloneChat.vue | 15 +- .../message_list_comps/ElicitationCard.vue | 352 +++++ .../MessagePartsRenderer.vue | 17 + dashboard/src/composables/useMessages.ts | 133 +- .../src/i18n/locales/en-US/features/chat.json | 15 + .../src/i18n/locales/zh-CN/features/chat.json | 15 + docs/en/use/mcp.md | 94 ++ docs/zh/use/mcp.md | 94 ++ tests/test_dashboard.py | 293 ++++ tests/unit/test_mcp_prompt_bridge.py | 212 +++ tests/unit/test_mcp_resource_bridge.py | 236 +++ tests/unit/test_mcp_stdio_client.py | 107 ++ tests/unit/test_mcp_subcapability_bridge.py | 1295 ++++++++++++++++ 28 files changed, 5602 insertions(+), 57 deletions(-) create mode 100644 astrbot/core/agent/mcp_elicitation_registry.py create mode 100644 astrbot/core/agent/mcp_prompt_bridge.py create mode 100644 astrbot/core/agent/mcp_resource_bridge.py create mode 100644 astrbot/core/agent/mcp_stdio_client.py create mode 100644 astrbot/core/agent/mcp_subcapability_bridge.py create mode 100644 dashboard/src/components/chat/message_list_comps/ElicitationCard.vue create mode 100644 tests/unit/test_mcp_prompt_bridge.py create mode 100644 tests/unit/test_mcp_resource_bridge.py create mode 100644 tests/unit/test_mcp_stdio_client.py create mode 100644 tests/unit/test_mcp_subcapability_bridge.py diff --git a/astrbot/core/agent/mcp_client.py b/astrbot/core/agent/mcp_client.py index a8ff0fdb90..10c14e66e3 100644 --- a/astrbot/core/agent/mcp_client.py +++ b/astrbot/core/agent/mcp_client.py @@ -2,7 +2,7 @@ import logging from contextlib import AsyncExitStack from datetime import timedelta -from typing import Generic +from typing import Any, Generic from tenacity import ( before_sleep_log, @@ -13,9 +13,16 @@ ) from astrbot import logger +from astrbot.core.agent.mcp_prompt_bridge import build_mcp_prompt_tool_names from astrbot.core.agent.run_context import ContextWrapper from astrbot.core.utils.log_pipe import LogPipe +from .mcp_resource_bridge import build_mcp_resource_tool_names +from .mcp_stdio_client import tolerant_stdio_client +from .mcp_subcapability_bridge import ( + MCPClientSubCapabilityBridge, + normalize_mcp_server_config, +) from .run_context import TContext from .tool import FunctionTool @@ -41,7 +48,10 @@ def _prepare_config(config: dict) -> dict: if config.get("mcpServers"): first_key = next(iter(config["mcpServers"])) config = config["mcpServers"][first_key] + config = normalize_mcp_server_config(config) config.pop("active", None) + config.pop("client_capabilities", None) + config.pop("provider", None) return config @@ -117,14 +127,22 @@ def __init__(self) -> None: self.name: str | None = None self.active: bool = True self.tools: list[mcp.Tool] = [] + self.prompts: list[mcp.types.Prompt] = [] + self.prompt_bridge_tool_names: list[str] = [] + self.resources: list[mcp.types.Resource] = [] + self.resource_templates: list[mcp.types.ResourceTemplate] = [] + self.resource_templates_supported: bool = False + self.resource_bridge_tool_names: list[str] = [] self.server_errlogs: list[str] = [] self.running_event = asyncio.Event() # Store connection config for reconnection self._mcp_server_config: dict | None = None self._server_name: str | None = None + self._server_capabilities: mcp.types.ServerCapabilities | None = None self._reconnect_lock = asyncio.Lock() # Lock for thread-safe reconnection self._reconnecting: bool = False # For logging and debugging + self.subcapability_bridge = MCPClientSubCapabilityBridge[Any]() async def connect_to_server(self, mcp_server_config: dict, name: str) -> None: """Connect to MCP server @@ -141,6 +159,8 @@ async def connect_to_server(self, mcp_server_config: dict, name: str) -> None: # Store config for reconnection self._mcp_server_config = mcp_server_config self._server_name = name + self.subcapability_bridge.set_server_name(name) + self.subcapability_bridge.configure_from_server_config(mcp_server_config) cfg = _prepare_config(mcp_server_config.copy()) @@ -184,6 +204,22 @@ def logging_callback( *streams, read_timeout_seconds=read_timeout, logging_callback=logging_callback, # type: ignore + sampling_callback=( + self.subcapability_bridge.handle_sampling + if self.subcapability_bridge.sampling_enabled + else None + ), + elicitation_callback=( + self.subcapability_bridge.handle_elicitation + if self.subcapability_bridge.elicitation_enabled + else None + ), + list_roots_callback=( + self.subcapability_bridge.handle_list_roots + if self.subcapability_bridge.roots_enabled + else None + ), + sampling_capabilities=self.subcapability_bridge.get_sampling_capabilities(), ), ) else: @@ -210,6 +246,22 @@ def logging_callback( write_stream=write_s, read_timeout_seconds=read_timeout, logging_callback=logging_callback, # type: ignore + sampling_callback=( + self.subcapability_bridge.handle_sampling + if self.subcapability_bridge.sampling_enabled + else None + ), + elicitation_callback=( + self.subcapability_bridge.handle_elicitation + if self.subcapability_bridge.elicitation_enabled + else None + ), + list_roots_callback=( + self.subcapability_bridge.handle_list_roots + if self.subcapability_bridge.roots_enabled + else None + ), + sampling_capabilities=self.subcapability_bridge.get_sampling_capabilities(), ), ) @@ -232,7 +284,7 @@ def callback(msg: str | mcp.types.LoggingMessageNotificationParams) -> None: self.server_errlogs.append(log_msg) stdio_transport = await self.exit_stack.enter_async_context( - mcp.stdio_client( + tolerant_stdio_client( server_params, errlog=LogPipe( level=logging.INFO, @@ -245,9 +297,41 @@ def callback(msg: str | mcp.types.LoggingMessageNotificationParams) -> None: # Create a new client session self.session = await self.exit_stack.enter_async_context( - mcp.ClientSession(*stdio_transport), + mcp.ClientSession( + *stdio_transport, + sampling_callback=( + self.subcapability_bridge.handle_sampling + if self.subcapability_bridge.sampling_enabled + else None + ), + elicitation_callback=( + self.subcapability_bridge.handle_elicitation + if self.subcapability_bridge.elicitation_enabled + else None + ), + list_roots_callback=( + self.subcapability_bridge.handle_list_roots + if self.subcapability_bridge.roots_enabled + else None + ), + sampling_capabilities=self.subcapability_bridge.get_sampling_capabilities(), + ), ) await self.session.initialize() + get_server_capabilities = getattr( + self.session, + "get_server_capabilities", + None, + ) + self._server_capabilities = ( + get_server_capabilities() if callable(get_server_capabilities) else None + ) + self.resources = [] + self.resource_templates = [] + self.resource_templates_supported = False + self.prompts = [] + self.prompt_bridge_tool_names = [] + self.resource_bridge_tool_names = [] async def list_tools_and_save(self) -> mcp.ListToolsResult: """List all tools from the server and save them to self.tools""" @@ -257,6 +341,120 @@ async def list_tools_and_save(self) -> mcp.ListToolsResult: self.tools = response.tools return response + @property + def supports_resources(self) -> bool: + return bool(self._server_capabilities and self._server_capabilities.resources) + + @property + def supports_prompts(self) -> bool: + return bool(self._server_capabilities and self._server_capabilities.prompts) + + async def load_resource_capabilities(self) -> None: + self.resources = [] + self.resource_templates = [] + self.resource_templates_supported = False + self.resource_bridge_tool_names = [] + + if not self._server_name or not self.supports_resources: + return + + try: + await self.list_resources_and_save() + except Exception as exc: # noqa: BLE001 + logger.warning( + "Failed to preload MCP resources for server %s: %s", + self._server_name, + exc, + ) + + try: + await self.list_resource_templates_and_save() + except Exception as exc: # noqa: BLE001 + logger.debug( + "Skipping MCP resource templates for server %s: %s", + self._server_name, + exc, + ) + + self.resource_bridge_tool_names = build_mcp_resource_tool_names( + self._server_name, + include_templates=self.resource_templates_supported, + ) + + async def load_prompt_capabilities(self) -> None: + self.prompts = [] + self.prompt_bridge_tool_names = [] + + if not self._server_name or not self.supports_prompts: + return + + try: + await self.list_prompts_and_save() + except Exception as exc: # noqa: BLE001 + logger.warning( + "Failed to preload MCP prompts for server %s: %s", + self._server_name, + exc, + ) + + self.prompt_bridge_tool_names = build_mcp_prompt_tool_names( + self._server_name, + ) + + async def list_prompts_and_save( + self, + cursor: str | None = None, + ) -> mcp.types.ListPromptsResult: + if not self.session: + raise ValueError("MCP session is not available for prompt listing.") + + params = ( + mcp.types.PaginatedRequestParams(cursor=cursor) + if cursor is not None + else None + ) + response = await self.session.list_prompts(params=params) + if cursor is None: + self.prompts = response.prompts + return response + + async def list_resources_and_save( + self, + cursor: str | None = None, + ) -> mcp.types.ListResourcesResult: + if not self.session: + raise ValueError("MCP session is not available for resource listing.") + + params = ( + mcp.types.PaginatedRequestParams(cursor=cursor) + if cursor is not None + else None + ) + response = await self.session.list_resources(params=params) + if cursor is None: + self.resources = response.resources + return response + + async def list_resource_templates_and_save( + self, + cursor: str | None = None, + ) -> mcp.types.ListResourceTemplatesResult: + if not self.session: + raise ValueError( + "MCP session is not available for resource template listing." + ) + + params = ( + mcp.types.PaginatedRequestParams(cursor=cursor) + if cursor is not None + else None + ) + response = await self.session.list_resource_templates(params=params) + self.resource_templates_supported = True + if cursor is None: + self.resource_templates = response.resourceTemplates + return response + async def _reconnect(self) -> None: """Reconnect to the MCP server using the stored configuration. @@ -281,6 +479,7 @@ async def _reconnect(self) -> None: logger.info( f"Attempting to reconnect to MCP server {self._server_name}..." ) + self.subcapability_bridge.clear_runtime_state() # Save old exit_stack for later cleanup (don't close it now to avoid cancel scope issues) if self.exit_stack: @@ -295,6 +494,8 @@ async def _reconnect(self) -> None: # Reconnect using stored config await self.connect_to_server(self._mcp_server_config, self._server_name) await self.list_tools_and_save() + await self.load_resource_capabilities() + await self.load_prompt_capabilities() logger.info( f"Successfully reconnected to MCP server {self._server_name}" @@ -312,6 +513,7 @@ async def call_tool_with_reconnect( tool_name: str, arguments: dict, read_timeout_seconds: timedelta, + run_context: ContextWrapper[Any] | None = None, ) -> mcp.types.CallToolResult: """Call MCP tool with automatic reconnection on failure, max 2 retries. @@ -336,28 +538,102 @@ async def call_tool_with_reconnect( reraise=True, ) async def _call_with_retry(): + async with self.subcapability_bridge.interactive_call(run_context): + if not self.session: + raise ValueError( + "MCP session is not available for MCP function tools." + ) + + try: + return await self.session.call_tool( + name=tool_name, + arguments=arguments, + read_timeout_seconds=read_timeout_seconds, + ) + except anyio.ClosedResourceError: + logger.warning( + f"MCP tool {tool_name} call failed (ClosedResourceError), attempting to reconnect..." + ) + # Attempt to reconnect + await self._reconnect() + # Reraise the exception to trigger tenacity retry + raise + + return await _call_with_retry() + + async def read_resource_with_reconnect( + self, + uri: str, + read_timeout_seconds: timedelta, + ) -> mcp.types.ReadResourceResult: + _ = read_timeout_seconds + + @retry( + retry=retry_if_exception_type(anyio.ClosedResourceError), + stop=stop_after_attempt(2), + wait=wait_exponential(multiplier=1, min=1, max=3), + before_sleep=before_sleep_log(logger, logging.WARNING), + reraise=True, + ) + async def _read_with_retry(): if not self.session: - raise ValueError("MCP session is not available for MCP function tools.") + raise ValueError("MCP session is not available for MCP resources.") try: - return await self.session.call_tool( - name=tool_name, + return await self.session.read_resource(uri=uri) + except anyio.ClosedResourceError: + logger.warning( + "MCP resource read for %s failed (ClosedResourceError), attempting to reconnect...", + uri, + ) + await self._reconnect() + raise + + return await _read_with_retry() + + async def get_prompt_with_reconnect( + self, + name: str, + arguments: dict[str, str] | None, + read_timeout_seconds: timedelta, + ) -> mcp.types.GetPromptResult: + @retry( + retry=retry_if_exception_type(anyio.ClosedResourceError), + stop=stop_after_attempt(2), + wait=wait_exponential(multiplier=1, min=1, max=3), + before_sleep=before_sleep_log(logger, logging.WARNING), + reraise=True, + ) + async def _get_with_retry(): + if not self.session: + raise ValueError("MCP session is not available for MCP prompts.") + + try: + return await self.session.get_prompt( + name=name, arguments=arguments, - read_timeout_seconds=read_timeout_seconds, ) except anyio.ClosedResourceError: logger.warning( - f"MCP tool {tool_name} call failed (ClosedResourceError), attempting to reconnect..." + "MCP prompt read for %s failed (ClosedResourceError), attempting to reconnect...", + name, ) - # Attempt to reconnect await self._reconnect() - # Reraise the exception to trigger tenacity retry raise - return await _call_with_retry() + _ = read_timeout_seconds + return await _get_with_retry() async def cleanup(self) -> None: """Clean up resources including old exit stacks from reconnections""" + self.subcapability_bridge.clear_runtime_state() + self._server_capabilities = None + self.prompts = [] + self.prompt_bridge_tool_names = [] + self.resources = [] + self.resource_templates = [] + self.resource_templates_supported = False + self.resource_bridge_tool_names = [] # Close current exit stack try: await self.exit_stack.aclose() @@ -395,4 +671,5 @@ async def call( tool_name=self.mcp_tool.name, arguments=kwargs, read_timeout_seconds=timedelta(seconds=context.tool_call_timeout), + run_context=context, ) diff --git a/astrbot/core/agent/mcp_elicitation_registry.py b/astrbot/core/agent/mcp_elicitation_registry.py new file mode 100644 index 0000000000..767c77bf73 --- /dev/null +++ b/astrbot/core/agent/mcp_elicitation_registry.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from astrbot.core.platform.astr_message_event import AstrMessageEvent + + +@dataclass(slots=True) +class MCPElicitationReply: + message_text: str + message_outline: str + + +@dataclass(slots=True) +class PendingMCPElicitation: + umo: str + sender_id: str + future: asyncio.Future[MCPElicitationReply] + + +_PENDING_MCP_ELICITATIONS: dict[str, PendingMCPElicitation] = {} + + +@asynccontextmanager +async def pending_mcp_elicitation( + umo: str, + sender_id: str, +) -> AsyncIterator[asyncio.Future[MCPElicitationReply]]: + loop = asyncio.get_running_loop() + future: asyncio.Future[MCPElicitationReply] = loop.create_future() + + current = _PENDING_MCP_ELICITATIONS.get(umo) + if current is not None and not current.future.done(): + raise RuntimeError( + f"Another MCP elicitation is already pending for session {umo}." + ) + + pending = PendingMCPElicitation( + umo=umo, + sender_id=sender_id, + future=future, + ) + _PENDING_MCP_ELICITATIONS[umo] = pending + + try: + yield future + finally: + current = _PENDING_MCP_ELICITATIONS.get(umo) + if current is pending: + _PENDING_MCP_ELICITATIONS.pop(umo, None) + if not future.done(): + future.cancel() + + +def try_capture_pending_mcp_elicitation(event: AstrMessageEvent) -> bool: + pending = _PENDING_MCP_ELICITATIONS.get(event.unified_msg_origin) + if pending is None: + return False + + sender_id = event.get_sender_id() + if not sender_id or sender_id != pending.sender_id: + return False + + if pending.future.done(): + _PENDING_MCP_ELICITATIONS.pop(event.unified_msg_origin, None) + return False + + pending.future.set_result( + MCPElicitationReply( + message_text=event.get_message_str() or "", + message_outline=event.get_message_outline(), + ) + ) + return True + + +def submit_pending_mcp_elicitation_reply( + umo: str, + sender_id: str, + reply_text: str, + *, + reply_outline: str | None = None, +) -> bool: + pending = _PENDING_MCP_ELICITATIONS.get(umo) + if pending is None or pending.sender_id != sender_id: + return False + + if pending.future.done(): + _PENDING_MCP_ELICITATIONS.pop(umo, None) + return False + + pending.future.set_result( + MCPElicitationReply( + message_text=reply_text, + message_outline=reply_outline or reply_text, + ) + ) + return True diff --git a/astrbot/core/agent/mcp_prompt_bridge.py b/astrbot/core/agent/mcp_prompt_bridge.py new file mode 100644 index 0000000000..b086accd5f --- /dev/null +++ b/astrbot/core/agent/mcp_prompt_bridge.py @@ -0,0 +1,289 @@ +from __future__ import annotations + +import json +import re +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Generic + +import mcp + +from astrbot.core.agent.run_context import ContextWrapper, TContext +from astrbot.core.agent.tool import FunctionTool + +if TYPE_CHECKING: + from .mcp_client import MCPClient + + +def build_mcp_prompt_tool_names(server_name: str) -> list[str]: + safe_server_name = _sanitize_tool_name_fragment(server_name) + return [ + f"mcp_{safe_server_name}_list_prompts", + f"mcp_{safe_server_name}_get_prompt", + ] + + +def build_mcp_prompt_tools( + mcp_client: MCPClient, + server_name: str, +) -> list[MCPPromptTool[TContext]]: + if not getattr(mcp_client, "supports_prompts", False): + return [] + + return [ + MCPListPromptsTool( + mcp_client=mcp_client, + mcp_server_name=server_name, + ), + MCPGetPromptTool( + mcp_client=mcp_client, + mcp_server_name=server_name, + ), + ] + + +class MCPPromptTool(FunctionTool, Generic[TContext]): + """Server-scoped synthetic tool for MCP prompts.""" + + def __init__(self, *, name: str, description: str, parameters: dict) -> None: + super().__init__( + name=name, + description=description, + parameters=parameters, + ) + self.mcp_client: MCPClient + self.mcp_server_name: str + + +class MCPListPromptsTool(MCPPromptTool[TContext]): + def __init__(self, *, mcp_client: MCPClient, mcp_server_name: str) -> None: + super().__init__( + name=build_mcp_prompt_tool_names(mcp_server_name)[0], + description=( + f"List MCP prompts exposed by server '{mcp_server_name}'. " + "Use this before getting a specific prompt template." + ), + parameters={ + "type": "object", + "properties": { + "cursor": { + "type": "string", + "description": ( + "Optional pagination cursor returned by a previous " + "prompt listing call." + ), + } + }, + }, + ) + self.mcp_client = mcp_client + self.mcp_server_name = mcp_server_name + + async def call( + self, + context: ContextWrapper[TContext], + **kwargs, + ) -> mcp.types.CallToolResult: + _ = context + response = await self.mcp_client.list_prompts_and_save( + cursor=kwargs.get("cursor"), + ) + return _text_result( + _format_prompts_listing( + server_name=self.mcp_server_name, + response=response, + ) + ) + + +class MCPGetPromptTool(MCPPromptTool[TContext]): + def __init__(self, *, mcp_client: MCPClient, mcp_server_name: str) -> None: + super().__init__( + name=build_mcp_prompt_tool_names(mcp_server_name)[1], + description=( + f"Get a specific MCP prompt from server '{mcp_server_name}' by " + "name, optionally providing prompt arguments." + ), + parameters={ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "The MCP prompt name to resolve.", + }, + "arguments": { + "type": "object", + "description": ( + "Optional prompt arguments. Keys and values are sent to " + "the MCP server as strings." + ), + "additionalProperties": { + "type": "string", + }, + }, + }, + "required": ["name"], + }, + ) + self.mcp_client = mcp_client + self.mcp_server_name = mcp_server_name + + async def call( + self, + context: ContextWrapper[TContext], + **kwargs, + ) -> mcp.types.CallToolResult: + read_timeout = timedelta(seconds=context.tool_call_timeout) + name = str(kwargs["name"]) + response = await self.mcp_client.get_prompt_with_reconnect( + name=name, + arguments=_normalize_prompt_arguments(kwargs.get("arguments")), + read_timeout_seconds=read_timeout, + ) + return _text_result( + shape_get_prompt_result( + server_name=self.mcp_server_name, + prompt_name=name, + response=response, + ) + ) + + +def shape_get_prompt_result( + *, + server_name: str, + prompt_name: str, + response: mcp.types.GetPromptResult, +) -> str: + lines = [ + f"MCP prompt from server '{server_name}':", + f"Prompt: {prompt_name}", + ] + if response.description: + lines.append(f"Description: {response.description}") + + if not response.messages: + lines.append("No prompt messages were returned.") + return "\n".join(lines) + + lines.append(f"Returned messages: {len(response.messages)}") + for idx, message in enumerate(response.messages, start=1): + lines.append("") + lines.append(f"Message {idx} ({message.role}):") + lines.extend(_format_prompt_message_content(message.content)) + return "\n".join(lines) + + +def _text_result(text: str) -> mcp.types.CallToolResult: + return mcp.types.CallToolResult( + content=[mcp.types.TextContent(type="text", text=text)] + ) + + +def _format_prompts_listing( + *, + server_name: str, + response: mcp.types.ListPromptsResult, +) -> str: + if not response.prompts: + text = f"No MCP prompts are currently exposed by server '{server_name}'." + if response.nextCursor: + text += f"\nNext cursor: {response.nextCursor}" + return text + + lines = [f"MCP prompts from server '{server_name}':"] + for idx, prompt in enumerate(response.prompts, start=1): + lines.extend(_format_prompt_metadata(idx, prompt)) + if response.nextCursor: + lines.append(f"Next cursor: {response.nextCursor}") + return "\n".join(lines) + + +def _format_prompt_metadata(index: int, prompt: mcp.types.Prompt) -> list[str]: + lines = [f"{index}. {prompt.name}"] + if prompt.title: + lines.append(f" Title: {prompt.title}") + if prompt.description: + lines.append(f" Description: {prompt.description}") + if prompt.arguments: + lines.append(" Arguments:") + for argument in prompt.arguments: + lines.append(_format_prompt_argument(argument)) + return lines + + +def _format_prompt_argument(argument: mcp.types.PromptArgument) -> str: + required_suffix = "required" if argument.required else "optional" + if argument.description: + return f" - {argument.name} ({required_suffix}): {argument.description}" + return f" - {argument.name} ({required_suffix})" + + +def _format_prompt_message_content( + content: mcp.types.ContentBlock, +) -> list[str]: + if isinstance(content, mcp.types.TextContent): + return content.text.splitlines() or [content.text] + if isinstance(content, mcp.types.ImageContent): + return [ + "Image block returned.", + f"MIME type: {content.mimeType}", + f"Base64 length: {len(content.data)}", + ] + if isinstance(content, mcp.types.AudioContent): + return [ + "Audio block returned.", + f"MIME type: {content.mimeType}", + f"Base64 length: {len(content.data)}", + ] + if isinstance(content, mcp.types.EmbeddedResource): + resource = content.resource + if isinstance(resource, mcp.types.TextResourceContents): + lines = [ + "Embedded text resource returned.", + f"URI: {resource.uri}", + ] + if resource.mimeType: + lines.append(f"MIME type: {resource.mimeType}") + lines.append("Text:") + lines.extend(resource.text.splitlines() or [resource.text]) + return lines + if isinstance(resource, mcp.types.BlobResourceContents): + lines = [ + "Embedded binary resource returned.", + f"URI: {resource.uri}", + ] + if resource.mimeType: + lines.append(f"MIME type: {resource.mimeType}") + lines.append(f"Base64 length: {len(resource.blob)}") + return lines + return [f"Unsupported prompt content block: {type(content).__name__}"] + + +def _normalize_prompt_arguments( + raw_arguments: Any, +) -> dict[str, str] | None: + if raw_arguments is None: + return None + if isinstance(raw_arguments, str): + stripped = raw_arguments.strip() + if not stripped: + return None + try: + parsed = json.loads(stripped) + except json.JSONDecodeError: + return None + raw_arguments = parsed + if not isinstance(raw_arguments, dict): + return None + normalized: dict[str, str] = {} + for key, value in raw_arguments.items(): + key_text = str(key).strip() + if not key_text: + continue + normalized[key_text] = "" if value is None else str(value) + return normalized or None + + +def _sanitize_tool_name_fragment(name: str) -> str: + sanitized = re.sub(r"[^a-zA-Z0-9]+", "_", name).strip("_").lower() + return sanitized or "server" diff --git a/astrbot/core/agent/mcp_resource_bridge.py b/astrbot/core/agent/mcp_resource_bridge.py new file mode 100644 index 0000000000..3663e44a60 --- /dev/null +++ b/astrbot/core/agent/mcp_resource_bridge.py @@ -0,0 +1,361 @@ +from __future__ import annotations + +import re +from datetime import timedelta +from typing import TYPE_CHECKING, Generic + +import mcp + +from astrbot.core.agent.run_context import ContextWrapper, TContext +from astrbot.core.agent.tool import FunctionTool + +if TYPE_CHECKING: + from .mcp_client import MCPClient + + +def build_mcp_resource_tool_names( + server_name: str, + *, + include_templates: bool, +) -> list[str]: + safe_server_name = _sanitize_tool_name_fragment(server_name) + names = [ + f"mcp_{safe_server_name}_list_resources", + f"mcp_{safe_server_name}_read_resource", + ] + if include_templates: + names.append(f"mcp_{safe_server_name}_list_resource_templates") + return names + + +def build_mcp_resource_tools( + mcp_client: MCPClient, + server_name: str, +) -> list[MCPResourceTool[TContext]]: + if not getattr(mcp_client, "supports_resources", False): + return [] + + tools: list[MCPResourceTool[TContext]] = [ + MCPListResourcesTool( + mcp_client=mcp_client, + mcp_server_name=server_name, + ), + MCPReadResourceTool( + mcp_client=mcp_client, + mcp_server_name=server_name, + ), + ] + if mcp_client.resource_templates_supported: + tools.append( + MCPListResourceTemplatesTool( + mcp_client=mcp_client, + mcp_server_name=server_name, + ) + ) + return tools + + +class MCPResourceTool(FunctionTool, Generic[TContext]): + """Server-scoped synthetic tool for MCP resources.""" + + def __init__(self, *, name: str, description: str, parameters: dict) -> None: + super().__init__( + name=name, + description=description, + parameters=parameters, + ) + self.mcp_client: MCPClient + self.mcp_server_name: str + + +class MCPListResourcesTool(MCPResourceTool[TContext]): + def __init__(self, *, mcp_client: MCPClient, mcp_server_name: str) -> None: + super().__init__( + name=build_mcp_resource_tool_names( + mcp_server_name, + include_templates=False, + )[0], + description=( + f"List readable MCP resources exposed by server '{mcp_server_name}'. " + "Use this before reading a specific resource URI." + ), + parameters={ + "type": "object", + "properties": { + "cursor": { + "type": "string", + "description": ( + "Optional pagination cursor returned by a previous " + "resource listing call." + ), + } + }, + }, + ) + self.mcp_client = mcp_client + self.mcp_server_name = mcp_server_name + + async def call( + self, + context: ContextWrapper[TContext], + **kwargs, + ) -> mcp.types.CallToolResult: + _ = context + response = await self.mcp_client.list_resources_and_save( + cursor=kwargs.get("cursor"), + ) + return _text_result( + _format_resources_listing( + server_name=self.mcp_server_name, + response=response, + ) + ) + + +class MCPListResourceTemplatesTool(MCPResourceTool[TContext]): + def __init__(self, *, mcp_client: MCPClient, mcp_server_name: str) -> None: + super().__init__( + name=build_mcp_resource_tool_names( + mcp_server_name, + include_templates=True, + )[2], + description=( + f"List MCP resource URI templates exposed by server " + f"'{mcp_server_name}'. Use the returned URI patterns to construct " + "resource URIs for read_resource." + ), + parameters={ + "type": "object", + "properties": { + "cursor": { + "type": "string", + "description": ( + "Optional pagination cursor returned by a previous " + "resource template listing call." + ), + } + }, + }, + ) + self.mcp_client = mcp_client + self.mcp_server_name = mcp_server_name + + async def call( + self, + context: ContextWrapper[TContext], + **kwargs, + ) -> mcp.types.CallToolResult: + _ = context + response = await self.mcp_client.list_resource_templates_and_save( + cursor=kwargs.get("cursor"), + ) + return _text_result( + _format_resource_templates_listing( + server_name=self.mcp_server_name, + response=response, + ) + ) + + +class MCPReadResourceTool(MCPResourceTool[TContext]): + def __init__(self, *, mcp_client: MCPClient, mcp_server_name: str) -> None: + super().__init__( + name=build_mcp_resource_tool_names( + mcp_server_name, + include_templates=False, + )[1], + description=( + f"Read a specific MCP resource from server '{mcp_server_name}' by " + "its URI." + ), + parameters={ + "type": "object", + "properties": { + "uri": { + "type": "string", + "description": "The MCP resource URI to read.", + } + }, + "required": ["uri"], + }, + ) + self.mcp_client = mcp_client + self.mcp_server_name = mcp_server_name + + async def call( + self, + context: ContextWrapper[TContext], + **kwargs, + ) -> mcp.types.CallToolResult: + read_timeout = timedelta(seconds=context.tool_call_timeout) + uri = str(kwargs["uri"]) + response = await self.mcp_client.read_resource_with_reconnect( + uri=uri, + read_timeout_seconds=read_timeout, + ) + return shape_read_resource_result( + server_name=self.mcp_server_name, + requested_uri=uri, + response=response, + ) + + +def shape_read_resource_result( + *, + server_name: str, + requested_uri: str, + response: mcp.types.ReadResourceResult, +) -> mcp.types.CallToolResult: + contents = response.contents + if not contents: + return _text_result( + f"MCP server '{server_name}' returned no contents for resource " + f"'{requested_uri}'." + ) + + if len(contents) == 1: + content = contents[0] + if isinstance(content, mcp.types.TextResourceContents): + return _text_result(_format_single_text_resource(server_name, content)) + if ( + isinstance(content, mcp.types.BlobResourceContents) + and content.mimeType + and content.mimeType.startswith("image/") + ): + return mcp.types.CallToolResult( + content=[ + mcp.types.EmbeddedResource( + type="resource", + resource=content, + ) + ] + ) + + return _text_result( + _format_multi_part_resource( + server_name=server_name, + requested_uri=requested_uri, + contents=contents, + ) + ) + + +def _text_result(text: str) -> mcp.types.CallToolResult: + return mcp.types.CallToolResult( + content=[mcp.types.TextContent(type="text", text=text)] + ) + + +def _format_resources_listing( + *, + server_name: str, + response: mcp.types.ListResourcesResult, +) -> str: + if not response.resources: + text = f"No MCP resources are currently exposed by server '{server_name}'." + if response.nextCursor: + text += f"\nNext cursor: {response.nextCursor}" + return text + + lines = [f"MCP resources from server '{server_name}':"] + for idx, resource in enumerate(response.resources, start=1): + lines.extend(_format_resource_metadata(idx, resource)) + if response.nextCursor: + lines.append(f"Next cursor: {response.nextCursor}") + return "\n".join(lines) + + +def _format_resource_templates_listing( + *, + server_name: str, + response: mcp.types.ListResourceTemplatesResult, +) -> str: + if not response.resourceTemplates: + text = ( + f"No MCP resource templates are currently exposed by server " + f"'{server_name}'." + ) + if response.nextCursor: + text += f"\nNext cursor: {response.nextCursor}" + return text + + lines = [f"MCP resource templates from server '{server_name}':"] + for idx, template in enumerate(response.resourceTemplates, start=1): + lines.extend(_format_resource_template_metadata(idx, template)) + if response.nextCursor: + lines.append(f"Next cursor: {response.nextCursor}") + return "\n".join(lines) + + +def _format_single_text_resource( + server_name: str, + content: mcp.types.TextResourceContents, +) -> str: + lines = [ + f"MCP text resource from server '{server_name}':", + f"URI: {content.uri}", + ] + if content.mimeType: + lines.append(f"MIME type: {content.mimeType}") + lines.extend(["", content.text]) + return "\n".join(lines) + + +def _format_multi_part_resource( + *, + server_name: str, + requested_uri: str, + contents: list[mcp.types.TextResourceContents | mcp.types.BlobResourceContents], +) -> str: + lines = [ + f"MCP resource read result from server '{server_name}':", + f"Requested URI: {requested_uri}", + f"Returned parts: {len(contents)}", + ] + for idx, content in enumerate(contents, start=1): + lines.append("") + lines.append(f"Part {idx}:") + lines.append(f"URI: {content.uri}") + if content.mimeType: + lines.append(f"MIME type: {content.mimeType}") + if isinstance(content, mcp.types.TextResourceContents): + lines.append("Text:") + lines.append(content.text) + else: + lines.append(f"Binary blob returned (base64 length: {len(content.blob)}).") + return "\n".join(lines) + + +def _format_resource_metadata( + index: int, + resource: mcp.types.Resource, +) -> list[str]: + lines = [f"{index}. {resource.name}", f" URI: {resource.uri}"] + if resource.title: + lines.append(f" Title: {resource.title}") + if resource.description: + lines.append(f" Description: {resource.description}") + if resource.mimeType: + lines.append(f" MIME type: {resource.mimeType}") + if resource.size is not None: + lines.append(f" Size: {resource.size} bytes") + return lines + + +def _format_resource_template_metadata( + index: int, + template: mcp.types.ResourceTemplate, +) -> list[str]: + lines = [f"{index}. {template.name}", f" URI template: {template.uriTemplate}"] + if template.title: + lines.append(f" Title: {template.title}") + if template.description: + lines.append(f" Description: {template.description}") + if template.mimeType: + lines.append(f" MIME type: {template.mimeType}") + return lines + + +def _sanitize_tool_name_fragment(name: str) -> str: + sanitized = re.sub(r"[^a-zA-Z0-9]+", "_", name).strip("_").lower() + return sanitized or "server" diff --git a/astrbot/core/agent/mcp_stdio_client.py b/astrbot/core/agent/mcp_stdio_client.py new file mode 100644 index 0000000000..e0433106d8 --- /dev/null +++ b/astrbot/core/agent/mcp_stdio_client.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import logging +import sys +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING, TextIO + +import anyio +import anyio.lowlevel +import mcp.types as types +from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream +from anyio.streams.text import TextReceiveStream +from mcp.client.stdio import ( + PROCESS_TERMINATION_TIMEOUT, + _create_platform_compatible_process, + _get_executable_command, + _terminate_process_tree, + get_default_environment, +) +from mcp.shared.message import SessionMessage + +from astrbot import logger + +if TYPE_CHECKING: + import mcp + + +def _normalize_stdout_line(line: str) -> str: + return line.rstrip("\r") + + +def _should_ignore_stdout_line(line: str) -> bool: + stripped = _normalize_stdout_line(line).strip() + if not stripped: + return True + + # JSON-RPC messages are serialized as JSON objects. Wrapper banners from + # tools such as npm/pnpm/yarn should not abort the session. + return not stripped.startswith("{") + + +@asynccontextmanager +async def tolerant_stdio_client( + server: mcp.StdioServerParameters, + errlog: TextIO = sys.stderr, +): + """A stdio MCP transport that ignores obvious non-protocol stdout noise.""" + + read_stream: MemoryObjectReceiveStream[SessionMessage | Exception] + read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception] + + write_stream: MemoryObjectSendStream[SessionMessage] + write_stream_reader: MemoryObjectReceiveStream[SessionMessage] + + read_stream_writer, read_stream = anyio.create_memory_object_stream(0) + write_stream, write_stream_reader = anyio.create_memory_object_stream(0) + + try: + command = _get_executable_command(server.command) + process = await _create_platform_compatible_process( + command=command, + args=server.args, + env=( + {**get_default_environment(), **server.env} + if server.env is not None + else get_default_environment() + ), + errlog=errlog, + cwd=server.cwd, + ) + except OSError: + await read_stream.aclose() + await write_stream.aclose() + await read_stream_writer.aclose() + await write_stream_reader.aclose() + raise + + async def stdout_reader(): + assert process.stdout, "Opened process is missing stdout" + + try: + async with read_stream_writer: + buffer = "" + async for chunk in TextReceiveStream( + process.stdout, + encoding=server.encoding, + errors=server.encoding_error_handler, + ): + lines = (buffer + chunk).split("\n") + buffer = lines.pop() + + for raw_line in lines: + line = _normalize_stdout_line(raw_line) + if _should_ignore_stdout_line(line): + if line.strip(): + logger.debug( + "Ignoring non-JSON stdout line from MCP stdio server: %s", + line.strip(), + ) + continue + + try: + message = types.JSONRPCMessage.model_validate_json( + line.strip() + ) + except Exception as exc: # pragma: no cover + logging.getLogger("mcp.client.stdio").exception( + "Failed to parse JSONRPC message from server" + ) + await read_stream_writer.send(exc) + continue + + await read_stream_writer.send(SessionMessage(message)) + except anyio.ClosedResourceError: # pragma: no cover + await anyio.lowlevel.checkpoint() + + async def stdin_writer(): + assert process.stdin, "Opened process is missing stdin" + + try: + async with write_stream_reader: + async for session_message in write_stream_reader: + payload = session_message.message.model_dump_json( + by_alias=True, + exclude_none=True, + ) + await process.stdin.send( + (payload + "\n").encode( + encoding=server.encoding, + errors=server.encoding_error_handler, + ) + ) + except anyio.ClosedResourceError: # pragma: no cover + await anyio.lowlevel.checkpoint() + + async with ( + anyio.create_task_group() as tg, + process, + ): + tg.start_soon(stdout_reader) + tg.start_soon(stdin_writer) + try: + yield read_stream, write_stream + finally: + if process.stdin: # pragma: no branch + try: + await process.stdin.aclose() + except Exception: # pragma: no cover + pass + + try: + with anyio.fail_after(PROCESS_TERMINATION_TIMEOUT): + await process.wait() + except TimeoutError: + await _terminate_process_tree(process) + except ProcessLookupError: # pragma: no cover + pass + + await read_stream.aclose() + await write_stream.aclose() + await read_stream_writer.aclose() + await write_stream_reader.aclose() diff --git a/astrbot/core/agent/mcp_subcapability_bridge.py b/astrbot/core/agent/mcp_subcapability_bridge.py new file mode 100644 index 0000000000..9b92b584fa --- /dev/null +++ b/astrbot/core/agent/mcp_subcapability_bridge.py @@ -0,0 +1,1302 @@ +from __future__ import annotations + +import asyncio +import copy +import json +import re +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any, Generic + +from astrbot import logger +from astrbot.core.agent.mcp_elicitation_registry import pending_mcp_elicitation +from astrbot.core.agent.run_context import ContextWrapper, TContext +from astrbot.core.message.components import Json +from astrbot.core.message.message_event_result import MessageChain +from astrbot.core.utils.astrbot_path import ( + get_astrbot_backups_path, + get_astrbot_config_path, + get_astrbot_data_path, + get_astrbot_knowledge_base_path, + get_astrbot_plugin_data_path, + get_astrbot_plugin_path, + get_astrbot_root, + get_astrbot_skills_path, + get_astrbot_temp_path, +) + +if TYPE_CHECKING: + import mcp + + +DEFAULT_MCP_CLIENT_CAPABILITIES = { + "elicitation": { + "enabled": False, + "timeout_seconds": 300, + }, + "sampling": { + "enabled": False, + }, + "roots": { + "enabled": False, + "paths": [], + }, +} + +DEFAULT_MCP_ROOT_PATHS = ("data", "temp") +DEFAULT_MCP_ELICITATION_TIMEOUT_SECONDS = 300 +MCP_ELICITATION_ACCEPT_KEYWORDS = { + "accept", + "done", + "ok", + "okay", + "yes", + "完成", + "已完成", + "同意", +} +MCP_ELICITATION_DECLINE_KEYWORDS = { + "decline", + "reject", + "refuse", + "no", + "拒绝", + "不同意", +} +MCP_ELICITATION_CANCEL_KEYWORDS = { + "cancel", + "stop", + "退出", + "取消", +} + + +def get_root_path_alias_resolvers(): + return { + "root": get_astrbot_root, + "data": get_astrbot_data_path, + "config": get_astrbot_config_path, + "plugins": get_astrbot_plugin_path, + "plugin_data": get_astrbot_plugin_data_path, + "temp": get_astrbot_temp_path, + "skills": get_astrbot_skills_path, + "knowledge_base": get_astrbot_knowledge_base_path, + "backups": get_astrbot_backups_path, + } + + +class UnsupportedSamplingRequestError(ValueError): + """Raised when a sampling request cannot be safely mapped.""" + + +class UnsupportedElicitationRequestError(ValueError): + """Raised when an elicitation request cannot be safely mapped.""" + + +@dataclass(slots=True) +class MCPElicitationCapabilityConfig: + enabled: bool = False + timeout_seconds: int = DEFAULT_MCP_ELICITATION_TIMEOUT_SECONDS + + +@dataclass(slots=True) +class MCPSamplingCapabilityConfig: + enabled: bool = False + + +@dataclass(slots=True) +class MCPRootsCapabilityConfig: + enabled: bool = False + paths: list[str] = field(default_factory=list) + + +@dataclass(slots=True) +class MCPClientCapabilitiesConfig: + elicitation: MCPElicitationCapabilityConfig + sampling: MCPSamplingCapabilityConfig + roots: MCPRootsCapabilityConfig + + @classmethod + def from_server_config( + cls, server_config: dict[str, Any] | None + ) -> MCPClientCapabilitiesConfig: + normalized = normalize_mcp_server_config(server_config or {}) + elicitation_cfg = normalized["client_capabilities"]["elicitation"] + sampling_cfg = normalized["client_capabilities"]["sampling"] + roots_cfg = normalized["client_capabilities"]["roots"] + return cls( + elicitation=MCPElicitationCapabilityConfig( + enabled=bool(elicitation_cfg.get("enabled", False)), + timeout_seconds=int( + elicitation_cfg.get( + "timeout_seconds", + DEFAULT_MCP_ELICITATION_TIMEOUT_SECONDS, + ) + ), + ), + sampling=MCPSamplingCapabilityConfig( + enabled=bool(sampling_cfg.get("enabled", False)), + ), + roots=MCPRootsCapabilityConfig( + enabled=bool(roots_cfg.get("enabled", False)), + paths=list(roots_cfg.get("paths", [])), + ), + ) + + +def normalize_mcp_server_config(server_config: dict[str, Any]) -> dict[str, Any]: + """Normalize persisted MCP server config fields for backward compatibility.""" + normalized = copy.deepcopy(server_config) + + client_capabilities = normalized.get("client_capabilities") + if not isinstance(client_capabilities, dict): + client_capabilities = {} + + elicitation_cfg = client_capabilities.get("elicitation") + if isinstance(elicitation_cfg, bool): + elicitation_cfg = {"enabled": elicitation_cfg} + elif not isinstance(elicitation_cfg, dict): + elicitation_cfg = {} + + sampling_cfg = client_capabilities.get("sampling") + if isinstance(sampling_cfg, bool): + sampling_cfg = {"enabled": sampling_cfg} + elif not isinstance(sampling_cfg, dict): + sampling_cfg = {} + + roots_cfg = client_capabilities.get("roots") + if isinstance(roots_cfg, bool): + roots_cfg = {"enabled": roots_cfg} + elif not isinstance(roots_cfg, dict): + roots_cfg = {} + + raw_root_paths = roots_cfg.get("paths", []) + if not isinstance(raw_root_paths, list): + raw_root_paths = [] + normalized_root_paths = [ + str(path).strip() + for path in raw_root_paths + if isinstance(path, str) and path.strip() + ] + + client_capabilities["elicitation"] = { + "enabled": bool(elicitation_cfg.get("enabled", False)), + "timeout_seconds": _normalize_positive_int( + elicitation_cfg.get( + "timeout_seconds", + DEFAULT_MCP_ELICITATION_TIMEOUT_SECONDS, + ), + DEFAULT_MCP_ELICITATION_TIMEOUT_SECONDS, + ), + } + client_capabilities["sampling"] = { + "enabled": bool(sampling_cfg.get("enabled", False)), + } + client_capabilities["roots"] = { + "enabled": bool(roots_cfg.get("enabled", False)), + "paths": normalized_root_paths, + } + normalized["client_capabilities"] = client_capabilities + return normalized + + +def _normalize_positive_int(value: Any, default: int) -> int: + if isinstance(value, bool): + return default + try: + normalized = int(value) + except (TypeError, ValueError): + return default + if normalized <= 0: + return default + return normalized + + +def normalize_mcp_config(config: dict[str, Any] | None) -> dict[str, Any]: + """Normalize the full MCP configuration file structure.""" + normalized = {"mcpServers": {}} + if not isinstance(config, dict): + return normalized + + raw_servers = config.get("mcpServers", {}) + if not isinstance(raw_servers, dict): + return normalized + + for name, server_config in raw_servers.items(): + if not isinstance(server_config, dict): + continue + normalized["mcpServers"][name] = normalize_mcp_server_config(server_config) + return normalized + + +class MCPClientSubCapabilityBridge(Generic[TContext]): + """Bridge MCP client sub-capability requests into AstrBot runtime calls.""" + + def __init__(self, server_name: str | None = None) -> None: + self._server_name = server_name or "" + self._capabilities = MCPClientCapabilitiesConfig.from_server_config({}) + self._interaction_lock = asyncio.Lock() + self._active_run_context: ContextWrapper[TContext] | None = None + + def configure_from_server_config(self, server_config: dict[str, Any]) -> None: + self._capabilities = MCPClientCapabilitiesConfig.from_server_config( + server_config + ) + + def set_server_name(self, server_name: str | None) -> None: + if server_name: + self._server_name = server_name + + @property + def sampling_enabled(self) -> bool: + return self._capabilities.sampling.enabled + + @property + def elicitation_enabled(self) -> bool: + return self._capabilities.elicitation.enabled + + @property + def elicitation_timeout_seconds(self) -> int: + return self._capabilities.elicitation.timeout_seconds + + @property + def roots_enabled(self) -> bool: + return self._capabilities.roots.enabled + + def get_sampling_capabilities(self) -> mcp.types.SamplingCapability | None: + if not self.sampling_enabled: + return None + + import mcp + + return mcp.types.SamplingCapability() + + async def handle_list_roots( + self, + _request_context: Any, + ) -> mcp.types.ListRootsResult | mcp.types.ErrorData: + import mcp + + if not self.roots_enabled: + return mcp.types.ErrorData( + code=mcp.types.INVALID_REQUEST, + message="Roots are not enabled for this MCP server.", + ) + + try: + return mcp.types.ListRootsResult(roots=self._build_root_entries()) + except Exception as exc: # noqa: BLE001 + logger.error( + "Roots request failed for MCP server %s: %s", + self._server_name, + exc, + exc_info=True, + ) + return mcp.types.ErrorData( + code=mcp.types.INTERNAL_ERROR, + message="Roots request failed inside AstrBot.", + data=str(exc), + ) + + def clear_runtime_state(self) -> None: + self._active_run_context = None + + @asynccontextmanager + async def interactive_call( + self, + run_context: ContextWrapper[TContext] | None, + ): + if not (self.sampling_enabled or self.elicitation_enabled): + yield + return + + async with self._interaction_lock: + self._active_run_context = run_context + try: + yield + finally: + self._active_run_context = None + + async def handle_sampling( + self, + _request_context: Any, + params: mcp.types.CreateMessageRequestParams, + ) -> ( + mcp.types.CreateMessageResult + | mcp.types.CreateMessageResultWithTools + | mcp.types.ErrorData + ): + import mcp + + if not self.sampling_enabled: + return mcp.types.ErrorData( + code=mcp.types.INVALID_REQUEST, + message="Sampling is not enabled for this MCP server.", + ) + + run_context = self._active_run_context + if run_context is None: + return mcp.types.ErrorData( + code=mcp.types.INVALID_REQUEST, + message=( + "Sampling requests are only supported during an active AstrBot " + "MCP interaction." + ), + ) + + try: + return await self._execute_sampling(run_context, params) + except UnsupportedSamplingRequestError as exc: + return mcp.types.ErrorData( + code=mcp.types.INVALID_REQUEST, + message=str(exc), + ) + except Exception as exc: # noqa: BLE001 + logger.error( + "Sampling request failed for MCP server %s: %s", + self._server_name, + exc, + exc_info=True, + ) + return mcp.types.ErrorData( + code=mcp.types.INTERNAL_ERROR, + message="Sampling request failed inside AstrBot.", + data=str(exc), + ) + + async def handle_elicitation( + self, + _request_context: Any, + params: mcp.types.ElicitRequestParams, + ) -> mcp.types.ElicitResult | mcp.types.ErrorData: + import mcp + + if not self.elicitation_enabled: + return mcp.types.ErrorData( + code=mcp.types.INVALID_REQUEST, + message="Elicitation is not enabled for this MCP server.", + ) + + run_context = self._active_run_context + if run_context is None: + return mcp.types.ErrorData( + code=mcp.types.INVALID_REQUEST, + message=( + "Elicitation requests are only supported during an active AstrBot " + "MCP interaction." + ), + ) + + try: + return await self._execute_elicitation(run_context, params) + except UnsupportedElicitationRequestError as exc: + return mcp.types.ErrorData( + code=mcp.types.INVALID_REQUEST, + message=str(exc), + ) + except Exception as exc: # noqa: BLE001 + logger.error( + "Elicitation request failed for MCP server %s: %s", + self._server_name, + exc, + exc_info=True, + ) + return mcp.types.ErrorData( + code=mcp.types.INTERNAL_ERROR, + message="Elicitation request failed inside AstrBot.", + data=str(exc), + ) + + async def _execute_sampling( + self, + run_context: ContextWrapper[TContext], + params: mcp.types.CreateMessageRequestParams, + ) -> mcp.types.CreateMessageResult: + import mcp + + plugin_context, event = self._extract_bound_runtime(run_context) + if plugin_context is None or event is None: + raise UnsupportedSamplingRequestError( + "Sampling requires an AstrBot agent context bound to the MCP tool call." + ) + + if params.includeContext not in (None, "none"): + raise UnsupportedSamplingRequestError( + "Sampling includeContext is not supported in the initial AstrBot integration." + ) + + if params.tools or params.toolChoice: + raise UnsupportedSamplingRequestError( + "Tool-assisted sampling is not supported in the initial AstrBot integration." + ) + + contexts = self._translate_sampling_messages(params.messages) + umo = getattr(event, "unified_msg_origin", None) + if not isinstance(umo, str) or not umo: + raise UnsupportedSamplingRequestError( + "Sampling requires a valid unified message origin." + ) + + provider_id = await plugin_context.get_current_chat_provider_id(umo) + provider = plugin_context.get_using_provider(umo) + if provider is None: + raise UnsupportedSamplingRequestError( + "Sampling requires an active chat provider." + ) + + provider_kwargs: dict[str, Any] = {"max_tokens": params.maxTokens} + if params.temperature is not None: + provider_kwargs["temperature"] = params.temperature + if params.stopSequences: + provider_kwargs["stop"] = params.stopSequences + provider_kwargs["stopSequences"] = params.stopSequences + if params.metadata: + provider_kwargs["metadata"] = params.metadata + + llm_resp = await plugin_context.llm_generate( + chat_provider_id=provider_id, + contexts=contexts, + system_prompt=params.systemPrompt or "", + **provider_kwargs, + ) + + if llm_resp.role == "err": + raise RuntimeError(llm_resp.completion_text or "Provider returned error") + if llm_resp.tools_call_args: + raise UnsupportedSamplingRequestError( + "Tool-assisted sampling responses are not supported in the initial AstrBot integration." + ) + + text = llm_resp.completion_text + if text is None: + raise RuntimeError("Provider returned no textual sampling result") + + model_name = provider.get_model() or provider.meta().model or provider.meta().id + return mcp.types.CreateMessageResult( + role="assistant", + content=mcp.types.TextContent(type="text", text=text), + model=model_name, + stopReason="endTurn", + ) + + @staticmethod + def _extract_bound_runtime( + run_context: ContextWrapper[TContext], + ) -> tuple[Any | None, Any | None]: + agent_context = getattr(run_context, "context", None) + plugin_context = getattr(agent_context, "context", None) + event = getattr(agent_context, "event", None) + return plugin_context, event + + async def _execute_elicitation( + self, + run_context: ContextWrapper[TContext], + params: mcp.types.ElicitRequestParams, + ) -> mcp.types.ElicitResult: + import mcp + + plugin_context, event = self._extract_bound_runtime(run_context) + if event is None: + raise UnsupportedElicitationRequestError( + "Elicitation requires an AstrBot event bound to the MCP tool call." + ) + + sender_id = event.get_sender_id() + if not sender_id: + raise UnsupportedElicitationRequestError( + "Elicitation requires a stable sender ID." + ) + + if isinstance(params, mcp.types.ElicitRequestFormParams): + return await self._execute_form_elicitation( + plugin_context, + event, + sender_id, + params, + ) + if isinstance(params, mcp.types.ElicitRequestURLParams): + return await self._execute_url_elicitation( + event, + sender_id, + params, + ) + raise UnsupportedElicitationRequestError( + f"Unsupported elicitation params type: {type(params).__name__}" + ) + + @staticmethod + def _translate_sampling_messages( + messages: list[mcp.types.SamplingMessage], + ) -> list[dict[str, str]]: + translated: list[dict[str, str]] = [] + for message in messages: + text = MCPClientSubCapabilityBridge._sampling_message_to_text(message) + translated.append( + { + "role": message.role, + "content": text, + } + ) + return translated + + @staticmethod + def _sampling_message_to_text(message: mcp.types.SamplingMessage) -> str: + import mcp + + text_parts: list[str] = [] + for block in message.content_as_list: + if isinstance(block, mcp.types.TextContent): + text_parts.append(block.text) + continue + + if isinstance(block, mcp.types.ImageContent): + raise UnsupportedSamplingRequestError( + "Image sampling inputs are not supported in the initial AstrBot integration." + ) + if isinstance(block, mcp.types.AudioContent): + raise UnsupportedSamplingRequestError( + "Audio sampling inputs are not supported in the initial AstrBot integration." + ) + + raise UnsupportedSamplingRequestError( + f"Sampling content block '{type(block).__name__}' is not supported in the initial AstrBot integration." + ) + + return "\n".join(text_parts) + + async def _execute_form_elicitation( + self, + plugin_context: Any, + event: Any, + sender_id: str, + params: mcp.types.ElicitRequestFormParams, + ) -> mcp.types.ElicitResult: + import mcp + + properties = self._get_elicitation_properties(params.requestedSchema) + deadline = asyncio.get_running_loop().time() + self.elicitation_timeout_seconds + await self._send_elicitation_message( + event, + self._build_form_elicitation_prompt(params, properties), + payload=self._build_form_elicitation_payload(params, properties), + ) + + while True: + reply_text = await self._wait_for_elicitation_reply( + event=event, + sender_id=sender_id, + deadline=deadline, + ) + if reply_text is None: + return mcp.types.ElicitResult(action="cancel") + + action = self._parse_cancel_or_decline_action(reply_text) + if action is not None: + return mcp.types.ElicitResult(action=action) + + try: + content = self._parse_form_elicitation_reply( + requested_schema=params.requestedSchema, + reply_text=reply_text, + ) + except UnsupportedElicitationRequestError as exc: + content = await self._try_llm_form_reply_fallback( + plugin_context=plugin_context, + event=event, + params=params, + reply_text=reply_text, + direct_parse_error=exc, + ) + if content is not None: + return mcp.types.ElicitResult( + action="accept", + content=content, + ) + await self._send_elicitation_message( + event, + self._build_form_retry_prompt(exc), + ) + continue + + return mcp.types.ElicitResult( + action="accept", + content=content, + ) + + async def _try_llm_form_reply_fallback( + self, + *, + plugin_context: Any, + event: Any, + params: mcp.types.ElicitRequestFormParams, + reply_text: str, + direct_parse_error: UnsupportedElicitationRequestError, + ) -> dict[str, str | int | float | bool | list[str] | None] | None: + if plugin_context is None: + return None + + umo = getattr(event, "unified_msg_origin", None) + if not isinstance(umo, str) or not umo: + return None + + try: + provider_id = await plugin_context.get_current_chat_provider_id(umo) + except Exception as exc: # noqa: BLE001 + logger.debug( + "Unable to resolve provider for MCP elicitation fallback on %s: %s", + self._server_name, + exc, + ) + return None + + prompt = self._build_elicitation_llm_fallback_prompt( + params=params, + reply_text=reply_text, + direct_parse_error=direct_parse_error, + ) + try: + llm_resp = await plugin_context.llm_generate( + chat_provider_id=provider_id, + prompt=prompt, + system_prompt=self._build_elicitation_llm_fallback_system_prompt(), + max_tokens=256, + ) + except Exception as exc: # noqa: BLE001 + logger.debug( + "LLM fallback failed during MCP elicitation for %s: %s", + self._server_name, + exc, + ) + return None + + if getattr(llm_resp, "role", None) == "err": + logger.debug( + "Provider returned error during MCP elicitation fallback for %s: %s", + self._server_name, + getattr(llm_resp, "completion_text", "") or "", + ) + return None + + raw_text = getattr(llm_resp, "completion_text", "") or "" + normalized = self._strip_code_fence(raw_text).strip() + if not normalized: + return None + + try: + payload = json.loads(normalized) + except json.JSONDecodeError: + logger.debug( + "LLM fallback returned non-JSON content during MCP elicitation for %s: %s", + self._server_name, + normalized, + ) + return None + + if not isinstance(payload, dict): + return None + + try: + return self._coerce_form_payload(payload, params.requestedSchema) + except UnsupportedElicitationRequestError as exc: + logger.debug( + "LLM fallback returned invalid MCP elicitation payload for %s: %s", + self._server_name, + exc, + ) + return None + + async def _execute_url_elicitation( + self, + event: Any, + sender_id: str, + params: mcp.types.ElicitRequestURLParams, + ) -> mcp.types.ElicitResult: + import mcp + + deadline = asyncio.get_running_loop().time() + self.elicitation_timeout_seconds + await self._send_elicitation_message( + event, + self._build_url_elicitation_prompt(params), + payload=self._build_url_elicitation_payload(params), + ) + + while True: + reply_text = await self._wait_for_elicitation_reply( + event=event, + sender_id=sender_id, + deadline=deadline, + ) + if reply_text is None: + return mcp.types.ElicitResult(action="cancel") + + action = self._parse_url_action(reply_text) + if action is not None: + return mcp.types.ElicitResult(action=action) + + await self._send_elicitation_message( + event, + "Please reply `done`, `decline`, or `cancel` to continue this MCP request.", + ) + + async def _send_elicitation_message( + self, + event: Any, + message: str, + *, + payload: dict[str, Any] | None = None, + ) -> None: + if payload and self._is_webchat_event(event): + try: + await event.send( + MessageChain( + chain=[Json(data=payload)], + type="elicitation", + ) + ) + return + except Exception as exc: # noqa: BLE001 + logger.debug( + "Falling back to plain-text MCP elicitation message for %s: %s", + self._server_name, + exc, + ) + + await event.send(MessageChain().message(message)) + + async def _wait_for_elicitation_reply( + self, + *, + event: Any, + sender_id: str, + deadline: float, + ) -> str | None: + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + return None + + try: + async with pending_mcp_elicitation( + event.unified_msg_origin, + sender_id, + ) as future: + reply = await asyncio.wait_for(future, timeout=remaining) + except asyncio.TimeoutError: + return None + + reply_text = reply.message_text.strip() + if reply_text: + return self._strip_code_fence(reply_text) + return reply.message_outline.strip() + + def _build_form_elicitation_prompt( + self, + params: mcp.types.ElicitRequestFormParams, + properties: dict[str, dict[str, Any]], + ) -> str: + required_fields = set( + self._get_required_elicitation_fields(params.requestedSchema) + ) + lines = [f"MCP server `{self._server_name}` needs more information."] + if params.message.strip(): + lines.append(params.message.strip()) + if properties: + lines.append("Requested fields:") + for field_name, schema in properties.items(): + field_type = self._get_elicitation_field_type(schema) + desc = str(schema.get("description", "")).strip() + suffix = " required" if field_name in required_fields else " optional" + if desc: + lines.append(f"- {field_name} ({field_type},{suffix}): {desc}") + else: + lines.append(f"- {field_name} ({field_type},{suffix})") + if len(properties) == 1: + lines.append("Reply with plain text or JSON.") + elif len(properties) > 1: + lines.append("Reply with JSON or `field: value` lines.") + else: + lines.append("Reply `accept` to continue.") + lines.append("Reply `decline` to refuse or `cancel` to stop.") + return "\n".join(lines) + + def _build_form_elicitation_payload( + self, + params: mcp.types.ElicitRequestFormParams, + properties: dict[str, dict[str, Any]], + ) -> dict[str, Any]: + required_fields = set( + self._get_required_elicitation_fields(params.requestedSchema) + ) + fields: list[dict[str, Any]] = [] + for field_name, schema in properties.items(): + enum_values = schema.get("enum") + fields.append( + { + "name": field_name, + "label": str(schema.get("title") or field_name), + "description": str(schema.get("description", "")).strip(), + "required": field_name in required_fields, + "type": self._get_elicitation_field_type(schema), + "enum": ( + [str(value) for value in enum_values] + if isinstance(enum_values, list) + else [] + ), + } + ) + return { + "kind": "form", + "server_name": self._server_name, + "message": params.message.strip(), + "prompt": self._build_form_elicitation_prompt(params, properties), + "fields": fields, + } + + @staticmethod + def _build_form_retry_prompt(exc: UnsupportedElicitationRequestError) -> str: + return ( + "I could not use that reply for the MCP elicitation.\n" + f"Reason: {exc}\n" + "Please try again, or reply `decline` / `cancel`." + ) + + def _build_url_elicitation_prompt( + self, + params: mcp.types.ElicitRequestURLParams, + ) -> str: + lines = [ + f"MCP server `{self._server_name}` needs an external confirmation step." + ] + if params.message.strip(): + lines.append(params.message.strip()) + lines.append(f"URL: {params.url}") + lines.append( + "Reply `done` after you finish, `decline` to refuse, or `cancel` to stop." + ) + return "\n".join(lines) + + def _build_url_elicitation_payload( + self, + params: mcp.types.ElicitRequestURLParams, + ) -> dict[str, Any]: + return { + "kind": "url", + "server_name": self._server_name, + "message": params.message.strip(), + "prompt": self._build_url_elicitation_prompt(params), + "url": params.url, + } + + @staticmethod + def _build_elicitation_llm_fallback_system_prompt() -> str: + return ( + "You extract structured MCP elicitation data from a user's natural-language reply.\n" + "Return only a JSON object.\n" + "Use only keys from the provided schema.\n" + "Do not invent facts. Omit fields that are not clearly supported.\n" + "Use proper JSON types for booleans, integers, numbers, and arrays.\n" + "Do not wrap the JSON in markdown fences." + ) + + def _build_elicitation_llm_fallback_prompt( + self, + *, + params: mcp.types.ElicitRequestFormParams, + reply_text: str, + direct_parse_error: UnsupportedElicitationRequestError, + ) -> str: + return ( + f"MCP server: {self._server_name}\n" + f"Original elicitation message:\n{params.message.strip() or ''}\n\n" + f"Requested JSON schema:\n" + f"{json.dumps(params.requestedSchema, ensure_ascii=False, indent=2)}\n\n" + f"User reply:\n{reply_text}\n\n" + f"Direct parser error:\n{direct_parse_error}\n\n" + "Produce the best possible JSON object that matches the schema." + ) + + def _parse_form_elicitation_reply( + self, + *, + requested_schema: dict[str, Any], + reply_text: str, + ) -> dict[str, str | int | float | bool | list[str] | None]: + properties = self._get_elicitation_properties(requested_schema) + if not properties: + return {} + + normalized_reply = reply_text.strip() + if not normalized_reply: + raise UnsupportedElicitationRequestError("The reply is empty.") + + if normalized_reply.startswith("{"): + try: + payload = json.loads(normalized_reply) + except json.JSONDecodeError as exc: + raise UnsupportedElicitationRequestError( + "The JSON reply could not be parsed." + ) from exc + if not isinstance(payload, dict): + raise UnsupportedElicitationRequestError( + "The JSON reply must be an object." + ) + elif len(properties) == 1: + field_name = next(iter(properties)) + payload = {field_name: normalized_reply} + else: + payload = self._parse_key_value_lines(normalized_reply, properties) + if not payload: + payload = self._parse_natural_language_form_reply( + reply_text=normalized_reply, + requested_schema=requested_schema, + ) + if not payload: + raise UnsupportedElicitationRequestError( + "Please reply with JSON, natural language, or `field: value` lines." + ) + + return self._coerce_form_payload(payload, requested_schema) + + def _parse_natural_language_form_reply( + self, + *, + reply_text: str, + requested_schema: dict[str, Any], + ) -> dict[str, Any]: + properties = self._get_elicitation_properties(requested_schema) + if not properties: + return {} + + parsed = self._parse_field_patterns(reply_text, properties) + parsed.update(self._match_enum_values(reply_text, properties, parsed.keys())) + if parsed: + return parsed + + target_fields = self._get_required_elicitation_fields(requested_schema) + if not target_fields: + target_fields = list(properties.keys()) + if len(target_fields) == 1: + return {target_fields[0]: reply_text} + + return {} + + def _coerce_form_payload( + self, + payload: dict[str, Any], + requested_schema: dict[str, Any], + ) -> dict[str, str | int | float | bool | list[str] | None]: + properties = self._get_elicitation_properties(requested_schema) + required_fields = self._get_required_elicitation_fields(requested_schema) + normalized_keys = { + field_name.casefold(): field_name for field_name in properties.keys() + } + + coerced: dict[str, str | int | float | bool | list[str] | None] = {} + for raw_key, raw_value in payload.items(): + normalized_key = str(raw_key).strip().casefold() + field_name = normalized_keys.get(normalized_key) + if field_name is None: + continue + coerced[field_name] = self._coerce_form_value( + field_name=field_name, + raw_value=raw_value, + schema=properties[field_name], + ) + + missing_required = [ + field_name for field_name in required_fields if field_name not in coerced + ] + if missing_required: + raise UnsupportedElicitationRequestError( + "Missing required field(s): " + ", ".join(missing_required) + ) + return coerced + + def _coerce_form_value( + self, + *, + field_name: str, + raw_value: Any, + schema: dict[str, Any], + ) -> str | int | float | bool | list[str] | None: + field_type = self._get_elicitation_field_type(schema) + if raw_value is None: + return None + + if field_type == "string": + value = str(raw_value).strip() + elif field_type == "integer": + if isinstance(raw_value, bool): + raise UnsupportedElicitationRequestError( + f"Field `{field_name}` must be an integer." + ) + try: + value = int(str(raw_value).strip()) + except (TypeError, ValueError) as exc: + raise UnsupportedElicitationRequestError( + f"Field `{field_name}` must be an integer." + ) from exc + elif field_type == "number": + if isinstance(raw_value, bool): + raise UnsupportedElicitationRequestError( + f"Field `{field_name}` must be a number." + ) + try: + value = float(str(raw_value).strip()) + except (TypeError, ValueError) as exc: + raise UnsupportedElicitationRequestError( + f"Field `{field_name}` must be a number." + ) from exc + elif field_type == "boolean": + value = self._coerce_boolean_value(field_name, raw_value) + elif field_type == "array": + value = self._coerce_string_array_value(field_name, raw_value) + else: + raise UnsupportedElicitationRequestError( + f"Field `{field_name}` uses unsupported type `{field_type}`." + ) + + enum_values = schema.get("enum") + if isinstance(enum_values, list) and value not in enum_values: + raise UnsupportedElicitationRequestError( + f"Field `{field_name}` must be one of: {', '.join(map(str, enum_values))}." + ) + return value + + @staticmethod + def _coerce_boolean_value(field_name: str, raw_value: Any) -> bool: + if isinstance(raw_value, bool): + return raw_value + + normalized = str(raw_value).strip().casefold() + truthy = {"true", "1", "yes", "y", "on", "是", "好的"} + falsy = {"false", "0", "no", "n", "off", "否", "不是"} + if normalized in truthy: + return True + if normalized in falsy: + return False + raise UnsupportedElicitationRequestError( + f"Field `{field_name}` must be a boolean." + ) + + @staticmethod + def _coerce_string_array_value(field_name: str, raw_value: Any) -> list[str]: + if isinstance(raw_value, list): + return [str(item).strip() for item in raw_value if str(item).strip()] + + normalized = str(raw_value).strip() + if not normalized: + return [] + parts = [ + part.strip() + for chunk in normalized.splitlines() + for part in chunk.split(",") + if part.strip() + ] + if not parts: + raise UnsupportedElicitationRequestError( + f"Field `{field_name}` must be a string array." + ) + return parts + + @staticmethod + def _parse_key_value_lines( + reply_text: str, + properties: dict[str, dict[str, Any]], + ) -> dict[str, str]: + normalized_keys = { + field_name.casefold(): field_name for field_name in properties.keys() + } + parsed: dict[str, str] = {} + for line in reply_text.splitlines(): + stripped = line.strip() + if not stripped: + continue + delimiter = ":" if ":" in stripped else ("：" if "：" in stripped else None) + if delimiter is None: + continue + raw_key, raw_value = stripped.split(delimiter, 1) + field_name = normalized_keys.get(raw_key.strip().casefold()) + if field_name is None: + continue + parsed[field_name] = raw_value.strip() + return parsed + + @staticmethod + def _parse_field_patterns( + reply_text: str, + properties: dict[str, dict[str, Any]], + ) -> dict[str, str]: + parsed: dict[str, str] = {} + separators = r"[:：=]|是|为" + boundaries = r"(?:[,，;；。]|$)" + for field_name in properties: + pattern = re.compile( + rf"{re.escape(field_name)}\s*(?:{separators})\s*(.+?)(?={boundaries})", + re.IGNORECASE, + ) + match = pattern.search(reply_text) + if match: + value = match.group(1).strip().strip("`'\"") + if value: + parsed[field_name] = value + return parsed + + @staticmethod + def _match_enum_values( + reply_text: str, + properties: dict[str, dict[str, Any]], + ignore_fields: set[str] | Any, + ) -> dict[str, str]: + normalized_reply = reply_text.casefold() + parsed: dict[str, str] = {} + ignored = set(ignore_fields) + for field_name, schema in properties.items(): + if field_name in ignored: + continue + enum_values = schema.get("enum") + if not isinstance(enum_values, list) or not enum_values: + continue + + matches = [ + str(enum_value) + for enum_value in enum_values + if str(enum_value).casefold() in normalized_reply + ] + if len(matches) == 1: + parsed[field_name] = matches[0] + return parsed + + @staticmethod + def _get_elicitation_properties( + requested_schema: dict[str, Any], + ) -> dict[str, dict[str, Any]]: + properties = requested_schema.get("properties", {}) + if not isinstance(properties, dict): + raise UnsupportedElicitationRequestError( + "Form-mode elicitation requires a top-level properties object." + ) + normalized_properties: dict[str, dict[str, Any]] = {} + for field_name, field_schema in properties.items(): + if isinstance(field_name, str) and isinstance(field_schema, dict): + normalized_properties[field_name] = field_schema + return normalized_properties + + @staticmethod + def _get_required_elicitation_fields( + requested_schema: dict[str, Any], + ) -> list[str]: + required_fields = requested_schema.get("required", []) + if not isinstance(required_fields, list): + return [] + return [field for field in required_fields if isinstance(field, str)] + + @staticmethod + def _get_elicitation_field_type(field_schema: dict[str, Any]) -> str: + field_type = field_schema.get("type", "string") + if isinstance(field_type, list): + for candidate in field_type: + if candidate in {"string", "integer", "number", "boolean", "array"}: + return candidate + raise UnsupportedElicitationRequestError( + "Unsupported multi-type elicitation field." + ) + if not isinstance(field_type, str): + return "string" + return field_type + + @staticmethod + def _parse_cancel_or_decline_action(reply_text: str) -> str | None: + normalized = reply_text.strip().casefold() + if normalized in MCP_ELICITATION_CANCEL_KEYWORDS: + return "cancel" + if normalized in MCP_ELICITATION_DECLINE_KEYWORDS: + return "decline" + return None + + @staticmethod + def _parse_url_action(reply_text: str) -> str | None: + normalized = reply_text.strip().casefold() + if normalized in MCP_ELICITATION_ACCEPT_KEYWORDS: + return "accept" + if normalized in MCP_ELICITATION_DECLINE_KEYWORDS: + return "decline" + if normalized in MCP_ELICITATION_CANCEL_KEYWORDS: + return "cancel" + return None + + @staticmethod + def _strip_code_fence(text: str) -> str: + stripped = text.strip() + if not stripped.startswith("```") or not stripped.endswith("```"): + return stripped + lines = stripped.splitlines() + if len(lines) <= 2: + return stripped.removeprefix("```").removesuffix("```").strip() + return "\n".join(lines[1:-1]).strip() + + @staticmethod + def _is_webchat_event(event: Any) -> bool: + platform_name = getattr(event, "get_platform_name", None) + if callable(platform_name): + try: + return platform_name() == "webchat" + except Exception: # noqa: BLE001 + return False + return False + + def _build_root_entries(self) -> list[mcp.types.Root]: + import mcp + + roots: list[mcp.types.Root] = [] + seen_paths: set[str] = set() + for name, path in self._iter_resolved_root_paths(): + normalized_path = str(path) + if normalized_path in seen_paths: + continue + seen_paths.add(normalized_path) + roots.append( + mcp.types.Root( + uri=path.as_uri(), + name=name, + ) + ) + return roots + + def _iter_resolved_root_paths(self) -> list[tuple[str, Path]]: + configured_paths = self._capabilities.roots.paths or list( + DEFAULT_MCP_ROOT_PATHS + ) + resolved_entries: list[tuple[str, Path]] = [] + for entry in configured_paths: + resolved = self._resolve_root_path_entry(entry) + if resolved is not None: + resolved_entries.append(resolved) + return resolved_entries + + def _resolve_root_path_entry(self, entry: str) -> tuple[str, Path] | None: + normalized_entry = entry.strip() + if not normalized_entry: + return None + + alias_key = normalized_entry.lower() + alias_resolvers = get_root_path_alias_resolvers() + if alias_key in alias_resolvers: + path = Path(alias_resolvers[alias_key]()).resolve() + display_name = alias_key + else: + candidate_path = Path(normalized_entry).expanduser() + if not candidate_path.is_absolute(): + candidate_path = Path(get_astrbot_root()) / candidate_path + path = candidate_path.resolve() + display_name = path.name or normalized_entry + + if not path.exists(): + logger.warning( + "Skipping missing MCP root path for server %s: %s", + self._server_name, + path, + ) + return None + + return display_name, path diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py index 523d758a0a..de3dafdd9a 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py @@ -6,6 +6,9 @@ from dataclasses import replace from astrbot.core import logger +from astrbot.core.agent.mcp_elicitation_registry import ( + try_capture_pending_mcp_elicitation, +) from astrbot.core.agent.message import Message from astrbot.core.agent.response import AgentStats from astrbot.core.astr_main_agent import ( @@ -164,6 +167,12 @@ async def process( return logger.debug("ready to request llm provider") + if try_capture_pending_mcp_elicitation(event): + logger.info( + "Captured MCP elicitation reply for active agent run, umo=%s", + event.unified_msg_origin, + ) + return follow_up_capture = try_capture_follow_up(event) if follow_up_capture: ( diff --git a/astrbot/core/platform/sources/webchat/webchat_event.py b/astrbot/core/platform/sources/webchat/webchat_event.py index b7da864aae..5ceededf10 100644 --- a/astrbot/core/platform/sources/webchat/webchat_event.py +++ b/astrbot/core/platform/sources/webchat/webchat_event.py @@ -66,6 +66,17 @@ async def _send( }, ) elif isinstance(comp, Json): + if message.type == "elicitation" and isinstance(comp.data, dict): + await web_chat_back_queue.put( + { + "type": "elicitation", + "data": comp.data, + "streaming": streaming, + "chain_type": message.type, + "message_id": message_id, + }, + ) + continue await web_chat_back_queue.put( { "type": "plain", diff --git a/astrbot/core/provider/func_tool_manager.py b/astrbot/core/provider/func_tool_manager.py index f950b00250..e2a9635fb7 100644 --- a/astrbot/core/provider/func_tool_manager.py +++ b/astrbot/core/provider/func_tool_manager.py @@ -8,6 +8,7 @@ import urllib.parse from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping from dataclasses import dataclass +from pathlib import Path from types import MappingProxyType from typing import Any @@ -16,6 +17,12 @@ from astrbot import logger from astrbot.core import sp from astrbot.core.agent.mcp_client import MCPClient, MCPTool +from astrbot.core.agent.mcp_prompt_bridge import build_mcp_prompt_tools +from astrbot.core.agent.mcp_resource_bridge import build_mcp_resource_tools +from astrbot.core.agent.mcp_subcapability_bridge import ( + normalize_mcp_config, + normalize_mcp_server_config, +) from astrbot.core.agent.tool import FunctionTool, ToolSet from astrbot.core.utils.astrbot_path import get_astrbot_data_path @@ -145,7 +152,10 @@ def _prepare_config(config: dict) -> dict: if config.get("mcpServers"): first_key = next(iter(config["mcpServers"])) config = config["mcpServers"][first_key] + config = normalize_mcp_server_config(config) config.pop("active", None) + config.pop("client_capabilities", None) + config.pop("provider", None) return config @@ -320,6 +330,13 @@ def get_full_tool_set(self) -> ToolSet: tool_set = ToolSet(self.func_list.copy()) return tool_set + def _remove_mcp_bound_tools(self, name: str) -> None: + self.func_list = [ + tool + for tool in self.func_list + if getattr(tool, "mcp_server_name", None) != name + ] + @staticmethod def _log_safe_mcp_debug_config(cfg: dict) -> None: # 仅记录脱敏后的摘要，避免泄露 command/args/url 中的敏感信息 @@ -370,18 +387,15 @@ async def init_mcp_clients( - 初始化超时使用环境变量 ASTRBOT_MCP_INIT_TIMEOUT 或默认值。 - 动态启用超时使用 ASTRBOT_MCP_ENABLE_TIMEOUT（独立于初始化超时）。 """ - data_dir = get_astrbot_data_path() - - mcp_json_file = os.path.join(data_dir, "mcp_server.json") - if not os.path.exists(mcp_json_file): + mcp_json_file = Path(get_astrbot_data_path()) / "mcp_server.json" + if not mcp_json_file.exists(): # 配置文件不存在错误处理 - with open(mcp_json_file, "w", encoding="utf-8") as f: + with mcp_json_file.open("w", encoding="utf-8") as f: json.dump(DEFAULT_MCP_CONFIG, f, ensure_ascii=False, indent=4) logger.info(f"未找到 MCP 服务配置文件，已创建默认配置文件 {mcp_json_file}") return MCPInitSummary(total=0, success=0, failed=[]) - with open(mcp_json_file, encoding="utf-8") as f: - mcp_server_json_obj: dict[str, dict] = json.load(f)["mcpServers"] + mcp_server_json_obj = self.load_mcp_config()["mcpServers"] init_timeout = self._init_timeout_default timeout_display = f"{init_timeout:g}" @@ -583,6 +597,8 @@ async def _init_mcp_client(self, name: str, config: dict) -> MCPClient: try: await mcp_client.connect_to_server(config, name) tools_res = await mcp_client.list_tools_and_save() + await mcp_client.load_resource_capabilities() + await mcp_client.load_prompt_capabilities() except asyncio.CancelledError: await self._cleanup_mcp_client_safely(mcp_client, name) raise @@ -591,13 +607,11 @@ async def _init_mcp_client(self, name: str, config: dict) -> MCPClient: raise logger.debug(f"MCP server {name} list tools response: {tools_res}") tool_names = [tool.name for tool in tools_res.tools] + tool_names.extend(getattr(mcp_client, "resource_bridge_tool_names", [])) + tool_names.extend(getattr(mcp_client, "prompt_bridge_tool_names", [])) # 移除该MCP服务之前的工具（如有） - self.func_list = [ - f - for f in self.func_list - if not (isinstance(f, MCPTool) and f.mcp_server_name == name) - ] + self._remove_mcp_bound_tools(name) # 将 MCP 工具转换为 FuncTool 并添加到 func_list for tool in mcp_client.tools: @@ -608,6 +622,11 @@ async def _init_mcp_client(self, name: str, config: dict) -> MCPClient: ) self.func_list.append(func_tool) + for resource_tool in build_mcp_resource_tools(mcp_client, name): + self.func_list.append(resource_tool) + for prompt_tool in build_mcp_prompt_tools(mcp_client, name): + self.func_list.append(prompt_tool) + logger.info(f"Connected to MCP server {name}, Tools: {tool_names}") return mcp_client @@ -620,11 +639,7 @@ async def _terminate_mcp_client(self, name: str) -> None: # 关闭MCP连接 await self._cleanup_mcp_client_safely(client, name) # 移除关联的FuncTool - self.func_list = [ - f - for f in self.func_list - if not (isinstance(f, MCPTool) and f.mcp_server_name == name) - ] + self._remove_mcp_bound_tools(name) async with self._runtime_lock: self._mcp_server_runtime.pop(name, None) self._mcp_starting.discard(name) @@ -632,11 +647,7 @@ async def _terminate_mcp_client(self, name: str) -> None: return # Runtime missing but stale tools may still exist after failed flows. - self.func_list = [ - f - for f in self.func_list - if not (isinstance(f, MCPTool) and f.mcp_server_name == name) - ] + self._remove_mcp_bound_tools(name) async with self._runtime_lock: self._mcp_starting.discard(name) @@ -652,7 +663,11 @@ async def test_mcp_server_connection(config: dict) -> list[str]: logger.debug(f"testing MCP server connection with config: {config}") await mcp_client.connect_to_server(config, "test") tools_res = await mcp_client.list_tools_and_save() + await mcp_client.load_resource_capabilities() + await mcp_client.load_prompt_capabilities() tool_names = [tool.name for tool in tools_res.tools] + tool_names.extend(getattr(mcp_client, "resource_bridge_tool_names", [])) + tool_names.extend(getattr(mcp_client, "prompt_bridge_tool_names", [])) finally: logger.debug("Cleaning up MCP client after testing connection.") await mcp_client.cleanup() @@ -820,28 +835,29 @@ def activate_llm_tool(self, name: str, star_map: dict) -> bool: @property def mcp_config_path(self): - data_dir = get_astrbot_data_path() - return os.path.join(data_dir, "mcp_server.json") + return Path(get_astrbot_data_path()) / "mcp_server.json" def load_mcp_config(self): - if not os.path.exists(self.mcp_config_path): + if not self.mcp_config_path.exists(): # 配置文件不存在，创建默认配置 - os.makedirs(os.path.dirname(self.mcp_config_path), exist_ok=True) - with open(self.mcp_config_path, "w", encoding="utf-8") as f: + self.mcp_config_path.parent.mkdir(parents=True, exist_ok=True) + with self.mcp_config_path.open("w", encoding="utf-8") as f: json.dump(DEFAULT_MCP_CONFIG, f, ensure_ascii=False, indent=4) - return DEFAULT_MCP_CONFIG + return normalize_mcp_config(DEFAULT_MCP_CONFIG) try: - with open(self.mcp_config_path, encoding="utf-8") as f: - return json.load(f) + with self.mcp_config_path.open(encoding="utf-8") as f: + return normalize_mcp_config(json.load(f)) except Exception as e: logger.error(f"加载 MCP 配置失败: {e}") - return DEFAULT_MCP_CONFIG + return normalize_mcp_config(DEFAULT_MCP_CONFIG) def save_mcp_config(self, config: dict) -> bool: try: - with open(self.mcp_config_path, "w", encoding="utf-8") as f: - json.dump(config, f, ensure_ascii=False, indent=4) + normalized = normalize_mcp_config(config) + self.mcp_config_path.parent.mkdir(parents=True, exist_ok=True) + with self.mcp_config_path.open("w", encoding="utf-8") as f: + json.dump(normalized, f, ensure_ascii=False, indent=4) return True except Exception as e: logger.error(f"保存 MCP 配置失败: {e}") diff --git a/astrbot/dashboard/routes/chat.py b/astrbot/dashboard/routes/chat.py index a914f3cbf0..84c6ccb515 100644 --- a/astrbot/dashboard/routes/chat.py +++ b/astrbot/dashboard/routes/chat.py @@ -10,6 +10,9 @@ from quart import g, make_response, request, send_file from astrbot.core import logger, sp +from astrbot.core.agent.mcp_elicitation_registry import ( + submit_pending_mcp_elicitation_reply, +) from astrbot.core.core_lifecycle import AstrBotCoreLifecycle from astrbot.core.db import BaseDatabase from astrbot.core.platform.message_type import MessageType @@ -50,6 +53,7 @@ def __init__( "/chat/sessions": ("GET", self.get_sessions), "/chat/get_session": ("GET", self.get_session), "/chat/stop": ("POST", self.stop_session), + "/chat/respond_elicitation": ("POST", self.respond_elicitation), "/chat/delete_session": ("GET", self.delete_webchat_session), "/chat/update_session_display_name": ( "POST", @@ -73,6 +77,18 @@ def __init__( self.running_convs: dict[str, bool] = {} + @staticmethod + def _build_webchat_session_umo(session) -> str: + message_type = ( + MessageType.GROUP_MESSAGE.value + if session.is_group + else MessageType.FRIEND_MESSAGE.value + ) + return ( + f"{session.platform_id}:{message_type}:" + f"{session.platform_id}!{session.creator}!{session.session_id}" + ) + async def get_file(self): filename = request.args.get("filename") if not filename: @@ -447,6 +463,19 @@ async def stream(): ) if part: accumulated_parts.append(part) + elif msg_type == "elicitation": + if accumulated_text: + accumulated_parts.append( + {"type": "plain", "text": accumulated_text} + ) + accumulated_text = "" + if isinstance(result_text, dict): + accumulated_parts.append( + { + "type": "elicitation", + "payload": result_text, + } + ) # 消息结束处理 if msg_type == "end": @@ -565,19 +594,72 @@ async def stop_session(self): if session.creator != username: return Response().error("Permission denied").__dict__ - message_type = ( - MessageType.GROUP_MESSAGE.value - if session.is_group - else MessageType.FRIEND_MESSAGE.value - ) - umo = ( - f"{session.platform_id}:{message_type}:" - f"{session.platform_id}!{username}!{session_id}" - ) + umo = self._build_webchat_session_umo(session) stopped_count = active_event_registry.request_agent_stop_all(umo) return Response().ok(data={"stopped_count": stopped_count}).__dict__ + async def respond_elicitation(self): + post_data = await request.json + if post_data is None: + return Response().error("Missing JSON body").__dict__ + + session_id = str(post_data.get("session_id", "")).strip() + reply_text = str(post_data.get("reply_text", "")).strip() + display_text = str(post_data.get("display_text", reply_text)).strip() + if not session_id: + return Response().error("Missing key: session_id").__dict__ + if not reply_text: + return Response().error("Missing key: reply_text").__dict__ + + username = g.get("username", "guest") + session = await self.db.get_platform_session_by_id(session_id) + if not session: + return Response().error(f"Session {session_id} not found").__dict__ + if session.creator != username: + return Response().error("Permission denied").__dict__ + + umo = self._build_webchat_session_umo(session) + if not submit_pending_mcp_elicitation_reply( + umo, + username, + reply_text, + reply_outline=display_text, + ): + return ( + Response().error("No pending MCP elicitation for this session").__dict__ + ) + + saved_record = await self.platform_history_mgr.insert( + platform_id=session.platform_id, + user_id=session_id, + content={ + "type": "user", + "message": [{"type": "plain", "text": display_text or reply_text}], + }, + sender_id=username, + sender_name=username, + ) + + return ( + Response() + .ok( + data={ + "saved_message": { + "id": saved_record.id, + "created_at": to_utc_isoformat(saved_record.created_at), + "content": { + "type": "user", + "message": [ + {"type": "plain", "text": display_text or reply_text} + ], + }, + } + } + ) + .__dict__ + ) + async def delete_webchat_session(self): """Delete a Platform session and all its related data.""" session_id = request.args.get("session_id") diff --git a/astrbot/dashboard/routes/live_chat.py b/astrbot/dashboard/routes/live_chat.py index 8d0af938d0..275c9afa85 100644 --- a/astrbot/dashboard/routes/live_chat.py +++ b/astrbot/dashboard/routes/live_chat.py @@ -576,6 +576,19 @@ async def _handle_chat_message( part = await self._create_attachment_from_file(filename, "video") if part: accumulated_parts.append(part) + elif msg_type == "elicitation": + if accumulated_text: + accumulated_parts.append( + {"type": "plain", "text": accumulated_text} + ) + accumulated_text = "" + if isinstance(result_text, dict): + accumulated_parts.append( + { + "type": "elicitation", + "payload": result_text, + } + ) should_save = False if msg_type == "end": diff --git a/astrbot/dashboard/routes/open_api.py b/astrbot/dashboard/routes/open_api.py index 9a736b1763..2f96c61b82 100644 --- a/astrbot/dashboard/routes/open_api.py +++ b/astrbot/dashboard/routes/open_api.py @@ -451,6 +451,19 @@ async def _handle_chat_ws_send(self, post_data: dict) -> None: ) if part: accumulated_parts.append(part) + elif msg_type == "elicitation": + if accumulated_text: + accumulated_parts.append( + {"type": "plain", "text": accumulated_text} + ) + accumulated_text = "" + if isinstance(result_text, dict): + accumulated_parts.append( + { + "type": "elicitation", + "payload": result_text, + } + ) if msg_type == "end": break diff --git a/astrbot/dashboard/routes/tools.py b/astrbot/dashboard/routes/tools.py index 84f8dcc6d7..19f6bf896c 100644 --- a/astrbot/dashboard/routes/tools.py +++ b/astrbot/dashboard/routes/tools.py @@ -104,7 +104,13 @@ async def get_mcp_servers(self): for name_key, runtime in self.tool_mgr.mcp_server_runtime_view.items(): if name_key == name: mcp_client = runtime.client - server_info["tools"] = [tool.name for tool in mcp_client.tools] + server_info["tools"] = ( + [tool.name for tool in mcp_client.tools] + + list( + getattr(mcp_client, "resource_bridge_tool_names", []) + ) + + list(getattr(mcp_client, "prompt_bridge_tool_names", [])) + ) server_info["errlogs"] = mcp_client.server_errlogs break else: @@ -431,7 +437,7 @@ async def get_tool_list(self): tools = self.tool_mgr.func_list tools_dict = [] for tool in tools: - if isinstance(tool, MCPTool): + if isinstance(tool, MCPTool) or getattr(tool, "mcp_server_name", None): origin = "mcp" origin_name = tool.mcp_server_name elif tool.handler_module_path and star_map.get( diff --git a/dashboard/src/components/chat/Chat.vue b/dashboard/src/components/chat/Chat.vue index 7c25e1bc3a..3d9c60ba69 100644 --- a/dashboard/src/components/chat/Chat.vue +++ b/dashboard/src/components/chat/Chat.vue @@ -51,6 +51,7 @@ { + messageList.value?.scrollToBottom(); + }); +} + // 路由变化监听 watch( () => route.path, diff --git a/dashboard/src/components/chat/MessageList.vue b/dashboard/src/components/chat/MessageList.vue index 42da95aa58..80041efd3a 100644 --- a/dashboard/src/components/chat/MessageList.vue +++ b/dashboard/src/components/chat/MessageList.vue @@ -97,6 +97,8 @@ @@ -221,6 +223,10 @@ export default { isLoadingMessages: { type: Boolean, default: false + }, + submitElicitation: { + type: Function, + default: null } }, emits: ['openImagePreview', 'replyMessage', 'replyWithText', 'openRefs'], @@ -417,6 +423,14 @@ export default { return messageParts.some(part => part.type === 'record' && part.embedded_url); }, + isActiveElicitationMessage(index, msg) { + if (!this.isStreaming || index !== this.messages.length - 1) { + return false; + } + return Array.isArray(msg?.content?.message) + && msg.content.message.some(part => part.type === 'elicitation' && part.payload); + }, + // 获取被引用消息的内容 getReplyContent(messageId) { const replyMsg = this.messages.find(m => m.id === messageId); diff --git a/dashboard/src/components/chat/StandaloneChat.vue b/dashboard/src/components/chat/StandaloneChat.vue index 69fac13f9b..5c96e56b42 100644 --- a/dashboard/src/components/chat/StandaloneChat.vue +++ b/dashboard/src/components/chat/StandaloneChat.vue @@ -5,7 +5,8 @@
@@ -158,6 +159,7 @@ const { enableStreaming, getSessionMessages: getSessionMsg, sendMessage: sendMsg, + submitElicitationResponse, stopMessage: stopMsg, toggleStreaming } = useMessages(currSessionId, getMediaFile, updateSessionTitle, getSessions); @@ -243,6 +245,17 @@ async function handleStopMessage() { await stopMsg(); } +async function handleSubmitElicitation(replyText: string, displayText: string) { + if (!currSessionId.value) { + return; + } + + await submitElicitationResponse(currSessionId.value, replyText, displayText); + nextTick(() => { + messageList.value?.scrollToBottom(); + }); +} + onMounted(async () => { // 独立模式在挂载时创建新会话 try { diff --git a/dashboard/src/components/chat/message_list_comps/ElicitationCard.vue b/dashboard/src/components/chat/message_list_comps/ElicitationCard.vue new file mode 100644 index 0000000000..a7e6f679af --- /dev/null +++ b/dashboard/src/components/chat/message_list_comps/ElicitationCard.vue @@ -0,0 +1,352 @@ + + + + + diff --git a/dashboard/src/components/chat/message_list_comps/MessagePartsRenderer.vue b/dashboard/src/components/chat/message_list_comps/MessagePartsRenderer.vue index 5fd7f59bed..4c9a953b28 100644 --- a/dashboard/src/components/chat/message_list_comps/MessagePartsRenderer.vue +++ b/dashboard/src/components/chat/message_list_comps/MessagePartsRenderer.vue @@ -60,6 +60,14 @@ + + import { useI18n, useModuleI18n } from '@/i18n/composables'; import { MarkdownRender } from 'markstream-vue'; +import ElicitationCard from './ElicitationCard.vue'; import IPythonToolBlock from './IPythonToolBlock.vue'; import ToolCallItem from './ToolCallItem.vue'; @@ -136,6 +145,14 @@ const props = defineProps({ downloadingFiles: { type: Object, default: () => new Set() + }, + interactiveElicitation: { + type: Boolean, + default: false + }, + submitElicitation: { + type: Function, + default: null } }); diff --git a/dashboard/src/composables/useMessages.ts b/dashboard/src/composables/useMessages.ts index c593fb283b..9c9d5ca6ba 100644 --- a/dashboard/src/composables/useMessages.ts +++ b/dashboard/src/composables/useMessages.ts @@ -34,14 +34,33 @@ export interface FileInfo { attachment_id?: string; // 用于按需下载 } +export interface ElicitationField { + name: string; + label: string; + description?: string; + required: boolean; + type: string; + enum?: string[]; +} + +export interface ElicitationPayload { + kind: 'form' | 'url'; + server_name: string; + message: string; + prompt: string; + url?: string; + fields?: ElicitationField[]; +} + // 消息部分的类型定义 export interface MessagePart { - type: 'plain' | 'image' | 'record' | 'file' | 'video' | 'reply' | 'tool_call'; + type: 'plain' | 'image' | 'record' | 'file' | 'video' | 'reply' | 'tool_call' | 'elicitation'; text?: string; // for plain attachment_id?: string; // for image, record, file, video filename?: string; // for file (filename from backend) message_id?: number; // for reply (PlatformSessionHistoryMessage.id) tool_calls?: ToolCall[]; // for tool_call + payload?: ElicitationPayload; // for elicitation // embedded fields - 加载后填充 embedded_url?: string; // blob URL for image, record embedded_file?: FileInfo; // for file (保留 attachment_id 用于按需下载) @@ -84,6 +103,35 @@ type StreamChunk = { [key: string]: any; }; +function normalizeElicitationPayload(payload: any): ElicitationPayload | null { + if (!payload || typeof payload !== 'object') { + return null; + } + + const kind = payload.kind === 'url' ? 'url' : 'form'; + const fields = Array.isArray(payload.fields) + ? payload.fields + .filter((field: any) => field && typeof field === 'object' && typeof field.name === 'string') + .map((field: any) => ({ + name: String(field.name), + label: String(field.label || field.name), + description: String(field.description || ''), + required: Boolean(field.required), + type: String(field.type || 'string'), + enum: Array.isArray(field.enum) ? field.enum.map((value: any) => String(value)) : [] + })) + : []; + + return { + kind, + server_name: String(payload.server_name || ''), + message: String(payload.message || ''), + prompt: String(payload.prompt || ''), + url: typeof payload.url === 'string' ? payload.url : undefined, + fields + }; +} + type WsStreamContext = { handleChunk: (payload: StreamChunk) => Promise; finish: (err?: unknown) => void; @@ -207,6 +255,24 @@ export function useMessages( return; } + if (payload.type === 'elicitation') { + const normalizedPayload = normalizeElicitationPayload(payload.data); + if (!normalizedPayload) { + return; + } + + messages.value.push({ + content: { + type: 'bot', + message: [{ + type: 'elicitation', + payload: normalizedPayload + }] + } + }); + return; + } + if (payload.type === 'plain') { const chainType = payload.chain_type || 'normal'; if (chainType === 'reasoning') { @@ -440,6 +506,24 @@ export function useMessages( return; } + if (chunkJson.type === 'elicitation') { + const normalizedPayload = normalizeElicitationPayload(chunkJson.data); + if (!normalizedPayload) { + return; + } + + messages.value.push({ + content: { + type: 'bot', + message: [{ + type: 'elicitation', + payload: normalizedPayload + }] + } + }); + return; + } + if (chunkJson.type === 'image') { const img = String(chunkJson.data || '').replace('[IMAGE]', ''); const imageUrl = await getMediaFile(img); @@ -671,6 +755,9 @@ export function useMessages( }; } // plain, reply, tool_call, video 保持原样 + if (part.type === 'elicitation') { + part.payload = normalizeElicitationPayload(part.payload) || undefined; + } } } @@ -1003,6 +1090,49 @@ export function useMessages( } } + async function submitElicitationResponse( + sessionId: string, + replyText: string, + displayText: string + ) { + const normalizedSessionId = sessionId.trim(); + const normalizedReplyText = replyText.trim(); + const normalizedDisplayText = (displayText || replyText).trim(); + if (!normalizedSessionId || !normalizedReplyText) { + throw new Error('Missing elicitation reply payload'); + } + + const response = await axios.post('/api/chat/respond_elicitation', { + session_id: normalizedSessionId, + reply_text: normalizedReplyText, + display_text: normalizedDisplayText + }); + if (response.data?.status !== 'ok') { + throw new Error(response.data?.message || 'Failed to submit elicitation reply'); + } + + const savedMessage = response.data?.data?.saved_message; + if (savedMessage?.content) { + await parseMessageContent(savedMessage.content); + messages.value.push({ + id: savedMessage.id, + created_at: savedMessage.created_at, + content: savedMessage.content + }); + return savedMessage; + } + + const fallbackMessage: MessageContent = { + type: 'user', + message: [{ + type: 'plain', + text: normalizedDisplayText + }] + }; + messages.value.push({ content: fallbackMessage }); + return null; + } + async function stopMessage() { const sessionId = currentRunningSessionId.value || currSessionId.value; if (!sessionId) { @@ -1057,6 +1187,7 @@ export function useMessages( currentSessionProject, getSessionMessages, sendMessage, + submitElicitationResponse, stopMessage, toggleStreaming, setTransportMode, diff --git a/dashboard/src/i18n/locales/en-US/features/chat.json b/dashboard/src/i18n/locales/en-US/features/chat.json index 53e98c991a..5d1e79eb6a 100644 --- a/dashboard/src/i18n/locales/en-US/features/chat.json +++ b/dashboard/src/i18n/locales/en-US/features/chat.json @@ -97,6 +97,21 @@ "replyTo": "Reply to", "notFound": "Message not found" }, + "elicitation": { + "title": "{server} needs more information", + "submit": "Submit", + "cancel": "Cancel", + "decline": "Decline", + "done": "Done", + "otherInput": "If none fit, enter a custom reply", + "arrayPlaceholder": "Use commas or new lines for multiple values", + "requiredField": "Please fill in {field}", + "emptyReply": "Please provide at least one field", + "submitted": "Submitted. Waiting for AstrBot to continue", + "submitFailed": "Failed to submit the elicitation reply", + "accepted": "Accepted", + "booleanLabel": "Turn on for true, off for false" + }, "project": { "title": "Projects", "create": "Create Project", diff --git a/dashboard/src/i18n/locales/zh-CN/features/chat.json b/dashboard/src/i18n/locales/zh-CN/features/chat.json index 9d60db1767..74b91e75d8 100644 --- a/dashboard/src/i18n/locales/zh-CN/features/chat.json +++ b/dashboard/src/i18n/locales/zh-CN/features/chat.json @@ -97,6 +97,21 @@ "replyTo": "引用", "notFound": "无法定位消息" }, + "elicitation": { + "title": "{server} 需要更多信息", + "submit": "提交", + "cancel": "取消", + "decline": "拒绝", + "done": "已完成", + "otherInput": "如果都不合适，可手动输入", + "arrayPlaceholder": "可使用逗号或换行分隔多个值", + "requiredField": "请填写 {field}", + "emptyReply": "请至少填写一个字段", + "submitted": "已提交，等待 AstrBot 继续处理", + "submitFailed": "提交 elicitation 回复失败", + "accepted": "已确认", + "booleanLabel": "打开以表示 true，关闭表示 false" + }, "project": { "title": "项目", "create": "创建项目", diff --git a/docs/en/use/mcp.md b/docs/en/use/mcp.md index 1681d31325..cebb5cde13 100644 --- a/docs/en/use/mcp.md +++ b/docs/en/use/mcp.md @@ -96,6 +96,100 @@ Configure it in the AstrBot WebUI: That's it. +## MCP Client Sub-Capabilities + +AstrBot now supports enabling MCP client sub-capabilities per server. The first integrated sub-capability is `sampling`, which allows an MCP server to send `sampling/createMessage` requests during a tool call and reuse the current bot's chat provider, persona context, and session origin. + +AstrBot also supports per-server `elicitation`. When enabled, an MCP server can ask the current user for missing input or require an external confirmation step while the current MCP tool call is still in progress. + +AstrBot also supports per-server `roots`. When enabled, an MCP server can call `roots/list` to learn which local file roots AstrBot is explicitly exposing to it. + +You can add the following to an MCP server configuration: + +```json +{ + "url": "https://example.com/mcp", + "transport": "sse", + "client_capabilities": { + "elicitation": { + "enabled": true, + "timeout_seconds": 300 + }, + "sampling": { + "enabled": true + }, + "roots": { + "enabled": true, + "paths": ["data", "temp"] + } + } +} +``` + +Current limitations: + +- `elicitation` is disabled by default and is only advertised for servers that explicitly set `enabled: true`. +- `elicitation.timeout_seconds` bounds the total time AstrBot waits for the user's reply; timeouts are returned as `cancel`. +- `elicitation` is only available while the MCP server is actively serving the current bot tool call; requests outside the active interaction context are rejected. +- `elicitation` currently uses plain chat messages instead of a dedicated form UI. +- Form-mode elicitation currently supports flat top-level fields with simple types: `string`, `integer`, `number`, `boolean`, and `array[string]`. +- For single-field form elicitation, the user can reply with plain text. For multi-field form elicitation, AstrBot accepts JSON or `field: value` lines. +- URL-mode elicitation currently sends the URL and instructions to the user, then waits for a chat reply such as `done`, `decline`, or `cancel`. +- Enabling `elicitation` only affects the configured MCP server and does not change other MCP servers or standard chat flows. +- `sampling` is disabled by default and is only advertised for servers that explicitly set `enabled: true`. +- `roots` is disabled by default and is only advertised for servers that explicitly set `enabled: true`. +- `roots.paths` can contain built-in aliases such as `data`, `temp`, `config`, `skills`, `plugins`, `plugin_data`, `knowledge_base`, `backups`, and `root`, as well as absolute paths or paths relative to AstrBot root. +- If `roots.enabled` is `true` and `paths` is omitted, AstrBot currently exposes `data` and `temp` as the default safe roots. +- `sampling` is only available while the MCP server is actively serving the current bot tool call; requests outside the active interaction context are rejected. +- The initial implementation only returns text sampling results. +- Tool-assisted sampling and multimodal sampling inputs such as image or audio are not supported yet. +- Enabling `sampling` only affects the configured MCP server and does not change other MCP servers or standard chat flows. +- Enabling `roots` only affects the configured MCP server and does not change other MCP servers or standard chat flows. + +## Notes for stdio servers + +When using an MCP server over stdio, the server should reserve stdout for JSON-RPC protocol messages and write logs to stderr. + +AstrBot now tolerates blank lines and common launcher banners such as `npm run` output to reduce noise during local testing, but that behavior is only a compatibility fallback. The robust setup is still to keep stdout protocol-only. + +## MCP Resources Bridge + +If an MCP server advertises the `resources` capability during initialization, AstrBot now registers a small set of bridge tools so the bot can interact with those resources through the existing tool loop. + +The first iteration exposes these server-scoped tools: + +- `mcp__list_resources` +- `mcp__read_resource` +- `mcp__list_resource_templates` (only when the server supports listing resource templates) + +This lets the bot discover available resources and read a specific resource URI without changing the normal chat flow or provider integration. + +Current limitations: + +- The bridge is read-only and does not support `resources/subscribe` push updates yet. +- AstrBot does not auto-inject MCP resources into prompt context; the bot still needs to read them explicitly through tools. +- A single text resource is returned as text. +- A single image blob resource is returned as an image-style tool result. +- Multi-part resources, mixed results, and non-image binary blobs are summarized into text in the first iteration. + +## MCP Prompts Bridge + +If an MCP server advertises the `prompts` capability during initialization, AstrBot also registers a small set of bridge tools so the bot can discover and fetch MCP prompts through the existing tool loop. + +The first iteration exposes these server-scoped tools: + +- `mcp__list_prompts` +- `mcp__get_prompt` + +This lets the bot inspect available prompt templates and resolve a specific prompt by name with optional arguments, without changing the normal chat flow or provider integration. + +Current limitations: + +- AstrBot does not auto-inject MCP prompts into the active chat context; the bot still needs to fetch them explicitly through tools. +- `get_prompt` results are currently summarized into text, preserving descriptions, message roles, and text blocks. +- Non-text prompt blocks such as images, audio, and embedded resources are summarized into text in the first iteration instead of being converted into multimodal context. +- The bridge does not support `prompts/list_changed` push updates yet, and it does not use MCP completions to auto-complete prompt arguments. + Reference links: 1. Learn how to use MCP here: [Model Context Protocol](https://modelcontextprotocol.io/introduction) diff --git a/docs/zh/use/mcp.md b/docs/zh/use/mcp.md index 79e3757fda..929c8c8176 100644 --- a/docs/zh/use/mcp.md +++ b/docs/zh/use/mcp.md @@ -95,6 +95,100 @@ npx -v 即可。 +## MCP Client 子能力 + +AstrBot 现在支持按服务开启 MCP Client 子能力。当前首个接入的子能力是 `sampling`，它允许 MCP Server 在工具调用过程中向 AstrBot 发起 `sampling/createMessage` 请求，并复用当前 Bot 的聊天模型、人格上下文与会话来源。 + +目前也支持按服务开启 `elicitation`。启用后，MCP Server 可以在工具调用过程中向当前用户追问缺失信息，或要求用户完成一个外部确认步骤，再继续当前 MCP 交互。 + +目前也支持按服务开启 `roots`。启用后，MCP Server 可以通过 `roots/list` 请求 AstrBot 暴露给它的文件根目录列表。 + +您可以在对应 MCP Server 配置中加入： + +```json +{ + "url": "https://example.com/mcp", + "transport": "sse", + "client_capabilities": { + "elicitation": { + "enabled": true, + "timeout_seconds": 300 + }, + "sampling": { + "enabled": true + }, + "roots": { + "enabled": true, + "paths": ["data", "temp"] + } + } +} +``` + +当前版本的限制： + +- `elicitation` 默认关闭，只有显式配置 `enabled: true` 的服务器才会声明该能力。 +- `elicitation.timeout_seconds` 用于限制等待用户回复的总时长；超时后当前 elicitation 会以 `cancel` 结束。 +- `elicitation` 当前只在对应 MCP Server 正在为当前 Bot 执行工具调用时可用；脱离当前交互上下文的请求会被拒绝。 +- `elicitation` 当前通过聊天消息交互，不提供独立表单 UI。 +- `elicitation` 的 form 模式当前支持顶层简单字段：`string`、`integer`、`number`、`boolean`、`array[string]`。 +- `elicitation` 的 form 模式在单字段时可直接回复纯文本，多字段时可回复 JSON 或 `field: value` 形式。 +- `elicitation` 的 url 模式当前会把 URL 和提示文本发送给用户，并等待用户在聊天中回复 `done` / `decline` / `cancel` 之类的确认。 +- 开启 `elicitation` 的影响范围仅限对应的 MCP Server，不会改变其他 MCP Server 或普通聊天流程。 +- `sampling` 默认关闭，只有显式配置 `enabled: true` 的服务器才会声明该能力。 +- `roots` 默认关闭，只有显式配置 `enabled: true` 的服务器才会声明该能力。 +- `roots.paths` 可填写内置别名（如 `data`、`temp`、`config`、`skills`、`plugins`、`plugin_data`、`knowledge_base`、`backups`、`root`）、绝对路径，或相对 AstrBot 根目录的路径。 +- 当 `roots.enabled` 为 `true` 且未显式填写 `paths` 时，AstrBot 当前默认暴露 `data` 和 `temp`。 +- `sampling` 仅在该 MCP Server 正在为当前 Bot 执行工具调用时可用；脱离当前交互上下文的请求会被拒绝。 +- 当前版本仅支持文本采样结果。 +- 当前版本不支持带工具的 sampling，也不支持图片、音频等多模态 sampling 输入。 +- 开启 `sampling` 的影响范围仅限对应的 MCP Server，不会改变其他 MCP Server 或普通聊天流程。 +- 开启 `roots` 的影响范围仅限对应的 MCP Server，不会改变其他 MCP Server 或普通聊天流程。 + +## stdio 服务输出说明 + +使用 stdio 方式接入 MCP Server 时，服务器应当只通过标准输出（stdout）发送 JSON-RPC 协议消息，并将日志写入标准错误（stderr）。 + +AstrBot 当前会尽量忽略空行以及 `npm run` 一类启动器输出的非协议横幅，减少测试时的噪声；但这只是兼容处理，不建议依赖。更稳妥的做法仍然是让 MCP Server 或启动脚本保持 stdout 干净。 + +## MCP Resources 桥接 + +如果某个 MCP Server 在初始化时声明了 `resources` 能力，AstrBot 会自动为它注册一组桥接工具，让 Bot 通过现有的工具调用流程与资源交互。 + +当前首版会按服务注册这些工具： + +- `mcp__list_resources` +- `mcp__read_resource` +- `mcp__list_resource_templates`（仅在服务支持模板列表时出现） + +这意味着 Bot 可以先列出资源，再读取某个具体 URI 的内容，而不需要修改普通聊天流程或额外配置 Provider。 + +当前版本的限制： + +- 当前只做只读桥接，不支持 `resources/subscribe` 订阅与推送刷新。 +- 当前不会自动把 MCP 资源注入提示词上下文，仍然需要 Bot 通过工具主动读取。 +- 单个文本资源会直接作为文本结果返回。 +- 单个图片 Blob 资源会按图片结果处理。 +- 多段资源、混合类型资源以及非图片二进制资源，当前会被整理为文本摘要返回。 + +## MCP Prompts 桥接 + +如果某个 MCP Server 在初始化时声明了 `prompts` 能力，AstrBot 也会自动为它注册一组桥接工具，让 Bot 通过现有的工具调用流程查看和获取 MCP prompt。 + +当前首版会按服务注册这些工具： + +- `mcp__list_prompts` +- `mcp__get_prompt` + +这意味着 Bot 可以先列出可用 prompt，再按名称和参数获取某个 prompt 的展开结果，而不需要修改普通聊天流程或额外配置 Provider。 + +当前版本的限制： + +- 当前不会自动把 MCP prompt 注入对话上下文，仍然需要 Bot 通过工具主动获取。 +- `get_prompt` 的结果当前会整理为文本摘要，保留描述、消息角色和文本内容。 +- 图片、音频、嵌入资源等非纯文本 prompt block，当前会被总结为文本说明，而不是直接转成多模态上下文。 +- 当前不支持 `prompts/list_changed` 推送刷新，也不支持通过 MCP completion 自动补全 prompt 参数。 + 参考链接： 1. 在这里了解如何使用 MCP: [Model Context Protocol](https://modelcontextprotocol.io/introduction) diff --git a/tests/test_dashboard.py b/tests/test_dashboard.py index 6c575910a0..83e27cc231 100644 --- a/tests/test_dashboard.py +++ b/tests/test_dashboard.py @@ -5,6 +5,7 @@ import zipfile from datetime import datetime from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest import pytest_asyncio @@ -275,6 +276,298 @@ async def test_commands_api(app: Quart, authenticated_header: dict): assert isinstance(data["data"], list) +@pytest.mark.asyncio +async def test_mcp_servers_api_exposes_client_capabilities( + app: Quart, + authenticated_header: dict, + core_lifecycle_td: AstrBotCoreLifecycle, + monkeypatch, +): + tool_mgr = core_lifecycle_td.provider_manager.llm_tools + monkeypatch.setattr( + tool_mgr, + "load_mcp_config", + lambda: { + "mcpServers": { + "demo": { + "url": "https://example.com/mcp", + "transport": "sse", + "active": True, + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 180, + }, + "sampling": { + "enabled": True, + } + }, + } + } + }, + ) + + test_client = app.test_client() + response = await test_client.get( + "/api/tools/mcp/servers", + headers=authenticated_header, + ) + assert response.status_code == 200 + data = await response.get_json() + assert data["status"] == "ok" + assert data["data"][0]["client_capabilities"]["elicitation"]["enabled"] is True + assert data["data"][0]["client_capabilities"]["elicitation"]["timeout_seconds"] == 180 + assert data["data"][0]["client_capabilities"]["sampling"]["enabled"] is True + + +@pytest.mark.asyncio +async def test_mcp_servers_api_includes_resource_bridge_tools( + app: Quart, + authenticated_header: dict, + core_lifecycle_td: AstrBotCoreLifecycle, + monkeypatch, +): + tool_mgr = core_lifecycle_td.provider_manager.llm_tools + monkeypatch.setattr( + tool_mgr, + "load_mcp_config", + lambda: { + "mcpServers": { + "demo": { + "command": "node", + "args": ["stdio.js"], + "active": True, + } + } + }, + ) + tool_mgr._mcp_server_runtime["demo"] = SimpleNamespace( + client=SimpleNamespace( + tools=[SimpleNamespace(name="demo_tool")], + resource_bridge_tool_names=["mcp_demo_list_resources"], + server_errlogs=[], + ) + ) + + test_client = app.test_client() + try: + response = await test_client.get( + "/api/tools/mcp/servers", + headers=authenticated_header, + ) + assert response.status_code == 200 + data = await response.get_json() + assert data["status"] == "ok" + assert data["data"][0]["tools"] == [ + "demo_tool", + "mcp_demo_list_resources", + ] + finally: + tool_mgr._mcp_server_runtime.pop("demo", None) + + +@pytest.mark.asyncio +async def test_mcp_servers_api_includes_prompt_bridge_tools( + app: Quart, + authenticated_header: dict, + core_lifecycle_td: AstrBotCoreLifecycle, + monkeypatch, +): + tool_mgr = core_lifecycle_td.provider_manager.llm_tools + monkeypatch.setattr( + tool_mgr, + "load_mcp_config", + lambda: { + "mcpServers": { + "demo": { + "command": "node", + "args": ["stdio.js"], + "active": True, + } + } + }, + ) + tool_mgr._mcp_server_runtime["demo"] = SimpleNamespace( + client=SimpleNamespace( + tools=[SimpleNamespace(name="demo_tool")], + resource_bridge_tool_names=[], + prompt_bridge_tool_names=["mcp_demo_list_prompts", "mcp_demo_get_prompt"], + server_errlogs=[], + ) + ) + + test_client = app.test_client() + try: + response = await test_client.get( + "/api/tools/mcp/servers", + headers=authenticated_header, + ) + assert response.status_code == 200 + data = await response.get_json() + assert data["status"] == "ok" + assert data["data"][0]["tools"] == [ + "demo_tool", + "mcp_demo_list_prompts", + "mcp_demo_get_prompt", + ] + finally: + tool_mgr._mcp_server_runtime.pop("demo", None) + + +@pytest.mark.asyncio +async def test_add_mcp_server_persists_client_capabilities( + app: Quart, + authenticated_header: dict, + core_lifecycle_td: AstrBotCoreLifecycle, + tmp_path, + monkeypatch, +): + monkeypatch.setattr( + "astrbot.core.provider.func_tool_manager.get_astrbot_data_path", + lambda: str(tmp_path), + ) + tool_mgr = core_lifecycle_td.provider_manager.llm_tools + monkeypatch.setattr( + tool_mgr, + "test_mcp_server_connection", + AsyncMock(return_value=["demo_tool"]), + ) + monkeypatch.setattr( + tool_mgr, + "enable_mcp_server", + AsyncMock(return_value=None), + ) + + test_client = app.test_client() + response = await test_client.post( + "/api/tools/mcp/add", + json={ + "name": "demo", + "url": "https://example.com/mcp", + "transport": "sse", + "active": True, + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 240, + }, + "sampling": { + "enabled": True, + } + }, + }, + headers=authenticated_header, + ) + assert response.status_code == 200 + data = await response.get_json() + assert data["status"] == "ok" + + persisted = tool_mgr.load_mcp_config() + assert ( + persisted["mcpServers"]["demo"]["client_capabilities"]["elicitation"][ + "enabled" + ] + is True + ) + assert ( + persisted["mcpServers"]["demo"]["client_capabilities"]["elicitation"][ + "timeout_seconds" + ] + == 240 + ) + assert ( + persisted["mcpServers"]["demo"]["client_capabilities"]["sampling"]["enabled"] + is True + ) + + +@pytest.mark.asyncio +async def test_chat_respond_elicitation_resolves_pending_reply_and_persists_message( + app: Quart, + authenticated_header: dict, + core_lifecycle_td: AstrBotCoreLifecycle, + monkeypatch, +): + captured: dict[str, str] = {} + + def _fake_submit(umo: str, sender_id: str, reply_text: str, *, reply_outline=None): + captured["umo"] = umo + captured["sender_id"] = sender_id + captured["reply_text"] = reply_text + captured["reply_outline"] = reply_outline or "" + return True + + monkeypatch.setattr( + "astrbot.dashboard.routes.chat.submit_pending_mcp_elicitation_reply", + _fake_submit, + ) + + test_client = app.test_client() + new_session_response = await test_client.get( + "/api/chat/new_session", + headers=authenticated_header, + ) + session_data = await new_session_response.get_json() + session_id = session_data["data"]["session_id"] + + response = await test_client.post( + "/api/chat/respond_elicitation", + json={ + "session_id": session_id, + "reply_text": '{"topic":"MCP 最小实现"}', + "display_text": "topic: MCP 最小实现", + }, + headers=authenticated_header, + ) + + assert response.status_code == 200 + data = await response.get_json() + assert data["status"] == "ok" + expected_username = core_lifecycle_td.astrbot_config["dashboard"]["username"] + assert captured["sender_id"] == expected_username + assert captured["reply_text"] == '{"topic":"MCP 最小实现"}' + assert captured["reply_outline"] == "topic: MCP 最小实现" + assert captured["umo"].endswith(f"webchat!{expected_username}!{session_id}") + assert data["data"]["saved_message"]["content"]["message"] == [ + {"type": "plain", "text": "topic: MCP 最小实现"} + ] + + +@pytest.mark.asyncio +async def test_chat_respond_elicitation_rejects_when_no_pending_request( + app: Quart, + authenticated_header: dict, + monkeypatch, +): + monkeypatch.setattr( + "astrbot.dashboard.routes.chat.submit_pending_mcp_elicitation_reply", + lambda *_args, **_kwargs: False, + ) + + test_client = app.test_client() + new_session_response = await test_client.get( + "/api/chat/new_session", + headers=authenticated_header, + ) + session_data = await new_session_response.get_json() + session_id = session_data["data"]["session_id"] + + response = await test_client.post( + "/api/chat/respond_elicitation", + json={ + "session_id": session_id, + "reply_text": "cancel", + "display_text": "cancel", + }, + headers=authenticated_header, + ) + + assert response.status_code == 200 + data = await response.get_json() + assert data["status"] == "error" + assert "No pending MCP elicitation" in data["message"] + + @pytest.mark.asyncio async def test_check_update( app: Quart, diff --git a/tests/unit/test_mcp_prompt_bridge.py b/tests/unit/test_mcp_prompt_bridge.py new file mode 100644 index 0000000000..14a6c57894 --- /dev/null +++ b/tests/unit/test_mcp_prompt_bridge.py @@ -0,0 +1,212 @@ +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import mcp +import pytest + +from astrbot.core.agent.mcp_prompt_bridge import ( + MCPGetPromptTool, + MCPListPromptsTool, + build_mcp_prompt_tool_names, + shape_get_prompt_result, +) +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.provider.func_tool_manager import ( + FunctionToolManager, + _MCPServerRuntime, +) + + +class _FakePromptCapableMCPClient: + def __init__(self) -> None: + self.name: str | None = None + self.tools = [ + mcp.types.Tool( + name="draft_brief", + description="Draft a short brief.", + inputSchema={"type": "object", "properties": {}}, + ) + ] + self.resource_bridge_tool_names: list[str] = [] + self.prompt_bridge_tool_names: list[str] = [] + self.prompts: list[mcp.types.Prompt] = [] + self.server_errlogs: list[str] = [] + self.received_arguments: list[dict[str, str] | None] = [] + + @property + def supports_prompts(self) -> bool: + return True + + @property + def supports_resources(self) -> bool: + return False + + async def connect_to_server(self, config: dict, name: str) -> None: + _ = config + self.name = name + + async def list_tools_and_save(self) -> mcp.types.ListToolsResult: + return mcp.types.ListToolsResult(tools=self.tools) + + async def load_resource_capabilities(self) -> None: + self.resource_bridge_tool_names = [] + + async def load_prompt_capabilities(self) -> None: + self.prompt_bridge_tool_names = build_mcp_prompt_tool_names( + self.name or "server" + ) + + async def list_prompts_and_save( + self, + cursor: str | None = None, + ) -> mcp.types.ListPromptsResult: + _ = cursor + self.prompts = [ + mcp.types.Prompt( + name="draft_brief", + description="Draft a short brief from a topic.", + arguments=[ + mcp.types.PromptArgument( + name="topic", + description="Topic to summarize", + required=True, + ) + ], + ) + ] + return mcp.types.ListPromptsResult(prompts=self.prompts) + + async def get_prompt_with_reconnect( + self, + name: str, + arguments: dict[str, str] | None, + read_timeout_seconds, + ) -> mcp.types.GetPromptResult: + _ = read_timeout_seconds + self.received_arguments.append(arguments) + return mcp.types.GetPromptResult( + description=f"Prompt '{name}' resolved.", + messages=[ + mcp.types.PromptMessage( + role="user", + content=mcp.types.TextContent( + type="text", + text=f"Write a concise brief about {arguments.get('topic', 'the topic') if arguments else 'the topic'}.", + ), + ) + ], + ) + + async def cleanup(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_prompt_bridge_tools_are_registered_and_removed( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + "astrbot.core.provider.func_tool_manager.MCPClient", + _FakePromptCapableMCPClient, + ) + + tool_mgr = FunctionToolManager() + client = await tool_mgr._init_mcp_client( + "demo-server", + {"command": "node", "args": ["stdio.js"]}, + ) + + tool_names = {tool.name for tool in tool_mgr.func_list} + assert "draft_brief" in tool_names + assert "mcp_demo_server_list_prompts" in tool_names + assert "mcp_demo_server_get_prompt" in tool_names + + completed_task = asyncio.create_task(asyncio.sleep(0)) + await completed_task + tool_mgr._mcp_server_runtime["demo-server"] = _MCPServerRuntime( + name="demo-server", + client=client, + shutdown_event=asyncio.Event(), + lifecycle_task=completed_task, + ) + + await tool_mgr._terminate_mcp_client("demo-server") + + assert not any( + getattr(tool, "mcp_server_name", None) == "demo-server" + for tool in tool_mgr.func_list + ) + + +@pytest.mark.asyncio +async def test_prompt_listing_tool_returns_text_summary(): + client = _FakePromptCapableMCPClient() + tool = MCPListPromptsTool( + mcp_client=client, + mcp_server_name="demo-server", + ) + context = ContextWrapper(context=SimpleNamespace()) + + result = await tool.call(context) + + assert isinstance(result.content[0], mcp.types.TextContent) + assert "draft_brief" in result.content[0].text + assert "topic" in result.content[0].text + + +@pytest.mark.asyncio +async def test_get_prompt_tool_passes_arguments_and_returns_text_summary(): + client = _FakePromptCapableMCPClient() + tool = MCPGetPromptTool( + mcp_client=client, + mcp_server_name="demo-server", + ) + context = ContextWrapper(context=SimpleNamespace(), tool_call_timeout=30) + + result = await tool.call( + context, + name="draft_brief", + arguments={"topic": "MCP 最小实现"}, + ) + + assert client.received_arguments == [{"topic": "MCP 最小实现"}] + assert isinstance(result.content[0], mcp.types.TextContent) + assert "Prompt: draft_brief" in result.content[0].text + assert "MCP 最小实现" in result.content[0].text + + +def test_shape_get_prompt_result_summarizes_non_text_blocks(): + result = shape_get_prompt_result( + server_name="demo-server", + prompt_name="rich_prompt", + response=mcp.types.GetPromptResult( + description="Rich prompt response.", + messages=[ + mcp.types.PromptMessage( + role="assistant", + content=mcp.types.ImageContent( + type="image", + data="ZmFrZQ==", + mimeType="image/png", + ), + ), + mcp.types.PromptMessage( + role="user", + content=mcp.types.EmbeddedResource( + type="resource", + resource=mcp.types.TextResourceContents( + uri="memo://prompt/context", + mimeType="text/plain", + text="embedded context", + ), + ), + ), + ], + ), + ) + + assert "Image block returned." in result + assert "Embedded text resource returned." in result + assert "embedded context" in result diff --git a/tests/unit/test_mcp_resource_bridge.py b/tests/unit/test_mcp_resource_bridge.py new file mode 100644 index 0000000000..e605ffcf15 --- /dev/null +++ b/tests/unit/test_mcp_resource_bridge.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +import mcp +from astrbot.core.agent.mcp_resource_bridge import ( + MCPListResourceTemplatesTool, + MCPListResourcesTool, + MCPReadResourceTool, + build_mcp_resource_tool_names, + shape_read_resource_result, +) +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.provider.func_tool_manager import ( + FunctionToolManager, + _MCPServerRuntime, +) + + +class _FakeResourceCapableMCPClient: + def __init__(self) -> None: + self.name: str | None = None + self.tools = [ + mcp.types.Tool( + name="draft_brief", + description="Draft a short brief.", + inputSchema={"type": "object", "properties": {}}, + ) + ] + self.resource_templates_supported = True + self.resource_bridge_tool_names: list[str] = [] + self.prompt_bridge_tool_names: list[str] = [] + self.server_errlogs: list[str] = [] + + @property + def supports_resources(self) -> bool: + return True + + @property + def supports_prompts(self) -> bool: + return False + + async def connect_to_server(self, config: dict, name: str) -> None: + _ = config + self.name = name + + async def list_tools_and_save(self) -> mcp.types.ListToolsResult: + return mcp.types.ListToolsResult(tools=self.tools) + + async def load_resource_capabilities(self) -> None: + self.resource_bridge_tool_names = build_mcp_resource_tool_names( + self.name or "server", + include_templates=True, + ) + + async def load_prompt_capabilities(self) -> None: + self.prompt_bridge_tool_names = [] + + async def list_resources_and_save( + self, + cursor: str | None = None, + ) -> mcp.types.ListResourcesResult: + _ = cursor + return mcp.types.ListResourcesResult( + resources=[ + mcp.types.Resource( + name="team_notes", + uri="memo://team/notes", + description="Shared team notes", + mimeType="text/plain", + ) + ] + ) + + async def list_resource_templates_and_save( + self, + cursor: str | None = None, + ) -> mcp.types.ListResourceTemplatesResult: + _ = cursor + self.resource_templates_supported = True + return mcp.types.ListResourceTemplatesResult( + resourceTemplates=[ + mcp.types.ResourceTemplate( + name="note_by_id", + uriTemplate="memo://notes/{id}", + description="Read a note by id", + mimeType="text/plain", + ) + ] + ) + + async def read_resource_with_reconnect( + self, + uri: str, + read_timeout_seconds, + ) -> mcp.types.ReadResourceResult: + _ = read_timeout_seconds + return mcp.types.ReadResourceResult( + contents=[ + mcp.types.TextResourceContents( + uri=uri, + mimeType="text/plain", + text="hello from resource", + ) + ] + ) + + async def cleanup(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_resource_bridge_tools_are_registered_and_removed( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + "astrbot.core.provider.func_tool_manager.MCPClient", + _FakeResourceCapableMCPClient, + ) + + tool_mgr = FunctionToolManager() + client = await tool_mgr._init_mcp_client( + "demo-server", + {"command": "node", "args": ["stdio.js"]}, + ) + + tool_names = {tool.name for tool in tool_mgr.func_list} + assert "draft_brief" in tool_names + assert "mcp_demo_server_list_resources" in tool_names + assert "mcp_demo_server_read_resource" in tool_names + assert "mcp_demo_server_list_resource_templates" in tool_names + + completed_task = asyncio.create_task(asyncio.sleep(0)) + await completed_task + tool_mgr._mcp_server_runtime["demo-server"] = _MCPServerRuntime( + name="demo-server", + client=client, + shutdown_event=asyncio.Event(), + lifecycle_task=completed_task, + ) + + await tool_mgr._terminate_mcp_client("demo-server") + + assert not any( + getattr(tool, "mcp_server_name", None) == "demo-server" + for tool in tool_mgr.func_list + ) + + +@pytest.mark.asyncio +async def test_resource_listing_tools_return_text_summaries(): + client = _FakeResourceCapableMCPClient() + list_tool = MCPListResourcesTool( + mcp_client=client, + mcp_server_name="demo-server", + ) + template_tool = MCPListResourceTemplatesTool( + mcp_client=client, + mcp_server_name="demo-server", + ) + context = ContextWrapper(context=SimpleNamespace()) + + list_result = await list_tool.call(context) + template_result = await template_tool.call(context) + + assert isinstance(list_result.content[0], mcp.types.TextContent) + assert "memo://team/notes" in list_result.content[0].text + assert isinstance(template_result.content[0], mcp.types.TextContent) + assert "memo://notes/{id}" in template_result.content[0].text + + +@pytest.mark.asyncio +async def test_read_resource_tool_returns_text_resource_summary(): + client = _FakeResourceCapableMCPClient() + tool = MCPReadResourceTool( + mcp_client=client, + mcp_server_name="demo-server", + ) + context = ContextWrapper(context=SimpleNamespace(), tool_call_timeout=30) + + result = await tool.call(context, uri="memo://team/notes") + + assert isinstance(result.content[0], mcp.types.TextContent) + assert "hello from resource" in result.content[0].text + assert "memo://team/notes" in result.content[0].text + + +def test_shape_read_resource_result_returns_embedded_image_for_single_image_blob(): + response = mcp.types.ReadResourceResult( + contents=[ + mcp.types.BlobResourceContents( + uri="memo://images/cover", + mimeType="image/png", + blob="ZmFrZQ==", + ) + ] + ) + + result = shape_read_resource_result( + server_name="demo-server", + requested_uri="memo://images/cover", + response=response, + ) + + assert isinstance(result.content[0], mcp.types.EmbeddedResource) + assert result.content[0].resource.mimeType == "image/png" + + +def test_shape_read_resource_result_summarizes_multi_part_content(): + response = mcp.types.ReadResourceResult( + contents=[ + mcp.types.TextResourceContents( + uri="memo://notes/1", + mimeType="text/plain", + text="first part", + ), + mcp.types.BlobResourceContents( + uri="memo://notes/1.bin", + mimeType="application/octet-stream", + blob="Zm9v", + ), + ] + ) + + result = shape_read_resource_result( + server_name="demo-server", + requested_uri="memo://notes/1", + response=response, + ) + + assert isinstance(result.content[0], mcp.types.TextContent) + assert "Returned parts: 2" in result.content[0].text + assert "Binary blob returned" in result.content[0].text diff --git a/tests/unit/test_mcp_stdio_client.py b/tests/unit/test_mcp_stdio_client.py new file mode 100644 index 0000000000..796d4ef659 --- /dev/null +++ b/tests/unit/test_mcp_stdio_client.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +import pytest + +from astrbot.core.agent.mcp_client import MCPClient +from astrbot.core.agent.mcp_stdio_client import _should_ignore_stdout_line + + +class _DummyAsyncContext: + def __init__(self, value): + self._value = value + + async def __aenter__(self): + return self._value + + async def __aexit__(self, exc_type, exc, tb): + return False + + +class _RecordingClientSession: + constructor_calls: list[dict] = [] + + def __init__(self, *args, **kwargs): + self.__class__.constructor_calls.append( + { + "args": args, + "kwargs": kwargs, + } + ) + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def initialize(self): + return None + + +@pytest.mark.parametrize( + "line", + [ + "", + " ", + "> mcp-minimal-server@1.0.0 start:stdio", + "> node stdio.js", + "npm notice", + ], +) +def test_should_ignore_non_protocol_stdio_stdout_lines(line: str): + assert _should_ignore_stdout_line(line) is True + + +@pytest.mark.parametrize( + "line", + [ + '{"jsonrpc":"2.0","id":1,"result":{}}', + '{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info"}}', + "{not-valid-json-yet}", + ], +) +def test_should_keep_json_like_stdio_stdout_lines_for_protocol_parsing(line: str): + assert _should_ignore_stdout_line(line) is False + + +@pytest.mark.asyncio +async def test_mcp_client_stdio_path_uses_tolerant_transport( + monkeypatch: pytest.MonkeyPatch, +): + _RecordingClientSession.constructor_calls.clear() + transport_calls: list[dict] = [] + + def _fake_tolerant_stdio_client(server_params, errlog): + transport_calls.append( + { + "server_params": server_params, + "errlog": errlog, + } + ) + return _DummyAsyncContext(("read", "write")) + + monkeypatch.setattr( + "astrbot.core.agent.mcp_client.tolerant_stdio_client", + _fake_tolerant_stdio_client, + ) + monkeypatch.setattr( + "astrbot.core.agent.mcp_client.mcp.ClientSession", + _RecordingClientSession, + ) + + client = MCPClient() + await client.connect_to_server( + { + "command": "node", + "args": ["stdio.js"], + }, + "demo", + ) + + assert len(transport_calls) == 1 + server_params = transport_calls[0]["server_params"] + assert server_params.command == "node" + assert server_params.args == ["stdio.js"] + assert _RecordingClientSession.constructor_calls + + await client.cleanup() diff --git a/tests/unit/test_mcp_subcapability_bridge.py b/tests/unit/test_mcp_subcapability_bridge.py new file mode 100644 index 0000000000..5ae573ed47 --- /dev/null +++ b/tests/unit/test_mcp_subcapability_bridge.py @@ -0,0 +1,1295 @@ +# ruff: noqa: ASYNC110 + +from __future__ import annotations + +import asyncio +import json +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import mcp +import pytest + +from astrbot.core.agent.mcp_client import MCPClient +from astrbot.core.agent.mcp_elicitation_registry import ( + pending_mcp_elicitation, + submit_pending_mcp_elicitation_reply, + try_capture_pending_mcp_elicitation, +) +from astrbot.core.agent.mcp_subcapability_bridge import MCPClientSubCapabilityBridge +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.provider.entities import LLMResponse +from astrbot.core.provider.func_tool_manager import FunctionToolManager + + +class _DummyAsyncContext: + def __init__(self, value): + self._value = value + + async def __aenter__(self): + return self._value + + async def __aexit__(self, exc_type, exc, tb): + return False + + +class _RecordingClientSession: + constructor_calls: list[dict] = [] + + def __init__(self, *args, **kwargs): + self.__class__.constructor_calls.append( + { + "args": args, + "kwargs": kwargs, + } + ) + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def initialize(self): + return None + + def get_server_capabilities(self): + return None + + +class _SamplingAwareSession: + def __init__(self, bridge, params): + self.bridge = bridge + self.params = params + self.calls: list[dict] = [] + + async def call_tool(self, *, name, arguments, read_timeout_seconds): + self.calls.append( + { + "name": name, + "arguments": arguments, + "read_timeout_seconds": read_timeout_seconds, + } + ) + sampling_result = await self.bridge.handle_sampling(None, self.params) + assert isinstance(sampling_result, mcp.types.CreateMessageResult) + return mcp.types.CallToolResult( + content=[ + mcp.types.TextContent( + type="text", + text=sampling_result.content.text, + ) + ] + ) + + +class _DummyProvider: + def __init__(self, model: str = "gpt-4o-mini"): + self._model = model + + def get_model(self) -> str: + return self._model + + def meta(self): + return SimpleNamespace(model=self._model, id="provider-1") + + +class _DummyPluginContext: + def __init__( + self, + *, + completion_text: str, + release_event: asyncio.Event | None = None, + entered_event: asyncio.Event | None = None, + ): + self.provider = _DummyProvider() + self.completion_text = completion_text + self.release_event = release_event + self.entered_event = entered_event + self.requests: list[dict] = [] + + async def get_current_chat_provider_id(self, umo: str) -> str: + assert umo + return "provider-1" + + def get_using_provider(self, umo: str | None = None): + return self.provider + + async def llm_generate(self, **kwargs): + self.requests.append(kwargs) + if self.entered_event is not None: + self.entered_event.set() + if self.release_event is not None: + await self.release_event.wait() + return LLMResponse(role="assistant", completion_text=self.completion_text) + + +class _DummyEvent: + def __init__( + self, + *, + umo: str = "test:umo", + sender_id: str = "user-1", + message_text: str = "", + outline: str = "", + platform_name: str = "test", + ) -> None: + self.unified_msg_origin = umo + self._sender_id = sender_id + self._message_text = message_text + self._outline = outline or message_text + self._platform_name = platform_name + self.sent_messages: list[str] = [] + self.sent_payloads: list[dict] = [] + + def get_sender_id(self) -> str: + return self._sender_id + + def get_message_str(self) -> str: + return self._message_text + + def get_message_outline(self) -> str: + return self._outline + + def get_platform_name(self) -> str: + return self._platform_name + + async def send(self, message_chain) -> None: + self.sent_messages.append( + message_chain.get_plain_text(with_other_comps_mark=True) + ) + if ( + getattr(message_chain, "type", None) == "elicitation" + and message_chain.chain + ): + first = message_chain.chain[0] + payload = getattr(first, "data", None) + if isinstance(payload, dict): + self.sent_payloads.append(payload) + + +def _build_run_context( + plugin_context: _DummyPluginContext, + *, + umo: str = "test:umo", + event: _DummyEvent | None = None, +): + event = event or _DummyEvent(umo=umo) + agent_context = SimpleNamespace( + context=plugin_context, + event=event, + ) + return ContextWrapper(context=agent_context) + + +def _build_sampling_params(*, tools=None, content=None): + if content is None: + content = mcp.types.TextContent(type="text", text="hello from server") + return mcp.types.CreateMessageRequestParams( + messages=[ + mcp.types.SamplingMessage( + role="user", + content=content, + ) + ], + maxTokens=64, + tools=tools, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("enabled", "expected_sampling_enabled"), + [ + pytest.param(False, False, id="sampling-disabled"), + pytest.param(True, True, id="sampling-enabled"), + ], +) +async def test_mcp_client_capability_advertisement_depends_on_config( + monkeypatch: pytest.MonkeyPatch, + enabled: bool, + expected_sampling_enabled: bool, +): + _RecordingClientSession.constructor_calls.clear() + + async def _fake_quick_test(_config): + return True, "" + + monkeypatch.setattr( + "astrbot.core.agent.mcp_client._quick_test_mcp_connection", + _fake_quick_test, + ) + monkeypatch.setattr( + "astrbot.core.agent.mcp_client.sse_client", + lambda **_kwargs: _DummyAsyncContext(("read", "write")), + ) + monkeypatch.setattr( + "astrbot.core.agent.mcp_client.mcp.ClientSession", + _RecordingClientSession, + ) + + client = MCPClient() + await client.connect_to_server( + { + "url": "https://example.com/mcp", + "transport": "sse", + "client_capabilities": { + "sampling": { + "enabled": enabled, + } + }, + }, + "demo", + ) + + kwargs = _RecordingClientSession.constructor_calls[0]["kwargs"] + sampling_callback = kwargs.get("sampling_callback") + sampling_capabilities = kwargs.get("sampling_capabilities") + + if expected_sampling_enabled: + assert callable(sampling_callback) + assert sampling_capabilities is not None + else: + assert sampling_callback is None + assert sampling_capabilities is None + + await client.cleanup() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("enabled", "expected_roots_enabled"), + [ + pytest.param(False, False, id="roots-disabled"), + pytest.param(True, True, id="roots-enabled"), + ], +) +async def test_mcp_client_roots_capability_advertisement_depends_on_config( + monkeypatch: pytest.MonkeyPatch, + enabled: bool, + expected_roots_enabled: bool, +): + _RecordingClientSession.constructor_calls.clear() + + async def _fake_quick_test(_config): + return True, "" + + monkeypatch.setattr( + "astrbot.core.agent.mcp_client._quick_test_mcp_connection", + _fake_quick_test, + ) + monkeypatch.setattr( + "astrbot.core.agent.mcp_client.sse_client", + lambda **_kwargs: _DummyAsyncContext(("read", "write")), + ) + monkeypatch.setattr( + "astrbot.core.agent.mcp_client.mcp.ClientSession", + _RecordingClientSession, + ) + + client = MCPClient() + await client.connect_to_server( + { + "url": "https://example.com/mcp", + "transport": "sse", + "client_capabilities": { + "roots": { + "enabled": enabled, + } + }, + }, + "demo", + ) + + kwargs = _RecordingClientSession.constructor_calls[0]["kwargs"] + roots_callback = kwargs.get("list_roots_callback") + + if expected_roots_enabled: + assert callable(roots_callback) + else: + assert roots_callback is None + + await client.cleanup() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("enabled", "expected_elicitation_enabled"), + [ + pytest.param(False, False, id="elicitation-disabled"), + pytest.param(True, True, id="elicitation-enabled"), + ], +) +async def test_mcp_client_elicitation_capability_advertisement_depends_on_config( + monkeypatch: pytest.MonkeyPatch, + enabled: bool, + expected_elicitation_enabled: bool, +): + _RecordingClientSession.constructor_calls.clear() + + async def _fake_quick_test(_config): + return True, "" + + monkeypatch.setattr( + "astrbot.core.agent.mcp_client._quick_test_mcp_connection", + _fake_quick_test, + ) + monkeypatch.setattr( + "astrbot.core.agent.mcp_client.sse_client", + lambda **_kwargs: _DummyAsyncContext(("read", "write")), + ) + monkeypatch.setattr( + "astrbot.core.agent.mcp_client.mcp.ClientSession", + _RecordingClientSession, + ) + + client = MCPClient() + await client.connect_to_server( + { + "url": "https://example.com/mcp", + "transport": "sse", + "client_capabilities": { + "elicitation": { + "enabled": enabled, + "timeout_seconds": 120, + } + }, + }, + "demo", + ) + + kwargs = _RecordingClientSession.constructor_calls[0]["kwargs"] + elicitation_callback = kwargs.get("elicitation_callback") + + if expected_elicitation_enabled: + assert callable(elicitation_callback) + else: + assert elicitation_callback is None + + await client.cleanup() + + +def test_load_and_save_mcp_config_normalizes_client_capabilities( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + "astrbot.core.provider.func_tool_manager.get_astrbot_data_path", + lambda: str(tmp_path), + ) + tool_mgr = FunctionToolManager() + + config = { + "mcpServers": { + "disabled-default": { + "command": "uv", + "args": ["run", "demo.py"], + }, + "enabled-sampling": { + "command": "uv", + "args": ["run", "demo.py"], + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 180, + }, + "sampling": { + "enabled": True, + }, + "roots": { + "enabled": True, + "paths": ["data", "temp"], + }, + }, + }, + } + } + + assert tool_mgr.save_mcp_config(config) is True + + saved_raw = json.loads((tmp_path / "mcp_server.json").read_text(encoding="utf-8")) + assert ( + saved_raw["mcpServers"]["disabled-default"]["client_capabilities"][ + "elicitation" + ]["enabled"] + is False + ) + assert ( + saved_raw["mcpServers"]["disabled-default"]["client_capabilities"][ + "elicitation" + ]["timeout_seconds"] + == 300 + ) + assert ( + saved_raw["mcpServers"]["disabled-default"]["client_capabilities"]["sampling"][ + "enabled" + ] + is False + ) + assert ( + saved_raw["mcpServers"]["disabled-default"]["client_capabilities"]["roots"][ + "enabled" + ] + is False + ) + assert ( + saved_raw["mcpServers"]["disabled-default"]["client_capabilities"]["roots"][ + "paths" + ] + == [] + ) + assert ( + saved_raw["mcpServers"]["enabled-sampling"]["client_capabilities"][ + "elicitation" + ]["enabled"] + is True + ) + assert ( + saved_raw["mcpServers"]["enabled-sampling"]["client_capabilities"][ + "elicitation" + ]["timeout_seconds"] + == 180 + ) + assert ( + saved_raw["mcpServers"]["enabled-sampling"]["client_capabilities"]["sampling"][ + "enabled" + ] + is True + ) + assert ( + saved_raw["mcpServers"]["enabled-sampling"]["client_capabilities"]["roots"][ + "enabled" + ] + is True + ) + assert saved_raw["mcpServers"]["enabled-sampling"]["client_capabilities"]["roots"][ + "paths" + ] == ["data", "temp"] + + loaded = tool_mgr.load_mcp_config() + assert ( + loaded["mcpServers"]["disabled-default"]["client_capabilities"]["elicitation"][ + "enabled" + ] + is False + ) + assert ( + loaded["mcpServers"]["disabled-default"]["client_capabilities"]["elicitation"][ + "timeout_seconds" + ] + == 300 + ) + assert ( + loaded["mcpServers"]["disabled-default"]["client_capabilities"]["sampling"][ + "enabled" + ] + is False + ) + assert ( + loaded["mcpServers"]["disabled-default"]["client_capabilities"]["roots"][ + "enabled" + ] + is False + ) + assert ( + loaded["mcpServers"]["enabled-sampling"]["client_capabilities"]["elicitation"][ + "enabled" + ] + is True + ) + assert ( + loaded["mcpServers"]["enabled-sampling"]["client_capabilities"]["elicitation"][ + "timeout_seconds" + ] + == 180 + ) + assert ( + loaded["mcpServers"]["enabled-sampling"]["client_capabilities"]["sampling"][ + "enabled" + ] + is True + ) + assert ( + loaded["mcpServers"]["enabled-sampling"]["client_capabilities"]["roots"][ + "enabled" + ] + is True + ) + assert loaded["mcpServers"]["enabled-sampling"]["client_capabilities"]["roots"][ + "paths" + ] == ["data", "temp"] + + +@pytest.mark.asyncio +async def test_roots_request_returns_default_roots_for_enabled_server( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +): + data_dir = tmp_path / "data" + temp_dir = data_dir / "temp" + data_dir.mkdir() + temp_dir.mkdir() + + monkeypatch.setattr( + "astrbot.core.agent.mcp_subcapability_bridge.get_astrbot_data_path", + lambda: str(data_dir), + ) + monkeypatch.setattr( + "astrbot.core.agent.mcp_subcapability_bridge.get_astrbot_temp_path", + lambda: str(temp_dir), + ) + + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "roots": { + "enabled": True, + } + } + } + ) + + result = await bridge.handle_list_roots(None) + + assert isinstance(result, mcp.types.ListRootsResult) + assert [root.name for root in result.roots] == ["data", "temp"] + assert str(result.roots[0].uri) == data_dir.resolve().as_uri() + assert str(result.roots[1].uri) == temp_dir.resolve().as_uri() + + +@pytest.mark.asyncio +async def test_roots_request_uses_explicit_paths_and_skips_missing_entries( + tmp_path, + monkeypatch: pytest.MonkeyPatch, +): + root_dir = tmp_path / "astrbot-root" + root_dir.mkdir() + explicit_dir = tmp_path / "explicit" + explicit_dir.mkdir() + nested_dir = root_dir / "workspace" + nested_dir.mkdir() + + monkeypatch.setattr( + "astrbot.core.agent.mcp_subcapability_bridge.get_astrbot_root", + lambda: str(root_dir), + ) + + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "roots": { + "enabled": True, + "paths": [ + str(explicit_dir), + "workspace", + "missing-dir", + ], + } + } + } + ) + + result = await bridge.handle_list_roots(None) + + assert isinstance(result, mcp.types.ListRootsResult) + assert [ + Path(str(root.uri).removeprefix("file:///")).name for root in result.roots + ] == [ + "explicit", + "workspace", + ] + + +@pytest.mark.asyncio +async def test_roots_request_is_rejected_when_disabled(): + bridge = MCPClientSubCapabilityBridge("demo") + + result = await bridge.handle_list_roots(None) + + assert isinstance(result, mcp.types.ErrorData) + assert result.code == mcp.types.INVALID_REQUEST + assert "Roots are not enabled" in result.message + + +@pytest.mark.asyncio +async def test_sampling_request_uses_bound_astrbot_context(): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "sampling": { + "enabled": True, + } + } + } + ) + plugin_context = _DummyPluginContext(completion_text="reply from astrbot") + run_context = _build_run_context(plugin_context) + params = _build_sampling_params() + + async with bridge.interactive_call(run_context): + result = await bridge.handle_sampling(None, params) + + assert isinstance(result, mcp.types.CreateMessageResult) + assert result.content.text == "reply from astrbot" + assert result.model == "gpt-4o-mini" + assert plugin_context.requests[0]["contexts"] == [ + { + "role": "user", + "content": "hello from server", + } + ] + assert plugin_context.requests[0]["max_tokens"] == 64 + + +@pytest.mark.asyncio +async def test_sampling_request_without_bound_context_is_rejected(): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "sampling": { + "enabled": True, + } + } + } + ) + + result = await bridge.handle_sampling(None, _build_sampling_params()) + + assert isinstance(result, mcp.types.ErrorData) + assert result.code == mcp.types.INVALID_REQUEST + assert "active AstrBot MCP interaction" in result.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "params", + [ + pytest.param( + _build_sampling_params( + tools=[ + mcp.types.Tool( + name="demo_tool", + description="demo", + inputSchema={"type": "object", "properties": {}}, + ) + ] + ), + id="tool-assisted-sampling", + ), + pytest.param( + _build_sampling_params( + content=mcp.types.ImageContent( + type="image", + data="ZmFrZQ==", + mimeType="image/png", + ) + ), + id="image-input", + ), + ], +) +async def test_sampling_request_rejects_unsupported_inputs(params): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "sampling": { + "enabled": True, + } + } + } + ) + plugin_context = _DummyPluginContext(completion_text="reply") + run_context = _build_run_context(plugin_context) + + async with bridge.interactive_call(run_context): + result = await bridge.handle_sampling(None, params) + + assert isinstance(result, mcp.types.ErrorData) + assert result.code == mcp.types.INVALID_REQUEST + + +@pytest.mark.asyncio +async def test_sampling_enabled_interactive_calls_are_serialized(): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "sampling": { + "enabled": True, + } + } + } + ) + + first_entered = asyncio.Event() + release_first = asyncio.Event() + first_plugin_context = _DummyPluginContext( + completion_text="first reply", + release_event=release_first, + entered_event=first_entered, + ) + second_plugin_context = _DummyPluginContext(completion_text="second reply") + order: list[str] = [] + params = _build_sampling_params() + + async def _first_call(): + async with bridge.interactive_call(_build_run_context(first_plugin_context)): + order.append("enter-1") + result = await bridge.handle_sampling(None, params) + order.append("exit-1") + return result + + async def _second_call(): + await first_entered.wait() + async with bridge.interactive_call(_build_run_context(second_plugin_context)): + order.append("enter-2") + result = await bridge.handle_sampling(None, params) + order.append("exit-2") + return result + + first_task = asyncio.create_task(_first_call()) + await first_entered.wait() + second_task = asyncio.create_task(_second_call()) + + await asyncio.sleep(0) + assert order == ["enter-1"] + + release_first.set() + first_result = await first_task + second_result = await second_task + + assert isinstance(first_result, mcp.types.CreateMessageResult) + assert isinstance(second_result, mcp.types.CreateMessageResult) + assert first_result.content.text == "first reply" + assert second_result.content.text == "second reply" + assert order == ["enter-1", "exit-1", "enter-2", "exit-2"] + + +@pytest.mark.asyncio +async def test_mcp_client_call_tool_with_reconnect_preserves_sampling_runtime_context(): + client = MCPClient() + client.subcapability_bridge.set_server_name("demo") + client.subcapability_bridge.configure_from_server_config( + { + "client_capabilities": { + "sampling": { + "enabled": True, + } + } + } + ) + client.session = _SamplingAwareSession( + client.subcapability_bridge, + _build_sampling_params(), + ) + + plugin_context = _DummyPluginContext(completion_text="reply from astrbot") + run_context = _build_run_context(plugin_context) + + result = await client.call_tool_with_reconnect( + tool_name="draft-brief", + arguments={"topic": "MCP 最小实现"}, + read_timeout_seconds=timedelta(seconds=60), + run_context=run_context, + ) + + assert isinstance(result, mcp.types.CallToolResult) + assert result.content[0].text == "reply from astrbot" + assert len(client.session.calls) == 1 + assert client.session.calls[0]["name"] == "draft-brief" + assert client.session.calls[0]["arguments"] == {"topic": "MCP 最小实现"} + assert plugin_context.requests[0]["contexts"] == [ + { + "role": "user", + "content": "hello from server", + } + ] + + +@pytest.mark.asyncio +async def test_pending_mcp_elicitation_captures_only_matching_sender(): + async with pending_mcp_elicitation("umo:1", "user-1") as future: + wrong_sender_event = _DummyEvent( + umo="umo:1", + sender_id="user-2", + message_text="ignored", + ) + assert try_capture_pending_mcp_elicitation(wrong_sender_event) is False + assert future.done() is False + + matching_event = _DummyEvent( + umo="umo:1", + sender_id="user-1", + message_text="accepted", + ) + assert try_capture_pending_mcp_elicitation(matching_event) is True + resolved_event = await future + assert resolved_event.message_text == "accepted" + assert resolved_event.message_outline == "accepted" + + +@pytest.mark.asyncio +async def test_pending_mcp_elicitation_accepts_direct_submission(): + async with pending_mcp_elicitation("umo:1", "user-1") as future: + assert ( + submit_pending_mcp_elicitation_reply( + "umo:1", + "user-1", + '{"topic":"MCP 最小实现"}', + reply_outline="topic: MCP 最小实现", + ) + is True + ) + resolved_reply = await future + + assert resolved_reply.message_text == '{"topic":"MCP 最小实现"}' + assert resolved_reply.message_outline == "topic: MCP 最小实现" + + +@pytest.mark.asyncio +async def test_elicitation_form_request_uses_next_matching_reply(): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 30, + } + } + } + ) + event = _DummyEvent(umo="test:umo", sender_id="user-1") + run_context = _build_run_context( + _DummyPluginContext(completion_text="unused"), + event=event, + ) + params = mcp.types.ElicitRequestFormParams( + message="Please provide the topic.", + requestedSchema={ + "type": "object", + "properties": { + "topic": { + "type": "string", + "description": "Brief topic name", + } + }, + "required": ["topic"], + }, + ) + + async def _resolve_reply(): + while not event.sent_messages: + await asyncio.sleep(0) + assert "topic" in event.sent_messages[0] + reply_event = _DummyEvent( + umo="test:umo", + sender_id="user-1", + message_text="MCP 最小实现", + ) + while not try_capture_pending_mcp_elicitation(reply_event): + await asyncio.sleep(0) + + async with bridge.interactive_call(run_context): + result_task = asyncio.create_task(bridge.handle_elicitation(None, params)) + await _resolve_reply() + result = await result_task + + assert isinstance(result, mcp.types.ElicitResult) + assert result.action == "accept" + assert result.content == {"topic": "MCP 最小实现"} + + +@pytest.mark.asyncio +async def test_elicitation_form_request_reprompts_after_invalid_reply(): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 30, + } + } + } + ) + event = _DummyEvent(umo="test:umo", sender_id="user-1") + run_context = _build_run_context( + _DummyPluginContext(completion_text="unused"), + event=event, + ) + params = mcp.types.ElicitRequestFormParams( + message="How many sections do you need?", + requestedSchema={ + "type": "object", + "properties": { + "count": { + "type": "integer", + } + }, + "required": ["count"], + }, + ) + + async def _resolve_replies(): + while len(event.sent_messages) < 1: + await asyncio.sleep(0) + first_reply = _DummyEvent( + umo="test:umo", + sender_id="user-1", + message_text="not-a-number", + ) + while not try_capture_pending_mcp_elicitation(first_reply): + await asyncio.sleep(0) + + while len(event.sent_messages) < 2: + await asyncio.sleep(0) + assert "could not use that reply" in event.sent_messages[1].lower() + second_reply = _DummyEvent( + umo="test:umo", + sender_id="user-1", + message_text="2", + ) + while not try_capture_pending_mcp_elicitation(second_reply): + await asyncio.sleep(0) + + async with bridge.interactive_call(run_context): + result_task = asyncio.create_task(bridge.handle_elicitation(None, params)) + await _resolve_replies() + result = await result_task + + assert isinstance(result, mcp.types.ElicitResult) + assert result.action == "accept" + assert result.content == {"count": 2} + + +@pytest.mark.asyncio +async def test_elicitation_form_request_accepts_natural_language_key_value_patterns(): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 30, + } + } + } + ) + event = _DummyEvent(umo="test:umo", sender_id="user-1") + run_context = _build_run_context( + _DummyPluginContext(completion_text="unused"), + event=event, + ) + params = mcp.types.ElicitRequestFormParams( + message="Collect plan details.", + requestedSchema={ + "type": "object", + "properties": { + "topic": {"type": "string"}, + "audience": {"type": "string"}, + }, + "required": ["topic", "audience"], + }, + ) + + async def _resolve_reply(): + while not event.sent_messages: + await asyncio.sleep(0) + reply_event = _DummyEvent( + umo="test:umo", + sender_id="user-1", + message_text="topic 是 MCP 最小实现，audience 为 新手", + ) + while not try_capture_pending_mcp_elicitation(reply_event): + await asyncio.sleep(0) + + async with bridge.interactive_call(run_context): + result_task = asyncio.create_task(bridge.handle_elicitation(None, params)) + await _resolve_reply() + result = await result_task + + assert isinstance(result, mcp.types.ElicitResult) + assert result.action == "accept" + assert result.content == { + "topic": "MCP 最小实现", + "audience": "新手", + } + + +@pytest.mark.asyncio +async def test_elicitation_form_request_uses_llm_fallback_for_bot_reply(): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 30, + } + } + } + ) + plugin_context = _DummyPluginContext( + completion_text='```json\n{"topic":"MCP 最小实现","audience":"新手"}\n```' + ) + event = _DummyEvent(umo="test:umo", sender_id="user-1") + run_context = _build_run_context( + plugin_context, + event=event, + ) + params = mcp.types.ElicitRequestFormParams( + message="Collect plan details.", + requestedSchema={ + "type": "object", + "properties": { + "topic": {"type": "string"}, + "audience": {"type": "string"}, + }, + "required": ["topic", "audience"], + }, + ) + + async def _resolve_reply(): + while not event.sent_messages: + await asyncio.sleep(0) + reply_event = _DummyEvent( + umo="test:umo", + sender_id="user-1", + message_text="面向新手，写一个关于 MCP 最小实现 的简要说明。", + ) + while not try_capture_pending_mcp_elicitation(reply_event): + await asyncio.sleep(0) + + async with bridge.interactive_call(run_context): + result_task = asyncio.create_task(bridge.handle_elicitation(None, params)) + await _resolve_reply() + result = await result_task + + assert isinstance(result, mcp.types.ElicitResult) + assert result.action == "accept" + assert result.content == { + "topic": "MCP 最小实现", + "audience": "新手", + } + assert len(plugin_context.requests) == 1 + assert plugin_context.requests[0]["chat_provider_id"] == "provider-1" + assert "Return only a JSON object." in plugin_context.requests[0]["system_prompt"] + + +@pytest.mark.asyncio +async def test_elicitation_form_request_retries_when_llm_fallback_returns_invalid_json(): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 30, + } + } + } + ) + plugin_context = _DummyPluginContext(completion_text="not-json") + event = _DummyEvent(umo="test:umo", sender_id="user-1") + run_context = _build_run_context( + plugin_context, + event=event, + ) + params = mcp.types.ElicitRequestFormParams( + message="Collect plan details.", + requestedSchema={ + "type": "object", + "properties": { + "topic": {"type": "string"}, + "audience": {"type": "string"}, + }, + "required": ["topic", "audience"], + }, + ) + + async def _resolve_replies(): + while not event.sent_messages: + await asyncio.sleep(0) + first_reply = _DummyEvent( + umo="test:umo", + sender_id="user-1", + message_text="帮我给新手准备一个关于 MCP 最小实现 的说明。", + ) + while not try_capture_pending_mcp_elicitation(first_reply): + await asyncio.sleep(0) + + while len(event.sent_messages) < 2: + await asyncio.sleep(0) + assert "could not use that reply" in event.sent_messages[1].lower() + + second_reply = _DummyEvent( + umo="test:umo", + sender_id="user-1", + message_text='{"topic":"MCP 最小实现","audience":"新手"}', + ) + while not try_capture_pending_mcp_elicitation(second_reply): + await asyncio.sleep(0) + + async with bridge.interactive_call(run_context): + result_task = asyncio.create_task(bridge.handle_elicitation(None, params)) + await _resolve_replies() + result = await result_task + + assert isinstance(result, mcp.types.ElicitResult) + assert result.action == "accept" + assert result.content == { + "topic": "MCP 最小实现", + "audience": "新手", + } + assert len(plugin_context.requests) == 1 + + +@pytest.mark.asyncio +async def test_webchat_elicitation_message_uses_structured_payload(): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 30, + } + } + } + ) + event = _DummyEvent( + umo="test:umo", + sender_id="user-1", + platform_name="webchat", + ) + run_context = _build_run_context( + _DummyPluginContext(completion_text="unused"), + event=event, + ) + params = mcp.types.ElicitRequestFormParams( + message="Choose a tone.", + requestedSchema={ + "type": "object", + "properties": { + "tone": { + "type": "string", + "enum": ["formal", "casual"], + } + }, + "required": ["tone"], + }, + ) + + async def _resolve_reply(): + while not event.sent_payloads: + await asyncio.sleep(0) + assert event.sent_payloads[0]["fields"][0]["enum"] == ["formal", "casual"] + submit_pending_mcp_elicitation_reply( + "test:umo", + "user-1", + "formal", + reply_outline="tone: formal", + ) + + async with bridge.interactive_call(run_context): + result_task = asyncio.create_task(bridge.handle_elicitation(None, params)) + await _resolve_reply() + result = await result_task + + assert isinstance(result, mcp.types.ElicitResult) + assert result.action == "accept" + assert result.content == {"tone": "formal"} + + +@pytest.mark.asyncio +async def test_elicitation_url_request_waits_for_confirmation(): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 30, + } + } + } + ) + event = _DummyEvent(umo="test:umo", sender_id="user-1") + run_context = _build_run_context( + _DummyPluginContext(completion_text="unused"), + event=event, + ) + params = mcp.types.ElicitRequestURLParams( + message="Authorize the test server.", + url="https://example.com/auth", + elicitationId="elic-1", + ) + + async def _resolve_reply(): + while not event.sent_messages: + await asyncio.sleep(0) + assert "https://example.com/auth" in event.sent_messages[0] + reply_event = _DummyEvent( + umo="test:umo", + sender_id="user-1", + message_text="done", + ) + while not try_capture_pending_mcp_elicitation(reply_event): + await asyncio.sleep(0) + + async with bridge.interactive_call(run_context): + result_task = asyncio.create_task(bridge.handle_elicitation(None, params)) + await _resolve_reply() + result = await result_task + + assert isinstance(result, mcp.types.ElicitResult) + assert result.action == "accept" + assert result.content is None + + +@pytest.mark.asyncio +async def test_elicitation_request_times_out_to_cancel(monkeypatch: pytest.MonkeyPatch): + bridge = MCPClientSubCapabilityBridge("demo") + bridge.configure_from_server_config( + { + "client_capabilities": { + "elicitation": { + "enabled": True, + "timeout_seconds": 30, + } + } + } + ) + event = _DummyEvent(umo="test:umo", sender_id="user-1") + run_context = _build_run_context( + _DummyPluginContext(completion_text="unused"), + event=event, + ) + params = mcp.types.ElicitRequestFormParams( + message="Please provide topic.", + requestedSchema={ + "type": "object", + "properties": { + "topic": {"type": "string"}, + }, + }, + ) + + async def _fake_wait_for_reply(*, event, sender_id, deadline): + del event, sender_id, deadline + return None + + monkeypatch.setattr(bridge, "_wait_for_elicitation_reply", _fake_wait_for_reply) + + async with bridge.interactive_call(run_context): + result = await bridge.handle_elicitation(None, params) + + assert isinstance(result, mcp.types.ElicitResult) + assert result.action == "cancel" From 87755aef5b6f3106b3015417a2f91b3a05a00ebb Mon Sep 17 00:00:00 2001 From: Last-emo-boy Date: Mon, 16 Mar 2026 20:23:53 +0800 Subject: [PATCH 2/5] merge: upstream/dev with MCP elicitation support --- astrbot/core/agent/mcp_client.py | 21 ++ .../core/agent/mcp_elicitation_registry.py | 57 +++++- astrbot/core/agent/mcp_resource_bridge.py | 13 +- .../core/agent/mcp_subcapability_bridge.py | 181 ++++++++++++++++-- 4 files changed, 250 insertions(+), 22 deletions(-) diff --git a/astrbot/core/agent/mcp_client.py b/astrbot/core/agent/mcp_client.py index ffe8e397b6..d77407fe61 100644 --- a/astrbot/core/agent/mcp_client.py +++ b/astrbot/core/agent/mcp_client.py @@ -23,6 +23,7 @@ MCPClientSubCapabilityBridge, normalize_mcp_server_config, ) +from .mcp_elicitation_registry import cleanup_elicitation_periodically from .run_context import TContext from .tool import FunctionTool @@ -143,6 +144,18 @@ def __init__(self) -> None: self._reconnect_lock = asyncio.Lock() # Lock for thread-safe reconnection self._reconnecting: bool = False # For logging and debugging self.subcapability_bridge = MCPClientSubCapabilityBridge[Any]() + + # Elicitation cleanup task + self._elicitation_cleanup_task: asyncio.Task[None] | None = None + self._start_elicitation_cleanup() + + def _start_elicitation_cleanup(self) -> None: + """启动后台 elicitation 清理任务。""" + self._elicitation_cleanup_task = asyncio.create_task( + cleanup_elicitation_periodically(interval=60), + name="mcp-elicitation-cleanup", + ) + logger.debug("已启动 MCP elicitation 后台清理任务") async def connect_to_server(self, mcp_server_config: dict, name: str) -> None: """Connect to MCP server @@ -640,6 +653,14 @@ async def cleanup(self) -> None: except Exception as e: logger.debug(f"Error closing current exit stack: {e}") + # Cancel elicitation cleanup task + if self._elicitation_cleanup_task: + self._elicitation_cleanup_task.cancel() + try: + await self._elicitation_cleanup_task + except asyncio.CancelledError: + logger.debug("Elicitation cleanup task cancelled") + # Don't close old exit stacks as they may be in different task contexts # They will be garbage collected naturally # Just clear the list to release references diff --git a/astrbot/core/agent/mcp_elicitation_registry.py b/astrbot/core/agent/mcp_elicitation_registry.py index 767c77bf73..6be61e1f77 100644 --- a/astrbot/core/agent/mcp_elicitation_registry.py +++ b/astrbot/core/agent/mcp_elicitation_registry.py @@ -1,11 +1,14 @@ from __future__ import annotations import asyncio +import time from collections.abc import AsyncIterator from contextlib import asynccontextmanager -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING +from astrbot import logger + if TYPE_CHECKING: from astrbot.core.platform.astr_message_event import AstrMessageEvent @@ -21,10 +24,18 @@ class PendingMCPElicitation: umo: str sender_id: str future: asyncio.Future[MCPElicitationReply] + created_at: float = field(default_factory=time.time) _PENDING_MCP_ELICITATIONS: dict[str, PendingMCPElicitation] = {} +# Elicitation 清理指标 +_cleanup_metrics = { + "total_cleaned": 0, + "last_cleanup_time": 0.0, + "last_cleanup_duration": 0.0, +} + @asynccontextmanager async def pending_mcp_elicitation( @@ -101,3 +112,47 @@ def submit_pending_mcp_elicitation_reply( ) ) return True + + +def cleanup_expired_elicitations() -> int: + """清理已完成的 elicitation 条目。 + + 返回清理的条目数量。 + """ + start_time = time.time() + expired = [ + umo for umo, p in _PENDING_MCP_ELICITATIONS.items() + if p.future.done() + ] + + for umo in expired: + _PENDING_MCP_ELICITATIONS.pop(umo, None) + + # 记录指标 + _cleanup_metrics["total_cleaned"] += len(expired) + _cleanup_metrics["last_cleanup_time"] = start_time + _cleanup_metrics["last_cleanup_duration"] = time.time() - start_time + + if expired: + logger.debug(f"清理了 {len(expired)} 个已完成的 elicitation 条目") + + return len(expired) + + +def get_cleanup_metrics() -> dict: + """获取清理指标。""" + return _cleanup_metrics.copy() + + +async def cleanup_elicitation_periodically(interval: int = 60) -> None: + """后台定期清理 elicitation 条目。 + + Args: + interval: 清理间隔秒数，默认 60 秒 + """ + while True: + await asyncio.sleep(interval) + try: + cleanup_expired_elicitations() + except Exception as e: + logger.error(f"Elicitation 清理任务出错：{e}", exc_info=True) diff --git a/astrbot/core/agent/mcp_resource_bridge.py b/astrbot/core/agent/mcp_resource_bridge.py index 3663e44a60..cf2cb3592a 100644 --- a/astrbot/core/agent/mcp_resource_bridge.py +++ b/astrbot/core/agent/mcp_resource_bridge.py @@ -356,6 +356,17 @@ def _format_resource_template_metadata( return lines -def _sanitize_tool_name_fragment(name: str) -> str: +def _sanitize_tool_name_fragment(name: str, server_config_hash: str | None = None) -> str: + """Sanitize server name to be used in tool names. + + Args: + name: Server name to sanitize + server_config_hash: Optional hash to append for uniqueness + + Returns: + Sanitized server name suitable for tool names + """ sanitized = re.sub(r"[^a-zA-Z0-9]+", "_", name).strip("_").lower() + if server_config_hash: + sanitized += f"_{server_config_hash[:8]}" return sanitized or "server" diff --git a/astrbot/core/agent/mcp_subcapability_bridge.py b/astrbot/core/agent/mcp_subcapability_bridge.py index 9b92b584fa..93fb5157fa 100644 --- a/astrbot/core/agent/mcp_subcapability_bridge.py +++ b/astrbot/core/agent/mcp_subcapability_bridge.py @@ -3,7 +3,10 @@ import asyncio import copy import json +import logging import re +import time +from collections import defaultdict from contextlib import asynccontextmanager from dataclasses import dataclass, field from pathlib import Path @@ -25,6 +28,13 @@ get_astrbot_skills_path, get_astrbot_temp_path, ) +from tenacity import ( + before_sleep_log, + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) if TYPE_CHECKING: import mcp @@ -94,6 +104,22 @@ class UnsupportedElicitationRequestError(ValueError): """Raised when an elicitation request cannot be safely mapped.""" +class MCPElicitationError(Exception): + """Base exception for elicitation failures.""" + + +class ElicitationParseError(MCPElicitationError): + """用户输入解析失败。""" + + +class ElicitationTimeoutError(MCPElicitationError): + """elicitation 超时。""" + + +class ElicitationValidationError(MCPElicitationError): + """schema 验证失败。""" + + @dataclass(slots=True) class MCPElicitationCapabilityConfig: enabled: bool = False @@ -236,7 +262,9 @@ class MCPClientSubCapabilityBridge(Generic[TContext]): def __init__(self, server_name: str | None = None) -> None: self._server_name = server_name or "" self._capabilities = MCPClientCapabilitiesConfig.from_server_config({}) - self._interaction_lock = asyncio.Lock() + # Per-UMO locks for better concurrency + self._interaction_locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock) + self._lock_last_used: dict[str, float] = {} self._active_run_context: ContextWrapper[TContext] | None = None def configure_from_server_config(self, server_config: dict[str, Any]) -> None: @@ -259,6 +287,35 @@ def elicitation_enabled(self) -> bool: @property def elicitation_timeout_seconds(self) -> int: return self._capabilities.elicitation.timeout_seconds + + def _compute_elicitation_timeout(self, properties: dict[str, dict[str, Any]]) -> int: + """根据表单复杂度动态计算超时时间。 + + 公式： + - 基础时间：60 秒 + - 普通字段：每个 30 秒 + - Enum 字段：每个 15 秒（用户只需选择） + + 如果用户配置了显式超时，则使用配置值。 + """ + # 如果用户配置了显式超时，优先使用配置 + configured = self._capabilities.elicitation.timeout_seconds + if configured != DEFAULT_MCP_ELICITATION_TIMEOUT_SECONDS: + return configured + + base = 60 # 基础 1 分钟 + per_field = 30 # 每个字段 30 秒 + per_enum_field = 15 # enum 字段减少时间 + + enum_count = sum( + 1 for f in properties.values() + if f.get("enum") and isinstance(f.get("enum"), list) + ) + field_count = len(properties) + + timeout = base + (field_count - enum_count) * per_field + enum_count * per_enum_field + # 限制最小 60 秒，最大 600 秒 + return max(60, min(600, timeout)) @property def roots_enabled(self) -> bool: @@ -301,17 +358,42 @@ async def handle_list_roots( def clear_runtime_state(self) -> None: self._active_run_context = None + + def _cleanup_unused_locks(self, max_age_seconds: int = 300) -> None: + """清理超过指定时间未使用的锁（LRU 清理）。""" + now = time.time() + expired = [ + umo for umo, last_used in self._lock_last_used.items() + if now - last_used > max_age_seconds + ] + for umo in expired: + self._interaction_locks.pop(umo, None) + self._lock_last_used.pop(umo, None) + if expired: + logger.debug(f"清理了 {len(expired)} 个未使用的 per-umo 锁") @asynccontextmanager async def interactive_call( self, run_context: ContextWrapper[TContext] | None, + umo: str | None = None, ): if not (self.sampling_enabled or self.elicitation_enabled): yield return - async with self._interaction_lock: + # 使用 per-umo 锁，如果没有提供 umo 则使用全局锁（向后兼容） + if umo: + lock = self._interaction_locks[umo] + self._lock_last_used[umo] = time.time() + # 定期清理未使用的锁 + if len(self._interaction_locks) > 100: + self._cleanup_unused_locks() + else: + # 向后兼容：如果没有 umo，创建一个临时锁 + lock = asyncio.Lock() + + async with lock: self._active_run_context = run_context try: yield @@ -575,7 +657,9 @@ async def _execute_form_elicitation( import mcp properties = self._get_elicitation_properties(params.requestedSchema) - deadline = asyncio.get_running_loop().time() + self.elicitation_timeout_seconds + # 使用动态超时计算 + timeout_seconds = self._compute_elicitation_timeout(properties) + deadline = asyncio.get_running_loop().time() + timeout_seconds await self._send_elicitation_message( event, self._build_form_elicitation_prompt(params, properties), @@ -624,6 +708,12 @@ async def _execute_form_elicitation( content=content, ) + @retry( + stop=stop_after_attempt(2), + wait=wait_exponential(multiplier=0.5, min=0.5, max=2), + before_sleep=before_sleep_log(logger, logging.DEBUG), + reraise=True, + ) async def _try_llm_form_reply_fallback( self, *, @@ -927,17 +1017,17 @@ def _parse_form_elicitation_reply( normalized_reply = reply_text.strip() if not normalized_reply: - raise UnsupportedElicitationRequestError("The reply is empty.") + raise ElicitationParseError("The reply is empty.") if normalized_reply.startswith("{"): try: payload = json.loads(normalized_reply) except json.JSONDecodeError as exc: - raise UnsupportedElicitationRequestError( + raise ElicitationParseError( "The JSON reply could not be parsed." ) from exc if not isinstance(payload, dict): - raise UnsupportedElicitationRequestError( + raise ElicitationParseError( "The JSON reply must be an object." ) elif len(properties) == 1: @@ -951,7 +1041,7 @@ def _parse_form_elicitation_reply( requested_schema=requested_schema, ) if not payload: - raise UnsupportedElicitationRequestError( + raise ElicitationParseError( "Please reply with JSON, natural language, or `field: value` lines." ) @@ -1007,7 +1097,7 @@ def _coerce_form_payload( field_name for field_name in required_fields if field_name not in coerced ] if missing_required: - raise UnsupportedElicitationRequestError( + raise ElicitationValidationError( "Missing required field(s): " + ", ".join(missing_required) ) return coerced @@ -1027,24 +1117,24 @@ def _coerce_form_value( value = str(raw_value).strip() elif field_type == "integer": if isinstance(raw_value, bool): - raise UnsupportedElicitationRequestError( + raise ElicitationValidationError( f"Field `{field_name}` must be an integer." ) try: value = int(str(raw_value).strip()) except (TypeError, ValueError) as exc: - raise UnsupportedElicitationRequestError( + raise ElicitationValidationError( f"Field `{field_name}` must be an integer." ) from exc elif field_type == "number": if isinstance(raw_value, bool): - raise UnsupportedElicitationRequestError( + raise ElicitationValidationError( f"Field `{field_name}` must be a number." ) try: value = float(str(raw_value).strip()) except (TypeError, ValueError) as exc: - raise UnsupportedElicitationRequestError( + raise ElicitationValidationError( f"Field `{field_name}` must be a number." ) from exc elif field_type == "boolean": @@ -1052,13 +1142,13 @@ def _coerce_form_value( elif field_type == "array": value = self._coerce_string_array_value(field_name, raw_value) else: - raise UnsupportedElicitationRequestError( + raise ElicitationValidationError( f"Field `{field_name}` uses unsupported type `{field_type}`." ) enum_values = schema.get("enum") if isinstance(enum_values, list) and value not in enum_values: - raise UnsupportedElicitationRequestError( + raise ElicitationValidationError( f"Field `{field_name}` must be one of: {', '.join(map(str, enum_values))}." ) return value @@ -1075,7 +1165,7 @@ def _coerce_boolean_value(field_name: str, raw_value: Any) -> bool: return True if normalized in falsy: return False - raise UnsupportedElicitationRequestError( + raise ElicitationValidationError( f"Field `{field_name}` must be a boolean." ) @@ -1094,7 +1184,7 @@ def _coerce_string_array_value(field_name: str, raw_value: Any) -> list[str]: if part.strip() ] if not parts: - raise UnsupportedElicitationRequestError( + raise ElicitationValidationError( f"Field `{field_name}` must be a string array." ) return parts @@ -1173,7 +1263,7 @@ def _get_elicitation_properties( ) -> dict[str, dict[str, Any]]: properties = requested_schema.get("properties", {}) if not isinstance(properties, dict): - raise UnsupportedElicitationRequestError( + raise ElicitationParseError( "Form-mode elicitation requires a top-level properties object." ) normalized_properties: dict[str, dict[str, Any]] = {} @@ -1198,7 +1288,7 @@ def _get_elicitation_field_type(field_schema: dict[str, Any]) -> str: for candidate in field_type: if candidate in {"string", "integer", "number", "boolean", "array"}: return candidate - raise UnsupportedElicitationRequestError( + raise ElicitationParseError( "Unsupported multi-type elicitation field." ) if not isinstance(field_type, str): @@ -1291,12 +1381,63 @@ def _resolve_root_path_entry(self, entry: str) -> tuple[str, Path] | None: path = candidate_path.resolve() display_name = path.name or normalized_entry - if not path.exists(): + # 安全检查 1: 符号链接检查 + if path.is_symlink(): + logger.warning( + "Skipping symlinked MCP root path for server %s: %s (symlinks are not allowed for security)", + self._server_name, + path, + ) + return None + + # 安全检查 2: 目录验证 + if not path.is_dir(): + logger.warning( + "Skipping non-directory MCP root path for server %s: %s (must be a directory)", + self._server_name, + path, + ) + return None + + # 安全检查 3: 白名单验证 + if not self._is_root_path_in_allowlist(path): logger.warning( - "Skipping missing MCP root path for server %s: %s", + "Skipping MCP root path for server %s: %s (not in allowlist)", self._server_name, path, ) return None return display_name, path + + def _is_root_path_in_allowlist(self, path: Path) -> bool: + """检查路径是否在允许的白名单内。 + + 白名单包括： + 1. AstrBot 标准目录（data, temp, config 等） + 2. 用户显式配置的其他目录 + """ + # 获取所有允许的根路径解析器 + alias_resolvers = get_root_path_alias_resolvers() + allowed_paths = set() + + # 添加所有别名解析的路径 + for resolver_func in alias_resolvers.values(): + try: + allowed_path = Path(resolver_func()).resolve() + allowed_paths.add(allowed_path) + except Exception: + continue + + # 检查路径是否是允许路径的子目录或本身 + for allowed_path in allowed_paths: + try: + # 使用 relative_to 检查路径关系 + path.relative_to(allowed_path) + return True + except ValueError: + # 不是子目录，继续检查 + continue + + # 检查是否完全匹配 + return path in allowed_paths From 19e154040618d0983e318b47b39d0c0c9ce4c3d0 Mon Sep 17 00:00:00 2001 From: w33d Date: Tue, 17 Mar 2026 12:17:05 +0800 Subject: [PATCH 3/5] fix(mcp): repair type annotations and pre-existing unit test failures Fixes type annotation issues (excessive use of Any) and three pre-existing failing tests in test_mcp_subcapability_bridge.py. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Type Annotation Fixes - Replace Any with precise types in mcp_subcapability_bridge.py - Change `_request_context: Any` → `_request_context: None` (unused params) - Change `event: Any` → `event: TContext` (elicitation methods) - Change `plugin_context: Any` → `plugin_context: ContextWrapper[TContext]` - Change `raw_value: Any` → `raw_value: str | int | float | bool | list | None` - Add SupportsEvent Protocol for type-safe event constraints - Add JsonValue type alias for JSON schema values - Fix MCPClient to use Generic[TContext] properly - Fix subcapability_bridge instantiation: `MCPClientSubCapabilityBridge[TContext]()` ## Unit Test Fixes - Fix roots allowlist logic to support user-configured absolute paths - Add _user_configured_root_paths set to track explicit user paths - Update _resolve_root_path_entry() to record absolute paths - Update _is_root_path_in_allowlist() to check temporary allowlist - Fix interactive_call lock logic for proper serialization - Auto-extract umo from run_context.event.unified_msg_origin when not provided - Ensures per-umo locking works correctly for concurrent requests - Fix elicitation exception handling to prevent test timeouts - Extend except clause to catch ElicitationValidationError and ElicitationParseError - These exceptions were not being caught, causing tests to hang --- astrbot/core/agent/mcp_client.py | 17 +-- .../core/agent/mcp_elicitation_registry.py | 17 +-- astrbot/core/agent/mcp_prompt_bridge.py | 2 +- .../core/agent/mcp_subcapability_bridge.py | 144 ++++++++++++------ 4 files changed, 111 insertions(+), 69 deletions(-) diff --git a/astrbot/core/agent/mcp_client.py b/astrbot/core/agent/mcp_client.py index d77407fe61..7303238b62 100644 --- a/astrbot/core/agent/mcp_client.py +++ b/astrbot/core/agent/mcp_client.py @@ -2,7 +2,7 @@ import logging from contextlib import AsyncExitStack from datetime import timedelta -from typing import Any, Generic +from typing import Generic from tenacity import ( before_sleep_log, @@ -14,17 +14,16 @@ from astrbot import logger from astrbot.core.agent.mcp_prompt_bridge import build_mcp_prompt_tool_names -from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.agent.run_context import ContextWrapper, TContext from astrbot.core.utils.log_pipe import LogPipe +from .mcp_elicitation_registry import cleanup_elicitation_periodically from .mcp_resource_bridge import build_mcp_resource_tool_names from .mcp_stdio_client import tolerant_stdio_client from .mcp_subcapability_bridge import ( MCPClientSubCapabilityBridge, normalize_mcp_server_config, ) -from .mcp_elicitation_registry import cleanup_elicitation_periodically -from .run_context import TContext from .tool import FunctionTool try: @@ -118,7 +117,7 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: return False, f"{e!s}" -class MCPClient: +class MCPClient(Generic[TContext]): def __init__(self) -> None: # Initialize session and client objects self.session: mcp.ClientSession | None = None @@ -143,12 +142,12 @@ def __init__(self) -> None: self._server_capabilities: mcp.types.ServerCapabilities | None = None self._reconnect_lock = asyncio.Lock() # Lock for thread-safe reconnection self._reconnecting: bool = False # For logging and debugging - self.subcapability_bridge = MCPClientSubCapabilityBridge[Any]() - + self.subcapability_bridge = MCPClientSubCapabilityBridge[TContext]() + # Elicitation cleanup task self._elicitation_cleanup_task: asyncio.Task[None] | None = None self._start_elicitation_cleanup() - + def _start_elicitation_cleanup(self) -> None: """启动后台 elicitation 清理任务。""" self._elicitation_cleanup_task = asyncio.create_task( @@ -526,7 +525,7 @@ async def call_tool_with_reconnect( tool_name: str, arguments: dict, read_timeout_seconds: timedelta, - run_context: ContextWrapper[Any] | None = None, + run_context: ContextWrapper[TContext] | None = None, ) -> mcp.types.CallToolResult: """Call MCP tool with automatic reconnection on failure, max 2 retries. diff --git a/astrbot/core/agent/mcp_elicitation_registry.py b/astrbot/core/agent/mcp_elicitation_registry.py index 6be61e1f77..f9f8ed1b2b 100644 --- a/astrbot/core/agent/mcp_elicitation_registry.py +++ b/astrbot/core/agent/mcp_elicitation_registry.py @@ -116,26 +116,23 @@ def submit_pending_mcp_elicitation_reply( def cleanup_expired_elicitations() -> int: """清理已完成的 elicitation 条目。 - + 返回清理的条目数量。 """ start_time = time.time() - expired = [ - umo for umo, p in _PENDING_MCP_ELICITATIONS.items() - if p.future.done() - ] - + expired = [umo for umo, p in _PENDING_MCP_ELICITATIONS.items() if p.future.done()] + for umo in expired: _PENDING_MCP_ELICITATIONS.pop(umo, None) - + # 记录指标 _cleanup_metrics["total_cleaned"] += len(expired) _cleanup_metrics["last_cleanup_time"] = start_time _cleanup_metrics["last_cleanup_duration"] = time.time() - start_time - + if expired: logger.debug(f"清理了 {len(expired)} 个已完成的 elicitation 条目") - + return len(expired) @@ -146,7 +143,7 @@ def get_cleanup_metrics() -> dict: async def cleanup_elicitation_periodically(interval: int = 60) -> None: """后台定期清理 elicitation 条目。 - + Args: interval: 清理间隔秒数，默认 60 秒 """ diff --git a/astrbot/core/agent/mcp_prompt_bridge.py b/astrbot/core/agent/mcp_prompt_bridge.py index b086accd5f..167d2b5d84 100644 --- a/astrbot/core/agent/mcp_prompt_bridge.py +++ b/astrbot/core/agent/mcp_prompt_bridge.py @@ -260,7 +260,7 @@ def _format_prompt_message_content( def _normalize_prompt_arguments( - raw_arguments: Any, + raw_arguments: str | dict[str, Any] | None, ) -> dict[str, str] | None: if raw_arguments is None: return None diff --git a/astrbot/core/agent/mcp_subcapability_bridge.py b/astrbot/core/agent/mcp_subcapability_bridge.py index 93fb5157fa..8f2db874ce 100644 --- a/astrbot/core/agent/mcp_subcapability_bridge.py +++ b/astrbot/core/agent/mcp_subcapability_bridge.py @@ -10,7 +10,14 @@ from contextlib import asynccontextmanager from dataclasses import dataclass, field from pathlib import Path -from typing import TYPE_CHECKING, Any, Generic +from typing import TYPE_CHECKING, Any, Generic, Protocol + +from tenacity import ( + before_sleep_log, + retry, + stop_after_attempt, + wait_exponential, +) from astrbot import logger from astrbot.core.agent.mcp_elicitation_registry import pending_mcp_elicitation @@ -28,13 +35,6 @@ get_astrbot_skills_path, get_astrbot_temp_path, ) -from tenacity import ( - before_sleep_log, - retry, - retry_if_exception_type, - stop_after_attempt, - wait_exponential, -) if TYPE_CHECKING: import mcp @@ -108,6 +108,23 @@ class MCPElicitationError(Exception): """Base exception for elicitation failures.""" +# Type definitions for improved type safety + + +class SupportsEvent(Protocol): + """Protocol for event objects that can receive MCP elicitation messages.""" + + unified_msg_origin: str | None + + async def send(self, message: MessageChain) -> None: + """Send a message to the user.""" + ... + + +# JSON Schema value types for MCP elicitation form fields +JsonValue = str | int | float | bool | list[str] | None + + class ElicitationParseError(MCPElicitationError): """用户输入解析失败。""" @@ -266,6 +283,8 @@ def __init__(self, server_name: str | None = None) -> None: self._interaction_locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock) self._lock_last_used: dict[str, float] = {} self._active_run_context: ContextWrapper[TContext] | None = None + # Temporary allowlist for user-configured root paths + self._user_configured_root_paths: set[Path] = set() def configure_from_server_config(self, server_config: dict[str, Any]) -> None: self._capabilities = MCPClientCapabilitiesConfig.from_server_config( @@ -287,33 +306,38 @@ def elicitation_enabled(self) -> bool: @property def elicitation_timeout_seconds(self) -> int: return self._capabilities.elicitation.timeout_seconds - - def _compute_elicitation_timeout(self, properties: dict[str, dict[str, Any]]) -> int: + + def _compute_elicitation_timeout( + self, properties: dict[str, dict[str, Any]] + ) -> int: """根据表单复杂度动态计算超时时间。 - + 公式： - 基础时间：60 秒 - 普通字段：每个 30 秒 - Enum 字段：每个 15 秒（用户只需选择） - + 如果用户配置了显式超时，则使用配置值。 """ # 如果用户配置了显式超时，优先使用配置 configured = self._capabilities.elicitation.timeout_seconds if configured != DEFAULT_MCP_ELICITATION_TIMEOUT_SECONDS: return configured - + base = 60 # 基础 1 分钟 per_field = 30 # 每个字段 30 秒 per_enum_field = 15 # enum 字段减少时间 - + enum_count = sum( - 1 for f in properties.values() + 1 + for f in properties.values() if f.get("enum") and isinstance(f.get("enum"), list) ) field_count = len(properties) - - timeout = base + (field_count - enum_count) * per_field + enum_count * per_enum_field + + timeout = ( + base + (field_count - enum_count) * per_field + enum_count * per_enum_field + ) # 限制最小 60 秒，最大 600 秒 return max(60, min(600, timeout)) @@ -331,7 +355,7 @@ def get_sampling_capabilities(self) -> mcp.types.SamplingCapability | None: async def handle_list_roots( self, - _request_context: Any, + _request_context: None, ) -> mcp.types.ListRootsResult | mcp.types.ErrorData: import mcp @@ -358,12 +382,13 @@ async def handle_list_roots( def clear_runtime_state(self) -> None: self._active_run_context = None - + def _cleanup_unused_locks(self, max_age_seconds: int = 300) -> None: """清理超过指定时间未使用的锁（LRU 清理）。""" now = time.time() expired = [ - umo for umo, last_used in self._lock_last_used.items() + umo + for umo, last_used in self._lock_last_used.items() if now - last_used > max_age_seconds ] for umo in expired: @@ -382,6 +407,18 @@ async def interactive_call( yield return + # 自动从 run_context 提取 umo（如果未显式提供） + if umo is None and run_context is not None: + event = getattr(run_context.context, "event", None) + if event is not None: + umo = getattr(event, "unified_msg_origin", None) + if umo: + logger.debug( + "Auto-extracted umo from run_context for server %s: %s", + self._server_name, + umo, + ) + # 使用 per-umo 锁，如果没有提供 umo 则使用全局锁（向后兼容） if umo: lock = self._interaction_locks[umo] @@ -392,7 +429,7 @@ async def interactive_call( else: # 向后兼容：如果没有 umo，创建一个临时锁 lock = asyncio.Lock() - + async with lock: self._active_run_context = run_context try: @@ -402,7 +439,7 @@ async def interactive_call( async def handle_sampling( self, - _request_context: Any, + _request_context: None, params: mcp.types.CreateMessageRequestParams, ) -> ( mcp.types.CreateMessageResult @@ -449,7 +486,7 @@ async def handle_sampling( async def handle_elicitation( self, - _request_context: Any, + _request_context: None, params: mcp.types.ElicitRequestParams, ) -> mcp.types.ElicitResult | mcp.types.ErrorData: import mcp @@ -649,8 +686,8 @@ def _sampling_message_to_text(message: mcp.types.SamplingMessage) -> str: async def _execute_form_elicitation( self, - plugin_context: Any, - event: Any, + plugin_context: ContextWrapper[TContext], + event: TContext, sender_id: str, params: mcp.types.ElicitRequestFormParams, ) -> mcp.types.ElicitResult: @@ -684,7 +721,11 @@ async def _execute_form_elicitation( requested_schema=params.requestedSchema, reply_text=reply_text, ) - except UnsupportedElicitationRequestError as exc: + except ( + UnsupportedElicitationRequestError, + ElicitationValidationError, + ElicitationParseError, + ) as exc: content = await self._try_llm_form_reply_fallback( plugin_context=plugin_context, event=event, @@ -717,8 +758,8 @@ async def _execute_form_elicitation( async def _try_llm_form_reply_fallback( self, *, - plugin_context: Any, - event: Any, + plugin_context: ContextWrapper[TContext], + event: TContext, params: mcp.types.ElicitRequestFormParams, reply_text: str, direct_parse_error: UnsupportedElicitationRequestError, @@ -798,7 +839,7 @@ async def _try_llm_form_reply_fallback( async def _execute_url_elicitation( self, - event: Any, + event: TContext, sender_id: str, params: mcp.types.ElicitRequestURLParams, ) -> mcp.types.ElicitResult: @@ -831,7 +872,7 @@ async def _execute_url_elicitation( async def _send_elicitation_message( self, - event: Any, + event: TContext, message: str, *, payload: dict[str, Any] | None = None, @@ -857,7 +898,7 @@ async def _send_elicitation_message( async def _wait_for_elicitation_reply( self, *, - event: Any, + event: TContext, sender_id: str, deadline: float, ) -> str | None: @@ -1027,9 +1068,7 @@ def _parse_form_elicitation_reply( "The JSON reply could not be parsed." ) from exc if not isinstance(payload, dict): - raise ElicitationParseError( - "The JSON reply must be an object." - ) + raise ElicitationParseError("The JSON reply must be an object.") elif len(properties) == 1: field_name = next(iter(properties)) payload = {field_name: normalized_reply} @@ -1106,7 +1145,7 @@ def _coerce_form_value( self, *, field_name: str, - raw_value: Any, + raw_value: str | int | float | bool | list | None, schema: dict[str, Any], ) -> str | int | float | bool | list[str] | None: field_type = self._get_elicitation_field_type(schema) @@ -1154,7 +1193,9 @@ def _coerce_form_value( return value @staticmethod - def _coerce_boolean_value(field_name: str, raw_value: Any) -> bool: + def _coerce_boolean_value( + field_name: str, raw_value: str | int | float | bool | None + ) -> bool: if isinstance(raw_value, bool): return raw_value @@ -1165,12 +1206,12 @@ def _coerce_boolean_value(field_name: str, raw_value: Any) -> bool: return True if normalized in falsy: return False - raise ElicitationValidationError( - f"Field `{field_name}` must be a boolean." - ) + raise ElicitationValidationError(f"Field `{field_name}` must be a boolean.") @staticmethod - def _coerce_string_array_value(field_name: str, raw_value: Any) -> list[str]: + def _coerce_string_array_value( + field_name: str, raw_value: str | list | None + ) -> list[str]: if isinstance(raw_value, list): return [str(item).strip() for item in raw_value if str(item).strip()] @@ -1288,9 +1329,7 @@ def _get_elicitation_field_type(field_schema: dict[str, Any]) -> str: for candidate in field_type: if candidate in {"string", "integer", "number", "boolean", "array"}: return candidate - raise ElicitationParseError( - "Unsupported multi-type elicitation field." - ) + raise ElicitationParseError("Unsupported multi-type elicitation field.") if not isinstance(field_type, str): return "string" return field_type @@ -1326,7 +1365,7 @@ def _strip_code_fence(text: str) -> str: return "\n".join(lines[1:-1]).strip() @staticmethod - def _is_webchat_event(event: Any) -> bool: + def _is_webchat_event(event: TContext) -> bool: platform_name = getattr(event, "get_platform_name", None) if callable(platform_name): try: @@ -1380,6 +1419,9 @@ def _resolve_root_path_entry(self, entry: str) -> tuple[str, Path] | None: candidate_path = Path(get_astrbot_root()) / candidate_path path = candidate_path.resolve() display_name = path.name or normalized_entry + # 将用户显式配置的绝对路径加入临时白名单 + if candidate_path.is_absolute(): + self._user_configured_root_paths.add(path) # 安全检查 1: 符号链接检查 if path.is_symlink(): @@ -1409,18 +1451,22 @@ def _resolve_root_path_entry(self, entry: str) -> tuple[str, Path] | None: return None return display_name, path - + def _is_root_path_in_allowlist(self, path: Path) -> bool: """检查路径是否在允许的白名单内。 - + 白名单包括： 1. AstrBot 标准目录（data, temp, config 等） 2. 用户显式配置的其他目录 """ + # 检查是否是用户显式配置的路径 + if path in self._user_configured_root_paths: + return True + # 获取所有允许的根路径解析器 alias_resolvers = get_root_path_alias_resolvers() allowed_paths = set() - + # 添加所有别名解析的路径 for resolver_func in alias_resolvers.values(): try: @@ -1428,7 +1474,7 @@ def _is_root_path_in_allowlist(self, path: Path) -> bool: allowed_paths.add(allowed_path) except Exception: continue - + # 检查路径是否是允许路径的子目录或本身 for allowed_path in allowed_paths: try: @@ -1438,6 +1484,6 @@ def _is_root_path_in_allowlist(self, path: Path) -> bool: except ValueError: # 不是子目录，继续检查 continue - + # 检查是否完全匹配 return path in allowed_paths From a090f4350166563d5460da6f915e48c52ff74c97 Mon Sep 17 00:00:00 2001 From: Last-emo-boy Date: Mon, 30 Mar 2026 23:48:38 +0800 Subject: [PATCH 4/5] Make AstrBotDashboard.run asynchronous Change AstrBotDashboard.run from a regular method to an async coroutine (def -> async def) so the dashboard can perform asynchronous startup and integrate with an asyncio event loop. Callers will need to await or schedule this method; no other logic changes were made in this diff. --- astrbot/dashboard/server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/astrbot/dashboard/server.py b/astrbot/dashboard/server.py index 5b5c6548c0..b335fd721a 100644 --- a/astrbot/dashboard/server.py +++ b/astrbot/dashboard/server.py @@ -533,7 +533,7 @@ def _resolve_dashboard_ssl_config( return True, resolved_ssl_config - def run(self): + async def run(self): ip_addr = [] dashboard_config = self.core_lifecycle.astrbot_config.get("dashboard", {}) host_value = ( From bdf743ffb82dfc2a20ea9ad749edb5044c66d6e5 Mon Sep 17 00:00:00 2001 From: Last-emo-boy Date: Tue, 31 Mar 2026 01:28:11 +0800 Subject: [PATCH 5/5] Expose and enforce elicitation field constraints Add schema constraint metadata (format, minimum, maximum, enumNames, default) to elicitation forms and include them in the human-readable field list. Validate numeric min/max on submitted values in the backend. Update the Vue elicitation card to display constraint hints, prefill defaults on mount, use enumNames for option labels, and add minor CSS. Also update the ElicitationField TypeScript interface to include the new properties. Files changed: mcp_subcapability_bridge.py, ElicitationCard.vue, useMessages.ts. --- .../core/agent/mcp_subcapability_bridge.py | 85 +++++++++++++++---- .../message_list_comps/ElicitationCard.vue | 52 +++++++++++- dashboard/src/composables/useMessages.ts | 5 ++ 3 files changed, 123 insertions(+), 19 deletions(-) diff --git a/astrbot/core/agent/mcp_subcapability_bridge.py b/astrbot/core/agent/mcp_subcapability_bridge.py index 8f2db874ce..2fac40571d 100644 --- a/astrbot/core/agent/mcp_subcapability_bridge.py +++ b/astrbot/core/agent/mcp_subcapability_bridge.py @@ -937,10 +937,38 @@ def _build_form_elicitation_prompt( field_type = self._get_elicitation_field_type(schema) desc = str(schema.get("description", "")).strip() suffix = " required" if field_name in required_fields else " optional" + constraints: list[str] = [] + if schema.get("format"): + constraints.append(f"format={schema['format']}") + if schema.get("minimum") is not None: + constraints.append(f"min={schema['minimum']}") + if schema.get("maximum") is not None: + constraints.append(f"max={schema['maximum']}") + enum_values = schema.get("enum") + if isinstance(enum_values, list) and enum_values: + enum_names = schema.get("enumNames") + if isinstance(enum_names, list) and len(enum_names) == len( + enum_values + ): + options = [ + f"{v}({n})" + for v, n in zip(enum_values, enum_names, strict=False) + ] + else: + options = [str(v) for v in enum_values] + constraints.append(f"options=[{', '.join(options)}]") + default_value = schema.get("default") + if default_value is not None: + constraints.append(f"default={default_value}") + constraint_str = f" [{', '.join(constraints)}]" if constraints else "" if desc: - lines.append(f"- {field_name} ({field_type},{suffix}): {desc}") + lines.append( + f"- {field_name} ({field_type},{suffix}{constraint_str}): {desc}" + ) else: - lines.append(f"- {field_name} ({field_type},{suffix})") + lines.append( + f"- {field_name} ({field_type},{suffix}{constraint_str})" + ) if len(properties) == 1: lines.append("Reply with plain text or JSON.") elif len(properties) > 1: @@ -961,20 +989,32 @@ def _build_form_elicitation_payload( fields: list[dict[str, Any]] = [] for field_name, schema in properties.items(): enum_values = schema.get("enum") - fields.append( - { - "name": field_name, - "label": str(schema.get("title") or field_name), - "description": str(schema.get("description", "")).strip(), - "required": field_name in required_fields, - "type": self._get_elicitation_field_type(schema), - "enum": ( - [str(value) for value in enum_values] - if isinstance(enum_values, list) - else [] - ), - } - ) + enum_names = schema.get("enumNames") + field_info: dict[str, Any] = { + "name": field_name, + "label": str(schema.get("title") or field_name), + "description": str(schema.get("description", "")).strip(), + "required": field_name in required_fields, + "type": self._get_elicitation_field_type(schema), + "enum": ( + [str(value) for value in enum_values] + if isinstance(enum_values, list) + else [] + ), + } + if isinstance(enum_names, list) and len(enum_names) == len( + field_info["enum"] + ): + field_info["enumNames"] = [str(n) for n in enum_names] + if "default" in schema: + field_info["default"] = schema["default"] + if "format" in schema and isinstance(schema["format"], str): + field_info["format"] = schema["format"] + if "minimum" in schema: + field_info["minimum"] = schema["minimum"] + if "maximum" in schema: + field_info["maximum"] = schema["maximum"] + fields.append(field_info) return { "kind": "form", "server_name": self._server_name, @@ -1190,6 +1230,19 @@ def _coerce_form_value( raise ElicitationValidationError( f"Field `{field_name}` must be one of: {', '.join(map(str, enum_values))}." ) + + if isinstance(value, int | float) and not isinstance(value, bool): + minimum = schema.get("minimum") + maximum = schema.get("maximum") + if minimum is not None and value < minimum: + raise ElicitationValidationError( + f"Field `{field_name}` must be >= {minimum}." + ) + if maximum is not None and value > maximum: + raise ElicitationValidationError( + f"Field `{field_name}` must be <= {maximum}." + ) + return value @staticmethod diff --git a/dashboard/src/components/chat/message_list_comps/ElicitationCard.vue b/dashboard/src/components/chat/message_list_comps/ElicitationCard.vue index a7e6f679af..e590f28185 100644 --- a/dashboard/src/components/chat/message_list_comps/ElicitationCard.vue +++ b/dashboard/src/components/chat/message_list_comps/ElicitationCard.vue @@ -65,10 +65,13 @@
{{ field.description }}
+
+ {{ getConstraintHint(field) }} +
- {{ option }} + {{ getEnumLabel(field, idx) }}
@@ -143,7 +146,7 @@