diff --git a/.codex-mergecheck-dev2 b/.codex-mergecheck-dev2 new file mode 160000 index 0000000000..97bcafc070 --- /dev/null +++ b/.codex-mergecheck-dev2 @@ -0,0 +1 @@ +Subproject commit 97bcafc070b6c53a096f52bebb12cc762867d6df diff --git a/astrbot/core/agent/mcp_client.py b/astrbot/core/agent/mcp_client.py index 0518bfa757..aabb0897b3 100644 --- a/astrbot/core/agent/mcp_client.py +++ b/astrbot/core/agent/mcp_client.py @@ -31,10 +31,18 @@ wait_exponential, ) -from astrbot.core.agent.run_context import ContextWrapper +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, TContext from astrbot.core.utils.log_pipe import LogPipe -from .run_context import TContext +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 .tool import FunctionTool logger = logging.getLogger("astrbot") @@ -87,7 +95,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 @@ -169,7 +180,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 @@ -179,6 +190,12 @@ 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() self.process_pid: int | None = None @@ -186,8 +203,22 @@ def __init__(self) -> None: # 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[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( + cleanup_elicitation_periodically(interval=60), + name="mcp-elicitation-cleanup", + ) + logger.debug("已启动 MCP elicitation 后台清理任务") @staticmethod def _extract_stdio_process_pid(streams_context: object) -> int | None: @@ -222,6 +253,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) self.process_pid = None cfg = _prepare_config(mcp_server_config.copy()) @@ -266,6 +299,22 @@ def logging_callback( *streams, read_timeout_seconds=read_timeout, logging_callback=logging_callback, # type: ignore[arg-type] + 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: @@ -292,6 +341,22 @@ def logging_callback( write_stream=write_s, read_timeout_seconds=read_timeout, logging_callback=logging_callback, # type: ignore[arg-type] + 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(), ), ) @@ -315,7 +380,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=cast( Any, @@ -332,9 +397,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""" @@ -344,6 +441,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. @@ -368,6 +579,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: @@ -382,6 +594,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}" @@ -399,6 +613,7 @@ async def call_tool_with_reconnect( tool_name: str, arguments: dict, read_timeout_seconds: timedelta, + run_context: ContextWrapper[TContext] | None = None, ) -> mcp.types.CallToolResult: """Call MCP tool with automatic reconnection on failure, max 2 retries. @@ -423,34 +638,116 @@ 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() 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 @@ -484,4 +781,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..f9f8ed1b2b --- /dev/null +++ b/astrbot/core/agent/mcp_elicitation_registry.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +import asyncio +import time +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +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 + + +@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] + 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( + 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 + + +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_prompt_bridge.py b/astrbot/core/agent/mcp_prompt_bridge.py new file mode 100644 index 0000000000..167d2b5d84 --- /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: str | dict[str, Any] | None, +) -> 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..cf2cb3592a --- /dev/null +++ b/astrbot/core/agent/mcp_resource_bridge.py @@ -0,0 +1,372 @@ +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, 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_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..2fac40571d --- /dev/null +++ b/astrbot/core/agent/mcp_subcapability_bridge.py @@ -0,0 +1,1542 @@ +from __future__ import annotations + +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 +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 +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.""" + + +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): + """用户输入解析失败。""" + + +class ElicitationTimeoutError(MCPElicitationError): + """elicitation 超时。""" + + +class ElicitationValidationError(MCPElicitationError): + """schema 验证失败。""" + + +@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({}) + # 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 + # 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( + 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 + + 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: + 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: None, + ) -> 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 + + 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 + + # 自动从 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] + 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 + finally: + self._active_run_context = None + + async def handle_sampling( + self, + _request_context: None, + 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: None, + 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: ContextWrapper[TContext], + event: TContext, + sender_id: str, + params: mcp.types.ElicitRequestFormParams, + ) -> mcp.types.ElicitResult: + import mcp + + properties = self._get_elicitation_properties(params.requestedSchema) + # 使用动态超时计算 + 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), + 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, + ElicitationValidationError, + ElicitationParseError, + ) 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, + ) + + @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, + *, + plugin_context: ContextWrapper[TContext], + event: TContext, + 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: TContext, + 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: TContext, + 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: TContext, + 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" + 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}{constraint_str}): {desc}" + ) + else: + 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: + 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") + 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, + "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 ElicitationParseError("The reply is empty.") + + if normalized_reply.startswith("{"): + try: + payload = json.loads(normalized_reply) + except json.JSONDecodeError as exc: + raise ElicitationParseError( + "The JSON reply could not be parsed." + ) from exc + if not isinstance(payload, dict): + raise ElicitationParseError("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 ElicitationParseError( + "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 ElicitationValidationError( + "Missing required field(s): " + ", ".join(missing_required) + ) + return coerced + + def _coerce_form_value( + self, + *, + field_name: str, + 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) + 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 ElicitationValidationError( + f"Field `{field_name}` must be an integer." + ) + try: + value = int(str(raw_value).strip()) + except (TypeError, ValueError) as exc: + raise ElicitationValidationError( + f"Field `{field_name}` must be an integer." + ) from exc + elif field_type == "number": + if isinstance(raw_value, bool): + raise ElicitationValidationError( + f"Field `{field_name}` must be a number." + ) + try: + value = float(str(raw_value).strip()) + except (TypeError, ValueError) as exc: + raise ElicitationValidationError( + 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 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 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 + def _coerce_boolean_value( + field_name: str, raw_value: str | int | float | bool | None + ) -> 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 ElicitationValidationError(f"Field `{field_name}` must be a boolean.") + + @staticmethod + 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()] + + 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 ElicitationValidationError( + 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 ElicitationParseError( + "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 ElicitationParseError("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: TContext) -> 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 candidate_path.is_absolute(): + self._user_configured_root_paths.add(path) + + # 安全检查 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 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. 用户显式配置的其他目录 + """ + # 检查是否是用户显式配置的路径 + 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: + 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 diff --git a/astrbot/core/astr_main_agent_resources.py b/astrbot/core/astr_main_agent_resources.py index 1b2d52f7aa..32e777fd2a 100644 --- a/astrbot/core/astr_main_agent_resources.py +++ b/astrbot/core/astr_main_agent_resources.py @@ -369,8 +369,7 @@ async def retrieve_knowledge_base( logger.debug("所配置的所有知识库全为空, 跳过检索过程") return - logger.debug(f"[知识库] 开始检索知识库,数量: {len(kb_names)}, top_k={top_k}") - + logger.debug(f"[知识库] 开始检索知识库，数量: {len(kb_names)}, top_k={top_k}") kb_context = await kb_mgr.retrieve( query=query, kb_names=kb_names, diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index c00dc563d9..b491fc758d 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -1,6 +1,10 @@ """如需修改配置，请在 `data/cmd_config.json` 中修改或者在管理面板中可视化修改。""" +import binascii +import hashlib import os +import secrets +from importlib import metadata from typing import Any, TypedDict from astrbot.builtin_stars.web_searcher.provider_constants import ( @@ -9,7 +13,20 @@ ) from astrbot.core.utils.astrbot_path import get_astrbot_data_path -VERSION = "4.25.0" + +def _generate_random_dashboard_password_hash() -> str: + iterations = 200_000 + salt = secrets.token_bytes(16) + secret = secrets.token_bytes(32) + dk = hashlib.pbkdf2_hmac("sha256", secret, salt, iterations) + return f"pbkdf2_sha256${iterations}${binascii.hexlify(salt).decode()}${dk.hex()}" + + +try: + __version__ = metadata.version("AstrBot") +except metadata.PackageNotFoundError: + __version__ = "unknown" +VERSION = __version__ DB_PATH = os.path.join(get_astrbot_data_path(), "data_v4.db") PERSONAL_WECHAT_CONFIG_METADATA = { "weixin_oc_base_url": { @@ -3492,6 +3509,7 @@ class ChatProviderTemplate(TypedDict): "description": "提供商可达性检测", "type": "bool", "hint": "/provider 命令列出模型时是否并发检测连通性。开启后会主动调用模型测试连通性，可能产生额外 token 消耗。", + "collapsed": True, }, "provider_settings.max_quoted_fallback_images": { "description": "引用图片回退解析上限", 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 4df0aad025..44b0481aa6 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 @@ -5,7 +5,10 @@ from collections.abc import AsyncGenerator from dataclasses import replace -from astrbot.core import logger +from astrbot.core import db_helper, 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_agent_run_util import AgentRunner, run_agent, run_live_agent @@ -179,6 +182,12 @@ async def process( # type: ignore[invalid-method-override] 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 da4759ecb1..f72afe7ecb 100644 --- a/astrbot/core/platform/sources/webchat/webchat_event.py +++ b/astrbot/core/platform/sources/webchat/webchat_event.py @@ -69,6 +69,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 45541e41c4..6b9e031b08 100644 --- a/astrbot/core/provider/func_tool_manager.py +++ b/astrbot/core/provider/func_tool_manager.py @@ -7,7 +7,131 @@ from __future__ import annotations -# Constants that are still imported by other modules +import asyncio +import copy +import json +import os +import threading +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 + +import aiohttp + +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 + +DEFAULT_MCP_CONFIG = {"mcpServers": {}} + +DEFAULT_MCP_INIT_TIMEOUT_SECONDS = 180.0 +DEFAULT_ENABLE_MCP_TIMEOUT_SECONDS = 180.0 +MCP_INIT_TIMEOUT_ENV = "ASTRBOT_MCP_INIT_TIMEOUT" +ENABLE_MCP_TIMEOUT_ENV = "ASTRBOT_MCP_ENABLE_TIMEOUT" +MAX_MCP_TIMEOUT_SECONDS = 300.0 + + +class MCPInitError(Exception): + """Base exception for MCP initialization failures.""" + + +class MCPInitTimeoutError(asyncio.TimeoutError, MCPInitError): + """Raised when MCP client initialization exceeds the configured timeout.""" + + +class MCPAllServicesFailedError(MCPInitError): + """Raised when all configured MCP services fail to initialize.""" + + +class MCPShutdownTimeoutError(asyncio.TimeoutError): + """Raised when MCP shutdown exceeds the configured timeout.""" + + def __init__(self, names: list[str], timeout: float) -> None: + self.names = names + self.timeout = timeout + message = f"MCP 服务关闭超时({timeout:g} 秒):{', '.join(names)}" + super().__init__(message) + + +@dataclass +class MCPInitSummary: + total: int + success: int + failed: list[str] + + +@dataclass +class _MCPServerRuntime: + name: str + client: MCPClient + shutdown_event: asyncio.Event + lifecycle_task: asyncio.Task[None] + + +class _MCPClientDictView(Mapping[str, MCPClient]): + """Read-only view of MCP clients derived from runtime state.""" + + def __init__(self, runtime: dict[str, _MCPServerRuntime]) -> None: + self._runtime = runtime + + def __getitem__(self, key: str) -> MCPClient: + return self._runtime[key].client + + def __iter__(self): + return iter(self._runtime) + + def __len__(self) -> int: + return len(self._runtime) + + +def _resolve_timeout( + timeout: float | str | None = None, + *, + env_name: str = MCP_INIT_TIMEOUT_ENV, + default: float = DEFAULT_MCP_INIT_TIMEOUT_SECONDS, +) -> float: + """Resolve timeout with precedence: explicit argument > env value > default.""" + source = f"环境变量 {env_name}" + if timeout is None: + timeout = os.getenv(env_name, str(default)) + else: + source = "显式参数 timeout" + + try: + timeout_value = float(timeout) + except (TypeError, ValueError): + logger.warning( + f"超时配置({source})={timeout!r} 无效,使用默认值 {default:g} 秒｡" + ) + return default + + if timeout_value <= 0: + logger.warning( + f"超时配置({source})={timeout_value:g} 必须大于 0,使用默认值 {default:g} 秒｡" + ) + return default + + if timeout_value > MAX_MCP_TIMEOUT_SECONDS: + logger.warning( + f"超时配置({source})={timeout_value:g} 过大,已限制为最大值 " + f"{MAX_MCP_TIMEOUT_SECONDS:g} 秒,以避免长时间等待｡" + ) + return MAX_MCP_TIMEOUT_SECONDS + + return timeout_value + + SUPPORTED_TYPES = [ "string", "number", @@ -27,40 +151,860 @@ "set": "array", } -# Re-export from _internal for backward compatibility - -from astrbot._internal.tools.registry import ( - DEFAULT_MCP_CONFIG, - ENABLE_MCP_TIMEOUT_ENV, - MCP_INIT_TIMEOUT_ENV, - FuncCall, - FunctionTool, - FunctionToolManager, - MCPAllServicesFailedError, - MCPInitError, - MCPInitSummary, - MCPInitTimeoutError, - MCPShutdownTimeoutError, - ToolSet, -) - -# For backward compatibility - alias FunctionTool as FuncTool +# Backward compatibility alias FuncTool = FunctionTool -__all__ = [ - "DEFAULT_MCP_CONFIG", - "ENABLE_MCP_TIMEOUT_ENV", - "MCP_INIT_TIMEOUT_ENV", - "PY_TO_JSON_TYPE", - "SUPPORTED_TYPES", - "FuncCall", - "FuncTool", - "FunctionTool", - "FunctionToolManager", - "MCPAllServicesFailedError", - "MCPInitError", - "MCPInitSummary", - "MCPInitTimeoutError", - "MCPShutdownTimeoutError", - "ToolSet", -] + +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 + + +async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: + """快速测试 MCP 服务器可达性""" + import aiohttp + + cfg = _prepare_config(config.copy()) + + url = cfg["url"] + headers = cfg.get("headers", {}) + timeout = cfg.get("timeout", 10) + + try: + async with aiohttp.ClientSession() as session: + if cfg.get("transport") == "streamable_http": + test_payload = { + "jsonrpc": "2.0", + "method": "initialize", + "id": 0, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.2.3"}, + }, + } + async with session.post( + url, + headers={ + **headers, + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + }, + json=test_payload, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + if response.status == 200: + return True, "" + return False, f"HTTP {response.status}: {response.reason}" + else: + async with session.get( + url, + headers={ + **headers, + "Accept": "application/json, text/event-stream", + }, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + if response.status == 200: + return True, "" + return False, f"HTTP {response.status}: {response.reason}" + + except asyncio.TimeoutError: + return False, f"连接超时: {timeout}秒" + except Exception as e: + return False, f"{e!s}" + + +class FunctionToolManager: + def __init__(self) -> None: + self.func_list: list[FuncTool] = [] + self._mcp_server_runtime: dict[str, _MCPServerRuntime] = {} + """MCP 服务运行时状态(唯一事实来源)""" + self._mcp_server_runtime_view = MappingProxyType(self._mcp_server_runtime) + self._mcp_client_dict_view = _MCPClientDictView(self._mcp_server_runtime) + self._timeout_mismatch_warned = False + self._timeout_warn_lock = threading.Lock() + self._runtime_lock = asyncio.Lock() + self._mcp_starting: set[str] = set() + self._init_timeout_default = _resolve_timeout( + timeout=None, + env_name=MCP_INIT_TIMEOUT_ENV, + default=DEFAULT_MCP_INIT_TIMEOUT_SECONDS, + ) + self._enable_timeout_default = _resolve_timeout( + timeout=None, + env_name=ENABLE_MCP_TIMEOUT_ENV, + default=DEFAULT_ENABLE_MCP_TIMEOUT_SECONDS, + ) + self._warn_on_timeout_mismatch( + self._init_timeout_default, + self._enable_timeout_default, + ) + + @property + def mcp_client_dict(self) -> Mapping[str, MCPClient]: + """Read-only compatibility view for external callers that still read mcp_client_dict. + + Note: Mutating this mapping is unsupported and will raise TypeError. + """ + return self._mcp_client_dict_view + + @property + def mcp_server_runtime_view(self) -> Mapping[str, _MCPServerRuntime]: + """Read-only view of MCP runtime metadata for external callers.""" + return self._mcp_server_runtime_view + + @property + def mcp_server_runtime(self) -> Mapping[str, _MCPServerRuntime]: + """Backward-compatible read-only view (deprecated). Do not mutate. + + Note: Mutations are not supported and will raise TypeError. + """ + return self._mcp_server_runtime_view + + def empty(self) -> bool: + return len(self.func_list) == 0 + + def spec_to_func( + self, + name: str, + func_args: list[dict], + desc: str, + handler: Callable[..., Awaitable[Any] | AsyncGenerator[Any]], + ) -> FuncTool: + params = { + "type": "object", # hard-coded here + "properties": {}, + } + for param in func_args: + p = copy.deepcopy(param) + p.pop("name", None) + params["properties"][param["name"]] = p + return FuncTool( + name=name, + parameters=params, + description=desc, + handler=handler, + ) + + def add_func( + self, + name: str, + func_args: list, + desc: str, + handler: Callable[..., Awaitable[Any] | AsyncGenerator[Any]], + ) -> None: + """添加函数调用工具 + + @param name: 函数名 + @param func_args: 函数参数列表,格式为 [{"type": "string", "name": "arg_name", "description": "arg_description"}, ...] + @param desc: 函数描述 + @param func_obj: 处理函数 + """ + # check if the tool has been added before + self.remove_func(name) + + self.func_list.append( + self.spec_to_func( + name=name, + func_args=func_args, + desc=desc, + handler=handler, + ), + ) + logger.info(f"添加函数调用工具: {name}") + + def remove_func(self, name: str) -> None: + """删除一个函数调用工具｡""" + for i, f in enumerate(self.func_list): + if f.name == name: + self.func_list.pop(i) + break + + def get_func(self, name) -> FuncTool | None: + # 优先返回已激活的工具(后加载的覆盖前面的,与 ToolSet.add_tool 保持一致) + # 使用 getattr(..., True) 与 ToolSet.add_tool 保持一致:没有 active 属性的工具视为已激活 + for f in reversed(self.func_list): + if f.name == name and getattr(f, "active", True): + return f + # 退化则拿最后一个同名工具 + for f in reversed(self.func_list): + if f.name == name: + return f + return None + + def get_full_tool_set(self) -> ToolSet: + """获取完整工具集 + + 使用 ToolSet.add_tool 进行填充｡对于同名工具,去重规则为: + - 优先保留 active=True 的工具; + - 当 active 状态相同时,后加载的工具会覆盖前面的工具｡ + + 因此,后加载的 inactive 工具不会覆盖已激活的工具; + 同时,MCP 工具在需要时仍可覆盖被禁用的内置工具｡ + """ + tool_set = ToolSet() + for tool in self.func_list: + tool_set.add_tool(tool) + 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 中的敏感信息 + if "command" in cfg: + cmd = cfg["command"] + executable = str(cmd[0] if isinstance(cmd, (list, tuple)) and cmd else cmd) + args_val = cfg.get("args", []) + args_count = ( + len(args_val) + if isinstance(args_val, (list, tuple)) + else (0 if args_val is None else 1) + ) + logger.debug(f" 命令可执行文件: {executable}, 参数数量: {args_count}") + return + + if "url" in cfg: + parsed = urllib.parse.urlparse(str(cfg["url"])) + host = parsed.hostname or "" + scheme = parsed.scheme or "unknown" + try: + port = f":{parsed.port}" if parsed.port else "" + except ValueError: + port = "" + logger.debug(f" 主机: {scheme}://{host}{port}") + + async def init_mcp_clients( + self, raise_on_all_failed: bool = False + ) -> MCPInitSummary: + """从项目根目录读取 mcp_server.json 文件,初始化 MCP 服务列表｡文件格式如下: + ``` + { + "mcpServers": { + "weather": { + "command": "uv", + "args": [ + "--directory", + "/ABSOLUTE/PATH/TO/PARENT/FOLDER/weather", + "run", + "weather.py" + ] + } + } + ... + } + ``` + + Timeout behavior: + - 初始化超时使用环境变量 ASTRBOT_MCP_INIT_TIMEOUT 或默认值｡ + - 动态启用超时使用 ASTRBOT_MCP_ENABLE_TIMEOUT(独立于初始化超时)｡ + """ + mcp_json_file = Path(get_astrbot_data_path()) / "mcp_server.json" + if not mcp_json_file.exists(): + # 配置文件不存在错误处理 + 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=[]) + + mcp_server_json_obj = self.load_mcp_config()["mcpServers"] + + init_timeout = self._init_timeout_default + timeout_display = f"{init_timeout:g}" + + active_configs: list[tuple[str, dict, asyncio.Event]] = [] + for name, cfg in mcp_server_json_obj.items(): + if cfg.get("active", True): + shutdown_event = asyncio.Event() + active_configs.append((name, cfg, shutdown_event)) + + if not active_configs: + return MCPInitSummary(total=0, success=0, failed=[]) + + logger.info(f"等待 {len(active_configs)} 个 MCP 服务初始化...") + + init_tasks = [ + asyncio.create_task( + self._start_mcp_server( + name=name, + cfg=cfg, + shutdown_event=shutdown_event, + init_timeout=init_timeout, + ), + name=f"mcp-init:{name}", + ) + for (name, cfg, shutdown_event) in active_configs + ] + results = await asyncio.gather(*init_tasks, return_exceptions=True) + + success_count = 0 + failed_services: list[str] = [] + + for (name, cfg, _), result in zip(active_configs, results, strict=False): + if isinstance(result, Exception): + if isinstance(result, MCPInitTimeoutError): + logger.error( + f"Connected to MCP server {name} timeout ({timeout_display} seconds)" + ) + else: + logger.error(f"Failed to initialize MCP server {name}: {result}") + self._log_safe_mcp_debug_config(cfg) + failed_services.append(name) + async with self._runtime_lock: + self._mcp_server_runtime.pop(name, None) + continue + + success_count += 1 + + if failed_services: + logger.warning( + f"The following MCP services failed to initialize: {', '.join(failed_services)}. " + f"Please check the mcp_server.json file and server availability." + ) + + summary = MCPInitSummary( + total=len(active_configs), success=success_count, failed=failed_services + ) + logger.info( + f"MCP services initialization completed: {summary.success}/{summary.total} successful, {len(summary.failed)} failed." + ) + if summary.total > 0 and summary.success == 0: + msg = "All MCP services failed to initialize, please check the mcp_server.json and server availability." + if raise_on_all_failed: + raise MCPAllServicesFailedError(msg) + logger.error(msg) + return summary + + async def _start_mcp_server( + self, + name: str, + cfg: dict, + *, + shutdown_event: asyncio.Event | None = None, + init_timeout: float, + ) -> None: + """Initialize MCP server with timeout and register task/event together. + + This method is idempotent. If the server is already running, the existing + runtime is kept and the new config is ignored. + """ + async with self._runtime_lock: + if name in self._mcp_server_runtime or name in self._mcp_starting: + logger.warning( + f"Connected to MCP server {name}, ignoring this startup request (timeout={init_timeout:g})." + ) + self._log_safe_mcp_debug_config(cfg) + return + self._mcp_starting.add(name) + + if shutdown_event is None: + shutdown_event = asyncio.Event() + + mcp_client: MCPClient | None = None + try: + mcp_client = await asyncio.wait_for( + self._init_mcp_client(name, cfg), + timeout=init_timeout, + ) + except asyncio.TimeoutError as exc: + raise MCPInitTimeoutError( + f"Connected to MCP server {name} timeout ({init_timeout:g} seconds)" + ) from exc + except Exception: + logger.error(f"Failed to initialize MCP client {name}", exc_info=True) + raise + finally: + if mcp_client is None: + async with self._runtime_lock: + self._mcp_starting.discard(name) + + async def lifecycle() -> None: + try: + await shutdown_event.wait() + logger.info(f"Received shutdown signal for MCP client {name}") + except asyncio.CancelledError: + logger.debug(f"MCP client {name} task was cancelled") + raise + finally: + await self._terminate_mcp_client(name) + + lifecycle_task = asyncio.create_task(lifecycle(), name=f"mcp-client:{name}") + async with self._runtime_lock: + self._mcp_server_runtime[name] = _MCPServerRuntime( + name=name, + client=mcp_client, + shutdown_event=shutdown_event, + lifecycle_task=lifecycle_task, + ) + self._mcp_starting.discard(name) + + async def _shutdown_runtimes( + self, + runtimes: list[_MCPServerRuntime], + shutdown_timeout: float, + *, + strict: bool = True, + ) -> list[str]: + """Shutdown runtimes and wait for lifecycle tasks to complete.""" + lifecycle_tasks = [ + runtime.lifecycle_task + for runtime in runtimes + if not runtime.lifecycle_task.done() + ] + if not lifecycle_tasks: + return [] + + for runtime in runtimes: + runtime.shutdown_event.set() + + try: + results = await asyncio.wait_for( + asyncio.gather(*lifecycle_tasks, return_exceptions=True), + timeout=shutdown_timeout, + ) + except asyncio.TimeoutError: + pending_names = [ + runtime.name + for runtime in runtimes + if not runtime.lifecycle_task.done() + ] + for task in lifecycle_tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*lifecycle_tasks, return_exceptions=True) + if strict: + raise MCPShutdownTimeoutError(pending_names, shutdown_timeout) + logger.warning( + "MCP server shutdown timeout (%s seconds), the following servers were not fully closed: %s", + f"{shutdown_timeout:g}", + ", ".join(pending_names), + ) + return pending_names + else: + for result in results: + if isinstance(result, asyncio.CancelledError): + logger.debug("MCP lifecycle task was cancelled during shutdown.") + elif isinstance(result, Exception): + logger.error( + "MCP lifecycle task failed during shutdown.", + exc_info=(type(result), result, result.__traceback__), + ) + return [] + + async def _cleanup_mcp_client_safely( + self, mcp_client: MCPClient, name: str + ) -> None: + """安全清理单个 MCP 客户端,避免清理异常中断主流程｡""" + try: + await mcp_client.cleanup() + except Exception as cleanup_exc: # only log here + logger.error( + f"Failed to cleanup MCP client resources {name}: {cleanup_exc}" + ) + + async def _init_mcp_client(self, name: str, config: dict) -> MCPClient: + """初始化单个MCP客户端""" + mcp_client = MCPClient() + mcp_client.name = name + 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 + except Exception: + await self._cleanup_mcp_client_safely(mcp_client, name) + 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._remove_mcp_bound_tools(name) + + # 将 MCP 工具转换为 FuncTool 并添加到 func_list + for tool in mcp_client.tools: + func_tool = MCPTool( + mcp_tool=tool, + mcp_client=mcp_client, + mcp_server_name=name, + ) + 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 + + async def _terminate_mcp_client(self, name: str) -> None: + """关闭并清理MCP客户端""" + async with self._runtime_lock: + runtime = self._mcp_server_runtime.get(name) + if runtime: + client = runtime.client + # 关闭MCP连接 + await self._cleanup_mcp_client_safely(client, name) + # 移除关联的FuncTool + self._remove_mcp_bound_tools(name) + async with self._runtime_lock: + self._mcp_server_runtime.pop(name, None) + self._mcp_starting.discard(name) + logger.info(f"Disconnected from MCP server {name}") + return + + # Runtime missing but stale tools may still exist after failed flows. + self._remove_mcp_bound_tools(name) + async with self._runtime_lock: + self._mcp_starting.discard(name) + + @staticmethod + async def test_mcp_server_connection(config: dict) -> list[str]: + if "url" in config: + success, error_msg = await _quick_test_mcp_connection(config) + if not success: + raise Exception(error_msg) + + mcp_client = MCPClient() + try: + 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() + return tool_names + + async def enable_mcp_server( + self, + name: str, + config: dict, + shutdown_event: asyncio.Event | None = None, + init_timeout: float | str | None = None, + ) -> None: + """Enable a new MCP server and initialize it. + + Args: + name: The name of the MCP server. + config: Configuration for the MCP server. + shutdown_event: Event to signal when the MCP client should shut down. + init_timeout: Timeout in seconds for initialization. + Uses ASTRBOT_MCP_ENABLE_TIMEOUT by default (separate from init timeout). + + Raises: + MCPInitTimeoutError: If initialization does not complete within timeout. + Exception: If there is an error during initialization. + """ + if init_timeout is None: + timeout_value = self._enable_timeout_default + else: + timeout_value = _resolve_timeout( + timeout=init_timeout, + env_name=ENABLE_MCP_TIMEOUT_ENV, + default=self._enable_timeout_default, + ) + await self._start_mcp_server( + name=name, + cfg=config, + shutdown_event=shutdown_event, + init_timeout=timeout_value, + ) + + async def disable_mcp_server( + self, + name: str | None = None, + shutdown_timeout: float = 10, + ) -> None: + """Disable an MCP server by its name. + + Args: + name (str): The name of the MCP server to disable. If None, ALL MCP servers will be disabled. + shutdown_timeout (int): Timeout. + + Raises: + MCPShutdownTimeoutError: If shutdown does not complete within timeout. + Only raised when disabling a specific server (name is not None). + + """ + if name: + async with self._runtime_lock: + runtime = self._mcp_server_runtime.get(name) + if runtime is None: + return + + await self._shutdown_runtimes([runtime], shutdown_timeout, strict=True) + else: + async with self._runtime_lock: + runtimes = list(self._mcp_server_runtime.values()) + await self._shutdown_runtimes(runtimes, shutdown_timeout, strict=False) + + def _warn_on_timeout_mismatch( + self, + init_timeout: float, + enable_timeout: float, + ) -> None: + if init_timeout == enable_timeout: + return + with self._timeout_warn_lock: + if self._timeout_mismatch_warned: + return + logger.info( + "检测到 MCP 初始化超时与动态启用超时配置不同:" + "初始化使用 %s 秒,动态启用使用 %s 秒｡如需一致,请设置相同值｡", + f"{init_timeout:g}", + f"{enable_timeout:g}", + ) + self._timeout_mismatch_warned = True + + def get_func_desc_openai_style(self, omit_empty_parameter_field=False) -> list: + """获得 OpenAI API 风格的**已经激活**的工具描述""" + tools = [f for f in self.func_list if f.active] + toolset = ToolSet(tools) + return toolset.openai_schema( + omit_empty_parameter_field=omit_empty_parameter_field, + ) + + def get_func_desc_anthropic_style(self) -> list: + """获得 Anthropic API 风格的**已经激活**的工具描述""" + tools = [f for f in self.func_list if f.active] + toolset = ToolSet(tools) + return toolset.anthropic_schema() + + def get_func_desc_google_genai_style(self) -> dict: + """获得 Google GenAI API 风格的**已经激活**的工具描述""" + tools = [f for f in self.func_list if f.active] + toolset = ToolSet(tools) + return toolset.google_schema() + + def deactivate_llm_tool(self, name: str) -> bool: + """停用一个已经注册的函数调用工具｡ + + Returns: + 如果没找到,会返回 False + + """ + func_tool = self.get_func(name) + if func_tool is not None: + func_tool.active = False + + inactivated_llm_tools: list = sp.get( + "inactivated_llm_tools", + [], + scope="global", + scope_id="global", + ) + if name not in inactivated_llm_tools: + inactivated_llm_tools.append(name) + sp.put( + "inactivated_llm_tools", + inactivated_llm_tools, + scope="global", + scope_id="global", + ) + + return True + return False + + # 因为不想解决循环引用,所以这里直接传入 star_map 先了... + def activate_llm_tool(self, name: str, star_map: dict) -> bool: + func_tool = self.get_func(name) + if func_tool is not None: + if func_tool.handler_module_path in star_map: + if not star_map[func_tool.handler_module_path].activated: + raise ValueError( + f"此函数调用工具所属的插件 {star_map[func_tool.handler_module_path].name} 已被禁用,请先在管理面板启用再激活此工具｡", + ) + + func_tool.active = True + + inactivated_llm_tools: list = sp.get( + "inactivated_llm_tools", + [], + scope="global", + scope_id="global", + ) + if name in inactivated_llm_tools: + inactivated_llm_tools.remove(name) + sp.put( + "inactivated_llm_tools", + inactivated_llm_tools, + scope="global", + scope_id="global", + ) + + return True + return False + + @property + def mcp_config_path(self): + return Path(get_astrbot_data_path()) / "mcp_server.json" + + def load_mcp_config(self): + if not self.mcp_config_path.exists(): + # 配置文件不存在，创建默认配置 + 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 normalize_mcp_config(DEFAULT_MCP_CONFIG) + + try: + 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 normalize_mcp_config(DEFAULT_MCP_CONFIG) + + def save_mcp_config(self, config: dict) -> bool: + try: + 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}") + return False + + async def sync_modelscope_mcp_servers(self, access_token: str) -> None: + """从 ModelScope 平台同步 MCP 服务器配置""" + base_url = "https://www.modelscope.cn/openapi/v1" + url = f"{base_url}/mcp/servers/operational" + headers = { + "Authorization": f"Bearer {access_token.strip()}", + "Content-Type": "application/json", + } + + try: + async with aiohttp.ClientSession() as session: + async with session.get(url, headers=headers) as response: + if response.status == 200: + data = await response.json() + mcp_server_list = data.get("data", {}).get( + "mcp_server_list", + [], + ) + local_mcp_config = self.load_mcp_config() + + synced_count = 0 + for server in mcp_server_list: + server_name = server["name"] + operational_urls = server.get("operational_urls", []) + if not operational_urls: + continue + url_info = operational_urls[0] + server_url = url_info.get("url") + if not server_url: + continue + # 添加到配置中(同名会覆盖) + local_mcp_config["mcpServers"][server_name] = { + "url": server_url, + "transport": "sse", + "active": True, + "provider": "modelscope", + } + synced_count += 1 + + if synced_count > 0: + self.save_mcp_config(local_mcp_config) + tasks = [] + for server in mcp_server_list: + name = server["name"] + tasks.append( + self.enable_mcp_server( + name=name, + config=local_mcp_config["mcpServers"][name], + ), + ) + await asyncio.gather(*tasks) + logger.info( + f"从 ModelScope 同步了 {synced_count} 个 MCP 服务器", + ) + else: + logger.warning("没有找到可用的 ModelScope MCP 服务器") + else: + raise Exception( + f"ModelScope API 请求失败: HTTP {response.status}", + ) + + except aiohttp.ClientError as e: + raise Exception(f"网络连接错误: {e!s}") + except Exception as e: + raise Exception(f"同步 ModelScope MCP 服务器时发生错误: {e!s}") + + # Module paths whose ``get_all_tools()`` function returns internal tools. + # To add a new internal-tool provider, simply append its module path here. + _INTERNAL_TOOL_PROVIDERS: list[str] = [ + "astrbot.core.tools.cron_tools", + "astrbot.core.tools.kb_query", + "astrbot.core.tools.send_message", + "astrbot.core.computer.computer_tool_provider", + ] + + def register_internal_tools(self) -> None: + """Register AstrBot built-in tools from all internal providers. + + Each provider module is expected to expose a ``get_all_tools()`` + function that returns a list of ``FunctionTool`` instances. + + Tools are marked with ``source='internal'`` so the WebUI can + distinguish them from plugin and MCP tools, and subagent + orchestrators can resolve them by name. + + Duplicate registration is idempotent (skips if name already present). + """ + import importlib + + existing_names = {t.name for t in self.func_list} + + for module_path in self._INTERNAL_TOOL_PROVIDERS: + try: + mod = importlib.import_module(module_path) + provider_tools = mod.get_all_tools() + except Exception as e: + logger.warning( + "Failed to load internal tool provider %s: %s", + module_path, + e, + ) + continue + + for tool in provider_tools: + tool.source = "internal" + if tool.name not in existing_names: + self.func_list.append(tool) + existing_names.add(tool.name) + logger.info("Registered internal tool: %s", tool.name) + + def __str__(self) -> str: + return str(self.func_list) + + def __repr__(self) -> str: + return str(self.func_list) + + +# alias +FuncCall = FunctionToolManager + diff --git a/astrbot/core/provider/sources/anthropic_source.py b/astrbot/core/provider/sources/anthropic_source.py index 41b85b6e98..6dfa094669 100644 --- a/astrbot/core/provider/sources/anthropic_source.py +++ b/astrbot/core/provider/sources/anthropic_source.py @@ -3,7 +3,6 @@ from collections.abc import AsyncGenerator from typing import Literal -import aiofiles import anthropic import httpx from anthropic import AsyncAnthropic @@ -136,10 +135,10 @@ def _prepare_payload(self, messages: list[dict]): """准备 Anthropic API 的请求 payload Args: - messages: OpenAI 格式的消息列表,包含用户输入和系统提示等信息 + messages: OpenAI 格式的消息列表，包含用户输入和系统提示等信息 Returns: system_prompt: 系统提示内容 - new_messages: 处理后的消息列表,去除系统提示 + new_messages: 处理后的消息列表，去除系统提示 """ system_prompt = "" @@ -174,7 +173,7 @@ def _prepare_payload(self, messages: list[dict]): if "tool_calls" in message and isinstance(message["tool_calls"], list): for tool_call in message["tool_calls"]: - blocks.append( + blocks.append( # noqa: PERF401 { "type": "tool_use", "name": tool_call["function"]["name"], @@ -405,7 +404,7 @@ async def _query_stream( # type: ignore[invalid-method-override] id=id, ) elif event.content_block.type == "tool_use": - # 工具使用块开始,初始化缓冲区 + # 工具使用块开始，初始化缓冲区 tool_use_buffer[event.index] = { "id": event.content_block.id, "name": event.content_block.name, @@ -477,7 +476,7 @@ async def _query_stream( # type: ignore[invalid-method-override] id=id, ) except json.JSONDecodeError: - # JSON 解析失败,跳过这个工具调用 + # JSON 解析失败，跳过这个工具调用 logger.warning(f"工具调用参数 JSON 解析失败: {tool_info}") # 清理缓冲区 @@ -644,7 +643,7 @@ async def assemble_context( # type: ignore[invalid-method-override] image_urls: list[str] | None = None, extra_user_content_parts: list[ContentPart] | None = None, ): - """组装上下文,支持文本和图片""" + """组装上下文，支持文本和图片""" async def resolve_image_url(image_url: str) -> dict | None: if image_url.startswith("http"): @@ -657,7 +656,7 @@ async def resolve_image_url(image_url: str) -> dict | None: image_data, mime_type = await self.encode_image_bs64(image_url) if not image_data: - logger.warning(f"图片 {image_url} 得到的结果为空,将忽略｡") + logger.warning(f"图片 {image_url} 得到的结果为空，将忽略。") return None return { @@ -675,17 +674,17 @@ async def resolve_image_url(image_url: str) -> dict | None: content = [] - # 1. 用户原始发言(OpenAI 建议:用户发言在前) + # 1. 用户原始发言（OpenAI 建议：用户发言在前） if text: content.append({"type": "text", "text": text}) elif image_urls: - # 如果没有文本但有图片,添加占位文本 + # 如果没有文本但有图片，添加占位文本 content.append({"type": "text", "text": "[图片]"}) elif extra_user_content_parts: - # 如果只有额外内容块,也需要添加占位文本 + # 如果只有额外内容块，也需要添加占位文本 content.append({"type": "text", "text": " "}) - # 2. 额外的内容块(系统提醒､指令等) + # 2. 额外的内容块（系统提醒、指令等） if extra_user_content_parts: for block in extra_user_content_parts: if isinstance(block, TextPart): @@ -704,7 +703,7 @@ async def resolve_image_url(image_url: str) -> dict | None: if image_dict: content.append(image_dict) - # 如果只有主文本且没有额外内容块和图片,返回简单格式以保持向后兼容 + # 如果只有主文本且没有额外内容块和图片，返回简单格式以保持向后兼容 if ( text and not extra_user_content_parts @@ -718,7 +717,7 @@ async def resolve_image_url(image_url: str) -> dict | None: return {"role": "user", "content": content} async def encode_image_bs64(self, image_url: str) -> tuple[str, str]: - """将图片转换为 base64,同时检测实际 MIME 类型""" + """将图片转换为 base64，同时检测实际 MIME 类型""" if image_url.startswith("base64://"): raw_base64 = image_url.replace("base64://", "") try: @@ -727,8 +726,8 @@ async def encode_image_bs64(self, image_url: str) -> tuple[str, str]: except Exception: mime_type = "image/jpeg" return f"data:{mime_type};base64,{raw_base64}", mime_type - async with aiofiles.open(image_url, "rb") as f: - image_bytes = await f.read() + with open(image_url, "rb") as f: + image_bytes = f.read() mime_type = self._detect_image_mime_type(image_bytes) image_bs64 = base64.b64encode(image_bytes).decode("utf-8") return f"data:{mime_type};base64,{image_bs64}", mime_type diff --git a/astrbot/core/provider/sources/gsvi_tts_source.py b/astrbot/core/provider/sources/gsvi_tts_source.py index f85c1b8fa4..9be718f979 100644 --- a/astrbot/core/provider/sources/gsvi_tts_source.py +++ b/astrbot/core/provider/sources/gsvi_tts_source.py @@ -1,75 +1,75 @@ -import uuid -from pathlib import Path - -import aiohttp - -from astrbot.core.provider.entities import ProviderType -from astrbot.core.provider.provider import TTSProvider -from astrbot.core.provider.register import register_provider_adapter -from astrbot.core.utils.astrbot_path import get_astrbot_temp_path - - -@register_provider_adapter( - "gsvi_tts_api", - "GSVI TTS API", - provider_type=ProviderType.TEXT_TO_SPEECH, -) -class ProviderGSVITTS(TTSProvider): - def __init__( - self, - provider_config: dict, - provider_settings: dict, - ) -> None: - super().__init__(provider_config, provider_settings) - self.api_key = provider_config.get("api_key", "") - self.api_base = provider_config.get("api_base", "http://127.0.0.1:8000") - self.api_base = self.api_base.removesuffix("/") - self.version = provider_config.get("version", "v4") - self.character = provider_config.get("character") - self.prompt_text_lang = provider_config.get("prompt_text_lang", "中文") - self.emotion = provider_config.get("emotion", "默认") - self.text_lang = provider_config.get("text_lang", "中文") - - async def get_audio(self, text: str) -> str: - temp_dir = get_astrbot_temp_path() - path = Path(temp_dir) / f"gsvi_tts_{uuid.uuid4()}.wav" - url = f"{self.api_base}/infer_single" - - headers = {"Content-Type": "application/json"} - if self.api_key: - headers["Authorization"] = f"Bearer {self.api_key}" - - data = { - "dl_url": self.api_base, - "version": self.version, - "model_name": self.character, - "prompt_text_lang": self.prompt_text_lang, - "emotion": self.emotion, - "text": text, - "text_lang": self.text_lang, - } - - async with aiohttp.ClientSession() as session: - async with session.post(url, json=data, headers=headers) as response: - if response.status == 200: - resp_json = await response.json() - msg = resp_json.get("msg") - audio_url = resp_json.get("audio_url") - if not msg or msg != "合成成功": - raise Exception(f"GSVI TTS API 合成失败: {msg}") - async with session.get(audio_url) as audio_response: - if audio_response.status == 200: - with open(path, "wb") as f: - f.write(await audio_response.read()) - else: - error_text = await audio_response.text() - raise Exception( - f"GSVI TTS API 下载音频失败，状态码: {audio_response.status}，错误: {error_text}", - ) - else: - error_text = await response.text() - raise Exception( - f"GSVI TTS API 请求失败,状态码: {response.status},错误: {error_text}", - ) - - return str(path) +import uuid +from pathlib import Path + +import aiohttp + +from astrbot.core.provider.entities import ProviderType +from astrbot.core.provider.provider import TTSProvider +from astrbot.core.provider.register import register_provider_adapter +from astrbot.core.utils.astrbot_path import get_astrbot_temp_path + + +@register_provider_adapter( + "gsvi_tts_api", + "GSVI TTS API", + provider_type=ProviderType.TEXT_TO_SPEECH, +) +class ProviderGSVITTS(TTSProvider): + def __init__( + self, + provider_config: dict, + provider_settings: dict, + ) -> None: + super().__init__(provider_config, provider_settings) + self.api_key = provider_config.get("api_key", "") + self.api_base = provider_config.get("api_base", "http://127.0.0.1:8000") + self.api_base = self.api_base.removesuffix("/") + self.version = provider_config.get("version", "v4") + self.character = provider_config.get("character") + self.prompt_text_lang = provider_config.get("prompt_text_lang", "中文") + self.emotion = provider_config.get("emotion", "默认") + self.text_lang = provider_config.get("text_lang", "中文") + + async def get_audio(self, text: str) -> str: + temp_dir = get_astrbot_temp_path() + path = Path(temp_dir) / f"gsvi_tts_{uuid.uuid4()}.wav" + url = f"{self.api_base}/infer_single" + + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + + data = { + "dl_url": self.api_base, + "version": self.version, + "model_name": self.character, + "prompt_text_lang": self.prompt_text_lang, + "emotion": self.emotion, + "text": text, + "text_lang": self.text_lang, + } + + async with aiohttp.ClientSession() as session: + async with session.post(url, json=data, headers=headers) as response: + if response.status == 200: + resp_json = await response.json() + msg = resp_json.get("msg") + audio_url = resp_json.get("audio_url") + if not msg or msg != "合成成功": + raise Exception(f"GSVI TTS API 合成失败: {msg}") + async with session.get(audio_url) as audio_response: + if audio_response.status == 200: + with open(path, "wb") as f: + f.write(await audio_response.read()) + else: + error_text = await audio_response.text() + raise Exception( + f"GSVI TTS API 下载音频失败，状态码: {audio_response.status}，错误: {error_text}", + ) + else: + error_text = await response.text() + raise Exception( + f"GSVI TTS API 请求失败,状态码: {response.status},错误: {error_text}", + ) + + return str(path) diff --git a/astrbot/dashboard/routes/chat.py b/astrbot/dashboard/routes/chat.py index c4a0556c6c..bb9528d850 100644 --- a/astrbot/dashboard/routes/chat.py +++ b/astrbot/dashboard/routes/chat.py @@ -12,6 +12,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 @@ -76,6 +79,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/batch_delete_sessions": ("POST", self.batch_delete_sessions), "/chat/update_session_display_name": ( @@ -102,6 +106,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: @@ -481,6 +497,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": @@ -599,19 +628,72 @@ async def stop_session(self): if session.creator != username: return Response().error("Permission denied").to_json() - 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}).to_json() + 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_session_internal(self, session, username: str) -> None: """Delete a single session and all its related data.""" session_id = session.session_id diff --git a/astrbot/dashboard/routes/live_chat.py b/astrbot/dashboard/routes/live_chat.py index 53153b1962..417d1e09e7 100644 --- a/astrbot/dashboard/routes/live_chat.py +++ b/astrbot/dashboard/routes/live_chat.py @@ -654,6 +654,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 8cd2d61def..9cfd54a078 100644 --- a/astrbot/dashboard/routes/open_api.py +++ b/astrbot/dashboard/routes/open_api.py @@ -473,6 +473,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/t2i.py b/astrbot/dashboard/routes/t2i.py index 3a63e7bcdc..7a95228744 100644 --- a/astrbot/dashboard/routes/t2i.py +++ b/astrbot/dashboard/routes/t2i.py @@ -19,7 +19,7 @@ def __init__( self.core_lifecycle = core_lifecycle self.config = core_lifecycle.astrbot_config self.manager = TemplateManager() - # 使用列表保证路由注册顺序,避免 / 路由优先匹配 /reset_default + # 使用列表保证路由注册顺序，避免 / 路由优先匹配 /reset_default self.routes = [ ("/t2i/templates", ("GET", self.list_templates)), ("/t2i/templates/active", ("GET", self.get_active_template)), @@ -148,7 +148,7 @@ async def update_template(self, name: str): self.manager.update_template(name, content) - # 检查更新的是否为当前激活的模板,如果是,则热重载 + # 检查更新的是否为当前激活的模板，如果是，则热重载 active_template = self.config.get("t2i_active_template", "base") if name == active_template: await self._reload_all_pipeline_schedulers() @@ -193,7 +193,7 @@ async def set_active_template(self): data = await request.json name = data.get("name") if not name: - response = jsonify(asdict(Response().error("模板名称(name)不能为空｡"))) + response = jsonify(asdict(Response().error("模板名称(name)不能为空。"))) response.status_code = 400 return response @@ -203,11 +203,11 @@ async def set_active_template(self): # 更新所有配置并热重载以应用更改 await self._sync_active_template_to_all_configs(name) - return jsonify(asdict(Response().ok(message=f"模板 '{name}' 已成功应用｡"))) + return jsonify(asdict(Response().ok(message=f"模板 '{name}' 已成功应用。"))) except FileNotFoundError: response = jsonify( - asdict(Response().error(f"模板 '{name}' 不存在,无法应用｡")), + asdict(Response().error(f"模板 '{name}' 不存在，无法应用。")), ) response.status_code = 404 return response diff --git a/astrbot/dashboard/routes/tools.py b/astrbot/dashboard/routes/tools.py index ee5eca22ef..672f58e535 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: diff --git a/dashboard/src/components/chat/MessageList.vue b/dashboard/src/components/chat/MessageList.vue index 8cbe70eade..6213dce985 100644 --- a/dashboard/src/components/chat/MessageList.vue +++ b/dashboard/src/components/chat/MessageList.vue @@ -208,6 +208,8 @@ :is-dark="isDark" :current-time="currentTime" :downloading-files="downloadingFiles" + :interactive-elicitation="isActiveElicitationMessage(index, msg)" + :submit-elicitation="submitElicitation" @open-image-preview="openImagePreview" @download-file="downloadFile" /> @@ -414,6 +416,10 @@ export default { type: Boolean, default: false, }, + submitElicitation: { + type: Function, + default: null, + }, }, emits: ["openImagePreview", "replyMessage", "replyWithText", "openRefs"], setup() { @@ -609,6 +615,18 @@ export default { ); }, + 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 b1dd2aa93d..fc159cba9b 100644 --- a/dashboard/src/components/chat/StandaloneChat.vue +++ b/dashboard/src/components/chat/StandaloneChat.vue @@ -10,6 +10,7 @@ :messages="messages" :is-dark="isDark" :is-streaming="isStreaming || isConvRunning" + :submit-elicitation="handleSubmitElicitation" @openImagePreview="openImagePreview" />
@@ -168,6 +169,7 @@ const { enableStreaming, getSessionMessages: getSessionMsg, sendMessage: sendMsg, + submitElicitationResponse, stopMessage: stopMsg, toggleStreaming, } = useMessages(currSessionId, getMediaFile, updateSessionTitle, getSessions); @@ -257,6 +259,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..e590f28185 --- /dev/null +++ b/dashboard/src/components/chat/message_list_comps/ElicitationCard.vue @@ -0,0 +1,398 @@ + + + + + diff --git a/dashboard/src/components/chat/message_list_comps/MessagePartsRenderer.vue b/dashboard/src/components/chat/message_list_comps/MessagePartsRenderer.vue index a61bb19b30..39cfa4f022 100644 --- a/dashboard/src/components/chat/message_list_comps/MessagePartsRenderer.vue +++ b/dashboard/src/components/chat/message_list_comps/MessagePartsRenderer.vue @@ -64,25 +64,16 @@
- -
- -
+ :payload="renderPart.part.payload" + :is-dark="isDark" + :interactive="interactiveElicitation" + :submit-elicitation="submitElicitation" + />
@@ -213,6 +205,7 @@ - - diff --git a/dashboard/src/views/dashboards/default/components/MemoryUsage.vue b/dashboard/src/views/dashboards/default/components/MemoryUsage.vue deleted file mode 100644 index dad27bf32f..0000000000 --- a/dashboard/src/views/dashboards/default/components/MemoryUsage.vue +++ /dev/null @@ -1,212 +0,0 @@ - - - - - diff --git a/dashboard/src/views/dashboards/default/components/MessageStat.vue b/dashboard/src/views/dashboards/default/components/MessageStat.vue deleted file mode 100644 index 3ecc2f824f..0000000000 --- a/dashboard/src/views/dashboards/default/components/MessageStat.vue +++ /dev/null @@ -1,428 +0,0 @@ - - - - - diff --git a/dashboard/src/views/dashboards/default/components/OnlinePlatform.vue b/dashboard/src/views/dashboards/default/components/OnlinePlatform.vue deleted file mode 100644 index 25e61c7019..0000000000 --- a/dashboard/src/views/dashboards/default/components/OnlinePlatform.vue +++ /dev/null @@ -1,179 +0,0 @@ - - - - - diff --git a/dashboard/src/views/dashboards/default/components/OnlineTime.vue b/dashboard/src/views/dashboards/default/components/OnlineTime.vue deleted file mode 100644 index 3b1564933b..0000000000 --- a/dashboard/src/views/dashboards/default/components/OnlineTime.vue +++ /dev/null @@ -1,260 +0,0 @@ - - - - - diff --git a/dashboard/src/views/dashboards/default/components/PlatformStat.vue b/dashboard/src/views/dashboards/default/components/PlatformStat.vue deleted file mode 100644 index eff60cb71e..0000000000 --- a/dashboard/src/views/dashboards/default/components/PlatformStat.vue +++ /dev/null @@ -1,299 +0,0 @@ - - - - - diff --git a/dashboard/src/views/dashboards/default/components/RunningTime.vue b/dashboard/src/views/dashboards/default/components/RunningTime.vue deleted file mode 100644 index 9ba079531e..0000000000 --- a/dashboard/src/views/dashboards/default/components/RunningTime.vue +++ /dev/null @@ -1,146 +0,0 @@ - - - - - diff --git a/dashboard/src/views/dashboards/default/components/TotalMessage.vue b/dashboard/src/views/dashboards/default/components/TotalMessage.vue deleted file mode 100644 index 97da00bbaa..0000000000 --- a/dashboard/src/views/dashboards/default/components/TotalMessage.vue +++ /dev/null @@ -1,116 +0,0 @@ - - - - - diff --git a/dashboard/src/views/persona/PersonaCard.vue b/dashboard/src/views/persona/PersonaCard.vue index cd6de044c6..dc915ac348 100644 --- a/dashboard/src/views/persona/PersonaCard.vue +++ b/dashboard/src/views/persona/PersonaCard.vue @@ -227,6 +227,10 @@ export default defineComponent({ transform: translateY(-2px); } +.persona-card:hover { + transform: translateY(-2px); +} + .system-prompt-preview { font-size: 14px; line-height: 1.4; 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/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"