mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 05:06:59 +00:00
refactor(mcp): expose per-server tool discovery results (#6734)
* refactor(mcp): expose per-server tool discovery results * refactor(mcp): clarify per-server discovery contracts
This commit is contained in:
spud
GitHub
parent
81f4d7707f
commit
3343522fa2
@ -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
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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]
|
||||
|
||||
713
backend/tests/test_mcp_per_server_discovery.py
Normal file
713
backend/tests/test_mcp_per_server_discovery.py
Normal file
@ -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"]
|
||||
@ -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`.
|
||||
|
||||
@ -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),
|
||||
|
||||
@ -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),
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user