mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-12 22:20:33 +00:00
2feec9cb09
* fix(mcp): degrade a blob without mimeType instead of failing the tool call A resource link or embedded resource without mimeType made the blob guard raise, so the whole tool call failed instead of degrading the attachment. Keep the guard for genuinely unusable payloads and let a missing media type fall through to the caller's default. * Modify test for text blocks and mime type validation Update test to check for text blocks and their mime types. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * test(mcp): remove MIME assertion from text content block --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
1509 lines
64 KiB
Python
1509 lines
64 KiB
Python
"""Load MCP tools using langchain-mcp-adapters with stdio session pooling."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import re
|
|
from collections.abc import Callable, Collection, Iterable, Mapping
|
|
from dataclasses import dataclass, field
|
|
from datetime import timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.parse import urlparse
|
|
from urllib.request import url2pathname
|
|
|
|
from langchain_core.tools import BaseTool, StructuredTool
|
|
from langgraph.config import get_config
|
|
|
|
from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig, resolve_effective_mcp_routing
|
|
from deerflow.config.paths import VIRTUAL_PATH_PREFIX, Paths, get_paths
|
|
from deerflow.constants import DEFAULT_MCP_SESSION_INIT_TIMEOUT, MCP_TMP_SUBDIR
|
|
from deerflow.mcp.client import build_servers_config
|
|
from deerflow.mcp.headers import apply_header_overrides
|
|
from deerflow.mcp.interceptors import build_mcp_tool_interceptors, compose_tool_interceptors
|
|
from deerflow.mcp.oauth import OAuthTokenManager, build_oauth_tool_interceptor, get_initial_oauth_headers
|
|
from deerflow.mcp.session_pool import (
|
|
MCPPoolDomain,
|
|
MCPSessionPool,
|
|
ServerBinding,
|
|
StaleMCPBindingError,
|
|
call_pooled_session_tool,
|
|
get_session_pool,
|
|
normalized_connection_fingerprint,
|
|
)
|
|
from deerflow.mcp.tasks import ORDINARY_MCP_TASK_DRIVER, TaskSubmitRequest
|
|
from deerflow.mcp.tasks.runtime import (
|
|
McpTaskConfigurationError,
|
|
get_mcp_task_oauth_token_manager,
|
|
get_mcp_task_submitter,
|
|
validate_mcp_task_config_snapshot,
|
|
)
|
|
from deerflow.mcp_scope import mcp_session_scope_key, runtime_thread_incarnation
|
|
from deerflow.reflection import resolve_variable
|
|
from deerflow.runtime.user_context import resolve_runtime_user_id
|
|
from deerflow.tools.mcp_metadata import tag_mcp_routing, tag_mcp_tool
|
|
from deerflow.tools.resource_placeholder import model_visible_location, resource_placeholder_text
|
|
from deerflow.tools.sync import make_sync_tool_wrapper
|
|
from deerflow.tools.types import Runtime
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# MCP tool names arrive verbatim from external (potentially hostile/compromised)
|
|
# servers. A tool name is only ever a function identifier: the provider's
|
|
# function-calling API validates it against this same charset at bind time. But
|
|
# deferred (tool_search) MCP tools are withheld from binding, so that provider
|
|
# check never runs on their names — they only ever live in the system-prompt
|
|
# string, where a crafted name (newlines, markdown, angle brackets) could forge
|
|
# framework prompt structure. Canonicalizing at the load boundary constrains
|
|
# both bound and deferred names to the same safe identifier charset, mirroring
|
|
# the load-time validation skill names get (skills/storage/skill_storage.py).
|
|
_VALID_MCP_TOOL_NAME = re.compile(r"^[A-Za-z0-9_-]+$")
|
|
|
|
# Matches local-file references embedded in free text returned by an MCP server.
|
|
# Some servers (notably Playwright's ``browser_take_screenshot``) report saved
|
|
# files only as text/markdown links rather than ``ResourceLink`` blocks. Those
|
|
# references may be absolute paths, ``file://`` URIs, or paths relative to the
|
|
# server process cwd (e.g. ``temp/page.yml``, ``./shot.png``). Each match is
|
|
# only rewritten when it resolves to an existing file inside the thread's
|
|
# user-data tree. URL spans are excluded before applying the path matchers so
|
|
# a URL query/fragment cannot be mistaken for a generated local file.
|
|
_LOCAL_PATH_IN_TEXT_RE = re.compile(
|
|
r"(?:file://)?/[^\s'\"<>|*?]+" # POSIX absolute path or file:// URI
|
|
r"|file://[A-Za-z]:[^\s'\"<>|*?]+" # file://C:/… — some Windows tools skip the third slash
|
|
# Windows drive-qualified absolute path; the lookbehind keeps a word
|
|
# character before the colon (file:/…, id:/…) on the earlier alternatives
|
|
r"|(?<![\w.-])[A-Za-z]:[\\/][^\s'\"<>|*?]+"
|
|
# path relative to the server cwd (Windows servers print "\" separators)
|
|
r"|(?:\.{0,2}[\\/]|[\w.-]+[\\/])[^\s'\"<>|*?]+"
|
|
)
|
|
|
|
# Require :// so labels such as "Saved:temp/page.yml" remain local references.
|
|
# At least two scheme characters leave Windows drive paths to local handling.
|
|
_URL_SCHEME_IN_TEXT_RE = re.compile(r"(?<![\w.+/\\-])(?P<scheme>[A-Za-z][A-Za-z0-9+.-]+)://")
|
|
|
|
# Trailing characters that are punctuation/markup rather than part of a path.
|
|
_TEXT_PATH_TRAILING_CHARS = ".,;:!?)]}>\"'`"
|
|
|
|
_FILE_SNAPSHOT = dict[Path, tuple[int, int]]
|
|
|
|
|
|
def _local_path_from_uri(uri: str, *, base_dir: Path | None = None) -> Path | None:
|
|
"""Return an absolute local filesystem ``Path`` if *uri* points to a local
|
|
file, otherwise ``None``.
|
|
|
|
Accepts bare paths and ``file://`` URIs. Remote URIs
|
|
(``http``/``https``/``data``/...) return ``None`` so the caller leaves them
|
|
untouched. Relative paths are resolved only when *base_dir* is supplied.
|
|
"""
|
|
if not uri:
|
|
return None
|
|
try:
|
|
parsed = urlparse(uri)
|
|
except ValueError:
|
|
return None
|
|
if parsed.scheme == "file":
|
|
# url2pathname converts the "/C:/..." form a file URI's path takes on
|
|
# Windows into a drive-qualified "C:\..." path; on POSIX it is identity.
|
|
# It already percent-decodes, so no extra unquote here, and it can
|
|
# reject odd Windows spellings with OSError — leave those untouched.
|
|
netloc = parsed.netloc
|
|
if netloc and netloc.lower() != "localhost":
|
|
# Some Windows tools emit file://C:/… (two slashes): the drive
|
|
# lands in the URI authority. Any other host is not a local file.
|
|
if len(netloc) != 2 or not netloc[0].isalpha() or netloc[1] != ":":
|
|
return None
|
|
url_path = f"/{netloc}{parsed.path}"
|
|
else:
|
|
url_path = parsed.path
|
|
try:
|
|
raw = url2pathname(url_path)
|
|
except OSError:
|
|
return None
|
|
elif len(parsed.scheme) == 1 and parsed.scheme.isalpha():
|
|
# urlparse reads a Windows drive prefix ("C:\...") as the URI scheme;
|
|
# the original string is a bare local path, not a remote URI.
|
|
raw = uri
|
|
elif parsed.scheme == "":
|
|
raw = uri
|
|
else:
|
|
return None
|
|
if not raw:
|
|
return None
|
|
path = Path(raw)
|
|
if not path.is_absolute():
|
|
if base_dir is None:
|
|
return None
|
|
path = base_dir / path
|
|
return path
|
|
|
|
|
|
def _local_uri_to_virtual_path(
|
|
uri: str,
|
|
*,
|
|
thread_id: str,
|
|
user_id: str,
|
|
source_base_dir: Path | None = None,
|
|
) -> str | None:
|
|
"""Translate a local file reference into its ``/mnt/user-data/...`` virtual path.
|
|
|
|
Stdio MCP servers run with their cwd and temp dir pinned inside the thread's
|
|
mounted user-data tree (see :func:`_make_session_pool_tool`), so the files
|
|
they produce already live somewhere the sandbox/artifact API can serve — the
|
|
only thing missing is the virtual prefix the rest of DeerFlow addresses them
|
|
by. This performs that purely deterministic host→virtual mapping: no copy, no
|
|
trusted-root list, and no exposure of files outside the thread's own tree.
|
|
|
|
Returns ``None`` (so the caller leaves the reference untouched) when the URI
|
|
is remote, cannot be resolved, points outside this thread's user-data tree,
|
|
or does not name an existing file. Relative references are resolved against
|
|
*source_base_dir* (the server's cwd).
|
|
"""
|
|
src = _local_path_from_uri(uri, base_dir=source_base_dir)
|
|
if src is None:
|
|
return None
|
|
|
|
try:
|
|
real = src.resolve()
|
|
if not real.is_file():
|
|
return None
|
|
except OSError:
|
|
return None
|
|
|
|
try:
|
|
user_data_root = get_paths().sandbox_user_data_dir(thread_id, user_id=user_id).resolve()
|
|
except OSError:
|
|
return None
|
|
|
|
try:
|
|
relative = real.relative_to(user_data_root)
|
|
except ValueError:
|
|
# The file lives outside this thread's user-data mount; we cannot
|
|
# express it as a virtual path, so leave the original reference as-is.
|
|
logger.debug("MCP path rewrite skipped outside user-data tree: %s", real)
|
|
return None
|
|
|
|
virtual_path = f"{VIRTUAL_PATH_PREFIX}/{relative.as_posix()}"
|
|
logger.debug("MCP path rewrite: %s -> %s", real, virtual_path)
|
|
return virtual_path
|
|
|
|
|
|
def _snapshot_workspace_files(root: Path) -> _FILE_SNAPSHOT:
|
|
"""Return a lightweight snapshot of regular files under *root*."""
|
|
snapshot: _FILE_SNAPSHOT = {}
|
|
if not root.exists():
|
|
return snapshot
|
|
|
|
try:
|
|
candidates = root.rglob("*")
|
|
for path in candidates:
|
|
try:
|
|
stat = path.stat()
|
|
except OSError:
|
|
continue
|
|
if path.is_file():
|
|
snapshot[path] = (stat.st_mtime_ns, stat.st_size)
|
|
except OSError:
|
|
return snapshot
|
|
return snapshot
|
|
|
|
|
|
def _changed_workspace_files(root: Path, before: _FILE_SNAPSHOT) -> list[Path]:
|
|
"""Return files under *root* that were created or modified since *before*."""
|
|
after = _snapshot_workspace_files(root)
|
|
return [path for path, signature in after.items() if before.get(path) != signature]
|
|
|
|
|
|
def _prepare_stdio_workspace(paths: Paths, *, thread_id: str, user_id: str) -> tuple[Path, Path, _FILE_SNAPSHOT]:
|
|
"""Prepare the thread workspace for a pinned stdio MCP subprocess.
|
|
|
|
Bundles all the synchronous filesystem work (dir creation, temp-dir prep,
|
|
and the pre-call snapshot) into one helper so the caller can run it off the
|
|
event loop via :func:`asyncio.to_thread`. Returns the workspace cwd, the
|
|
pinned temp dir, and the pre-call file snapshot.
|
|
"""
|
|
paths.ensure_thread_dirs(thread_id, user_id=user_id)
|
|
source_base_dir = paths.sandbox_work_dir(thread_id, user_id=user_id)
|
|
tmp_dir = source_base_dir / MCP_TMP_SUBDIR
|
|
try:
|
|
tmp_dir.mkdir(parents=True, exist_ok=True)
|
|
tmp_dir.chmod(0o700)
|
|
except OSError:
|
|
logger.warning("Failed to prepare MCP temp dir: %s", tmp_dir, exc_info=True)
|
|
before_files = _snapshot_workspace_files(source_base_dir)
|
|
return source_base_dir, tmp_dir, before_files
|
|
|
|
|
|
def _result_has_text_content(call_tool_result: Any) -> bool:
|
|
"""Return ``True`` when the MCP result carries any text content.
|
|
|
|
The after-call snapshot diff only feeds bare-filename correlation in free
|
|
text. When the result has no text blocks there is nothing to rewrite, so the
|
|
caller can skip the second recursive walk entirely.
|
|
"""
|
|
from mcp.types import EmbeddedResource, TextContent, TextResourceContents
|
|
|
|
content = getattr(call_tool_result, "content", None)
|
|
if not content:
|
|
return False
|
|
for item in content:
|
|
if isinstance(item, TextContent):
|
|
return True
|
|
if isinstance(item, EmbeddedResource) and isinstance(item.resource, TextResourceContents):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _sub_outside_remote_urls(pattern: re.Pattern[str], replacement: Callable[[re.Match[str]], str], text: str) -> tuple[str, int]:
|
|
"""Apply a local-file substitution outside explicit non-file URLs."""
|
|
parts: list[str] = []
|
|
previous_end = 0
|
|
count = 0
|
|
search_from = 0
|
|
while match := _URL_SCHEME_IN_TEXT_RE.search(text, search_from):
|
|
end = match.end()
|
|
prefix_start = match.start()
|
|
single_quoted = prefix_start > 0 and text[prefix_start - 1] == "'"
|
|
while prefix_start > 0 and text[prefix_start - 1].isspace():
|
|
prefix_start -= 1
|
|
in_markdown_link = text[max(0, prefix_start - 2) : prefix_start] == "]("
|
|
depth = 0
|
|
escaped = False
|
|
while end < len(text):
|
|
char = text[end]
|
|
if char.isspace() or (char in '<>"`|' and not escaped):
|
|
break
|
|
# An apostrophe is a delimiter only for a single-quoted URL.
|
|
if single_quoted and char == "'" and not escaped:
|
|
break
|
|
# Stop at the Markdown destination's closing paren, not at balanced
|
|
# or escaped parentheses inside the URL. Adjacent links stay local.
|
|
if in_markdown_link and not escaped:
|
|
if char == "(":
|
|
depth += 1
|
|
elif char == ")":
|
|
if depth == 0:
|
|
break
|
|
depth -= 1
|
|
escaped = char == "\\" and not escaped
|
|
end += 1
|
|
search_from = end
|
|
if match.group("scheme").lower() == "file":
|
|
continue
|
|
rewritten, substitutions = pattern.subn(replacement, text[previous_end : match.start()])
|
|
parts.extend((rewritten, text[match.start() : end]))
|
|
count += substitutions
|
|
previous_end = end
|
|
rewritten, substitutions = pattern.subn(replacement, text[previous_end:])
|
|
parts.append(rewritten)
|
|
return "".join(parts), count + substitutions
|
|
|
|
|
|
def _rewrite_unique_bare_filenames(
|
|
text: str,
|
|
*,
|
|
changed_files: Iterable[Path],
|
|
thread_id: str,
|
|
user_id: str,
|
|
source_base_dir: Path | None = None,
|
|
) -> str:
|
|
"""Rewrite bare filenames only when this call produced a unique match.
|
|
|
|
A response like ``Saved as page-2026.yml`` is not structurally a path. The
|
|
only safe way to interpret it is to correlate the filename with files
|
|
created/modified by this exact tool call, and rewrite only when the basename
|
|
maps to exactly one file inside this thread's mounted user-data tree.
|
|
"""
|
|
candidates: dict[str, list[str]] = {}
|
|
for path in changed_files:
|
|
virtual_path = _local_uri_to_virtual_path(
|
|
str(path),
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
source_base_dir=source_base_dir,
|
|
)
|
|
if virtual_path is None:
|
|
continue
|
|
candidates.setdefault(path.name, []).append(virtual_path)
|
|
|
|
unique = {name: paths[0] for name, paths in candidates.items() if len(set(paths)) == 1}
|
|
if not unique:
|
|
if candidates:
|
|
logger.debug("MCP bare filename rewrite skipped: no unique candidate in %s", sorted(candidates))
|
|
else:
|
|
logger.debug("MCP bare filename rewrite skipped: no snapshot candidates")
|
|
return text
|
|
|
|
rewritten = text
|
|
for name in sorted(unique, key=len, reverse=True):
|
|
# Do not rewrite inside longer paths/words, with either path separator.
|
|
# A final sentence period is allowed, but ".bak" or another segment is not.
|
|
pattern = re.compile(rf"(?<![\w./\\-]){re.escape(name)}(?!(?:[\w/\\-]|\.[\w]))")
|
|
# A callable replacement, not a template: the virtual path is built from
|
|
# the real file's relative path, where a backslash is an ordinary
|
|
# character, so it must never be read as a regex escape.
|
|
rewritten_text, count = _sub_outside_remote_urls(pattern, lambda _match: unique[name], rewritten)
|
|
if count:
|
|
logger.debug("MCP bare filename rewrite: %s -> %s", name, unique[name])
|
|
rewritten = rewritten_text
|
|
return rewritten
|
|
|
|
|
|
def _rewrite_changed_paths_with_spaces(
|
|
text: str,
|
|
*,
|
|
changed_files: Iterable[Path],
|
|
thread_id: str,
|
|
user_id: str,
|
|
source_base_dir: Path | None,
|
|
) -> str:
|
|
"""Rewrite changed paths with spaces only at unambiguous text boundaries."""
|
|
candidates: dict[str, set[str]] = {}
|
|
for path in changed_files:
|
|
spellings = [str(path)]
|
|
if path.anchor == "/":
|
|
# Some servers print file URIs without percent-encoding spaces.
|
|
spellings.extend((f"file:{path}", f"file://{path}"))
|
|
if source_base_dir is not None:
|
|
try:
|
|
relative = path.relative_to(source_base_dir).as_posix()
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
spellings.extend((relative, f"./{relative}"))
|
|
matching_spellings = [spelling for spelling in spellings if any(char.isspace() for char in spelling) and spelling in text]
|
|
if not matching_spellings:
|
|
continue
|
|
virtual_path = _local_uri_to_virtual_path(str(path), thread_id=thread_id, user_id=user_id)
|
|
if virtual_path is None:
|
|
continue
|
|
for spelling in matching_spellings:
|
|
candidates.setdefault(spelling, set()).add(virtual_path)
|
|
|
|
rewritten = text
|
|
for spelling in sorted(candidates, key=len, reverse=True):
|
|
destinations = candidates[spelling]
|
|
if len(destinations) != 1 or spelling not in rewritten:
|
|
continue
|
|
# Whitespace after an unquoted spelling may continue a longer filename.
|
|
# Require punctuation or the end of the text instead of rewriting a prefix.
|
|
pattern = re.compile(rf"(?<![\w./\\-]){re.escape(spelling)}(?!(?:[\w\s/\\-]|\.[\w]))")
|
|
replacement = next(iter(destinations))
|
|
rewritten, _ = _sub_outside_remote_urls(pattern, lambda _match: replacement, rewritten)
|
|
return rewritten
|
|
|
|
|
|
def _rewrite_local_paths_in_text(
|
|
text: str,
|
|
*,
|
|
thread_id: str,
|
|
user_id: str,
|
|
source_base_dir: Path | None = None,
|
|
changed_files: Iterable[Path] | None = None,
|
|
) -> str:
|
|
"""Best-effort rewrite of local file references found in free text.
|
|
|
|
Some MCP servers (notably Playwright's ``browser_take_screenshot``) report
|
|
the saved file only as free text — e.g. ``Took the screenshot and saved it
|
|
as temp/page-2026.png`` — instead of a ``ResourceLink``. Free text is not a
|
|
reliable protocol, so this is deliberately conservative: every candidate
|
|
token is handed to :func:`_local_uri_to_virtual_path`, which only rewrites
|
|
it when it resolves to an existing file inside this thread's user-data tree.
|
|
Tokens that are not real paths (or point elsewhere) are left exactly as they
|
|
were. Explicit non-file URLs are preserved in every pass, even if a query or
|
|
fragment happens to contain a path or filename that exists locally.
|
|
"""
|
|
translated_by_source: dict[str, str | None] = {}
|
|
|
|
def _replace(match: re.Match[str]) -> str:
|
|
token = match.group(0)
|
|
# A path can end a sentence ("saved as temp/a.png."); strip trailing
|
|
# punctuation and restore it after the (possibly rewritten) path.
|
|
stripped = token.rstrip(_TEXT_PATH_TRAILING_CHARS)
|
|
trailing = token[len(stripped) :]
|
|
if stripped not in translated_by_source:
|
|
translated_by_source[stripped] = _local_uri_to_virtual_path(
|
|
stripped,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
source_base_dir=source_base_dir,
|
|
)
|
|
rewritten = translated_by_source[stripped]
|
|
if rewritten is None:
|
|
return token
|
|
return f"{rewritten}{trailing}"
|
|
|
|
if changed_files is not None:
|
|
changed_files = list(changed_files)
|
|
text = _rewrite_changed_paths_with_spaces(
|
|
text,
|
|
changed_files=changed_files,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
source_base_dir=source_base_dir,
|
|
)
|
|
|
|
rewritten, _ = _sub_outside_remote_urls(_LOCAL_PATH_IN_TEXT_RE, _replace, text)
|
|
if changed_files is None:
|
|
return rewritten
|
|
return _rewrite_unique_bare_filenames(
|
|
rewritten,
|
|
changed_files=changed_files,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
source_base_dir=source_base_dir,
|
|
)
|
|
|
|
|
|
def _extract_thread_id(runtime: Runtime | None) -> str:
|
|
"""Extract thread_id from the injected tool runtime or LangGraph config."""
|
|
if runtime is not None:
|
|
tid = runtime.context.get("thread_id") if runtime.context else None
|
|
if tid is not None:
|
|
return str(tid)
|
|
config = runtime.config or {}
|
|
tid = config.get("configurable", {}).get("thread_id")
|
|
if tid is not None:
|
|
return str(tid)
|
|
|
|
try:
|
|
tid = get_config().get("configurable", {}).get("thread_id")
|
|
return str(tid) if tid is not None else "default"
|
|
except RuntimeError:
|
|
return "default"
|
|
|
|
|
|
def _convert_call_tool_result(
|
|
call_tool_result: Any,
|
|
*,
|
|
thread_id: str | None = None,
|
|
user_id: str | None = None,
|
|
source_base_dir: Path | None = None,
|
|
changed_files: Iterable[Path] | None = None,
|
|
) -> Any:
|
|
"""Convert an MCP CallToolResult to the LangChain ``content_and_artifact`` format.
|
|
|
|
Implements the same conversion logic as the adapter without relying on
|
|
the private ``langchain_mcp_adapters.tools._convert_call_tool_result`` symbol.
|
|
|
|
When ``thread_id`` and ``user_id`` are provided, local files referenced by
|
|
``ResourceLink`` blocks or plain text (e.g. screenshots saved by Playwright
|
|
MCP) have their references translated from the host path to the
|
|
``/mnt/user-data/...`` virtual path so they can be resolved by the sandbox
|
|
and artifact API. The files themselves are not copied — stdio servers run
|
|
with their cwd/temp pinned inside the mounted tree, so they already live in
|
|
a servable location. Remote URIs and files outside the thread's user-data
|
|
tree are left untouched.
|
|
|
|
``ResourceLink`` items never become URL-sourced ``file`` blocks: Chat
|
|
Completions message serialization rejects those, and once the result is
|
|
checkpointed every later turn of the thread fails. An ``http(s)`` image
|
|
link stays an image block; every other link (``ui://`` cards, local files
|
|
at virtual paths, remote non-image links, unknown schemes) becomes a short
|
|
text placeholder, and the link is preserved as structured data in the
|
|
artifact under ``resource_links``. Placeholders never embed host
|
|
filesystem paths: a link that stays unresolved (outside the thread's
|
|
user-data tree, or no thread context) omits the location segment rather
|
|
than leaking a raw ``file://`` URI or bare path, and non-referenceable
|
|
``data:``/``blob:`` URIs are dropped from the artifact channel too.
|
|
"""
|
|
from langchain_core.messages import ToolMessage
|
|
from langchain_core.messages.content import create_file_block, create_image_block, create_text_block
|
|
from langchain_core.tools import ToolException
|
|
from mcp.types import AudioContent, EmbeddedResource, ImageContent, ResourceLink, TextContent, TextResourceContents
|
|
|
|
# Pass ToolMessage through directly (interceptor short-circuit).
|
|
if isinstance(call_tool_result, ToolMessage):
|
|
return call_tool_result, None
|
|
|
|
# Pass LangGraph Command through directly when langgraph is installed.
|
|
try:
|
|
from langgraph.types import Command
|
|
|
|
if isinstance(call_tool_result, Command):
|
|
return call_tool_result, None
|
|
except ImportError:
|
|
# langgraph is optional; if unavailable, continue with standard MCP content conversion.
|
|
pass
|
|
|
|
def _resolve_link_url(uri: str) -> str:
|
|
if thread_id is None or user_id is None:
|
|
return uri
|
|
rewritten = _local_uri_to_virtual_path(uri, thread_id=thread_id, user_id=user_id, source_base_dir=source_base_dir)
|
|
return rewritten if rewritten is not None else uri
|
|
|
|
def _resolve_text(text: str) -> str:
|
|
# Servers like Playwright report saved files only as plain text, with no
|
|
# ResourceLink to hook into. Scan the text for local paths and translate
|
|
# them so the produced files are readable through the sandbox/artifact API.
|
|
if thread_id is None or user_id is None:
|
|
return text
|
|
return _rewrite_local_paths_in_text(
|
|
text,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
source_base_dir=source_base_dir,
|
|
changed_files=changed_files,
|
|
)
|
|
|
|
# Convert MCP content blocks to LangChain content blocks.
|
|
lc_content = []
|
|
resource_links: list[dict[str, Any]] = []
|
|
for item in call_tool_result.content:
|
|
if isinstance(item, TextContent):
|
|
lc_content.append(create_text_block(text=_resolve_text(item.text)))
|
|
elif isinstance(item, ImageContent):
|
|
lc_content.append(create_image_block(base64=item.data, mime_type=item.mimeType))
|
|
elif isinstance(item, AudioContent):
|
|
raise NotImplementedError(f"AudioContent conversion to LangChain content blocks is not yet supported. Received audio with mime type: {item.mimeType}")
|
|
elif isinstance(item, ResourceLink):
|
|
mime = item.mimeType or None
|
|
url = _resolve_link_url(str(item.uri))
|
|
if mime and mime.startswith("image/") and url.lower().startswith(("http://", "https://")):
|
|
lc_content.append(create_image_block(url=url, mime_type=mime))
|
|
elif url.lower().startswith(("data:", "blob:")):
|
|
# ``data:``/``blob:`` URIs never enter state: a ``data:`` URI
|
|
# embeds the whole payload, so inlining it would put megabytes
|
|
# of base64 into model-visible text and into checkpointed
|
|
# ``resource_links``; a ``blob:`` URI names a browser-local
|
|
# object nothing else can dereference. ``image/*`` links are
|
|
# downgraded here too, by design: the conversion layer keeps
|
|
# NEW inline payloads out of state, while the read-time
|
|
# middleware passes persisted ``data:`` image blocks through
|
|
# (it only heals blocks already in checkpoints). Neither is
|
|
# referenceable, so the raw URI is dropped from the artifact
|
|
# channel and the placeholder omits the location segment.
|
|
lc_content.append(create_text_block(text=resource_placeholder_text(name=item.name or "unnamed", mime_type=mime)))
|
|
else:
|
|
# URL-sourced file blocks are rejected by Chat Completions
|
|
# serialization (langchain-core), which bricks the thread once
|
|
# the result is checkpointed. Downgrade every other link to a
|
|
# text placeholder and keep the structured link in the artifact
|
|
# channel. The placeholder shows the location only when the
|
|
# resolved URL carries no host-path risk — a raw ``file://``
|
|
# URI or bare host path is withheld from model-visible text
|
|
# but still recorded in the artifact channel.
|
|
lc_content.append(
|
|
create_text_block(
|
|
text=resource_placeholder_text(
|
|
name=item.name or "unnamed",
|
|
mime_type=mime,
|
|
url=model_visible_location(url),
|
|
)
|
|
)
|
|
)
|
|
resource_links.append({"name": item.name, "uri": url, "mime_type": mime})
|
|
elif isinstance(item, EmbeddedResource):
|
|
from mcp.types import BlobResourceContents
|
|
|
|
res = item.resource
|
|
if isinstance(res, TextResourceContents):
|
|
lc_content.append(create_text_block(text=_resolve_text(res.text)))
|
|
elif isinstance(res, BlobResourceContents):
|
|
mime = res.mimeType or None
|
|
if mime and mime.startswith("image/"):
|
|
lc_content.append(create_image_block(base64=res.blob, mime_type=mime))
|
|
else:
|
|
# ``BlobResourceContents.mimeType`` is optional in the MCP spec, but
|
|
# langchain-core refuses a base64 block without one. Degrade to a
|
|
# generic binary type instead of raising, so the rest of the result —
|
|
# including the text blocks the server returned — survives. The
|
|
# sibling ``ResourceLink`` branch already degrades the same way.
|
|
lc_content.append(
|
|
create_file_block(
|
|
base64=res.blob,
|
|
mime_type=mime or "application/octet-stream",
|
|
)
|
|
)
|
|
else:
|
|
lc_content.append(create_text_block(text=str(res)))
|
|
else:
|
|
lc_content.append(create_text_block(text=str(item)))
|
|
|
|
if call_tool_result.isError:
|
|
error_parts = [item["text"] for item in lc_content if isinstance(item, dict) and item.get("type") == "text"]
|
|
raise ToolException("\n".join(error_parts) if error_parts else str(lc_content))
|
|
|
|
artifact = None
|
|
if call_tool_result.structuredContent is not None:
|
|
artifact = {"structured_content": call_tool_result.structuredContent}
|
|
if resource_links:
|
|
if artifact is None:
|
|
artifact = {}
|
|
artifact["resource_links"] = resource_links
|
|
|
|
return lc_content, artifact
|
|
|
|
|
|
def _resolve_session_init_timeout(server_cfg: Any) -> float | None:
|
|
"""Return the effective session-init timeout for *server_cfg*.
|
|
|
|
``None`` (an explicit opt-out) stays ``None``. Any other non-numeric value
|
|
falls back to the default rather than being passed to ``asyncio.wait_for``
|
|
(which would raise on it) or silently disabling the bound: pydantic
|
|
guarantees a float for real configs, but configs built with mocks in tests
|
|
can supply anything, and the fallback keeps the hang-protection in place.
|
|
"""
|
|
value = server_cfg.session_init_timeout if server_cfg is not None else DEFAULT_MCP_SESSION_INIT_TIMEOUT
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
|
return DEFAULT_MCP_SESSION_INIT_TIMEOUT
|
|
return float(value)
|
|
|
|
|
|
def _make_session_pool_tool(
|
|
tool: BaseTool,
|
|
server_name: str,
|
|
connection: dict[str, Any],
|
|
tool_interceptors: list[Any] | None = None,
|
|
tool_call_timeout: float | None = None,
|
|
session_init_timeout: float | None = None,
|
|
tool_name_prefix: bool = True,
|
|
ownership_domain: MCPPoolDomain = "deployment",
|
|
*,
|
|
pool: MCPSessionPool | None = None,
|
|
binding: ServerBinding | None = None,
|
|
) -> BaseTool:
|
|
"""Wrap an MCP tool so it reuses a persistent session from the pool.
|
|
|
|
Replaces the per-call session creation with pool-managed sessions scoped
|
|
by ``(server_name, user/thread/incarnation, ownership_domain)``. This ensures stateful MCP servers
|
|
(e.g. Playwright) keep their state across tool calls within the same thread
|
|
while staying isolated per user.
|
|
|
|
The configured ``tool_interceptors`` (OAuth, custom) are preserved and
|
|
applied on every call before invoking the pooled session.
|
|
"""
|
|
# Strip only prefixes added by the adapter. An unprefixed server may expose
|
|
# a tool whose own name happens to start with ``<server_name>_``.
|
|
original_name = tool.name
|
|
prefix = f"{server_name}_"
|
|
if tool_name_prefix and original_name.startswith(prefix):
|
|
original_name = original_name[len(prefix) :]
|
|
|
|
if pool is None:
|
|
pool = get_session_pool()
|
|
if binding is None and isinstance(pool, MCPSessionPool):
|
|
binding = pool.ensure_binding(
|
|
server_name,
|
|
normalized_connection_fingerprint(connection),
|
|
domain=ownership_domain,
|
|
)
|
|
effective_domain = binding.domain if binding is not None else ownership_domain
|
|
|
|
async def call_with_persistent_session(
|
|
deerflow_runtime: Runtime | None = None,
|
|
**arguments: Any,
|
|
) -> Any:
|
|
thread_id = _extract_thread_id(deerflow_runtime)
|
|
user_id = resolve_runtime_user_id(deerflow_runtime)
|
|
# Scope the pooled session by user *and* thread. Filesystem isolation is
|
|
# per-(user_id, thread_id), so a thread_id alone could otherwise let two
|
|
# users with a colliding thread_id share one stateful MCP session.
|
|
thread_incarnation = runtime_thread_incarnation(deerflow_runtime)
|
|
scope_key = mcp_session_scope_key(
|
|
user_id=user_id,
|
|
thread_id=thread_id,
|
|
thread_incarnation=thread_incarnation,
|
|
)
|
|
session_connection = dict(connection)
|
|
# cwd/temp pinning and the workspace snapshot only matter for stdio
|
|
# servers, which run as local subprocesses writing to a real filesystem.
|
|
# SSE/HTTP servers have no local cwd to pin, so skip the filesystem work
|
|
# entirely for them (avoids needless dir creation and recursive walks).
|
|
is_stdio = session_connection.get("transport", "stdio") == "stdio"
|
|
source_base_dir: Path | None = None
|
|
process_cwd: Path | None = None
|
|
before_files: _FILE_SNAPSHOT | None = None
|
|
if is_stdio:
|
|
paths = get_paths()
|
|
# Bundle the synchronous filesystem prep (dir creation, temp-dir
|
|
# setup, pre-call snapshot) and run it off the event loop — the
|
|
# snapshot walks the whole workspace and would otherwise block.
|
|
source_base_dir, tmp_dir, before_files = await asyncio.to_thread(_prepare_stdio_workspace, paths, thread_id=thread_id, user_id=user_id)
|
|
# Stdio MCP servers resolve relative output links against their
|
|
# process cwd. Keep that cwd inside the thread's mounted user-data
|
|
# tree so files produced by tools like Playwright land where the
|
|
# sandbox/artifact API can serve them and their references can be
|
|
# translated to virtual paths.
|
|
configured_cwd = session_connection.get("cwd", str(source_base_dir))
|
|
session_connection["cwd"] = str(configured_cwd)
|
|
process_cwd = Path(configured_cwd)
|
|
# Pin the subprocess temp dir under the same mounted tree. Tools that
|
|
# default to the OS temp dir (Node's os.tmpdir(), Python's tempfile,
|
|
# many CLIs) then write inside user-data instead of an unreachable
|
|
# host path — the tool-agnostic counterpart to fixing the cwd. Merge
|
|
# rather than replace any operator-provided env.
|
|
session_env = dict(session_connection.get("env") or {})
|
|
session_env.setdefault("TMPDIR", str(tmp_dir))
|
|
session_env.setdefault("TMP", str(tmp_dir))
|
|
session_env.setdefault("TEMP", str(tmp_dir))
|
|
session_connection["env"] = session_env
|
|
if binding is not None:
|
|
session_request = pool.get_session(
|
|
server_name,
|
|
scope_key,
|
|
session_connection,
|
|
binding=binding,
|
|
)
|
|
elif effective_domain == "deployment":
|
|
# Compatibility for direct unit-test fakes that construct this
|
|
# wrapper without going through get_mcp_tools(). Production
|
|
# discovery always supplies the explicit binding above.
|
|
session_request = pool.get_session(server_name, scope_key, session_connection)
|
|
else:
|
|
session_request = pool.get_session(
|
|
server_name,
|
|
scope_key,
|
|
session_connection,
|
|
domain=effective_domain,
|
|
)
|
|
if session_init_timeout is not None:
|
|
# Cancellation here is safe: MCPSessionPool.get_session owns the
|
|
# teardown of a session stuck mid-creation (it signals close and
|
|
# waits for the owner task's __aexit__ to run in its own task),
|
|
# so a hung server cannot leak a session or block the turn.
|
|
try:
|
|
session = await asyncio.wait_for(
|
|
session_request,
|
|
timeout=session_init_timeout,
|
|
)
|
|
except TimeoutError:
|
|
# Surface the timeout at the same log level as discovery
|
|
# timeouts: the tool call still fails with a TimeoutError the
|
|
# model can react to, but operators need the WARNING to
|
|
# diagnose tool-call failures caused by hung MCP sessions.
|
|
logger.warning(
|
|
"MCP session initialization for server '%s' timed out after %.1fs",
|
|
server_name,
|
|
session_init_timeout,
|
|
)
|
|
raise
|
|
else:
|
|
session = await session_request
|
|
|
|
domain_kwargs = {"domain": effective_domain} if effective_domain != "deployment" else {}
|
|
|
|
# Build common call_tool kwargs once — only add keys when needed so
|
|
# existing call-sites that assert on exact arguments are not affected.
|
|
call_kwargs: dict[str, Any] = {}
|
|
if tool_call_timeout:
|
|
call_kwargs["read_timeout_seconds"] = timedelta(seconds=tool_call_timeout)
|
|
|
|
if tool_interceptors:
|
|
from langchain_mcp_adapters.interceptors import MCPToolCallRequest
|
|
|
|
async def base_handler(request: MCPToolCallRequest) -> Any:
|
|
# Preserve interceptor-injected headers for stdio MCP calls by
|
|
# forwarding them through MCP call meta.
|
|
kwargs = dict(call_kwargs)
|
|
if request.headers:
|
|
if isinstance(request.headers, Mapping):
|
|
kwargs["meta"] = {"headers": dict(request.headers)}
|
|
else:
|
|
logger.warning("Ignoring MCP interceptor headers with unsupported type: %s", type(request.headers).__name__)
|
|
return await call_pooled_session_tool(
|
|
session,
|
|
pool,
|
|
server_name=server_name,
|
|
scope_key=scope_key,
|
|
tool_name=request.name,
|
|
arguments=request.args,
|
|
call_kwargs=kwargs,
|
|
**domain_kwargs,
|
|
)
|
|
|
|
handler = compose_tool_interceptors(tool_interceptors, base_handler)
|
|
|
|
request = MCPToolCallRequest(
|
|
name=original_name,
|
|
args=arguments,
|
|
server_name=server_name,
|
|
runtime=deerflow_runtime,
|
|
)
|
|
call_tool_result = await handler(request)
|
|
else:
|
|
call_tool_result = await call_pooled_session_tool(
|
|
session,
|
|
pool,
|
|
server_name=server_name,
|
|
scope_key=scope_key,
|
|
tool_name=original_name,
|
|
arguments=arguments,
|
|
call_kwargs=call_kwargs,
|
|
**domain_kwargs,
|
|
)
|
|
|
|
# The after-call snapshot diff only feeds bare-filename correlation in
|
|
# free text, so skip the second recursive walk when there is no text
|
|
# content to rewrite. Both the diff and the per-token path resolution
|
|
# inside _convert_call_tool_result touch the filesystem, so run them off
|
|
# the event loop.
|
|
changed_files: list[Path] | None = None
|
|
if is_stdio and before_files is not None and _result_has_text_content(call_tool_result):
|
|
changed_files = await asyncio.to_thread(_changed_workspace_files, source_base_dir, before_files)
|
|
return await asyncio.to_thread(
|
|
_convert_call_tool_result,
|
|
call_tool_result,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
source_base_dir=process_cwd,
|
|
changed_files=changed_files,
|
|
)
|
|
|
|
return StructuredTool(
|
|
name=tool.name,
|
|
description=tool.description,
|
|
args_schema=tool.args_schema,
|
|
coroutine=call_with_persistent_session,
|
|
response_format="content_and_artifact",
|
|
metadata=tool.metadata,
|
|
)
|
|
|
|
|
|
def _raw_mcp_tool_name(
|
|
tool: BaseTool,
|
|
*,
|
|
server_name: str,
|
|
tool_name_prefix: bool,
|
|
) -> str:
|
|
prefix = f"{server_name}_"
|
|
if tool_name_prefix and tool.name.startswith(prefix):
|
|
return tool.name[len(prefix) :]
|
|
return tool.name
|
|
|
|
|
|
def _make_background_submit_tool(
|
|
tool: BaseTool,
|
|
*,
|
|
server_name: str,
|
|
task_name: str,
|
|
submit_tool: str,
|
|
status_tool: str,
|
|
cancel_tool: str,
|
|
connection_scope: str,
|
|
) -> BaseTool:
|
|
background_contract = f"Submitted as durable background task {task_name!r}; returns a DeerFlow task ID immediately and status polling is handled automatically."
|
|
|
|
async def submit_in_background(
|
|
deerflow_runtime: Runtime | None = None,
|
|
**arguments: Any,
|
|
) -> dict[str, Any]:
|
|
submitter = get_mcp_task_submitter()
|
|
thread_id = _extract_thread_id(deerflow_runtime)
|
|
user_id = resolve_runtime_user_id(deerflow_runtime)
|
|
thread_incarnation = runtime_thread_incarnation(deerflow_runtime)
|
|
context = deerflow_runtime.context if deerflow_runtime is not None and deerflow_runtime.context else {}
|
|
run_id = context.get("run_id")
|
|
tool_call_id = getattr(deerflow_runtime, "tool_call_id", None) if deerflow_runtime is not None else None
|
|
created = await submitter.submit(
|
|
driver_name=ORDINARY_MCP_TASK_DRIVER,
|
|
request=TaskSubmitRequest(
|
|
user_id=user_id,
|
|
thread_id=thread_id,
|
|
thread_incarnation=thread_incarnation,
|
|
run_id=str(run_id) if run_id is not None else None,
|
|
tool_call_id=str(tool_call_id) if tool_call_id is not None else None,
|
|
server_name=server_name,
|
|
task_name=task_name,
|
|
arguments=arguments,
|
|
driver_data={
|
|
"submit_tool": submit_tool,
|
|
"status_tool": status_tool,
|
|
"cancel_tool": cancel_tool,
|
|
"connection_scope": connection_scope,
|
|
},
|
|
),
|
|
)
|
|
return {
|
|
"task_id": created["id"],
|
|
"task_name": task_name,
|
|
"status": created["status"],
|
|
"message": "Task is running in the background.",
|
|
}
|
|
|
|
return StructuredTool(
|
|
name=tool.name,
|
|
description=(f"{tool.description}\n\n{background_contract}" if tool.description else background_contract),
|
|
args_schema=tool.args_schema,
|
|
coroutine=submit_in_background,
|
|
metadata=tool.metadata,
|
|
)
|
|
|
|
|
|
def _configure_task_tools_for_server(
|
|
tools: list[BaseTool],
|
|
*,
|
|
server_name: str,
|
|
server_config: McpServerConfig,
|
|
tool_name_prefix: bool,
|
|
connection_scope: str = "deployment",
|
|
) -> list[BaseTool]:
|
|
"""Hide driver-only tools and replace submit with a durable wrapper."""
|
|
if not server_config.task_toolsets:
|
|
return tools
|
|
|
|
by_raw_name = {
|
|
_raw_mcp_tool_name(
|
|
tool,
|
|
server_name=server_name,
|
|
tool_name_prefix=tool_name_prefix,
|
|
): tool
|
|
for tool in tools
|
|
}
|
|
expected = {
|
|
raw_name
|
|
for toolset in server_config.task_toolsets
|
|
for raw_name in (
|
|
toolset.submit_tool,
|
|
toolset.status_tool,
|
|
toolset.cancel_tool,
|
|
)
|
|
}
|
|
missing = sorted(expected - by_raw_name.keys())
|
|
if missing:
|
|
raise McpTaskConfigurationError(f"MCP server {server_name!r} task_toolsets reference missing raw tool(s): {', '.join(missing)}")
|
|
|
|
hidden = {raw_name for toolset in server_config.task_toolsets for raw_name in (toolset.status_tool, toolset.cancel_tool)}
|
|
submit_by_name = {toolset.submit_tool: toolset for toolset in server_config.task_toolsets}
|
|
configured: list[BaseTool] = []
|
|
for tool in tools:
|
|
raw_name = _raw_mcp_tool_name(
|
|
tool,
|
|
server_name=server_name,
|
|
tool_name_prefix=tool_name_prefix,
|
|
)
|
|
if raw_name in hidden:
|
|
continue
|
|
toolset = submit_by_name.get(raw_name)
|
|
if toolset is None:
|
|
configured.append(tool)
|
|
continue
|
|
configured.append(
|
|
_make_background_submit_tool(
|
|
tool,
|
|
server_name=server_name,
|
|
task_name=toolset.name,
|
|
submit_tool=toolset.submit_tool,
|
|
status_tool=toolset.status_tool,
|
|
cancel_tool=toolset.cancel_tool,
|
|
connection_scope=connection_scope,
|
|
)
|
|
)
|
|
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
|
|
personal_user_id: str | None = None
|
|
|
|
|
|
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:
|
|
# Never interpolate ``e`` nor attach its traceback: ``ExtensionsConfig``
|
|
# resolves ``$VAR`` before validation, so a discovery error message can
|
|
# embed a resolved-credential request URL. Keep the server name and the
|
|
# exception type for diagnosis.
|
|
logger.warning(
|
|
"Skipping MCP server '%s' after tool discovery failed (%s)",
|
|
server_name,
|
|
type(e).__name__,
|
|
)
|
|
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",
|
|
personal_user_id=personal_user_id,
|
|
)
|
|
|
|
|
|
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, owner_user_id=prepared.personal_user_id) 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,
|
|
*,
|
|
personal_user_id: str | None = None,
|
|
session_pool: MCPSessionPool | None = None,
|
|
) -> list[BaseTool]:
|
|
"""Get all tools from enabled MCP servers.
|
|
|
|
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 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:
|
|
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.
|
|
"""
|
|
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=None,
|
|
personal_user_id=personal_user_id,
|
|
session_pool=session_pool,
|
|
)
|
|
if prepared is None:
|
|
return []
|
|
|
|
try:
|
|
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:
|
|
# Only the exception type is logged: the message/traceback can carry a
|
|
# resolved-credential request URL (see ``_discover_server_tools``).
|
|
logger.error("Failed to load MCP tools (%s)", type(e).__name__)
|
|
return []
|