From 3343522fa2e62b1a1bb7df430affcba647103daf Mon Sep 17 00:00:00 2001 From: spud <92900806+jamespud@users.noreply.github.com> Date: Tue, 13 Oct 2026 08:38:36 +0800 Subject: [PATCH] refactor(mcp): expose per-server tool discovery results (#6734) * refactor(mcp): expose per-server tool discovery results * refactor(mcp): clarify per-server discovery contracts --- backend/docs/MCP_SERVER.md | 8 + .../packages/harness/deerflow/mcp/oauth.py | 5 + .../packages/harness/deerflow/mcp/tools.py | 716 ++++++++++++------ backend/tests/test_mcp_custom_interceptors.py | 69 +- .../tests/test_mcp_per_server_discovery.py | 713 +++++++++++++++++ backend/tests/test_mcp_session_pool.py | 20 +- backend/tests/test_mcp_sync_wrapper.py | 49 +- .../tests/test_mcp_tool_name_validation.py | 15 +- 8 files changed, 1347 insertions(+), 248 deletions(-) create mode 100644 backend/tests/test_mcp_per_server_discovery.py diff --git a/backend/docs/MCP_SERVER.md b/backend/docs/MCP_SERVER.md index efa9879b66..7a5e1aabc8 100644 --- a/backend/docs/MCP_SERVER.md +++ b/backend/docs/MCP_SERVER.md @@ -576,6 +576,14 @@ Declare interceptors in `extensions_config.json` using the `mcpInterceptors` fie Each entry is a Python import path in `module:variable` format (resolved via `resolve_variable`). The variable must be a **no-arg builder function** that returns an async interceptor compatible with `MultiServerMCPClient`’s `tool_interceptors` interface, or `None` to skip. +**Lifecycle.** Builders run once per enabled MCP server, not once per discovery. +Tool discovery is server-scoped, so every server gets its own interceptor instance built +from a single-server configuration view, and a stateful interceptor (a counter, a cache, +a rate limiter) only ever observes that server's calls. Do not rely on one instance being +shared across servers. The built-in OAuth, per-user and context-header interceptors are +unaffected: they resolve everything per `request.server_name`. Servers that are not +discovered never run their builder. + Example interceptor that injects an authorization header from the request-scoped LangGraph secret context. For a plain header mapping prefer the declarative `headers_from_context` block above; write an interceptor when the header value diff --git a/backend/packages/harness/deerflow/mcp/oauth.py b/backend/packages/harness/deerflow/mcp/oauth.py index 38d3f9f215..16ee5eafc0 100644 --- a/backend/packages/harness/deerflow/mcp/oauth.py +++ b/backend/packages/harness/deerflow/mcp/oauth.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio import logging import threading +from collections.abc import Collection from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta from typing import Any @@ -281,6 +282,7 @@ async def get_initial_oauth_headers( extensions_config: ExtensionsConfig, *, token_manager: OAuthTokenManager | None = None, + server_names: Collection[str] | None = None, ) -> dict[str, str]: """Get initial OAuth Authorization headers for MCP server connections.""" token_manager = token_manager or OAuthTokenManager.from_extensions_config(extensions_config) @@ -288,7 +290,10 @@ async def get_initial_oauth_headers( return {} headers: dict[str, str] = {} + requested = set(server_names) if server_names is not None else None for server_name in token_manager.oauth_server_names(): + if requested is not None and server_name not in requested: + continue try: value = await token_manager.get_authorization_header(server_name) except Exception: diff --git a/backend/packages/harness/deerflow/mcp/tools.py b/backend/packages/harness/deerflow/mcp/tools.py index a85f436b7c..0d24aa563a 100644 --- a/backend/packages/harness/deerflow/mcp/tools.py +++ b/backend/packages/harness/deerflow/mcp/tools.py @@ -5,7 +5,8 @@ from __future__ import annotations import asyncio import logging import re -from collections.abc import Iterable, Mapping +from collections.abc import Collection, Iterable, Mapping +from dataclasses import dataclass, field from datetime import timedelta from pathlib import Path from typing import Any @@ -26,6 +27,7 @@ from deerflow.mcp.session_pool import ( MCPPoolDomain, MCPSessionPool, ServerBinding, + StaleMCPBindingError, call_pooled_session_tool, get_session_pool, normalized_connection_fingerprint, @@ -929,6 +931,457 @@ def _configure_task_tools_for_server( return configured +@dataclass(frozen=True, slots=True) +class ServerDiscoveryResult: + """Successful tool discovery for one MCP server. + + A present entry with ``tools=()`` is a successful discovery of zero tools; + a failed discovery is represented by the server's absence from the grouped + mapping. ``pool``/``binding`` identify the exact stdio session owner a + result was produced against and stay out of ``repr`` because they are + secret-adjacent ownership handles. + """ + + tools: tuple[BaseTool, ...] + pool: MCPSessionPool | None = field(default=None, repr=False) + binding: ServerBinding | None = field(default=None, repr=False) + + +def _flatten_server_tool_groups( + enabled_order: Collection[str], + groups: Mapping[str, ServerDiscoveryResult], +) -> list[BaseTool]: + """Flatten grouped results in enabled-server declaration order.""" + return [tool for name in enabled_order if name in groups for tool in groups[name].tools] + + +def _load_mcp_adapter_symbols() -> tuple[Any, Any] | None: + """Load the adapter symbols needed for discovery, or fail closed.""" + try: + from langchain_mcp_adapters.client import MultiServerMCPClient + from langchain_mcp_adapters.tools import load_mcp_tools + except ImportError: + logger.warning("langchain-mcp-adapters not installed. Install it to enable MCP tools: pip install langchain-mcp-adapters") + return None + return MultiServerMCPClient, load_mcp_tools + + +@dataclass(frozen=True, slots=True) +class _RawServerDiscovery: + tools: list[BaseTool] + tool_interceptors: list[Any] + + +@dataclass(frozen=True, slots=True) +class _PreparedServerClient: + server_name: str + connection: dict[str, Any] + client: Any + tool_interceptors: list[Any] + tool_name_prefix: bool + session_init_timeout: float | None + + +@dataclass(frozen=True, slots=True) +class _PreparedServerDiscovery: + config: ExtensionsConfig + server_names: tuple[str, ...] + connections: dict[str, dict[str, Any]] + pool: MCPSessionPool + bindings: dict[str, ServerBinding] + ownership_domain: MCPPoolDomain + connection_scope: str + + +def _prepare_server_client( + extensions_config: ExtensionsConfig, + server_name: str, + connection: dict[str, Any], + *, + client_cls: Any, + initial_oauth_header: str | None, + oauth_token_manager: OAuthTokenManager, +) -> _PreparedServerClient: + """Build the client/interceptors for one selected server.""" + server_cfg = extensions_config.mcp_servers.get(server_name) + transport = connection.get("transport", "stdio") + tool_name_prefix = server_cfg.tool_name_prefix if server_cfg is not None else True + session_init_timeout = _resolve_session_init_timeout(server_cfg) + server_view = extensions_config.model_copy( + update={"mcp_servers": {server_name: extensions_config.mcp_servers[server_name]}}, + ) + tool_interceptors = build_mcp_tool_interceptors( + server_view, + oauth_builder=lambda config: build_oauth_tool_interceptor(config, token_manager=oauth_token_manager), + resolver=resolve_variable, + target_logger=logger, + ) + if initial_oauth_header and transport in ("sse", "http"): + # Case-insensitive write: a static header spelled 'authorization' + # must be replaced, not joined on the wire by a second field. + connection["headers"] = apply_header_overrides( + connection.get("headers", {}), + {"Authorization": initial_oauth_header}, + ) + + client = client_cls( + {server_name: connection}, + tool_interceptors=tool_interceptors, + tool_name_prefix=True, + ) + return _PreparedServerClient( + server_name=server_name, + connection=connection, + client=client, + tool_interceptors=tool_interceptors, + tool_name_prefix=tool_name_prefix, + session_init_timeout=session_init_timeout, + ) + + +async def _discover_server_tools( + prepared_client: _PreparedServerClient, + *, + load_mcp_tools_fn: Any, +) -> tuple[str, _RawServerDiscovery | None]: + """Run only the discovery await for one selected MCP server.""" + server_name = prepared_client.server_name + connection = prepared_client.connection + client = prepared_client.client + + try: + if prepared_client.tool_name_prefix: + discovery = client.get_tools(server_name=server_name) + else: + discovery = load_mcp_tools_fn( + None, + connection=connection, + callbacks=client.callbacks, + server_name=server_name, + tool_interceptors=client.tool_interceptors, + tool_name_prefix=False, + ) + + if prepared_client.session_init_timeout is not None: + # Timeout tool discovery (subprocess spawn + initialize + + # tools/list) so a hung stdio server cannot block agent + # construction indefinitely. This runs per server so one slow + # server does not prevent healthy siblings from contributing. + try: + server_tools = await asyncio.wait_for(discovery, timeout=prepared_client.session_init_timeout) + except TimeoutError: + logger.warning( + "Skipping MCP server '%s' after tool discovery timed out (%.1fs)", + server_name, + prepared_client.session_init_timeout, + ) + return server_name, None + else: + server_tools = await discovery + + return server_name, _RawServerDiscovery( + tools=list(server_tools), + tool_interceptors=prepared_client.tool_interceptors, + ) + except (asyncio.CancelledError, StaleMCPBindingError, McpTaskConfigurationError): + raise + except Exception as e: + logger.warning( + f"Skipping MCP server '{server_name}' after tool discovery failed: {e}", + exc_info=True, + ) + return server_name, None + + +def _wrap_discovered_server_tools( + extensions_config: ExtensionsConfig, + server_name: str, + connection: dict[str, Any], + raw_discovery: _RawServerDiscovery, + *, + tool_name_prefix: bool, + session_init_timeout: float | None, + pool: MCPSessionPool | None, + binding: ServerBinding | None, + ownership_domain: MCPPoolDomain, + connection_scope: str, +) -> ServerDiscoveryResult: + """Validate, tag, wrap, and task-configure one successful discovery. + + ``tool_name_prefix`` and ``session_init_timeout`` are supplied by the + prepared client so discovery and wrapping share one derivation of the + server config instead of two that can drift apart. + """ + server_cfg = extensions_config.mcp_servers.get(server_name) + transport = connection.get("transport", "stdio") + current_server_tools: list[BaseTool] = [] + + for tool in raw_discovery.tools: + if not _VALID_MCP_TOOL_NAME.fullmatch(tool.name or ""): + logger.warning( + "Dropping MCP tool from server '%s' with invalid name %r: tool names must match %s. A name outside this charset cannot be bound as a function tool and could forge prompt structure when listed as a deferred tool.", + server_name, + tool.name, + _VALID_MCP_TOOL_NAME.pattern, + ) + continue + tag_mcp_tool(tool, server_name=server_name, transport=transport) + prefix = f"{server_name}_" + original_name = tool.name[len(prefix) :] if tool_name_prefix and tool.name.startswith(prefix) else tool.name + routing = resolve_effective_mcp_routing(server_cfg, original_name) + if routing.get("mode") != "off": + tag_mcp_routing(tool, routing) + if transport == "stdio": + _timeout = server_cfg.tool_call_timeout if server_cfg else None + current_server_tools.append( + _make_session_pool_tool( + tool, + server_name, + connection, + raw_discovery.tool_interceptors, + tool_call_timeout=_timeout, + session_init_timeout=session_init_timeout, + tool_name_prefix=tool_name_prefix, + ownership_domain=ownership_domain, + pool=pool, + binding=binding, + ) + ) + else: + if server_cfg and server_cfg.tool_call_timeout is not None: + logger.warning( + "Ignoring tool_call_timeout for MCP server '%s' because transport '%s' is not stdio; configure HTTP/SSE transport-level timeouts instead.", + server_name, + transport, + ) + current_server_tools.append(tool) + + if server_cfg is not None: + current_server_tools = _configure_task_tools_for_server( + current_server_tools, + server_name=server_name, + server_config=server_cfg, + tool_name_prefix=tool_name_prefix, + connection_scope=connection_scope, + ) + + # Grouped results are directly usable by the cache/adapters that consume + # them, so sync invocation is patched inside each server's group. + for tool in current_server_tools: + if getattr(tool, "func", None) is None and getattr(tool, "coroutine", None) is not None: + tool.func = make_sync_tool_wrapper(tool.coroutine, tool.name) + + result_pool = pool if transport == "stdio" else None + result_binding = binding if transport == "stdio" else None + return ServerDiscoveryResult( + tools=tuple(current_server_tools), + pool=result_pool, + binding=result_binding, + ) + + +async def _resolve_discovery_config( + extensions_config: ExtensionsConfig | None, + personal_user_id: str | None, +) -> ExtensionsConfig: + """Resolve and authorize config outside the public discovery error handler.""" + if extensions_config is None: + # NOTE: We use ExtensionsConfig.from_file() instead of get_extensions_config() + # to always read the latest configuration from disk. This ensures that changes + # made through the Gateway API (which runs in a separate process) are immediately + # reflected when initializing MCP tools. Callers that need to prove which + # revision produced these tools pass the instance they snapshotted instead. + extensions_config = ExtensionsConfig.from_file() + + if personal_user_id is None: + validate_mcp_task_config_snapshot(extensions_config) + return extensions_config + + from deerflow.mcp.personal_access import authorized_personal_config + from deerflow.mcp.tasks.runtime import is_mcp_task_runtime_available + from deerflow.mcp.user_config import load_user_mcp_config + + current = await asyncio.to_thread(load_user_mcp_config, personal_user_id) + if current != extensions_config: + raise McpTaskConfigurationError("Personal MCP configuration changed during discovery; retry the run") + extensions_config = current = await authorized_personal_config(personal_user_id, current) + if any(server.task_toolsets for server in current.mcp_servers.values()) and not is_mcp_task_runtime_available(): + raise McpTaskConfigurationError("Personal MCP task toolsets require the platform's durable task runtime") + return extensions_config + + +def _prepare_server_discovery( + extensions_config: ExtensionsConfig, + *, + server_names: Collection[str] | None, + personal_user_id: str | None, + session_pool: MCPSessionPool | None, +) -> _PreparedServerDiscovery | None: + """Build selected connections and capture bindings before discovery.""" + all_enabled = extensions_config.get_enabled_mcp_servers() + requested = set(all_enabled) if server_names is None else set(server_names) & set(all_enabled) + selected_names = [name for name in all_enabled if name in requested] + if not selected_names: + # Distinguish "nothing is enabled" from "the caller's selection matched + # nothing". Unknown or disabled names are dropped per the documented + # selection contract, but that must not be reported as an empty config. + if server_names is None: + logger.info("No enabled MCP servers configured") + else: + logger.info( + "No enabled MCP servers matched the requested selection: %s", + sorted(set(server_names)), + ) + return None + + selected_view = extensions_config.model_copy( + update={"mcp_servers": {name: all_enabled[name] for name in selected_names}}, + ) + connections = build_servers_config(selected_view) + attempted_names = list(selected_names) + selected_names = [name for name in selected_names if name in connections] + if not connections: + logger.info("No valid MCP connections for selected servers: %s", sorted(attempted_names)) + return None + + pool = session_pool + if pool is None: + pool = get_session_pool() + ownership_domain: MCPPoolDomain = "personal" if personal_user_id is not None else "deployment" + server_bindings: dict[str, ServerBinding] = {} + for server_name, server_connection in connections.items(): + if server_connection.get("transport", "stdio") == "stdio": + server_bindings[server_name] = pool.ensure_binding( + server_name, + normalized_connection_fingerprint(server_connection), + domain=ownership_domain, + ) + + return _PreparedServerDiscovery( + config=extensions_config, + server_names=tuple(selected_names), + connections=connections, + pool=pool, + bindings=server_bindings, + ownership_domain=ownership_domain, + connection_scope="personal" if personal_user_id is not None else "deployment", + ) + + +async def _discover_and_wrap_prepared( + adapter_symbols: tuple[Any, Any], + prepared: _PreparedServerDiscovery, +) -> dict[str, ServerDiscoveryResult]: + """Run the OAuth/discovery phase, then wrap successful results in order.""" + client_cls, load_mcp_tools_fn = adapter_symbols + extensions_config = prepared.config + oauth_token_manager = OAuthTokenManager.from_extensions_config(extensions_config) if prepared.connection_scope == "personal" else get_mcp_task_oauth_token_manager(extensions_config) + + # Only selected, valid connections reach this point. Selected-but-invalid + # servers are intentionally not contacted for OAuth headers (F1 contract). + initial_oauth_headers = await get_initial_oauth_headers( + extensions_config, + token_manager=oauth_token_manager, + server_names=set(prepared.connections), + ) + + prepared_clients = [ + _prepare_server_client( + extensions_config, + server_name, + prepared.connections[server_name], + client_cls=client_cls, + initial_oauth_header=initial_oauth_headers.get(server_name), + oauth_token_manager=oauth_token_manager, + ) + for server_name in prepared.server_names + ] + + pairs = await asyncio.gather( + *( + _discover_server_tools( + prepared_client, + load_mcp_tools_fn=load_mcp_tools_fn, + ) + for prepared_client in prepared_clients + ) + ) + raw_by_server = {server_name: raw for server_name, raw in pairs if raw is not None} + prepared_by_server = {prepared_client.server_name: prepared_client for prepared_client in prepared_clients} + + groups: dict[str, ServerDiscoveryResult] = {} + for server_name in prepared.server_names: + raw_discovery = raw_by_server.get(server_name) + if raw_discovery is None: + continue + prepared_client = prepared_by_server[server_name] + groups[server_name] = _wrap_discovered_server_tools( + extensions_config, + server_name, + prepared_client.connection, + raw_discovery, + tool_name_prefix=prepared_client.tool_name_prefix, + session_init_timeout=prepared_client.session_init_timeout, + pool=prepared.pool, + binding=prepared.bindings.get(server_name), + ownership_domain=prepared.ownership_domain, + connection_scope=prepared.connection_scope, + ) + + logger.info( + "Discovered %d tool(s) from %d MCP server(s)", + sum(len(result.tools) for result in groups.values()), + len(groups), + ) + return groups + + +async def get_mcp_tools_by_server( + extensions_config: ExtensionsConfig | None = None, + *, + server_names: Collection[str] | None = None, + personal_user_id: str | None = None, + session_pool: MCPSessionPool | None = None, +) -> dict[str, ServerDiscoveryResult]: + """Discover selected enabled MCP servers independently. + + Args: + extensions_config: Optional pre-loaded config. ``None`` reads the latest + revision from disk; callers that must prove which revision produced + the tools pass the exact snapshotted instance. + server_names: Restrict discovery to these server names. Names that are + unknown or currently disabled are ignored, because a cache caller's + selection can legitimately lag the latest config; an effective + selection of nothing returns ``{}`` without contacting any server. + ``None`` selects every enabled server. + personal_user_id: When set, the caller's personal MCP configuration is + authorized and used instead of the deployment configuration. + session_pool: Optional session pool captured by the caller. Cache + initializers pass the exact pool their generation owns so a + superseded claim cannot install bindings into a replacement pool; + ``None`` resolves the current singleton. + + Returns: + One entry per server whose discovery succeeded, keyed by server name. + A server whose discovery failed or timed out is **absent**; a server + that succeeded with zero tools is **present** with ``tools=()``. + """ + adapter_symbols = _load_mcp_adapter_symbols() + if adapter_symbols is None: + return {} + + extensions_config = await _resolve_discovery_config(extensions_config, personal_user_id) + prepared = _prepare_server_discovery( + extensions_config, + server_names=server_names, + personal_user_id=personal_user_id, + session_pool=session_pool, + ) + if prepared is None: + return {} + return await _discover_and_wrap_prepared(adapter_symbols, prepared) + + async def get_mcp_tools( extensions_config: ExtensionsConfig | None = None, *, @@ -937,247 +1390,46 @@ async def get_mcp_tools( ) -> list[BaseTool]: """Get all tools from enabled MCP servers. - Tools using stdio transport are wrapped with persistent-session logic so - consecutive calls within the same thread reuse the same MCP session. - HTTP/SSE tools are returned unwrapped to avoid cross-task TaskGroup - cleanup errors. + Config resolution, personal authorization, connection construction, and + binding capture stay outside the broad discovery-phase handler so they keep + the base API's exception semantics. Only OAuth/discovery/wrapping is + softened by that handler. Args: - extensions_config: Optional pre-loaded extensions config. Callers that - must prove which config revision produced these tools pass the exact - instance they snapshotted; ``None`` loads the latest config from disk. - session_pool: Optional session pool captured by the caller. Cache - initializers pass the exact pool their generation owns so a - superseded claim cannot install bindings into a replacement pool; - ``None`` resolves the current singleton. + extensions_config: Optional pre-loaded config. ``None`` reads the latest + revision from disk, so Gateway API edits are picked up without a + restart. Callers that must prove which revision produced these tools + (the cache initializer) pass the exact instance they snapshotted. + personal_user_id: When set, that user's authorized personal MCP + configuration is discovered instead of the deployment one. + session_pool: Optional session pool pinned by the caller. Cache + initializers pass the exact pool their initialization generation + owns, so a superseded claim cannot install bindings into a + replacement pool; ``None`` resolves the current singleton. Returns: - List of LangChain tools from all enabled MCP servers. + Wrapped tools from every enabled server, in enabled-server declaration + order. Returns an empty list when the adapter is unavailable or when a + discovery-phase failure is caught by the handler described above. """ - try: - from langchain_mcp_adapters.client import MultiServerMCPClient - from langchain_mcp_adapters.tools import load_mcp_tools - except ImportError: - logger.warning("langchain-mcp-adapters not installed. Install it to enable MCP tools: pip install langchain-mcp-adapters") + adapter_symbols = _load_mcp_adapter_symbols() + if adapter_symbols is None: return [] - if extensions_config is None: - # NOTE: We use ExtensionsConfig.from_file() instead of get_extensions_config() - # to always read the latest configuration from disk. This ensures that changes - # made through the Gateway API (which runs in a separate process) are immediately - # reflected when initializing MCP tools. Callers that need to prove which - # revision produced these tools pass the instance they snapshotted instead. - extensions_config = ExtensionsConfig.from_file() - if personal_user_id is None: - validate_mcp_task_config_snapshot(extensions_config) - else: - from deerflow.mcp.personal_access import authorized_personal_config - from deerflow.mcp.tasks.runtime import is_mcp_task_runtime_available - from deerflow.mcp.user_config import load_user_mcp_config - - current = await asyncio.to_thread(load_user_mcp_config, personal_user_id) - if current != extensions_config: - raise McpTaskConfigurationError("Personal MCP configuration changed during discovery; retry the run") - extensions_config = current = await authorized_personal_config(personal_user_id, current) - if any(server.task_toolsets for server in current.mcp_servers.values()) and not is_mcp_task_runtime_available(): - raise McpTaskConfigurationError("Personal MCP task toolsets require the platform's durable task runtime") - servers_config = build_servers_config(extensions_config) - - if not servers_config: - logger.info("No enabled MCP servers configured") + extensions_config = await _resolve_discovery_config(extensions_config, personal_user_id) + prepared = _prepare_server_discovery( + extensions_config, + server_names=None, + personal_user_id=personal_user_id, + session_pool=session_pool, + ) + if prepared is None: return [] - # Capture the exact pool and stdio binding before the first discovery await. - # A later reset/reconcile can supersede this capability, but a stale wrapper - # can never silently pair its old connection with a replacement pool/epoch. - pool = session_pool - if pool is None: - pool = get_session_pool() - ownership_domain: MCPPoolDomain = "personal" if personal_user_id is not None else "deployment" - server_bindings: dict[str, ServerBinding] = {} - for server_name, server_connection in servers_config.items(): - if server_connection.get("transport", "stdio") == "stdio": - server_bindings[server_name] = pool.ensure_binding( - server_name, - normalized_connection_fingerprint(server_connection), - domain=ownership_domain, - ) - try: - # Create the multi-server MCP client - logger.info(f"Initializing MCP client with {len(servers_config)} server(s)") - - # Inject initial OAuth headers for server connections (tool discovery/session init) - oauth_token_manager = OAuthTokenManager.from_extensions_config(extensions_config) if personal_user_id is not None else get_mcp_task_oauth_token_manager(extensions_config) - initial_oauth_headers = await get_initial_oauth_headers(extensions_config, token_manager=oauth_token_manager) - for server_name, auth_header in initial_oauth_headers.items(): - if server_name not in servers_config: - continue - if servers_config[server_name].get("transport") in ("sse", "http"): - # Case-insensitive write: a static header spelled 'authorization' - # must be replaced, not joined on the wire by a second field. - servers_config[server_name]["headers"] = apply_header_overrides( - servers_config[server_name].get("headers", {}), - {"Authorization": auth_header}, - ) - - tool_interceptors = build_mcp_tool_interceptors( - extensions_config, - oauth_builder=lambda config: build_oauth_tool_interceptor(config, token_manager=oauth_token_manager), - resolver=resolve_variable, - target_logger=logger, - ) - - client = MultiServerMCPClient( - servers_config, - tool_interceptors=tool_interceptors, - tool_name_prefix=True, - ) - - async def load_server_tools(server_name: str) -> list[BaseTool] | None: - # ``None`` marks a server whose discovery failed, which is distinct - # from a server that answered with zero tools: only the latter is - # evidence that configured task_toolsets name missing tools. - try: - server_cfg = extensions_config.mcp_servers.get(server_name) - tool_name_prefix = server_cfg.tool_name_prefix if server_cfg is not None else True - session_init_timeout = _resolve_session_init_timeout(server_cfg) - if tool_name_prefix: - discovery = client.get_tools(server_name=server_name) - else: - discovery = load_mcp_tools( - None, - connection=servers_config[server_name], - callbacks=client.callbacks, - server_name=server_name, - tool_interceptors=client.tool_interceptors, - tool_name_prefix=False, - ) - if session_init_timeout is not None: - # Timeout tool discovery (subprocess spawn + initialize + - # tools/list) so a hung stdio server cannot block agent - # construction indefinitely. Per-server because the gather - # below runs each server independently — one slow server - # must not prevent the others from contributing tools. - # - # Cancellation here is safe: discovery runs inside the - # adapter's nested async context managers (load_mcp_tools → - # create_session → _create_stdio_session → stdio_client), - # and wait_for's CancelledError unwinds them. stdio_client's - # finally closes stdin, waits for a graceful exit, then - # escalates to _terminate_process_tree (SIGTERM→SIGKILL on - # POSIX, process-tree termination on Windows), so the npx - # subprocess and any children it spawned are reaped — no - # orphan processes accumulate across repeated timeouts. - try: - return await asyncio.wait_for(discovery, timeout=session_init_timeout) - except TimeoutError: - # Only our own bound is logged as "timed out": the - # branch condition guarantees the value is not None, so - # the %.1f format cannot fail. A TimeoutError raised by - # discovery itself (e.g. an internal SDK timeout on the - # opted-out path) falls through to the generic failure - # handler below instead. - logger.warning( - "Skipping MCP server '%s' after tool discovery timed out (%.1fs)", - server_name, - session_init_timeout, - ) - return None - return await discovery - except Exception as e: - logger.warning( - f"Skipping MCP server '{server_name}' after tool discovery failed: {e}", - exc_info=True, - ) - return None - - # Get tools from each server independently so one broken MCP server does - # not prevent healthy servers from contributing their tools. - tools_by_server = await asyncio.gather(*(load_server_tools(name) for name in servers_config)) - tools = [tool for server_tools in tools_by_server if server_tools is not None for tool in server_tools] - logger.info(f"Successfully loaded {len(tools)} tool(s) from MCP servers") - - # Wrap each tool with persistent-session logic. - # Only pool stdio sessions. HTTP/SSE transports use anyio TaskGroups - # internally which cannot be closed from a different async task, so - # pooling them causes RuntimeError on cleanup (see #3203). - wrapped_tools: list[BaseTool] = [] - # Route each tool by the server that actually produced it: tools_by_server[i] - # corresponds to the i-th server in servers_config. Inferring the source server by - # scanning servers_config for a name prefix is ambiguous when one server name is a - # prefix of another (e.g. "web" vs "web_scraper" → "web_scraper_search".startswith( - # "web_") matches "web" first), which pools the tool under the wrong server. Using the - # source grouping makes routing exact even when a server opts out of name prefixing. - for source_name, server_tools in zip(servers_config.keys(), tools_by_server, strict=True): - if server_tools is None: - # Discovery already logged the skip. Validating task_toolsets - # against a server that never answered would report every - # raw tool as missing and discard the healthy servers' tools. - continue - transport = servers_config[source_name].get("transport", "stdio") - server_cfg = extensions_config.mcp_servers.get(source_name) - tool_name_prefix = server_cfg.tool_name_prefix if server_cfg is not None else True - current_server_tools: list[BaseTool] = [] - for tool in server_tools: - if not _VALID_MCP_TOOL_NAME.fullmatch(tool.name or ""): - logger.warning( - "Dropping MCP tool from server '%s' with invalid name %r: tool names must match %s. A name outside this charset cannot be bound as a function tool and could forge prompt structure when listed as a deferred tool.", - source_name, - tool.name, - _VALID_MCP_TOOL_NAME.pattern, - ) - continue - tag_mcp_tool(tool, server_name=source_name, transport=transport) - prefix = f"{source_name}_" - original_name = tool.name[len(prefix) :] if tool_name_prefix and tool.name.startswith(prefix) else tool.name - routing = resolve_effective_mcp_routing(server_cfg, original_name) - if routing.get("mode") != "off": - tag_mcp_routing(tool, routing) - if transport == "stdio": - _timeout = server_cfg.tool_call_timeout if server_cfg else None - _init_timeout = _resolve_session_init_timeout(server_cfg) - current_server_tools.append( - _make_session_pool_tool( - tool, - source_name, - servers_config[source_name], - tool_interceptors, - tool_call_timeout=_timeout, - session_init_timeout=_init_timeout, - tool_name_prefix=tool_name_prefix, - ownership_domain=ownership_domain, - pool=pool, - binding=server_bindings[source_name], - ) - ) - else: - if transport != "stdio" and server_cfg and server_cfg.tool_call_timeout is not None: - logger.warning( - "Ignoring tool_call_timeout for MCP server '%s' because transport '%s' is not stdio; configure HTTP/SSE transport-level timeouts instead.", - source_name, - transport, - ) - current_server_tools.append(tool) - - if server_cfg is not None: - current_server_tools = _configure_task_tools_for_server( - current_server_tools, - server_name=source_name, - server_config=server_cfg, - tool_name_prefix=tool_name_prefix, - connection_scope="personal" if personal_user_id is not None else "deployment", - ) - wrapped_tools.extend(current_server_tools) - - # Patch tools to support sync invocation, as deerflow client streams synchronously - for tool in wrapped_tools: - if getattr(tool, "func", None) is None and getattr(tool, "coroutine", None) is not None: - tool.func = make_sync_tool_wrapper(tool.coroutine, tool.name) - - return wrapped_tools - - except McpTaskConfigurationError: + groups = await _discover_and_wrap_prepared(adapter_symbols, prepared) + return _flatten_server_tool_groups(extensions_config.get_enabled_mcp_servers(), groups) + except (McpTaskConfigurationError, StaleMCPBindingError): raise except Exception as e: logger.error(f"Failed to load MCP tools: {e}", exc_info=True) diff --git a/backend/tests/test_mcp_custom_interceptors.py b/backend/tests/test_mcp_custom_interceptors.py index 08432de98d..1c5236f13b 100644 --- a/backend/tests/test_mcp_custom_interceptors.py +++ b/backend/tests/test_mcp_custom_interceptors.py @@ -3,6 +3,7 @@ import asyncio from unittest.mock import AsyncMock, MagicMock, patch +from deerflow.config.extensions_config import ExtensionsConfig from deerflow.mcp.tools import get_mcp_tools @@ -14,9 +15,17 @@ def _make_patches(*, interceptor_paths=None): mock_client = MagicMock() mock_client.get_tools = AsyncMock(return_value=[]) - extra = {} + config_payload = { + "mcpServers": { + "test-server": { + "type": "http", + "url": "https://example.test/mcp", + } + } + } if interceptor_paths is not None: - extra["mcpInterceptors"] = interceptor_paths + config_payload["mcpInterceptors"] = interceptor_paths + extensions_config = ExtensionsConfig.model_validate(config_payload) return { "client_cls": patch( @@ -25,10 +34,7 @@ def _make_patches(*, interceptor_paths=None): ), "from_file": patch( "deerflow.config.extensions_config.ExtensionsConfig.from_file", - return_value=MagicMock( - model_extra=extra, - get_enabled_mcp_servers=MagicMock(return_value={}), - ), + return_value=extensions_config, ), "build_servers": patch( "deerflow.mcp.tools.build_servers_config", @@ -272,3 +278,54 @@ def test_custom_interceptor_non_callable_return_logs_warning(): assert len(_get_interceptors(mock_cls)) == 0 mock_warn.assert_called_once() assert "non-callable" in mock_warn.call_args[0][0] + + +def test_custom_interceptor_builder_runs_once_per_enabled_server(): + """Builders are per server, so instances are never shared across servers. + + Locks the F1 lifecycle contract (see ``docs/MCP_SERVER.md``): server-scoped + discovery builds interceptors from a single-server config view. If this ever + regresses to one shared instance per discovery, the F2 server-scoped cache + would hand a server a foreign interceptor instance. + """ + built: list = [] + + def fake_builder(): + async def interceptor(request, handler): + return await handler(request) + + built.append(interceptor) + return interceptor + + config = ExtensionsConfig.model_validate( + { + "mcpInterceptors": ["my_package.auth:build_interceptor"], + "mcpServers": { + "alpha": {"type": "http", "url": "https://alpha.example/mcp"}, + "beta": {"type": "http", "url": "https://beta.example/mcp"}, + }, + } + ) + clients: list[tuple[set[str], list]] = [] + + def make_client(connections, **kwargs): + client = MagicMock() + client.get_tools = AsyncMock(return_value=[]) + clients.append((set(connections), kwargs.get("tool_interceptors") or [])) + return client + + with ( + patch("deerflow.config.extensions_config.ExtensionsConfig.from_file", return_value=config), + patch("langchain_mcp_adapters.client.MultiServerMCPClient", side_effect=make_client), + patch("deerflow.mcp.tools.get_initial_oauth_headers", new_callable=AsyncMock, return_value={}), + patch("deerflow.mcp.tools.build_oauth_tool_interceptor", return_value=None), + patch("deerflow.mcp.tools.resolve_variable", return_value=fake_builder), + ): + asyncio.run(get_mcp_tools()) + + assert len(built) == 2, f"builder ran {len(built)} times, expected once per server" + assert len(clients) == 2 + per_client = [interceptors[0] for _names, interceptors in clients] + assert per_client[0] is built[0] + assert per_client[1] is built[1] + assert per_client[0] is not per_client[1] diff --git a/backend/tests/test_mcp_per_server_discovery.py b/backend/tests/test_mcp_per_server_discovery.py new file mode 100644 index 0000000000..1f0ed69e28 --- /dev/null +++ b/backend/tests/test_mcp_per_server_discovery.py @@ -0,0 +1,713 @@ +"""F1 coverage for grouped, per-server MCP tool discovery.""" + +from __future__ import annotations + +import asyncio +import builtins +import logging +from unittest.mock import AsyncMock, MagicMock + +import pytest +from langchain_core.tools import StructuredTool + +from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.mcp.oauth import get_initial_oauth_headers +from deerflow.mcp.session_pool import ServerBinding, StaleMCPBindingError, normalized_connection_fingerprint +from deerflow.mcp.tasks.runtime import McpTaskConfigurationError +from deerflow.mcp.tools import ( + ServerDiscoveryResult, + _flatten_server_tool_groups, + get_mcp_tools, + get_mcp_tools_by_server, +) + + +def _tool(name: str) -> StructuredTool: + async def _call() -> str: + return name + + return StructuredTool.from_function(coroutine=_call, name=name, description=name) + + +def _config(servers: dict[str, dict]) -> ExtensionsConfig: + return ExtensionsConfig.model_validate({"mcpServers": servers}) + + +def _patch_common(monkeypatch, *, oauth_headers=None, interceptors=None): + headers = oauth_headers or AsyncMock(return_value={}) + monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", headers) + monkeypatch.setattr( + "deerflow.mcp.tools.build_mcp_tool_interceptors", + interceptors or (lambda *args, **kwargs: []), + ) + return headers + + +def test_flatten_server_tool_groups_preserves_enabled_order_and_empty_success(): + a, b = _tool("a"), _tool("b") + groups = { + "A": ServerDiscoveryResult(tools=(a,)), + "B": ServerDiscoveryResult(tools=(b,)), + "EMPTY": ServerDiscoveryResult(tools=()), + } + + flattened = _flatten_server_tool_groups(("B", "EMPTY", "A", "ABSENT"), groups) + + assert flattened == [b, a] + assert flattened[0] is b + assert flattened[1] is a + assert "EMPTY" in groups + assert groups["EMPTY"].tools == () + + +def test_server_discovery_result_repr_hides_ownership_handles(): + pool = MagicMock() + binding = ServerBinding( + domain="deployment", + server_name="A", + epoch=7, + fingerprint="secret-fingerprint", + ) + + rendered = repr(ServerDiscoveryResult(tools=(), pool=pool, binding=binding)) + + assert "pool" not in rendered + assert "binding" not in rendered + assert "secret-fingerprint" not in rendered + + +@pytest.mark.asyncio +async def test_get_mcp_tools_config_load_errors_propagate(monkeypatch): + from_file = MagicMock(side_effect=FileNotFoundError("missing config")) + logger_error = MagicMock() + monkeypatch.setattr("deerflow.mcp.tools.ExtensionsConfig.from_file", from_file) + monkeypatch.setattr("deerflow.mcp.tools.logger.error", logger_error) + + with pytest.raises(FileNotFoundError, match="missing config"): + await get_mcp_tools() + + from_file.assert_called_once() + logger_error.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_mcp_tools_missing_adapter_returns_before_side_effects(monkeypatch): + from_file = MagicMock(side_effect=AssertionError("config must not load")) + get_session_pool = MagicMock(side_effect=AssertionError("pool must not resolve")) + real_import = builtins.__import__ + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + if name.startswith("langchain_mcp_adapters"): + raise ImportError("missing adapter") + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr("deerflow.mcp.tools.ExtensionsConfig.from_file", from_file) + monkeypatch.setattr("deerflow.mcp.tools.get_session_pool", get_session_pool) + monkeypatch.setattr(builtins, "__import__", fake_import) + + tools = await get_mcp_tools() + + assert tools == [] + from_file.assert_not_called() + get_session_pool.assert_not_called() + + +@pytest.mark.asyncio +async def test_wrap_failure_returns_empty_instead_of_partial_tools(monkeypatch): + config = _config( + { + "A": {"enabled": True, "type": "stdio", "command": "a-server"}, + "B": {"enabled": True, "type": "stdio", "command": "b-server"}, + } + ) + pool = MagicMock() + pool.ensure_binding.side_effect = lambda name, *_args, **_kwargs: ServerBinding( + domain="deployment", + server_name=name, + epoch=1, + fingerprint=f"fp-{name}", + ) + + class FakeClient: + def __init__(self, connections, **kwargs): + self.callbacks = None + self.tool_interceptors = kwargs.get("tool_interceptors") or [] + + async def get_tools(self, *, server_name=None): + return [_tool(f"{server_name}_search")] + + def fake_wrap(tool, server_name, *args, **kwargs): + if server_name == "A": + raise RuntimeError("wrap failed") + return tool + + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + _patch_common(monkeypatch) + monkeypatch.setattr("deerflow.mcp.tools._make_session_pool_tool", fake_wrap) + + tools = await get_mcp_tools(config, session_pool=pool) + + assert tools == [] + + +@pytest.mark.asyncio +async def test_task_configuration_error_waits_for_sibling_discovery(monkeypatch): + config = _config( + { + "A": { + "enabled": True, + "type": "http", + "url": "https://a.example/mcp", + "task_toolsets": [ + { + "name": "reports", + "submit_tool": "submit", + "status_tool": "status", + "cancel_tool": "cancel", + } + ], + }, + "B": {"enabled": True, "type": "http", "url": "https://b.example/mcp"}, + } + ) + sibling_finished = asyncio.Event() + + class FakeClient: + def __init__(self, connections, **kwargs): + self.callbacks = None + self.tool_interceptors = kwargs.get("tool_interceptors") or [] + + async def get_tools(self, *, server_name=None): + if server_name == "B": + await asyncio.sleep(0.05) + sibling_finished.set() + return [_tool("B_search")] + return [] + + manager = MagicMock() + manager.has_oauth_servers.return_value = False + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_task_oauth_token_manager", lambda config: manager) + monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", AsyncMock(return_value={})) + monkeypatch.setattr("deerflow.mcp.tools.build_mcp_tool_interceptors", lambda *args, **kwargs: []) + + with pytest.raises(McpTaskConfigurationError, match="missing raw tool"): + await get_mcp_tools_by_server(config, server_names={"A", "B"}) + + assert sibling_finished.is_set() + + +@pytest.mark.asyncio +async def test_selected_stdio_isolation_builds_one_client_and_only_its_binding(monkeypatch): + config = _config( + { + "A": {"enabled": True, "type": "stdio", "command": "a-server"}, + "B": {"enabled": True, "type": "stdio", "command": "b-server"}, + } + ) + created: list[set[str]] = [] + wrapped: list[str] = [] + oauth_selection: list[set[str]] = [] + binding = ServerBinding( + domain="deployment", + server_name="A", + epoch=1, + fingerprint="fp-a", + ) + pool = MagicMock() + pool.ensure_binding.return_value = binding + + class FakeClient: + def __init__(self, connections, **kwargs): + created.append(set(connections)) + self.connections = connections + self.callbacks = None + self.tool_interceptors = kwargs.get("tool_interceptors") or [] + + async def get_tools(self, *, server_name=None): + assert server_name in self.connections + return [_tool(f"{server_name}_search")] + + async def fake_oauth_headers(_config, *, token_manager=None, server_names=None): + oauth_selection.append(set(server_names or ())) + return {} + + def fake_wrapper(tool, server_name, *args, **kwargs): + wrapped.append(server_name) + return tool + + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", fake_oauth_headers) + monkeypatch.setattr("deerflow.mcp.tools.build_mcp_tool_interceptors", lambda *args, **kwargs: []) + monkeypatch.setattr("deerflow.mcp.tools._make_session_pool_tool", fake_wrapper) + monkeypatch.setattr("deerflow.mcp.tools.get_session_pool", MagicMock(side_effect=AssertionError("caller pool must win"))) + + result = await get_mcp_tools_by_server( + config, + server_names={"A"}, + session_pool=pool, + ) + + assert created == [{"A"}] + assert list(result) == ["A"] + assert [tool.name for tool in result["A"].tools] == ["A_search"] + assert wrapped == ["A"] + assert oauth_selection == [{"A"}] + pool.ensure_binding.assert_called_once() + assert pool.ensure_binding.call_args.args[0] == "A" + assert pool.ensure_binding.call_args.args[1] == normalized_connection_fingerprint( + {"transport": "stdio", "command": "a-server", "args": []}, + ) + assert pool.ensure_binding.call_args.kwargs == {"domain": "deployment"} + assert result["A"].pool is pool + assert result["A"].binding is binding + + +@pytest.mark.asyncio +async def test_each_selected_server_gets_its_own_single_server_client(monkeypatch): + config = _config( + { + "A": {"enabled": True, "type": "http", "url": "https://a.example/mcp"}, + "B": {"enabled": True, "type": "http", "url": "https://b.example/mcp"}, + "C": {"enabled": True, "type": "http", "url": "https://c.example/mcp"}, + } + ) + created: list[set[str]] = [] + + class FakeClient: + def __init__(self, connections, **kwargs): + created.append(set(connections)) + self.callbacks = None + self.tool_interceptors = kwargs.get("tool_interceptors") or [] + + async def get_tools(self, *, server_name=None): + return [_tool(f"{server_name}_search")] + + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + _patch_common(monkeypatch) + + result = await get_mcp_tools_by_server(config, server_names={"A", "B"}) + + assert created == [{"A"}, {"B"}] + assert list(result) == ["A", "B"] + assert [tool.name for tool in result["A"].tools] == ["A_search"] + assert [tool.name for tool in result["B"].tools] == ["B_search"] + + +@pytest.mark.asyncio +async def test_each_client_gets_interceptors_built_from_only_its_server(monkeypatch): + config = _config( + { + "A": {"enabled": True, "type": "http", "url": "https://a.example/mcp"}, + "B": {"enabled": True, "type": "http", "url": "https://b.example/mcp"}, + } + ) + clients: list[tuple[set[str], list[object]]] = [] + built_for: list[tuple[str, ...]] = [] + + class FakeClient: + def __init__(self, connections, **kwargs): + interceptors = kwargs.get("tool_interceptors") or [] + clients.append((set(connections), interceptors)) + self.callbacks = None + self.tool_interceptors = interceptors + + async def get_tools(self, *, server_name=None): + return [_tool(f"{server_name}_search")] + + def fake_interceptors(server_config, *args, **kwargs): + built_for.append(tuple(server_config.get_enabled_mcp_servers())) + return [object()] + + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", AsyncMock(return_value={})) + monkeypatch.setattr("deerflow.mcp.tools.build_mcp_tool_interceptors", fake_interceptors) + + result = await get_mcp_tools_by_server(config, server_names={"A", "B"}) + + assert list(result) == ["A", "B"] + assert [connections for connections, _ in clients] == [{"A"}, {"B"}] + assert sorted(built_for) == [("A",), ("B",)] + assert clients[0][1] is not clients[1][1] + + +@pytest.mark.asyncio +async def test_failed_discovery_is_absent_but_successful_empty_is_present(monkeypatch): + config = _config( + { + "EMPTY": {"enabled": True, "type": "http", "url": "https://empty.example/mcp"}, + "FAIL": {"enabled": True, "type": "http", "url": "https://fail.example/mcp"}, + } + ) + + class FakeClient: + def __init__(self, connections, **kwargs): + self.callbacks = None + self.tool_interceptors = kwargs.get("tool_interceptors") or [] + + async def get_tools(self, *, server_name=None): + if server_name == "FAIL": + raise RuntimeError("offline") + return [] + + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + _patch_common(monkeypatch) + + result = await get_mcp_tools_by_server(config) + + assert "EMPTY" in result + assert result["EMPTY"].tools == () + assert "FAIL" not in result + + +@pytest.mark.asyncio +async def test_initial_oauth_headers_only_fetch_selected_names(monkeypatch): + config = _config( + { + "A": { + "enabled": True, + "type": "http", + "url": "https://a.example/mcp", + "oauth": {"enabled": True, "token_url": "https://auth.example/a"}, + }, + "B": { + "enabled": True, + "type": "http", + "url": "https://b.example/mcp", + "oauth": {"enabled": True, "token_url": "https://auth.example/b"}, + }, + } + ) + manager = MagicMock() + manager.has_oauth_servers.return_value = True + manager.oauth_server_names.return_value = ["A", "B"] + manager.get_authorization_header = AsyncMock(side_effect=lambda name: f"Bearer {name}") + monkeypatch.setattr( + "deerflow.mcp.oauth.OAuthTokenManager.from_extensions_config", + lambda config: manager, + ) + + assert await get_initial_oauth_headers(config, server_names={"A"}) == {"A": "Bearer A"} + manager.get_authorization_header.assert_awaited_once_with("A") + + manager.get_authorization_header.reset_mock() + assert await get_initial_oauth_headers(config, server_names=set()) == {} + manager.get_authorization_header.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_get_mcp_tools_flattens_grouped_discovery_in_enabled_order(monkeypatch): + config = _config( + { + "A": {"enabled": True, "type": "http", "url": "https://a.example/mcp"}, + "B": {"enabled": True, "type": "http", "url": "https://b.example/mcp"}, + } + ) + a, b = _tool("A_search"), _tool("B_search") + calls: list[object] = [] + + async def fake_grouped(_adapter_symbols, prepared): + calls.append(prepared) + return { + "B": ServerDiscoveryResult(tools=(b,)), + "A": ServerDiscoveryResult(tools=(a,)), + } + + monkeypatch.setattr("deerflow.mcp.tools._discover_and_wrap_prepared", fake_grouped) + + tools = await get_mcp_tools(config) + + assert tools == [a, b] + assert len(calls) == 1 + assert calls[0].config is config + + +@pytest.mark.asyncio +async def test_get_mcp_tools_full_discovery_preserves_name_order_and_loads_config_once(monkeypatch): + config = _config( + { + "A": {"enabled": True, "type": "http", "url": "https://a.example/mcp"}, + "B": {"enabled": True, "type": "http", "url": "https://b.example/mcp"}, + } + ) + + a_first = _tool("A_first") + b_first = _tool("B_first") + b_second = _tool("B_second") + + class FakeClient: + def __init__(self, connections, **kwargs): + self.callbacks = None + self.tool_interceptors = kwargs.get("tool_interceptors") or [] + + async def get_tools(self, *, server_name=None): + if server_name == "A": + return [a_first] + return [b_first, b_second] + + from_file = MagicMock(return_value=config) + monkeypatch.setattr("deerflow.mcp.tools.ExtensionsConfig.from_file", from_file) + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + _patch_common(monkeypatch) + + tools = await get_mcp_tools() + + assert [tool.name for tool in tools] == ["A_first", "B_first", "B_second"] + assert tools == [a_first, b_first, b_second] + assert tools[0] is a_first + assert tools[1] is b_first + assert tools[2] is b_second + assert from_file.call_count == 1 + + +@pytest.mark.asyncio +async def test_bindings_are_captured_before_the_first_discovery_await(monkeypatch): + config = _config( + { + "A": {"enabled": True, "type": "stdio", "command": "a-server"}, + } + ) + order: list[str] = [] + binding = ServerBinding(domain="deployment", server_name="A", epoch=1, fingerprint="fp-a") + pool = MagicMock() + + def ensure_binding(*args, **kwargs): + order.append("bind") + return binding + + pool.ensure_binding.side_effect = ensure_binding + + async def fake_oauth_headers(*args, **kwargs): + assert order == ["bind"] + order.append("oauth") + return {} + + class FakeClient: + def __init__(self, connections, **kwargs): + self.callbacks = None + self.tool_interceptors = kwargs.get("tool_interceptors") or [] + + async def get_tools(self, *, server_name=None): + assert order == ["bind", "oauth"] + order.append("discover") + return [] + + monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", fake_oauth_headers) + monkeypatch.setattr("deerflow.mcp.tools.build_mcp_tool_interceptors", lambda *args, **kwargs: []) + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + + result = await get_mcp_tools_by_server(config, session_pool=pool) + + assert result["A"].tools == () + assert result["A"].pool is pool + assert result["A"].binding is binding + assert order == ["bind", "oauth", "discover"] + pool.ensure_binding.assert_called_once() + + +@pytest.mark.asyncio +async def test_personal_discovery_uses_personal_pool_domain_and_task_scope(monkeypatch): + config = _config( + { + "A": { + "enabled": True, + "type": "stdio", + "command": "a-server", + "task_toolsets": [ + { + "name": "reports", + "submit_tool": "submit", + "status_tool": "status", + "cancel_tool": "cancel", + } + ], + } + } + ) + pool = MagicMock() + binding = ServerBinding(domain="personal", server_name="A", epoch=1, fingerprint="fp-a") + pool.ensure_binding.return_value = binding + ownership_domains: list[str] = [] + connection_scopes: list[str] = [] + deployment_manager = MagicMock(side_effect=AssertionError("deployment manager used for personal discovery")) + personal_manager = MagicMock() + personal_manager.has_oauth_servers.return_value = False + + class FakeClient: + def __init__(self, connections, **kwargs): + self.callbacks = None + self.tool_interceptors = kwargs.get("tool_interceptors") or [] + + async def get_tools(self, *, server_name=None): + return [_tool("A_submit")] + + def fake_wrapper(tool, server_name, *args, **kwargs): + ownership_domains.append(kwargs["ownership_domain"]) + return tool + + def fake_configure(tools, *, server_name, server_config, tool_name_prefix, connection_scope="deployment"): + connection_scopes.append(connection_scope) + return tools + + monkeypatch.setattr("deerflow.mcp.user_config.load_user_mcp_config", lambda owner: config) + monkeypatch.setattr( + "deerflow.mcp.personal_access.authorized_personal_config", + AsyncMock(return_value=config), + ) + monkeypatch.setattr("deerflow.mcp.tasks.runtime.is_mcp_task_runtime_available", lambda: True) + monkeypatch.setattr("deerflow.mcp.tools.OAuthTokenManager.from_extensions_config", lambda config: personal_manager) + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_task_oauth_token_manager", deployment_manager) + monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", AsyncMock(return_value={})) + monkeypatch.setattr("deerflow.mcp.tools.build_mcp_tool_interceptors", lambda *args, **kwargs: []) + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + monkeypatch.setattr("deerflow.mcp.tools._make_session_pool_tool", fake_wrapper) + monkeypatch.setattr("deerflow.mcp.tools._configure_task_tools_for_server", fake_configure) + + result = await get_mcp_tools_by_server( + config, + personal_user_id="alice", + session_pool=pool, + ) + + assert list(result) == ["A"] + assert pool.ensure_binding.call_args.kwargs == {"domain": "personal"} + assert ownership_domains == ["personal"] + assert connection_scopes == ["personal"] + deployment_manager.assert_not_called() + + +@pytest.mark.asyncio +async def test_stale_binding_propagates_before_any_client_is_built(monkeypatch): + config = _config( + { + "A": {"enabled": True, "type": "stdio", "command": "a-server"}, + } + ) + pool = MagicMock() + pool.ensure_binding.side_effect = StaleMCPBindingError("A") + clients: list[set[str]] = [] + + class FakeClient: + def __init__(self, connections, **kwargs): + clients.append(set(connections)) + + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + + with pytest.raises(StaleMCPBindingError): + await get_mcp_tools_by_server(config, session_pool=pool) + + assert clients == [] + + +@pytest.mark.asyncio +async def test_task_configuration_error_propagates_from_grouped_discovery(monkeypatch): + config = _config( + { + "A": { + "enabled": True, + "type": "http", + "url": "https://a.example/mcp", + "task_toolsets": [ + { + "name": "reports", + "submit_tool": "submit", + "status_tool": "status", + "cancel_tool": "cancel", + } + ], + } + } + ) + + class FakeClient: + def __init__(self, connections, **kwargs): + self.callbacks = None + self.tool_interceptors = kwargs.get("tool_interceptors") or [] + + async def get_tools(self, *, server_name=None): + return [] + + manager = MagicMock() + manager.has_oauth_servers.return_value = False + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + monkeypatch.setattr("deerflow.mcp.tools.get_mcp_task_oauth_token_manager", lambda config: manager) + monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", AsyncMock(return_value={})) + monkeypatch.setattr("deerflow.mcp.tools.build_mcp_tool_interceptors", lambda *args, **kwargs: []) + + with pytest.raises(McpTaskConfigurationError, match="missing raw tool"): + await get_mcp_tools_by_server(config, server_names={"A"}) + + +@pytest.mark.asyncio +async def test_selection_matching_only_unknown_or_disabled_names(monkeypatch, caplog): + """A selection that matches nothing is not "nothing is configured".""" + config = _config( + { + "A": {"enabled": True, "type": "http", "url": "https://a.example/mcp"}, + "B": {"enabled": False, "type": "http", "url": "https://b.example/mcp"}, + } + ) + client_calls: list = [] + monkeypatch.setattr( + "langchain_mcp_adapters.client.MultiServerMCPClient", + lambda *args, **kwargs: client_calls.append((args, kwargs)) or MagicMock(), + ) + headers = AsyncMock(return_value={}) + monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", headers) + + with caplog.at_level(logging.INFO, logger="deerflow.mcp.tools"): + assert await get_mcp_tools_by_server(config, server_names={"B", "UNKNOWN"}) == {} + + assert client_calls == [] + headers.assert_not_awaited() + assert "No enabled MCP servers matched the requested selection" in caplog.text + assert "No enabled MCP servers configured" not in caplog.text + + +@pytest.mark.asyncio +async def test_discovery_and_wrapper_use_the_same_session_init_timeout(monkeypatch): + """Preparation derives the timeout once; discovery and wrapping share it.""" + config = _config( + { + "A": { + "enabled": True, + "type": "stdio", + "command": "npx", + "args": [], + "session_init_timeout": 5, + } + } + ) + observed: dict[str, float | None] = {} + + class FakeClient: + def __init__(self, connections, **kwargs): + pass + + async def get_tools(self, *, server_name=None): + return [_tool("A_search")] + + real_wait_for = asyncio.wait_for + + async def spy_wait_for(awaitable, timeout): + observed["discovery"] = timeout + return await real_wait_for(awaitable, timeout) + + monkeypatch.setattr("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient) + monkeypatch.setattr("deerflow.mcp.tools.asyncio.wait_for", spy_wait_for) + monkeypatch.setattr("deerflow.mcp.tools.get_initial_oauth_headers", AsyncMock(return_value={})) + monkeypatch.setattr("deerflow.mcp.tools.build_mcp_tool_interceptors", lambda *a, **kw: []) + + real_make_tool = __import__("deerflow.mcp.tools", fromlist=["_make_session_pool_tool"])._make_session_pool_tool + + def spy_make_tool(*args, **kwargs): + observed["wrapper"] = kwargs.get("session_init_timeout") + return real_make_tool(*args, **kwargs) + + monkeypatch.setattr("deerflow.mcp.tools._make_session_pool_tool", spy_make_tool) + + groups = await get_mcp_tools_by_server(config, server_names={"A"}) + + assert len(groups["A"].tools) == 1 + assert observed["discovery"] == 5.0 + assert observed["wrapper"] == observed["discovery"] diff --git a/backend/tests/test_mcp_session_pool.py b/backend/tests/test_mcp_session_pool.py index 78e7129f43..344e360ef5 100644 --- a/backend/tests/test_mcp_session_pool.py +++ b/backend/tests/test_mcp_session_pool.py @@ -3092,8 +3092,24 @@ async def test_mcp_tools_routed_to_source_server_with_prefix_overlap(): response_format="content_and_artifact", ) - extensions_config = MagicMock() - extensions_config.model_extra = {} + from deerflow.config.extensions_config import ExtensionsConfig + + extensions_config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "web": { + "type": "stdio", + "command": "npx", + "args": ["web"], + }, + "web_scraper": { + "type": "stdio", + "command": "npx", + "args": ["scraper"], + }, + } + } + ) # `web` is inserted before `web_scraper`, so a first-prefix-match mis-routes # `web_scraper_search` to `web`. diff --git a/backend/tests/test_mcp_sync_wrapper.py b/backend/tests/test_mcp_sync_wrapper.py index 68281fe595..30197b8219 100644 --- a/backend/tests/test_mcp_sync_wrapper.py +++ b/backend/tests/test_mcp_sync_wrapper.py @@ -8,6 +8,7 @@ from langchain_core.runnables import RunnableConfig from langchain_core.tools import InjectedToolArg, StructuredTool from pydantic import BaseModel, Field +from deerflow.config.extensions_config import ExtensionsConfig from deerflow.mcp.tools import get_mcp_tools from deerflow.tools.sync import make_sync_tool_wrapper @@ -33,10 +34,20 @@ def test_mcp_tool_sync_wrapper_generation(): mock_client_instance = MagicMock() # Use AsyncMock for get_tools as it's awaited (Fix for Comment 5) mock_client_instance.get_tools = AsyncMock(return_value=[mock_tool]) + extensions_config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "test-server": { + "type": "http", + "url": "https://example.test/mcp", + } + } + } + ) with ( patch("langchain_mcp_adapters.client.MultiServerMCPClient", return_value=mock_client_instance), - patch("deerflow.config.extensions_config.ExtensionsConfig.from_file"), + patch("deerflow.config.extensions_config.ExtensionsConfig.from_file", return_value=extensions_config), patch( "deerflow.mcp.tools.build_servers_config", return_value={ @@ -85,11 +96,31 @@ def test_mcp_tool_loading_skips_failed_server(): mock_client_instance = MagicMock() mock_client_instance.get_tools = AsyncMock(side_effect=get_tools_for_server) + extensions_config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "good-server": { + "type": "http", + "url": "https://good.example/mcp", + }, + "bad-server": { + "type": "http", + "url": "https://bad.example/mcp", + }, + } + } + ) with ( patch("langchain_mcp_adapters.client.MultiServerMCPClient", return_value=mock_client_instance), - patch("deerflow.config.extensions_config.ExtensionsConfig.from_file", return_value=MagicMock(model_extra={})), - patch("deerflow.mcp.tools.build_servers_config", return_value={"good-server": {}, "bad-server": {}}), + patch("deerflow.config.extensions_config.ExtensionsConfig.from_file", return_value=extensions_config), + patch( + "deerflow.mcp.tools.build_servers_config", + return_value={ + "good-server": {"transport": "http", "url": "https://good.example/mcp"}, + "bad-server": {"transport": "http", "url": "https://bad.example/mcp"}, + }, + ), patch("deerflow.mcp.tools.get_initial_oauth_headers", new_callable=AsyncMock, return_value={}), patch("deerflow.mcp.tools.build_oauth_tool_interceptor", return_value=None), patch("deerflow.mcp.tools.logger.warning") as mock_warning, @@ -229,8 +260,16 @@ def test_func_patched_mcp_tool_keeps_toolnode_runtime_injection(tmp_path): client.get_tools = AsyncMock(return_value=[discovered]) client.tool_interceptors = [] client.callbacks = None - cfg = MagicMock() - cfg.mcp_servers = {} + cfg = ExtensionsConfig.model_validate( + { + "mcpServers": { + "pw": { + "type": "stdio", + "command": "x", + } + } + } + ) with ( patch("langchain_mcp_adapters.client.MultiServerMCPClient", return_value=client), diff --git a/backend/tests/test_mcp_tool_name_validation.py b/backend/tests/test_mcp_tool_name_validation.py index 1f6db7685c..1fa983de7a 100644 --- a/backend/tests/test_mcp_tool_name_validation.py +++ b/backend/tests/test_mcp_tool_name_validation.py @@ -18,6 +18,7 @@ from unittest.mock import AsyncMock, MagicMock, patch from langchain_core.tools import StructuredTool from pydantic import BaseModel, Field +from deerflow.config.extensions_config import ExtensionsConfig from deerflow.mcp.tools import get_mcp_tools @@ -37,12 +38,20 @@ def _load(server_tools: list[StructuredTool]) -> tuple[list, MagicMock]: mock_client = MagicMock() mock_client.get_tools = AsyncMock(return_value=server_tools) - ext = MagicMock(model_extra={}) - ext.mcp_servers.get = MagicMock(return_value=None) # no per-tool routing config + extensions_config = ExtensionsConfig.model_validate( + { + "mcpServers": { + "srv": { + "type": "sse", + "url": "https://example.test/mcp", + } + } + } + ) with ( patch("langchain_mcp_adapters.client.MultiServerMCPClient", return_value=mock_client), - patch("deerflow.config.extensions_config.ExtensionsConfig.from_file", return_value=ext), + patch("deerflow.config.extensions_config.ExtensionsConfig.from_file", return_value=extensions_config), patch("deerflow.mcp.tools.build_servers_config", return_value={"srv": {"transport": "sse"}}), patch("deerflow.mcp.tools.get_initial_oauth_headers", new_callable=AsyncMock, return_value={}), patch("deerflow.mcp.tools.build_oauth_tool_interceptor", return_value=None),