mirror of
https://github.com/langchain-ai/deepagents.git
synced 2026-07-22 01:35:28 -04:00
perf(code): load MCP servers concurrently during graph build (#4659)
Deep Agents Code now discovers MCP servers concurrently, cutting initial graph-load time when multiple MCP servers are configured. --- Initial Deep Agents Code graph load scaled linearly with the number of configured MCP servers because `_load_tools_from_config` ran per-server preflight and tool discovery sequentially — the dominant contributor to startup time. This runs independent per-server preflight/connection-build and discovery concurrently under a bounded semaphore (`_MCP_LOAD_CONCURRENCY = 8`) while keeping deterministic server/tool ordering, per-server error isolation, auth, cleanup, and session-manager semantics unchanged. The adapter-import warmup was moved into `_load_tools_from_config`, so MCP adapters are imported only when at least one active MCP server exists. Measured (simulated, 0.4s preflight + 0.6s discovery per server): 6 servers 6.02s → 1.01s (~6x); 10 servers 10.03s → 2.01s (limit=8). Tool ordering and `server_infos` ordering verified identical to the sequential loader. The slow-load warning is untouched and no public APIs changed. Made by [Open SWE](https://openswe.vercel.app/agents/8dfe87f3-154b-050b-93a0-4a2a9af39fee) --------- Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
This commit is contained in:
Mason Daugherty
@@ -11,6 +11,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import copy
|
||||
import fnmatch
|
||||
import functools
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
@@ -18,10 +19,10 @@ import shutil
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Literal, cast, overload
|
||||
from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast, overload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
|
||||
from langchain_core.tools import BaseTool
|
||||
from langchain_mcp_adapters.client import Connection
|
||||
@@ -32,6 +33,8 @@ if TYPE_CHECKING:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
# Maintainer note: `deepagents-talon` imports `MCPConfigError`,
|
||||
# `MCPServerInfo`, and `get_mcp_tools` from this module, and its tests construct
|
||||
# `MCPToolInfo`. Keep those symbols' names, signatures, and return/dataclass
|
||||
@@ -1474,6 +1477,84 @@ def _apply_tool_filter(
|
||||
return [t for t in tools if not _any_entry_matches(t.name, entries)]
|
||||
|
||||
|
||||
_MCP_LOAD_CONCURRENCY = 8
|
||||
"""Upper bound on MCP servers preflighted/discovered concurrently.
|
||||
|
||||
Independent servers are probed in parallel so graph load no longer scales
|
||||
linearly with server count, but the fan-out is capped so a large config cannot
|
||||
spawn an unbounded number of simultaneous socket/subprocess handshakes (or
|
||||
`asyncio.to_thread` `shutil.which` workers).
|
||||
"""
|
||||
|
||||
|
||||
def _warm_mcp_adapter_imports() -> None:
|
||||
"""Eagerly import MCP adapter modules whose first import may block.
|
||||
|
||||
Run via `asyncio.to_thread` before adapter symbols are used, so the initial
|
||||
(potentially blocking) package-resource scan stays off the server event
|
||||
loop. Because this runs inside `_load_tools_from_config`, it happens only
|
||||
when at least one active MCP server exists — a config with no MCP servers
|
||||
never imports the adapters.
|
||||
"""
|
||||
from langchain_mcp_adapters import (
|
||||
sessions as _sessions, # noqa: F401
|
||||
tools as _tools, # noqa: F401
|
||||
)
|
||||
|
||||
|
||||
async def _gather_bounded(
|
||||
factories: Sequence[Callable[[], Awaitable[_T]]],
|
||||
*,
|
||||
limit: int,
|
||||
) -> list[_T]:
|
||||
"""Await coroutine factories with bounded concurrency, preserving order.
|
||||
|
||||
Results are returned in submission order (not completion order), so callers
|
||||
can zip them back against their inputs to keep deterministic ordering. If a
|
||||
factory raises (including a cancellation/shutdown signal), the remaining
|
||||
tasks are cancelled and awaited before the exception propagates, so no
|
||||
background work is left running.
|
||||
|
||||
`asyncio.gather` propagates only the *first* task to finish with an
|
||||
exception; when several tasks fail concurrently the rest are cancelled
|
||||
during teardown and their exceptions would otherwise be discarded silently.
|
||||
To keep concurrent failures debuggable, each dropped (non-cancellation)
|
||||
sibling exception is logged at debug level before the first one propagates.
|
||||
|
||||
Args:
|
||||
factories: Zero-arg callables each returning an awaitable to run.
|
||||
limit: Maximum number of awaitables in flight at once. Values below 1
|
||||
are clamped to 1.
|
||||
|
||||
Returns:
|
||||
The awaited results in the same order as `factories`.
|
||||
"""
|
||||
semaphore = asyncio.Semaphore(max(1, limit))
|
||||
|
||||
async def _run(factory: Callable[[], Awaitable[_T]]) -> _T:
|
||||
async with semaphore:
|
||||
return await factory()
|
||||
|
||||
tasks = [asyncio.create_task(_run(factory)) for factory in factories]
|
||||
try:
|
||||
return await asyncio.gather(*tasks)
|
||||
except BaseException:
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
for result in results:
|
||||
if isinstance(result, BaseException) and not isinstance(
|
||||
result, asyncio.CancelledError
|
||||
):
|
||||
logger.debug(
|
||||
"MCP concurrent load: a sibling task failed while another "
|
||||
"failure was already propagating; logging the dropped "
|
||||
"exception for debugging",
|
||||
exc_info=result,
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
async def _load_tools_from_config(
|
||||
config: dict[str, Any],
|
||||
*,
|
||||
@@ -1504,7 +1585,11 @@ async def _load_tools_from_config(
|
||||
Raises:
|
||||
RuntimeError: If `session_manager` is reconfigured incompatibly with
|
||||
sessions already active on it.
|
||||
""" # noqa: DOC501, DOC502 - `RuntimeError` surfaces via `MCPSessionManager.configure`; `KeyboardInterrupt` / `SystemExit` / `CancelledError` are re-raised pass-throughs
|
||||
""" # noqa: DOC502 - `RuntimeError` surfaces via `MCPSessionManager.configure`
|
||||
# Warm the adapter imports off the event loop *here* (rather than in the
|
||||
# caller) so a config with no active MCP servers — which returns before
|
||||
# ever reaching this function — never pays the adapter-import cost.
|
||||
await asyncio.to_thread(_warm_mcp_adapter_imports)
|
||||
from langchain_mcp_adapters.sessions import (
|
||||
SSEConnection,
|
||||
StdioConnection,
|
||||
@@ -1513,10 +1598,26 @@ async def _load_tools_from_config(
|
||||
)
|
||||
from langchain_mcp_adapters.tools import convert_mcp_tool_to_langchain_tool
|
||||
|
||||
skipped: dict[str, tuple[MCPServerStatus, str]] = {}
|
||||
server_items = list(config["mcpServers"].items())
|
||||
# Resolve each server's transport once, up front. `_resolve_server_type` is
|
||||
# pure, so this is a readability/DRY win over recomputing it in preflight,
|
||||
# discovery, and the final fold-in loop below.
|
||||
transports = {name: _resolve_server_type(cfg) for name, cfg in server_items}
|
||||
|
||||
for server_name, server_config in config["mcpServers"].items():
|
||||
server_type = _resolve_server_type(server_config)
|
||||
async def _preflight_and_connect(
|
||||
server_name: str,
|
||||
server_config: dict[str, Any],
|
||||
) -> tuple[MCPServerStatus, str] | Connection:
|
||||
"""Preflight one server and build its connection config.
|
||||
|
||||
Per-server preflight/config failures are captured here so one bad
|
||||
server never aborts loading the others.
|
||||
|
||||
Returns:
|
||||
A `(status, error)` tuple when the server must be skipped, or a
|
||||
ready `Connection` otherwise.
|
||||
"""
|
||||
server_type = transports[server_name]
|
||||
try:
|
||||
if server_type in _SUPPORTED_REMOTE_TYPES:
|
||||
await _check_remote_server(server_name, server_config)
|
||||
@@ -1530,13 +1631,8 @@ async def _load_tools_from_config(
|
||||
server_name,
|
||||
exc,
|
||||
)
|
||||
skipped[server_name] = ("error", str(exc))
|
||||
return ("error", str(exc))
|
||||
|
||||
connections: dict[str, Connection] = {}
|
||||
for server_name, server_config in config["mcpServers"].items():
|
||||
if server_name in skipped:
|
||||
continue
|
||||
server_type = _resolve_server_type(server_config)
|
||||
try:
|
||||
if server_type in _SUPPORTED_REMOTE_TYPES:
|
||||
if server_type == "http":
|
||||
@@ -1582,8 +1678,7 @@ async def _load_tools_from_config(
|
||||
"MCP server '%s' skipped: not authenticated.",
|
||||
server_name,
|
||||
)
|
||||
skipped[server_name] = ("unauthenticated", auth_msg)
|
||||
continue
|
||||
return ("unauthenticated", auth_msg)
|
||||
|
||||
if explicit_oauth or (
|
||||
stored_tokens is not None and not has_authorization_header
|
||||
@@ -1599,21 +1694,42 @@ async def _load_tools_from_config(
|
||||
interactive=False,
|
||||
)
|
||||
|
||||
connections[server_name] = conn
|
||||
else:
|
||||
connections[server_name] = StdioConnection(
|
||||
command=server_config["command"],
|
||||
args=server_config.get("args", []),
|
||||
env=server_config.get("env") or None,
|
||||
transport="stdio",
|
||||
)
|
||||
return conn
|
||||
return StdioConnection(
|
||||
command=server_config["command"],
|
||||
args=server_config.get("args", []),
|
||||
env=server_config.get("env") or None,
|
||||
transport="stdio",
|
||||
)
|
||||
except (OSError, RuntimeError, TypeError, ValueError) as exc:
|
||||
logger.warning(
|
||||
"MCP server '%s' skipped: config/setup failed: %s",
|
||||
server_name,
|
||||
exc,
|
||||
)
|
||||
skipped[server_name] = ("error", str(exc))
|
||||
return ("error", str(exc))
|
||||
|
||||
# Preflight + connection build runs concurrently across servers (bounded).
|
||||
# Results come back in submission order, so `skipped`/`connections` are
|
||||
# assembled in config order and stay deterministic regardless of which
|
||||
# server's probe finished first.
|
||||
preflight_results = await _gather_bounded(
|
||||
[
|
||||
functools.partial(_preflight_and_connect, name, cfg)
|
||||
for name, cfg in server_items
|
||||
],
|
||||
limit=_MCP_LOAD_CONCURRENCY,
|
||||
)
|
||||
|
||||
skipped: dict[str, tuple[MCPServerStatus, str]] = {}
|
||||
connections: dict[str, Connection] = {}
|
||||
for (server_name, _server_config), result in zip(
|
||||
server_items, preflight_results, strict=True
|
||||
):
|
||||
if isinstance(result, tuple):
|
||||
skipped[server_name] = result
|
||||
else:
|
||||
connections[server_name] = result
|
||||
|
||||
runtime_manager: MCPSessionManager | None = session_manager
|
||||
if runtime_manager is not None:
|
||||
@@ -1621,23 +1737,22 @@ async def _load_tools_from_config(
|
||||
elif not stateless:
|
||||
runtime_manager = MCPSessionManager(connections=connections)
|
||||
|
||||
all_tools: list[BaseTool] = []
|
||||
server_infos: list[MCPServerInfo] = []
|
||||
async def _discover_server(
|
||||
server_name: str,
|
||||
server_config: dict[str, Any],
|
||||
transport: str,
|
||||
) -> tuple[list[BaseTool], MCPServerInfo]:
|
||||
"""Discover one server's tools and build its `MCPServerInfo`.
|
||||
|
||||
for server_name, server_config in config["mcpServers"].items():
|
||||
transport = _resolve_server_type(server_config)
|
||||
if server_name in skipped:
|
||||
status, error = skipped[server_name]
|
||||
server_infos.append(
|
||||
MCPServerInfo(
|
||||
name=server_name,
|
||||
transport=transport,
|
||||
status=status,
|
||||
error=error,
|
||||
),
|
||||
)
|
||||
continue
|
||||
Both discovery failures (classified as auth vs. generic error) and
|
||||
post-discovery tool-construction failures are captured as a non-`ok`
|
||||
`MCPServerInfo` with no tools, so a single failing server never aborts
|
||||
the load for the others. Cancellation/shutdown signals are re-raised so
|
||||
the bounded runner can tear the whole load down.
|
||||
|
||||
Returns:
|
||||
The server's LangChain tools plus its `MCPServerInfo` entry.
|
||||
""" # noqa: DOC501 - CancelledError/KeyboardInterrupt/SystemExit are re-raised pass-throughs
|
||||
try:
|
||||
async with create_session(connections[server_name]) as discover_session:
|
||||
await discover_session.initialize()
|
||||
@@ -1721,93 +1836,152 @@ async def _load_tools_from_config(
|
||||
server_name,
|
||||
exc_info=True,
|
||||
)
|
||||
return [], MCPServerInfo(
|
||||
name=server_name,
|
||||
transport=transport,
|
||||
status=status,
|
||||
error=error,
|
||||
)
|
||||
|
||||
# Tool construction and filtering run after the discovery session has
|
||||
# closed and can still fail (schema conversion, custom tool filters).
|
||||
# Isolate them too so a construction error degrades this one server to
|
||||
# an error entry instead of aborting the whole concurrent load — the
|
||||
# same guarantee the discovery `try` above provides. Cancellation and
|
||||
# shutdown signals still propagate so the bounded runner can tear down.
|
||||
try:
|
||||
if runtime_manager is None:
|
||||
server_tools: list[BaseTool] = [
|
||||
convert_mcp_tool_to_langchain_tool(
|
||||
None,
|
||||
mcp_tool,
|
||||
connection=connections[server_name],
|
||||
server_name=server_name,
|
||||
tool_name_prefix=True,
|
||||
)
|
||||
for mcp_tool in mcp_tools
|
||||
]
|
||||
else:
|
||||
server_tools = [
|
||||
_build_cached_mcp_tool(
|
||||
mcp_tool=mcp_tool,
|
||||
server_name=server_name,
|
||||
session_manager=runtime_manager,
|
||||
tool_name_prefix=True,
|
||||
)
|
||||
for mcp_tool in mcp_tools
|
||||
]
|
||||
|
||||
server_tools = _apply_tool_filter(server_tools, server_name, server_config)
|
||||
|
||||
# Pair each tool's input_schema by its LangChain (server-prefixed)
|
||||
# name — the same form `server_tools` carries — so the lookup needs
|
||||
# no string surgery and stays correct if `tool_name_prefix` ever
|
||||
# changes. Deep-copy the raw dict because `MCPToolInfo` is `frozen`
|
||||
# but Python's `frozen=True` does not freeze nested mutables; a
|
||||
# shared reference would let one holder mutate every other's view.
|
||||
schemas: dict[str, dict[str, Any] | None] = {}
|
||||
for mcp_tool in mcp_tools:
|
||||
tool_name = getattr(mcp_tool, "name", "")
|
||||
try:
|
||||
raw_schema = getattr(mcp_tool, "inputSchema", None)
|
||||
schema_copy = (
|
||||
copy.deepcopy(raw_schema) if raw_schema is not None else None
|
||||
)
|
||||
except (AttributeError, TypeError, RecursionError) as exc:
|
||||
logger.warning(
|
||||
"MCP tool %r on server %r: inputSchema access raised "
|
||||
"%s: %s; rendering with no parameters",
|
||||
tool_name,
|
||||
server_name,
|
||||
exc.__class__.__name__,
|
||||
exc,
|
||||
)
|
||||
schema_copy = None
|
||||
lc_name = f"{server_name}_{tool_name}"
|
||||
schemas[lc_name] = schema_copy
|
||||
|
||||
tool_infos: list[MCPToolInfo] = []
|
||||
for tool in server_tools:
|
||||
schema = schemas.get(tool.name)
|
||||
if schema is None and schemas:
|
||||
logger.debug(
|
||||
"MCP tool %r on server %r: no schema matched in lookup "
|
||||
"(available keys: %s); rendering with no parameters",
|
||||
tool.name,
|
||||
server_name,
|
||||
list(schemas.keys())[:5],
|
||||
)
|
||||
tool_infos.append(
|
||||
MCPToolInfo(
|
||||
name=tool.name,
|
||||
description=tool.description or "",
|
||||
input_schema=schema,
|
||||
),
|
||||
)
|
||||
except (asyncio.CancelledError, KeyboardInterrupt, SystemExit):
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"MCP server '%s' skipped: tool construction failed",
|
||||
server_name,
|
||||
exc_info=True,
|
||||
)
|
||||
return [], MCPServerInfo(
|
||||
name=server_name,
|
||||
transport=transport,
|
||||
status="error",
|
||||
error=str(exc),
|
||||
)
|
||||
|
||||
return server_tools, MCPServerInfo(
|
||||
name=server_name,
|
||||
transport=transport,
|
||||
tools=tuple(tool_infos),
|
||||
)
|
||||
|
||||
# Discovery also runs concurrently (bounded) across the servers that
|
||||
# survived preflight. Because `_gather_bounded` returns results in
|
||||
# submission order and skipped servers are folded back in below by
|
||||
# iterating `server_items` in config order, `server_infos` stays in config
|
||||
# order and the returned tools stay sorted by tool name — regardless of
|
||||
# which server's probe finished first.
|
||||
discover_items = [
|
||||
(server_name, server_config, transports[server_name])
|
||||
for server_name, server_config in server_items
|
||||
if server_name not in skipped
|
||||
]
|
||||
discovery_results = await _gather_bounded(
|
||||
[
|
||||
functools.partial(_discover_server, name, cfg, transport)
|
||||
for name, cfg, transport in discover_items
|
||||
],
|
||||
limit=_MCP_LOAD_CONCURRENCY,
|
||||
)
|
||||
discovered: dict[str, tuple[list[BaseTool], MCPServerInfo]] = {
|
||||
server_name: result
|
||||
for (server_name, _cfg, _transport), result in zip(
|
||||
discover_items, discovery_results, strict=True
|
||||
)
|
||||
}
|
||||
|
||||
all_tools: list[BaseTool] = []
|
||||
server_infos: list[MCPServerInfo] = []
|
||||
for server_name, _server_config in server_items:
|
||||
if server_name in skipped:
|
||||
status, error = skipped[server_name]
|
||||
server_infos.append(
|
||||
MCPServerInfo(
|
||||
name=server_name,
|
||||
transport=transport,
|
||||
transport=transports[server_name],
|
||||
status=status,
|
||||
error=error,
|
||||
),
|
||||
)
|
||||
continue
|
||||
|
||||
if runtime_manager is None:
|
||||
server_tools = [
|
||||
convert_mcp_tool_to_langchain_tool(
|
||||
None,
|
||||
mcp_tool,
|
||||
connection=connections[server_name],
|
||||
server_name=server_name,
|
||||
tool_name_prefix=True,
|
||||
)
|
||||
for mcp_tool in mcp_tools
|
||||
]
|
||||
else:
|
||||
server_tools = [
|
||||
_build_cached_mcp_tool(
|
||||
mcp_tool=mcp_tool,
|
||||
server_name=server_name,
|
||||
session_manager=runtime_manager,
|
||||
tool_name_prefix=True,
|
||||
)
|
||||
for mcp_tool in mcp_tools
|
||||
]
|
||||
|
||||
server_tools = _apply_tool_filter(server_tools, server_name, server_config)
|
||||
server_tools, server_info = discovered[server_name]
|
||||
all_tools.extend(server_tools)
|
||||
|
||||
# Pair each tool's input_schema by its LangChain (server-prefixed)
|
||||
# name — the same form `server_tools` carries — so the lookup needs
|
||||
# no string surgery and stays correct if `tool_name_prefix` ever
|
||||
# changes. Deep-copy the raw dict because `MCPToolInfo` is `frozen`
|
||||
# but Python's `frozen=True` does not freeze nested mutables; a
|
||||
# shared reference would let one holder mutate every other's view.
|
||||
schemas: dict[str, dict[str, Any] | None] = {}
|
||||
for mcp_tool in mcp_tools:
|
||||
tool_name = getattr(mcp_tool, "name", "")
|
||||
try:
|
||||
raw_schema = getattr(mcp_tool, "inputSchema", None)
|
||||
schema_copy = (
|
||||
copy.deepcopy(raw_schema) if raw_schema is not None else None
|
||||
)
|
||||
except (AttributeError, TypeError, RecursionError) as exc:
|
||||
logger.warning(
|
||||
"MCP tool %r on server %r: inputSchema access raised %s: %s; "
|
||||
"rendering with no parameters",
|
||||
tool_name,
|
||||
server_name,
|
||||
exc.__class__.__name__,
|
||||
exc,
|
||||
)
|
||||
schema_copy = None
|
||||
lc_name = f"{server_name}_{tool_name}"
|
||||
schemas[lc_name] = schema_copy
|
||||
|
||||
tool_infos: list[MCPToolInfo] = []
|
||||
for tool in server_tools:
|
||||
schema = schemas.get(tool.name)
|
||||
if schema is None and schemas:
|
||||
logger.debug(
|
||||
"MCP tool %r on server %r: no schema matched in lookup "
|
||||
"(available keys: %s); rendering with no parameters",
|
||||
tool.name,
|
||||
server_name,
|
||||
list(schemas.keys())[:5],
|
||||
)
|
||||
tool_infos.append(
|
||||
MCPToolInfo(
|
||||
name=tool.name,
|
||||
description=tool.description or "",
|
||||
input_schema=schema,
|
||||
),
|
||||
)
|
||||
server_infos.append(
|
||||
MCPServerInfo(
|
||||
name=server_name,
|
||||
transport=transport,
|
||||
tools=tuple(tool_infos),
|
||||
),
|
||||
)
|
||||
server_infos.append(server_info)
|
||||
|
||||
all_tools.sort(key=lambda tool: tool.name)
|
||||
return all_tools, None if stateless else runtime_manager, server_infos
|
||||
|
||||
@@ -47,18 +47,6 @@ def _print_startup_error(message: str) -> None:
|
||||
)
|
||||
|
||||
|
||||
def _warm_mcp_adapter_imports() -> None:
|
||||
"""Eagerly import MCP adapter modules whose first import may block.
|
||||
|
||||
Called via `asyncio.to_thread` before MCP loading; kept as a small helper so
|
||||
tests can patch it and assert the warmup stays off the server event loop.
|
||||
"""
|
||||
from langchain_mcp_adapters import (
|
||||
sessions as _sessions, # noqa: F401
|
||||
tools as _tools, # noqa: F401
|
||||
)
|
||||
|
||||
|
||||
def _get_mcp_session_manager() -> Any: # noqa: ANN401
|
||||
"""Return the process-wide MCP session manager singleton.
|
||||
|
||||
@@ -92,8 +80,9 @@ async def _build_tools(
|
||||
`asyncio.run` (which raises inside a running loop). `stateless=True` ensures
|
||||
discovery only uses throwaway sessions, while the shared runtime session
|
||||
manager binds real sessions lazily inside the server loop on first tool
|
||||
invocation. MCP adapter imports are warmed in a worker thread first because
|
||||
first import can perform blocking package-resource scans.
|
||||
invocation. MCP adapter imports are warmed in a worker thread inside
|
||||
`_load_tools_from_config` (only when active servers exist) because first
|
||||
import can perform blocking package-resource scans.
|
||||
|
||||
Args:
|
||||
config: Deserialized server configuration.
|
||||
@@ -117,8 +106,6 @@ async def _build_tools(
|
||||
if not config.no_mcp:
|
||||
from deepagents_code.mcp_tools import resolve_and_load_mcp_tools
|
||||
|
||||
await asyncio.to_thread(_warm_mcp_adapter_imports)
|
||||
|
||||
try:
|
||||
mcp_tools, _, mcp_server_info = await resolve_and_load_mcp_tools(
|
||||
explicit_config_path=config.mcp_config_path,
|
||||
|
||||
@@ -6,6 +6,7 @@ import asyncio
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
@@ -27,6 +28,7 @@ from deepagents_code.mcp_tools import (
|
||||
_apply_tool_filter,
|
||||
_check_remote_server,
|
||||
_check_stdio_server,
|
||||
_gather_bounded,
|
||||
_json_error_snippet,
|
||||
_load_tools_from_config,
|
||||
_normalize_mcp_arguments,
|
||||
@@ -1628,6 +1630,28 @@ class TestResolveAndLoadMcpTools:
|
||||
assert manager is None
|
||||
assert infos == []
|
||||
|
||||
@patch("deepagents_code.mcp_tools._warm_mcp_adapter_imports")
|
||||
@patch("deepagents_code.mcp_tools.discover_mcp_configs")
|
||||
async def test_no_adapter_warmup_when_no_active_servers(
|
||||
self,
|
||||
mock_discover: MagicMock,
|
||||
mock_warm: MagicMock,
|
||||
) -> None:
|
||||
"""With no configured servers, MCP adapters are never imported.
|
||||
|
||||
`_warm_mcp_adapter_imports` (and the adapter imports that follow it)
|
||||
live inside `_load_tools_from_config`, which the resolver never reaches
|
||||
when discovery yields no servers — so the warmup must not run.
|
||||
"""
|
||||
mock_discover.return_value = []
|
||||
|
||||
tools, manager, infos = await resolve_and_load_mcp_tools(no_mcp=False)
|
||||
|
||||
assert tools == []
|
||||
assert manager is None
|
||||
assert infos == []
|
||||
mock_warm.assert_not_called()
|
||||
|
||||
@patch("deepagents_code.mcp_tools._load_tools_from_config")
|
||||
@patch("deepagents_code.mcp_tools.discover_mcp_configs")
|
||||
async def test_explicit_path_merges_with_discovery(
|
||||
@@ -2177,6 +2201,528 @@ class TestToolOrdering:
|
||||
await manager.cleanup()
|
||||
|
||||
|
||||
class TestLoadToolsConcurrency:
|
||||
"""`_load_tools_from_config` probes independent servers concurrently.
|
||||
|
||||
These tests pin the new behavior: bounded-concurrency preflight and
|
||||
discovery, with per-server error isolation and cancellation semantics
|
||||
matching the previous sequential loader. The load-bearing ordering
|
||||
guarantee is that `server_infos` follows config order regardless of
|
||||
completion order; the returned tool list is always sorted by tool name
|
||||
(via the terminal sort in the loader), so tool-name assertions here are
|
||||
content checks rather than ordering proofs.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _bypass_stdio_health_check(self) -> Generator[None]:
|
||||
"""Bypass stdio pre-flight so tests focus on discovery concurrency."""
|
||||
with patch("deepagents_code.mcp_tools._check_stdio_server"):
|
||||
yield
|
||||
|
||||
@staticmethod
|
||||
def _config(*names: str) -> dict[str, Any]:
|
||||
return {
|
||||
"mcpServers": {
|
||||
name: {"command": "node", "args": [f"{name}.js"]} for name in names
|
||||
}
|
||||
}
|
||||
|
||||
def _tracking_session_factory(
|
||||
self,
|
||||
*,
|
||||
tool_by_server: dict[str, str],
|
||||
hold: asyncio.Event | None = None,
|
||||
sleep_s: float = 0.05,
|
||||
) -> tuple[Any, dict[str, int]]:
|
||||
"""Build a `create_session` fake that records concurrency.
|
||||
|
||||
Returns the async context-manager fake plus a mutable stats dict with
|
||||
`max_inflight` (peak simultaneous open discovery sessions).
|
||||
"""
|
||||
stats = {"inflight": 0, "max_inflight": 0}
|
||||
|
||||
@asynccontextmanager
|
||||
async def _fake(
|
||||
connection: dict[str, Any],
|
||||
*,
|
||||
_mcp_callbacks: object | None = None,
|
||||
) -> AsyncIterator[AsyncMock]:
|
||||
stats["inflight"] += 1
|
||||
stats["max_inflight"] = max(stats["max_inflight"], stats["inflight"])
|
||||
try:
|
||||
# Derive the server name from the recorded command arg so each
|
||||
# session yields that server's tool.
|
||||
args = connection.get("args") or []
|
||||
server = args[0].removesuffix(".js") if args else "srv"
|
||||
session = AsyncMock()
|
||||
session.initialize = AsyncMock()
|
||||
session.list_tools = AsyncMock(
|
||||
return_value=_make_tool_page(
|
||||
[_make_mcp_tool(tool_by_server[server])]
|
||||
)
|
||||
)
|
||||
if hold is not None:
|
||||
await hold.wait()
|
||||
else:
|
||||
await asyncio.sleep(sleep_s)
|
||||
yield session
|
||||
finally:
|
||||
stats["inflight"] -= 1
|
||||
|
||||
return _fake, stats
|
||||
|
||||
async def test_discovery_runs_concurrently(self) -> None:
|
||||
"""All servers' discovery sessions are open at the same time."""
|
||||
names = ["a", "b", "c", "d"]
|
||||
tool_by_server = {n: f"tool_{n}" for n in names}
|
||||
hold = asyncio.Event()
|
||||
fake, stats = self._tracking_session_factory(
|
||||
tool_by_server=tool_by_server, hold=hold
|
||||
)
|
||||
|
||||
async def _release_when_all_open() -> None:
|
||||
for _ in range(200):
|
||||
if stats["inflight"] >= len(names):
|
||||
break
|
||||
await asyncio.sleep(0.005)
|
||||
hold.set()
|
||||
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", fake):
|
||||
releaser = asyncio.create_task(_release_when_all_open())
|
||||
tools, manager, infos = await _load_tools_from_config(self._config(*names))
|
||||
await releaser
|
||||
|
||||
assert stats["max_inflight"] == len(names)
|
||||
assert [t.name for t in tools] == [
|
||||
"a_tool_a",
|
||||
"b_tool_b",
|
||||
"c_tool_c",
|
||||
"d_tool_d",
|
||||
]
|
||||
assert [i.name for i in infos] == names
|
||||
assert manager is not None
|
||||
await manager.cleanup()
|
||||
|
||||
async def test_discovery_concurrency_is_bounded(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""No more than `_MCP_LOAD_CONCURRENCY` sessions are open at once."""
|
||||
monkeypatch.setattr(
|
||||
"deepagents_code.mcp_tools._MCP_LOAD_CONCURRENCY", 2, raising=True
|
||||
)
|
||||
names = ["s1", "s2", "s3", "s4", "s5"]
|
||||
tool_by_server = {n: f"t_{n}" for n in names}
|
||||
fake, stats = self._tracking_session_factory(
|
||||
tool_by_server=tool_by_server, sleep_s=0.03
|
||||
)
|
||||
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", fake):
|
||||
tools, manager, infos = await _load_tools_from_config(self._config(*names))
|
||||
|
||||
assert stats["max_inflight"] == 2
|
||||
assert [i.name for i in infos] == names
|
||||
assert len(tools) == len(names)
|
||||
assert manager is not None
|
||||
await manager.cleanup()
|
||||
|
||||
async def test_order_preserved_when_later_servers_finish_first(self) -> None:
|
||||
"""`server_infos` follows config order regardless of completion order."""
|
||||
names = ["first", "second", "third"]
|
||||
# Later servers sleep less, so they finish discovery before earlier ones.
|
||||
delays = {"first": 0.09, "second": 0.05, "third": 0.01}
|
||||
finish_order: list[str] = []
|
||||
|
||||
@asynccontextmanager
|
||||
async def _fake(
|
||||
connection: dict[str, Any],
|
||||
*,
|
||||
_mcp_callbacks: object | None = None,
|
||||
) -> AsyncIterator[AsyncMock]:
|
||||
server = (connection.get("args") or ["x"])[0].removesuffix(".js")
|
||||
session = AsyncMock()
|
||||
session.initialize = AsyncMock()
|
||||
session.list_tools = AsyncMock(
|
||||
return_value=_make_tool_page([_make_mcp_tool(f"tool_{server}")])
|
||||
)
|
||||
await asyncio.sleep(delays[server])
|
||||
finish_order.append(server)
|
||||
yield session
|
||||
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", _fake):
|
||||
tools, manager, infos = await _load_tools_from_config(self._config(*names))
|
||||
|
||||
assert finish_order == ["third", "second", "first"]
|
||||
assert [i.name for i in infos] == names
|
||||
assert [t.name for t in tools] == [
|
||||
"first_tool_first",
|
||||
"second_tool_second",
|
||||
"third_tool_third",
|
||||
]
|
||||
assert manager is not None
|
||||
await manager.cleanup()
|
||||
|
||||
async def test_one_server_failure_isolated_from_others(self) -> None:
|
||||
"""A single discovery failure does not abort the other servers."""
|
||||
names = ["ok1", "boom", "ok2"]
|
||||
|
||||
@asynccontextmanager
|
||||
async def _fake(
|
||||
connection: dict[str, Any],
|
||||
*,
|
||||
_mcp_callbacks: object | None = None,
|
||||
) -> AsyncIterator[AsyncMock]:
|
||||
server = (connection.get("args") or ["x"])[0].removesuffix(".js")
|
||||
await asyncio.sleep(0.01)
|
||||
if server == "boom":
|
||||
msg = "discovery exploded"
|
||||
raise RuntimeError(msg)
|
||||
session = AsyncMock()
|
||||
session.initialize = AsyncMock()
|
||||
session.list_tools = AsyncMock(
|
||||
return_value=_make_tool_page([_make_mcp_tool(f"tool_{server}")])
|
||||
)
|
||||
yield session
|
||||
|
||||
with patch("langchain_mcp_adapters.sessions.create_session", _fake):
|
||||
tools, manager, infos = await _load_tools_from_config(self._config(*names))
|
||||
|
||||
by_name = {i.name: i for i in infos}
|
||||
assert [i.name for i in infos] == names
|
||||
assert by_name["ok1"].status == "ok"
|
||||
assert by_name["ok2"].status == "ok"
|
||||
assert by_name["boom"].status == "error"
|
||||
assert "discovery exploded" in (by_name["boom"].error or "")
|
||||
assert [t.name for t in tools] == ["ok1_tool_ok1", "ok2_tool_ok2"]
|
||||
assert manager is not None
|
||||
await manager.cleanup()
|
||||
|
||||
async def test_preflight_failure_isolated_and_order_preserved(self) -> None:
|
||||
"""A mid-config preflight failure is skipped; survivors keep order.
|
||||
|
||||
Exercises the preflight-error path and the fold-in loop that
|
||||
interleaves a skipped server *between* discovered ones in config order,
|
||||
with discovery finishing out of order so the ordering cannot be an
|
||||
accident of completion timing.
|
||||
"""
|
||||
names = ["ok_a", "pf_fail", "ok_b"]
|
||||
# `ok_b` discovers faster than `ok_a`, so completion order is reversed.
|
||||
delays = {"ok_a": 0.06, "ok_b": 0.01}
|
||||
|
||||
def _check(name: str, _cfg: dict[str, Any]) -> None:
|
||||
if name == "pf_fail":
|
||||
msg = "preflight boom"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _fake(
|
||||
connection: dict[str, Any],
|
||||
*,
|
||||
_mcp_callbacks: object | None = None,
|
||||
) -> AsyncIterator[AsyncMock]:
|
||||
server = (connection.get("args") or ["x"])[0].removesuffix(".js")
|
||||
session = AsyncMock()
|
||||
session.initialize = AsyncMock()
|
||||
session.list_tools = AsyncMock(
|
||||
return_value=_make_tool_page([_make_mcp_tool(f"tool_{server}")])
|
||||
)
|
||||
await asyncio.sleep(delays[server])
|
||||
yield session
|
||||
|
||||
with (
|
||||
patch("deepagents_code.mcp_tools._check_stdio_server", _check),
|
||||
patch("langchain_mcp_adapters.sessions.create_session", _fake),
|
||||
):
|
||||
tools, manager, infos = await _load_tools_from_config(self._config(*names))
|
||||
|
||||
by_name = {i.name: i for i in infos}
|
||||
# server_infos follows config order, with the skipped server in place.
|
||||
assert [i.name for i in infos] == names
|
||||
assert by_name["ok_a"].status == "ok"
|
||||
assert by_name["ok_b"].status == "ok"
|
||||
assert by_name["pf_fail"].status == "error"
|
||||
assert "preflight boom" in (by_name["pf_fail"].error or "")
|
||||
assert [t.name for t in tools] == ["ok_a_tool_ok_a", "ok_b_tool_ok_b"]
|
||||
assert manager is not None
|
||||
await manager.cleanup()
|
||||
|
||||
async def test_all_servers_fail_preflight_yields_empty(self) -> None:
|
||||
"""Every server failing preflight yields no tools, infos in order.
|
||||
|
||||
Drives the empty-discovery `_gather_bounded([], ...)` path and asserts a
|
||||
non-`None` (empty) session manager is still returned.
|
||||
"""
|
||||
names = ["x1", "x2", "x3"]
|
||||
|
||||
def _check(_name: str, _cfg: dict[str, Any]) -> None:
|
||||
msg = "nope"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
with patch("deepagents_code.mcp_tools._check_stdio_server", _check):
|
||||
tools, manager, infos = await _load_tools_from_config(self._config(*names))
|
||||
|
||||
assert tools == []
|
||||
assert [i.name for i in infos] == names
|
||||
assert all(i.status == "error" for i in infos)
|
||||
assert manager is not None
|
||||
await manager.cleanup()
|
||||
|
||||
async def test_tool_construction_failure_isolated(self) -> None:
|
||||
"""A post-discovery construction failure degrades one server only.
|
||||
|
||||
Discovery succeeds for every server, but tool filtering raises for one.
|
||||
That server must become an `error` info while its siblings load
|
||||
normally — proving isolation now covers the post-discovery construction
|
||||
block, not just the discovery session. Without that guard the failure
|
||||
would abort the entire concurrent load.
|
||||
"""
|
||||
names = ["good", "bad_build"]
|
||||
fake, _ = self._tracking_session_factory(
|
||||
tool_by_server={n: f"t_{n}" for n in names}, sleep_s=0.0
|
||||
)
|
||||
|
||||
def _filter(
|
||||
server_tools: list[Any], server_name: str, _server_config: dict[str, Any]
|
||||
) -> list[Any]:
|
||||
if server_name == "bad_build":
|
||||
msg = "build boom"
|
||||
raise RuntimeError(msg)
|
||||
return server_tools
|
||||
|
||||
with (
|
||||
patch("langchain_mcp_adapters.sessions.create_session", fake),
|
||||
patch("deepagents_code.mcp_tools._apply_tool_filter", _filter),
|
||||
):
|
||||
tools, manager, infos = await _load_tools_from_config(self._config(*names))
|
||||
|
||||
by_name = {i.name: i for i in infos}
|
||||
assert [i.name for i in infos] == names
|
||||
assert by_name["good"].status == "ok"
|
||||
assert by_name["bad_build"].status == "error"
|
||||
assert "build boom" in (by_name["bad_build"].error or "")
|
||||
assert [t.name for t in tools] == ["good_t_good"]
|
||||
assert manager is not None
|
||||
await manager.cleanup()
|
||||
|
||||
async def test_cancellation_propagates_and_cancels_siblings(self) -> None:
|
||||
"""A cancelled worker propagates and tears down its siblings."""
|
||||
names = ["cancel", "sibling"]
|
||||
sibling_cancelled = asyncio.Event()
|
||||
|
||||
@asynccontextmanager
|
||||
async def _fake(
|
||||
connection: dict[str, Any],
|
||||
*,
|
||||
_mcp_callbacks: object | None = None,
|
||||
) -> AsyncIterator[AsyncMock]:
|
||||
server = (connection.get("args") or ["x"])[0].removesuffix(".js")
|
||||
if server == "cancel":
|
||||
await asyncio.sleep(0.01)
|
||||
raise asyncio.CancelledError
|
||||
try:
|
||||
await asyncio.sleep(10)
|
||||
except asyncio.CancelledError:
|
||||
sibling_cancelled.set()
|
||||
raise
|
||||
yield AsyncMock() # pragma: no cover - never reached
|
||||
|
||||
with (
|
||||
patch("langchain_mcp_adapters.sessions.create_session", _fake),
|
||||
pytest.raises(asyncio.CancelledError),
|
||||
):
|
||||
await _load_tools_from_config(self._config(*names))
|
||||
|
||||
assert sibling_cancelled.is_set()
|
||||
|
||||
async def test_preflight_runs_concurrently(self) -> None:
|
||||
"""Stdio pre-flight checks run concurrently across servers."""
|
||||
names = ["p1", "p2", "p3"]
|
||||
tool_by_server = {n: f"pt_{n}" for n in names}
|
||||
stats = {"inflight": 0, "max_inflight": 0}
|
||||
# `_slow_check` runs in `asyncio.to_thread` worker threads, so the shared
|
||||
# counters must be guarded with a real lock (`+=` is not atomic across
|
||||
# threads) and the barrier must be a thread-safe primitive.
|
||||
stats_lock = threading.Lock()
|
||||
barrier = threading.Event()
|
||||
|
||||
def _slow_check(_name: str, _cfg: dict[str, Any]) -> None:
|
||||
# `_check_stdio_server` is sync and invoked via asyncio.to_thread,
|
||||
# so bump the counter and block until every worker is in-flight.
|
||||
with stats_lock:
|
||||
stats["inflight"] += 1
|
||||
stats["max_inflight"] = max(stats["max_inflight"], stats["inflight"])
|
||||
barrier.wait()
|
||||
with stats_lock:
|
||||
stats["inflight"] -= 1
|
||||
|
||||
async def _release() -> None:
|
||||
for _ in range(200):
|
||||
with stats_lock:
|
||||
peak = stats["max_inflight"]
|
||||
if peak >= len(names):
|
||||
break
|
||||
await asyncio.sleep(0.005)
|
||||
barrier.set()
|
||||
|
||||
fake, _ = self._tracking_session_factory(
|
||||
tool_by_server=tool_by_server, sleep_s=0.0
|
||||
)
|
||||
with (
|
||||
patch("deepagents_code.mcp_tools._check_stdio_server", _slow_check),
|
||||
patch("langchain_mcp_adapters.sessions.create_session", fake),
|
||||
):
|
||||
releaser = asyncio.create_task(_release())
|
||||
_tools, manager, infos = await _load_tools_from_config(self._config(*names))
|
||||
await releaser
|
||||
|
||||
assert stats["max_inflight"] == len(names)
|
||||
assert [i.name for i in infos] == names
|
||||
assert manager is not None
|
||||
await manager.cleanup()
|
||||
|
||||
async def test_warmup_runs_off_loop_before_discovery(self) -> None:
|
||||
"""Adapter warmup runs, off the event loop, before any discovery."""
|
||||
loop_thread_id = threading.get_ident()
|
||||
events: list[tuple[str, int]] = []
|
||||
|
||||
def _warm() -> None:
|
||||
events.append(("warm", threading.get_ident()))
|
||||
|
||||
@asynccontextmanager
|
||||
async def _fake(
|
||||
connection: dict[str, Any],
|
||||
*,
|
||||
_mcp_callbacks: object | None = None,
|
||||
) -> AsyncIterator[AsyncMock]:
|
||||
events.append(("discover", threading.get_ident()))
|
||||
session = AsyncMock()
|
||||
session.initialize = AsyncMock()
|
||||
session.list_tools = AsyncMock(return_value=_make_tool_page([]))
|
||||
yield session
|
||||
|
||||
with (
|
||||
patch("deepagents_code.mcp_tools._warm_mcp_adapter_imports", _warm),
|
||||
patch("langchain_mcp_adapters.sessions.create_session", _fake),
|
||||
):
|
||||
_tools, manager, _infos = await _load_tools_from_config(
|
||||
self._config("only")
|
||||
)
|
||||
|
||||
assert events[0][0] == "warm"
|
||||
assert events[0][1] != loop_thread_id
|
||||
assert any(kind == "discover" for kind, _ in events)
|
||||
assert manager is not None
|
||||
await manager.cleanup()
|
||||
|
||||
|
||||
class TestGatherBounded:
|
||||
"""Direct tests for the `_gather_bounded` concurrency helper.
|
||||
|
||||
These pin the helper's contract independently of MCP loading: submission
|
||||
(not completion) ordering, the empty and clamped-limit edge cases, and the
|
||||
failure path that cancels + awaits siblings and never silently drops a
|
||||
concurrent failure.
|
||||
"""
|
||||
|
||||
async def test_results_follow_submission_order(self) -> None:
|
||||
"""Results zip back to submission order even when completion differs."""
|
||||
completed: list[int] = []
|
||||
|
||||
def _factory(idx: int, delay: float) -> Callable[[], Any]:
|
||||
async def _run() -> int:
|
||||
await asyncio.sleep(delay)
|
||||
completed.append(idx)
|
||||
return idx
|
||||
|
||||
return _run
|
||||
|
||||
# Index 0 finishes last, index 2 finishes first.
|
||||
factories = [_factory(0, 0.03), _factory(1, 0.02), _factory(2, 0.001)]
|
||||
results = await _gather_bounded(factories, limit=8)
|
||||
|
||||
assert results == [0, 1, 2]
|
||||
assert completed == [2, 1, 0]
|
||||
|
||||
async def test_empty_returns_empty(self) -> None:
|
||||
"""Zero factories return an empty list without touching the loop."""
|
||||
assert await _gather_bounded([], limit=8) == []
|
||||
|
||||
async def test_limit_below_one_is_clamped_to_serial(self) -> None:
|
||||
"""A limit < 1 is clamped to 1, so factories run strictly serially."""
|
||||
active = {"n": 0, "max": 0}
|
||||
|
||||
def _factory() -> Callable[[], Any]:
|
||||
async def _run() -> None:
|
||||
active["n"] += 1
|
||||
active["max"] = max(active["max"], active["n"])
|
||||
await asyncio.sleep(0.01)
|
||||
active["n"] -= 1
|
||||
|
||||
return _run
|
||||
|
||||
await _gather_bounded([_factory(), _factory(), _factory()], limit=0)
|
||||
assert active["max"] == 1
|
||||
|
||||
async def test_failure_cancels_and_awaits_siblings(self) -> None:
|
||||
"""A raising factory cancels the rest and awaits them before raising."""
|
||||
sibling = {"cancelled": False, "completed": False}
|
||||
started = asyncio.Event()
|
||||
|
||||
def _failing() -> Callable[[], Any]:
|
||||
async def _run() -> None:
|
||||
await started.wait()
|
||||
msg = "boom"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
return _run
|
||||
|
||||
def _sibling() -> Callable[[], Any]:
|
||||
async def _run() -> None:
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.sleep(10)
|
||||
except asyncio.CancelledError:
|
||||
sibling["cancelled"] = True
|
||||
raise
|
||||
sibling["completed"] = True # pragma: no cover - never reached
|
||||
|
||||
return _run
|
||||
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
await _gather_bounded([_failing(), _sibling()], limit=8)
|
||||
|
||||
assert sibling["cancelled"] is True
|
||||
assert sibling["completed"] is False
|
||||
|
||||
async def test_concurrent_failures_are_logged_not_lost(
|
||||
self, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""When several factories fail, no failure vanishes silently.
|
||||
|
||||
`asyncio.gather` propagates only the first exception; the others are
|
||||
logged at debug so a concurrent failure is never dropped without trace.
|
||||
"""
|
||||
|
||||
def _failing(message: str) -> Callable[[], Any]:
|
||||
async def _run() -> None:
|
||||
raise RuntimeError(message)
|
||||
|
||||
return _run
|
||||
|
||||
with (
|
||||
caplog.at_level(logging.DEBUG, logger="deepagents_code.mcp_tools"),
|
||||
pytest.raises(RuntimeError),
|
||||
):
|
||||
await _gather_bounded(
|
||||
[_failing("first_failure"), _failing("second_failure")], limit=8
|
||||
)
|
||||
|
||||
assert "sibling task failed" in caplog.text
|
||||
# Both failures are represented in the captured logs (the propagated one
|
||||
# plus the logged sibling), so neither is lost.
|
||||
assert "first_failure" in caplog.text
|
||||
assert "second_failure" in caplog.text
|
||||
|
||||
|
||||
class TestCachedSessionProxy:
|
||||
"""Runtime tool wrappers use lazy cached sessions with retry semantics."""
|
||||
|
||||
|
||||
@@ -78,7 +78,6 @@ class TestServerGraph:
|
||||
loop_thread_id = threading.get_ident()
|
||||
create_cli_agent_thread_ids: list[int] = []
|
||||
create_model_thread_ids: list[int] = []
|
||||
warm_import_thread_ids: list[int] = []
|
||||
|
||||
def create_cli_agent_side_effect(**_: object) -> tuple[object, object]:
|
||||
create_cli_agent_thread_ids.append(threading.get_ident())
|
||||
@@ -88,9 +87,6 @@ class TestServerGraph:
|
||||
create_model_thread_ids.append(threading.get_ident())
|
||||
return model_result
|
||||
|
||||
def warm_import_side_effect() -> None:
|
||||
warm_import_thread_ids.append(threading.get_ident())
|
||||
|
||||
create_cli_agent = MagicMock(side_effect=create_cli_agent_side_effect)
|
||||
agent_module = _module_with_attrs(
|
||||
"deepagents_code.agent",
|
||||
@@ -164,17 +160,10 @@ class TestServerGraph:
|
||||
|
||||
module = _import_fresh_server_graph()
|
||||
resolve_mcp_tools.assert_not_awaited()
|
||||
with patch.object(
|
||||
module,
|
||||
"_warm_mcp_adapter_imports",
|
||||
side_effect=warm_import_side_effect,
|
||||
):
|
||||
assert await module.make_graph() is graph_obj
|
||||
assert await module.make_graph() is graph_obj
|
||||
|
||||
configure_redaction.assert_called_once_with()
|
||||
resolve_mcp_tools.assert_awaited_once()
|
||||
assert warm_import_thread_ids
|
||||
assert warm_import_thread_ids[0] != loop_thread_id
|
||||
assert create_cli_agent_thread_ids
|
||||
assert create_cli_agent_thread_ids[0] != loop_thread_id
|
||||
# `create_model` must run off the loop thread: it does blocking disk IO
|
||||
@@ -214,7 +203,7 @@ class TestServerGraph:
|
||||
)
|
||||
|
||||
async def test_build_tools_skips_mcp_when_disabled(self) -> None:
|
||||
"""`no_mcp=True` should not warm imports or call the MCP resolver."""
|
||||
"""`no_mcp=True` should not call the MCP resolver at all."""
|
||||
fetch_tool = object()
|
||||
thread_tool = object()
|
||||
resolve_mcp_tools = AsyncMock()
|
||||
@@ -242,16 +231,13 @@ class TestServerGraph:
|
||||
},
|
||||
):
|
||||
module = _import_fresh_server_graph()
|
||||
warm_imports = MagicMock(side_effect=AssertionError("MCP warmup ran"))
|
||||
with patch.object(module, "_warm_mcp_adapter_imports", warm_imports):
|
||||
tools, mcp_server_info = await module._build_tools(
|
||||
ServerConfig(no_mcp=True),
|
||||
None,
|
||||
)
|
||||
tools, mcp_server_info = await module._build_tools(
|
||||
ServerConfig(no_mcp=True),
|
||||
None,
|
||||
)
|
||||
|
||||
assert tools == [fetch_tool, thread_tool]
|
||||
assert mcp_server_info is None
|
||||
warm_imports.assert_not_called()
|
||||
resolve_mcp_tools.assert_not_awaited()
|
||||
|
||||
async def test_interpreter_settings_apply_before_agent_construction(self) -> None:
|
||||
@@ -332,18 +318,20 @@ class TestServerGraph:
|
||||
"enable_interpreter": True,
|
||||
}
|
||||
|
||||
async def test_mcp_adapter_warmup_runs_before_mcp_resolver(self) -> None:
|
||||
"""MCP imports must be warmed before resolver imports adapter modules."""
|
||||
events: list[str] = []
|
||||
async def test_build_tools_delegates_mcp_loading_to_resolver(self) -> None:
|
||||
"""`_build_tools` should defer all MCP work to the resolver.
|
||||
|
||||
Adapter warmup now lives inside `_load_tools_from_config` (gated on
|
||||
active servers existing), so `_build_tools` no longer warms imports
|
||||
itself — it just calls the resolver and appends the returned tools.
|
||||
"""
|
||||
fetch_tool = object()
|
||||
thread_tool = object()
|
||||
|
||||
def resolve_mcp_tools(
|
||||
**_: object,
|
||||
) -> tuple[list[object], None, list[object]]:
|
||||
events.append("resolver")
|
||||
return [], None, []
|
||||
class FakeSessionManager:
|
||||
pass
|
||||
|
||||
resolve_mcp_tools = AsyncMock(return_value=([], None, []))
|
||||
config_module = _module_with_attrs(
|
||||
"deepagents_code.config",
|
||||
settings=SimpleNamespace(has_tavily=False),
|
||||
@@ -354,14 +342,10 @@ class TestServerGraph:
|
||||
get_current_thread_id=thread_tool,
|
||||
web_search=object(),
|
||||
)
|
||||
|
||||
class FakeSessionManager:
|
||||
pass
|
||||
|
||||
mcp_module = _module_with_attrs(
|
||||
"deepagents_code.mcp_tools",
|
||||
MCPSessionManager=FakeSessionManager,
|
||||
resolve_and_load_mcp_tools=AsyncMock(side_effect=resolve_mcp_tools),
|
||||
resolve_and_load_mcp_tools=resolve_mcp_tools,
|
||||
)
|
||||
|
||||
with patch.dict(
|
||||
@@ -373,17 +357,13 @@ class TestServerGraph:
|
||||
},
|
||||
):
|
||||
module = _import_fresh_server_graph()
|
||||
assert not hasattr(module, "_warm_mcp_adapter_imports")
|
||||
tools, mcp_server_info = await module._build_tools(
|
||||
ServerConfig(no_mcp=False),
|
||||
None,
|
||||
)
|
||||
|
||||
def warm_imports() -> None:
|
||||
events.append("warmup")
|
||||
|
||||
with patch.object(module, "_warm_mcp_adapter_imports", warm_imports):
|
||||
tools, mcp_server_info = await module._build_tools(
|
||||
ServerConfig(no_mcp=False),
|
||||
None,
|
||||
)
|
||||
|
||||
assert events == ["warmup", "resolver"]
|
||||
resolve_mcp_tools.assert_awaited_once()
|
||||
assert tools == [fetch_tool, thread_tool]
|
||||
assert mcp_server_info == []
|
||||
|
||||
|
||||
Reference in New Issue
Block a user