"""Registering a connected (or schema-cached) MCP server's tools into the tool registry: include/exclude filtering, trust-tier metadata capture, utility-tool selection, name-collision resolution and the schema-cache write-through.""" import logging from types import SimpleNamespace from typing import TYPE_CHECKING, Any, Callable, Dict, List from tools.mcp_tool_common import _parse_boolish, _core from tools.mcp_tool_handlers import _make_check_fn, _make_get_prompt_handler, _make_list_prompts_handler, _make_list_resources_handler, _make_read_resource_handler from tools.mcp_tool_common import _resolve_tool_timeout from tools.mcp_tool_schema import _UTILITY_CAPABILITY_ATTRS, _UTILITY_CAPABILITY_METHODS, _build_utility_schemas, _normalize_name_filter, matches_name_filter if TYPE_CHECKING: # pragma: no cover from tools.mcp_tool import MCPServerTask logger = logging.getLogger("tools.mcp_tool") def _normalize_server_trust(value: Any) -> str: """Normalize a config ``trust`` value: None -> ``full`` (backward-compatible default); any unrecognized string -> ``untrusted``, so a misspelled tier fails closed rather than silently disabling gating.""" if value is None: return _core._TRUST_FULL text = str(value).strip().lower() if text == _core._TRUST_FULL: return _core._TRUST_FULL if text == _core._TRUST_UNTRUSTED: return _core._TRUST_UNTRUSTED logger.warning( "MCP trust: unrecognized trust value %r — treating as 'untrusted' " "(valid values: full, untrusted)", value, ) return _core._TRUST_UNTRUSTED def _annotation_read_only_hint(mcp_tool: Any) -> bool: """True only when annotations (SDK object or schema-cache dict) carry ``readOnlyHint is True``; anything else is False — unknown metadata means the tool must be treated as write-capable.""" annotations = getattr(mcp_tool, "annotations", None) if annotations is None: return False if isinstance(annotations, dict): hint = annotations.get("readOnlyHint") else: hint = getattr(annotations, "readOnlyHint", None) return hint is True def _record_tool_trust_metadata( server_name: str, config: dict, tools: List[Any] ) -> None: """Capture per-server trust and per-tool readOnlyHint at discovery.""" with _core._lock: _core._server_trust_levels[server_name] = _normalize_server_trust( (config or {}).get("trust") ) hints = _core._tool_read_only_hints.setdefault(server_name, {}) for tool in tools: name = getattr(tool, "name", None) if name: hints[name] = _annotation_read_only_hint(tool) def _track_mcp_tool_server(tool_name: str, server_name: str) -> None: """Remember the exact raw MCP server that registered *tool_name*.""" with _core._lock: _core._mcp_tool_server_names[tool_name] = server_name def _forget_mcp_tool_server(tool_name: str) -> None: """Forget MCP server provenance for a deregistered tool.""" with _core._lock: _core._mcp_tool_server_names.pop(tool_name, None) def _select_utility_schemas(server_name: str, server: "MCPServerTask", config: dict) -> List[dict]: """Select utility schemas based on config and server capabilities.""" tools_filter = config.get("tools") or {} resources_enabled = _parse_boolish(tools_filter.get("resources"), default=True) prompts_enabled = _parse_boolish(tools_filter.get("prompts"), default=True) # ``initialize_result.capabilities`` is the source of truth: its sub-objects # are non-None iff the server advertises that request family. The old # ``hasattr(server.session, ...)`` gate never filtered anything because # ClientSession defines all four methods on the class. advertised_caps = None init_result = getattr(server, "initialize_result", None) if init_result is not None: advertised_caps = getattr(init_result, "capabilities", None) selected: List[dict] = [] for entry in _build_utility_schemas(server_name): handler_key = entry["handler_key"] if handler_key in {"list_resources", "read_resource"} and not resources_enabled: logger.debug("MCP server '%s': skipping utility '%s' (resources disabled)", server_name, handler_key) continue if handler_key in {"list_prompts", "get_prompt"} and not prompts_enabled: logger.debug("MCP server '%s': skipping utility '%s' (prompts disabled)", server_name, handler_key) continue if advertised_caps is not None: cap_attr = _UTILITY_CAPABILITY_ATTRS[handler_key] if getattr(advertised_caps, cap_attr, None) is None: logger.debug( "MCP server '%s': skipping utility '%s' " "(server does not advertise '%s' capability)", server_name, handler_key, cap_attr, ) continue else: # Legacy fallback when initialize_result wasn't captured (test # fixtures, older paths): register every stub, as before. required_method = _UTILITY_CAPABILITY_METHODS[handler_key] if not hasattr(server.session, required_method): logger.debug( "MCP server '%s': skipping utility '%s' (session lacks %s)", server_name, handler_key, required_method, ) continue selected.append(entry) return selected def _existing_tool_names() -> List[str]: """Return tool names for all currently connected servers.""" names: List[str] = [] for _sname, server in _core._servers.items(): if hasattr(server, "_registered_tool_names"): names.extend(server._registered_tool_names) continue for mcp_tool in server._tools: schema = _core._convert_mcp_schema(server.name, mcp_tool) names.append(schema["name"]) # Lazy servers registered from the schema cache have no MCPServerTask yet — # their tools live only in the registry. with _core._lock: lazy_names = [ n for sname, tool_names in _core._lazy_server_tool_names.items() if sname not in _core._servers for n in tool_names ] names.extend(lazy_names) return names # Utility tool key -> handler factory; each takes (server_name, tool_timeout). _UTILITY_HANDLER_FACTORIES = { "list_resources": _make_list_resources_handler, "read_resource": _make_read_resource_handler, "list_prompts": _make_list_prompts_handler, "get_prompt": _make_get_prompt_handler, } def _make_tool_filter(name: str, config: dict) -> Callable[[str], bool]: """Build the include/exclude predicate for a server's tool names. Rules: ``tools.include`` is a whitelist, ``tools.exclude`` a blacklist; entries may be exact names or fnmatch globs; include wins over exclude; ``include: []`` is an explicit empty whitelist (register nothing); neither set registers everything. """ tools_filter = config.get("tools") or {} include_raw = tools_filter.get("include") include_set = _normalize_name_filter(include_raw, f"mcp_servers.{name}.tools.include") include_active = isinstance(include_raw, (str, list, tuple, set)) exclude_set = _normalize_name_filter( tools_filter.get("exclude"), f"mcp_servers.{name}.tools.exclude" ) def _should_register(tool_name: str) -> bool: if include_active: return matches_name_filter(tool_name, include_set) if exclude_set: return not matches_name_filter(tool_name, exclude_set) return True return _should_register def _resolve_name_collisions(name: str, candidates: List[dict]): """Preflight registry-name collisions among one server's candidates. Returns ``(unique_candidates, ambiguous_names, shadowed_utilities)``. Exact duplicate rows (same name + origin) are dropped silently; a generated utility that normalizes onto a server-native tool's name is shadowed (the native tool wins); any other multi-origin collision is ambiguous and every colliding entry is skipped (fail closed). """ unique_candidates: List[dict] = [] seen_candidates: set[tuple[str, str]] = set() origins_by_name: Dict[str, set[str]] = {} for candidate in candidates: key = (candidate["registry_name"], candidate["origin"]) if key in seen_candidates: logger.debug( "MCP server '%s': duplicate registration candidate %s for '%s'; " "keeping one", name, candidate["origin"], candidate["registry_name"], ) continue seen_candidates.add(key) unique_candidates.append(candidate) origins_by_name.setdefault(candidate["registry_name"], set()).add( candidate["origin"] ) ambiguous_names: Dict[str, List[str]] = {} shadowed_utilities: set[tuple[str, str]] = set() for registry_name, origins in origins_by_name.items(): if len(origins) <= 1: continue utility_origins = sorted( o for o in origins if o.startswith("generated utility ") ) native_origins = sorted(origins - set(utility_origins)) if len(native_origins) == 1 and utility_origins: for util_origin in utility_origins: shadowed_utilities.add((registry_name, util_origin)) logger.info( "MCP server '%s': generated utility %s normalizes onto " "server-native %s — keeping the native tool and dropping the " "utility (the utility only applies when the server has no such " "tool of its own)", name, ", ".join(utility_origins), native_origins[0], ) continue ambiguous_names[registry_name] = sorted(origins) for registry_name, origins in sorted(ambiguous_names.items()): logger.error( "MCP server '%s': name normalization collision for '%s' from %s; " "skipping every colliding entry instead of choosing an arbitrary " "handler", name, registry_name, ", ".join(origins), ) return unique_candidates, ambiguous_names, shadowed_utilities def _write_schema_cache(name: str, server: "MCPServerTask", config: dict, should_register) -> None: """Write-through: persist the manifest so the next startup can register this server lazily without spawning it. Never raises.""" try: from tools.mcp_schema_cache import config_fingerprint, write_cache_entry tools_payload: List[dict] = [] for mcp_tool in server._tools: if not should_register(mcp_tool.name): continue schema_obj = getattr(mcp_tool, "inputSchema", None) tools_payload.append({ "name": mcp_tool.name, "description": mcp_tool.description or "", "inputSchema": schema_obj if isinstance(schema_obj, dict) else {}, # Persisted so the lazy path trust-gates identically next startup. "annotations": { "readOnlyHint": _annotation_read_only_hint(mcp_tool), }, }) utility_payload = [ {"schema": entry["schema"], "handler_key": entry["handler_key"]} for entry in _select_utility_schemas(name, server, config) ] cache_meta = getattr(server, "_list_cache_meta", None) or {} write_cache_entry( name, config_fingerprint(config), tools=tools_payload, utility_tools=utility_payload, ttl_ms=cache_meta.get("ttl_ms"), cache_scope=cache_meta.get("cache_scope"), ) except Exception as exc: logger.debug("MCP schema cache write failed for '%s': %s", name, exc) def _register_server_tools(name: str, server: "MCPServerTask", config: dict) -> List[str]: """Register an already-connected server's tools (plus utility tools) into the registry; used by initial discovery and list_changed refresh. Toolset resolution for ``mcp-{server}`` / raw-name aliases derives from the live registry rather than mutating ``toolsets.TOOLSETS``. Lossy name normalization can map distinct raw names (``read-file``/``read_file``) to one registry name; such collisions fail closed — every ambiguous entry is skipped. Returns the registered prefixed names. """ from tools.registry import registry registered_names: List[str] = [] toolset_name = f"mcp-{name}" _should_register = _make_tool_filter(name, config) check_fn = _make_check_fn(name) candidates: List[dict] = [] # Security boundary: capture trust tier and readOnlyHint NOW, at discovery, # so the call-time gate classifies from data we control, not re-read # server-supplied state. _record_tool_trust_metadata(name, config, server._tools) for mcp_tool in server._tools: if not _should_register(mcp_tool.name): logger.debug( "MCP server '%s': skipping tool '%s' (filtered by config)", name, mcp_tool.name, ) continue _core._scan_mcp_description(name, mcp_tool.name, mcp_tool.description or "") schema = _core._convert_mcp_schema(name, mcp_tool) candidates.append( { "registry_name": schema["name"], "origin": f"tool {mcp_tool.name!r}", "schema": schema, "handler": _core._make_tool_handler( name, mcp_tool.name, server.tool_timeout ), "check_fn": check_fn, } ) # Generated resource/prompt utility tools share the same namespace as raw # MCP tools, so they must participate in the same collision preflight. for entry in _select_utility_schemas(name, server, config): schema = entry["schema"] handler_key = entry["handler_key"] candidates.append( { "registry_name": schema["name"], "origin": f"generated utility {handler_key!r}", "schema": schema, "handler": _UTILITY_HANDLER_FACTORIES[handler_key]( name, server.tool_timeout ), "check_fn": check_fn, } ) unique_candidates, ambiguous_names, shadowed_utilities = _resolve_name_collisions(name, candidates) for candidate in unique_candidates: registry_name = candidate["registry_name"] if registry_name in ambiguous_names: continue if (registry_name, candidate["origin"]) in shadowed_utilities: continue existing_toolset = registry.get_toolset_for_tool(registry_name) if existing_toolset and existing_toolset != toolset_name: if existing_toolset.startswith("mcp-"): logger.error( "MCP server '%s': %s normalizes to '%s', already owned by " "MCP toolset '%s' — skipping to preserve the existing owner", name, candidate["origin"], registry_name, existing_toolset, ) else: logger.warning( "MCP server '%s': %s (→ '%s') collides with built-in tool " "in toolset '%s' — skipping to preserve built-in", name, candidate["origin"], registry_name, existing_toolset, ) continue registry.register( name=registry_name, toolset=toolset_name, schema=candidate["schema"], handler=candidate["handler"], check_fn=candidate["check_fn"], is_async=False, description=candidate["schema"]["description"], scope=_core._server_registry_scope(name), ) # The pre-check above is advisory only. Multiple servers connect in # parallel, so ToolRegistry.register() is the atomic ownership gate. if registry.get_toolset_for_tool(registry_name) != toolset_name: logger.error( "MCP server '%s': registration of %s as '%s' was rejected by " "the registry; skipping provenance/count updates", name, candidate["origin"], registry_name, ) continue _core._track_mcp_tool_server(registry_name, name) registered_names.append(registry_name) if registered_names: registry.register_toolset_alias(name, toolset_name) _write_schema_cache(name, server, config, _should_register) return registered_names class _CachedMCPTool: """Minimal stand-in for MCP Tool objects loaded from the schema cache.""" __slots__ = ("name", "description", "inputSchema") def __init__(self, name: str, description: str, inputSchema: dict): self.name = name self.description = description self.inputSchema = inputSchema or {} def _register_from_cache_sync(name: str, config: dict, entry: dict) -> List[str]: """Lazy startup: register a server's tools from a cached manifest with no child process. The first real call routes through ``_get_connected_server_for_call`` -> ``_ensure_lazy_server_connected``.""" from tools.registry import registry from tools.mcp_schema_cache import ( config_fingerprint, tools_from_cache_entry, utility_tools_from_cache_entry, ) registered_names: List[str] = [] toolset_name = f"mcp-{name}" fingerprint = config_fingerprint(config) tool_timeout = _resolve_tool_timeout(config) _should_register = _make_tool_filter(name, config) check_fn = _make_check_fn(name) # Record trust metadata before registration so the call-time gate is # identical whether the server was spawned live or registered from cache. # Missing "annotations" in older cache files fails closed to write-capable. cached_tool_objs = [ SimpleNamespace( name=raw.get("name"), annotations=raw.get("annotations") if isinstance(raw.get("annotations"), dict) else None, ) for raw in tools_from_cache_entry(entry) if isinstance(raw, dict) and raw.get("name") ] _record_tool_trust_metadata(name, config, cached_tool_objs) for raw in tools_from_cache_entry(entry): if not isinstance(raw, dict): continue raw_name = raw.get("name") if not raw_name or not _should_register(raw_name): continue raw_schema = raw.get("inputSchema") mcp_tool = _CachedMCPTool( raw_name, raw.get("description") or "", raw_schema if isinstance(raw_schema, dict) else {}, ) # Defense-in-depth: the cache file is user-writable JSON, so apply the # same injection scan as eager discovery. _core._scan_mcp_description(name, mcp_tool.name, mcp_tool.description or "") schema = _core._convert_mcp_schema(name, mcp_tool) registry_name = schema["name"] existing_toolset = registry.get_toolset_for_tool(registry_name) if existing_toolset and existing_toolset != toolset_name: logger.warning( "MCP server '%s' (lazy): cached tool '%s' collides with " "toolset '%s' — skipping", name, registry_name, existing_toolset, ) continue registry.register( name=registry_name, toolset=toolset_name, schema=schema, handler=_core._make_tool_handler(name, raw_name, tool_timeout), check_fn=check_fn, is_async=False, description=schema["description"], scope=_core._mcp_registry_scope(), ) if registry.get_toolset_for_tool(registry_name) != toolset_name: continue _core._track_mcp_tool_server(registry_name, name) registered_names.append(registry_name) for raw in utility_tools_from_cache_entry(entry): if not isinstance(raw, dict): continue schema = raw.get("schema") handler_key = raw.get("handler_key") if not isinstance(schema, dict) or handler_key not in _UTILITY_HANDLER_FACTORIES: continue util_name = schema.get("name") or "" if not util_name: continue existing_toolset = registry.get_toolset_for_tool(util_name) if existing_toolset and existing_toolset != toolset_name: continue registry.register( name=util_name, toolset=toolset_name, schema=schema, handler=_UTILITY_HANDLER_FACTORIES[handler_key](name, tool_timeout), check_fn=check_fn, is_async=False, description=schema.get("description") or "", scope=_core._mcp_registry_scope(), ) if registry.get_toolset_for_tool(util_name) != toolset_name: continue _core._track_mcp_tool_server(util_name, name) registered_names.append(util_name) if registered_names: registry.register_toolset_alias(name, toolset_name) with _core._lock: _core._lazy_server_configs[name] = dict(config) _core._lazy_server_fingerprints[name] = fingerprint _core._lazy_server_tool_names[name] = list(registered_names) logger.info( "MCP server '%s' (lazy): registered %d tool(s) from schema cache", name, len(registered_names), ) return registered_names