from __future__ import annotations import asyncio from collections.abc import Callable, Iterator import hashlib import importlib.util import inspect from pathlib import Path import re import sys import threading from typing import TYPE_CHECKING, Any from vibe.core.config.harness_files import get_harness_files_manager from vibe.core.logger import logger from vibe.core.paths import DEFAULT_TOOL_DIR from vibe.core.tools.base import BaseTool, BaseToolConfig from vibe.core.tools.connectors import ConnectorRegistry from vibe.core.tools.mcp import MCPRegistry from vibe.core.tools.mcp.tools import MCPTool from vibe.core.utils import name_matches, run_sync if TYPE_CHECKING: from vibe.core.config import VibeConfig def _try_canonical_module_name(path: Path) -> str | None: """Extract canonical module name for vibe package files. Prevents Pydantic class identity mismatches when the same module is imported via dynamic discovery and regular imports. """ try: parts = path.resolve().parts except (OSError, ValueError): return None try: vibe_idx = parts.index("vibe") except ValueError: return None if vibe_idx + 1 >= len(parts): return None module_parts = [p.removesuffix(".py") for p in parts[vibe_idx:]] return ".".join(module_parts) def _compute_module_name(path: Path) -> str: """Return canonical module name for vibe files, hash-based synthetic name otherwise.""" if canonical := _try_canonical_module_name(path): return canonical resolved = path.resolve() path_hash = hashlib.md5(str(resolved).encode()).hexdigest()[:8] stem = re.sub(r"[^0-9A-Za-z_]", "_", path.stem) or "mod" return f"vibe_tools_discovered_{stem}_{path_hash}" class NoSuchToolError(Exception): """Exception raised when a tool is not found.""" class ToolManager: """Manages tool discovery and instantiation for an Agent. Discovers available tools from the provided search paths. Each Agent should have its own ToolManager instance. """ def __init__( self, config_getter: Callable[[], VibeConfig], mcp_registry: MCPRegistry | None = None, connector_registry: ConnectorRegistry | None = None, *, defer_mcp: bool = False, ) -> None: self._config_getter = config_getter self._mcp_registry = mcp_registry or MCPRegistry() self._connector_registry = connector_registry self._instances: dict[str, BaseTool] = {} self._search_paths: list[Path] = self._compute_search_paths(self._config) self._lock = threading.Lock() self._mcp_integrated = False self._available: dict[str, type[BaseTool]] = { cls.get_name(): cls for cls in self._iter_tool_classes(self._search_paths) } if not defer_mcp: self.integrate_all() @property def _config(self) -> VibeConfig: return self._config_getter() @staticmethod def _compute_search_paths(config: VibeConfig) -> list[Path]: paths: list[Path] = [DEFAULT_TOOL_DIR.path] paths.extend(config.tool_paths) mgr = get_harness_files_manager() paths.extend(mgr.project_tools_dirs) paths.extend(mgr.user_tools_dirs) unique: list[Path] = [] seen: set[Path] = set() for p in paths: rp = p.resolve() if rp not in seen: seen.add(rp) unique.append(rp) return unique @staticmethod def _iter_tool_classes(search_paths: list[Path]) -> Iterator[type[BaseTool]]: """Iterate over all search_paths to find tool classes. Note: if a search path is not a directory, it is treated as a single tool file. """ for base in search_paths: if not base.is_dir() and base.name.endswith(".py"): if tools := ToolManager._load_tools_from_file(base): for tool in tools: yield tool for path in base.rglob("*.py"): if tools := ToolManager._load_tools_from_file(path): for tool in tools: yield tool @staticmethod def _load_tools_from_file(file_path: Path) -> list[type[BaseTool]] | None: if not file_path.is_file(): return name = file_path.name if name.startswith("_"): return module_name = _compute_module_name(file_path) if module_name in sys.modules: module = sys.modules[module_name] else: spec = importlib.util.spec_from_file_location(module_name, file_path) if spec is None or spec.loader is None: return module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module try: spec.loader.exec_module(module) except Exception: return tools = [] for tool_obj in vars(module).values(): if not inspect.isclass(tool_obj): continue if not issubclass(tool_obj, BaseTool) or tool_obj is BaseTool: continue if inspect.isabstract(tool_obj): continue tools.append(tool_obj) return tools @staticmethod def discover_tool_defaults( search_paths: list[Path] | None = None, ) -> dict[str, dict[str, Any]]: if search_paths is None: search_paths = [DEFAULT_TOOL_DIR.path] defaults: dict[str, dict[str, Any]] = {} for cls in ToolManager._iter_tool_classes(search_paths): try: tool_name = cls.get_name() config_class = cls._get_tool_config_class() defaults[tool_name] = config_class().model_dump(exclude_none=True) except Exception as e: logger.warning( "Failed to get defaults for tool %s: %s", cls.__name__, e ) continue return defaults @property def registered_tools(self) -> dict[str, type[BaseTool]]: with self._lock: return dict(self._available) @property def available_tools(self) -> dict[str, type[BaseTool]]: with self._lock: runtime_available = { name: cls for name, cls in self._available.items() if cls.is_available() } # Per-source filtering first (MCP server/connector disabled flags). result = self._apply_per_source_filtering(runtime_available) # Global overrides take precedence. if self._config.enabled_tools: return { name: cls for name, cls in result.items() if name_matches(name, self._config.enabled_tools) } if self._config.disabled_tools: return { name: cls for name, cls in result.items() if not name_matches(name, self._config.disabled_tools) } return result def _apply_per_source_filtering( self, tools: dict[str, type[BaseTool]] ) -> dict[str, type[BaseTool]]: """Filter out MCP/connector tools disabled at the server or connector level.""" disabled_sources, per_source_disabled = self._build_source_disable_index() if not disabled_sources and not per_source_disabled: return tools return { name: cls for name, cls in tools.items() if not self._is_source_disabled(cls, disabled_sources, per_source_disabled) } def _build_source_disable_index( self, ) -> tuple[set[tuple[str, bool]], dict[tuple[str, bool], set[str]]]: """Return (fully_disabled, per_tool_disabled) keyed by (source_name, is_connector).""" disabled_sources: set[tuple[str, bool]] = set() per_source_disabled: dict[tuple[str, bool], set[str]] = {} for srv in self._config.mcp_servers: key = (srv.name, False) if srv.disabled: disabled_sources.add(key) elif srv.disabled_tools: per_source_disabled[key] = set(srv.disabled_tools) for cfg in self._config.connectors: key = (cfg.name, True) if cfg.disabled: disabled_sources.add(key) elif cfg.disabled_tools: per_source_disabled[key] = set(cfg.disabled_tools) return disabled_sources, per_source_disabled @staticmethod def _is_source_disabled( tool_cls: type[BaseTool], disabled_sources: set[tuple[str, bool]], per_source_disabled: dict[tuple[str, bool], set[str]], ) -> bool: if not issubclass(tool_cls, MCPTool): return False server_name = tool_cls.get_server_name() if server_name is None: return False key = (server_name, tool_cls.is_connector()) if key in disabled_sources: return True return tool_cls.get_remote_name() in per_source_disabled.get(key, set()) def integrate_mcp(self, *, raise_on_failure: bool = False) -> None: """Discover and register MCP tools (sync wrapper). Idempotent: subsequent calls after a successful integration are no-ops to avoid redundant MCP discovery. """ run_sync(self._integrate_mcp_async(raise_on_failure=raise_on_failure)) async def _integrate_mcp_async(self, *, raise_on_failure: bool = False) -> None: """Async MCP discovery — canonical implementation.""" if self._mcp_integrated: return if not self._config.mcp_servers: return try: mcp_tools = await self._mcp_registry.get_tools_async( self._config.mcp_servers ) except Exception as exc: logger.warning("MCP integration failed: %s", exc) if raise_on_failure: raise return with self._lock: self._available = {**self._available, **mcp_tools} self._mcp_integrated = True logger.info( "MCP integration registered %d tools (via registry)", len(mcp_tools) ) def _purge_connector_state(self) -> None: """Remove stale connector tool classes and cached instances.""" stale_keys = [ name for name, cls in self._available.items() if issubclass(cls, MCPTool) and cls.is_connector() ] for key in stale_keys: self._available.pop(key, None) self._instances.pop(key, None) def _purge_mcp_state(self) -> None: """Remove stale MCP tool classes and cached instances.""" stale_keys = [ name for name, cls in self._available.items() if issubclass(cls, MCPTool) and not cls.is_connector() ] for key in stale_keys: self._available.pop(key, None) self._instances.pop(key, None) def integrate_connectors(self) -> None: """Discover and register connector tools (sync wrapper).""" run_sync(self.integrate_connectors_async()) async def integrate_connectors_async(self) -> None: """Discover and register connector tools — canonical implementation. Thread-safe: can be called from the deferred-init background thread. """ if self._connector_registry is None: return try: connector_tools = await self._connector_registry.get_tools_async() except Exception as exc: logger.warning(f"Connector integration failed: {exc}") with self._lock: self._purge_connector_state() return with self._lock: self._purge_connector_state() self._available.update(connector_tools) logger.info(f"Connector integration registered {len(connector_tools)} tools") async def refresh_remote_tools_async(self) -> None: """Force MCP and connector re-discovery for the current config.""" with self._lock: self._mcp_registry.clear() self._purge_mcp_state() self._mcp_integrated = False self._purge_connector_state() if self._connector_registry is not None: self._connector_registry.clear() await self._integrate_all_async() def refresh_remote_tools(self) -> None: """Sync wrapper for :meth:`refresh_remote_tools_async`.""" run_sync(self.refresh_remote_tools_async()) def integrate_all(self, *, raise_on_mcp_failure: bool = False) -> None: """Discover MCP and connector tools in parallel. Runs both async discovery paths concurrently via ``asyncio.gather`` inside a single ``run_sync`` call. """ run_sync(self._integrate_all_async(raise_on_mcp_failure=raise_on_mcp_failure)) async def _integrate_all_async(self, *, raise_on_mcp_failure: bool = False) -> None: """Run MCP and connector discovery concurrently. Uses ``return_exceptions=True`` so that a failing MCP server does not cancel in-flight connector discovery (or vice-versa). """ mcp_result, connector_result = await asyncio.gather( self._integrate_mcp_async(raise_on_failure=raise_on_mcp_failure), self.integrate_connectors_async(), return_exceptions=True, ) # Re-raise MCP errors when the caller asked for them. if isinstance(mcp_result, BaseException): if raise_on_mcp_failure: raise mcp_result logger.warning(f"MCP integration failed: {mcp_result}") if isinstance(connector_result, BaseException): logger.warning(f"Connector integration failed: {connector_result}") def get_tool_config(self, tool_name: str) -> BaseToolConfig: with self._lock: tool_class = self._available.get(tool_name) if tool_class: config_class = tool_class._get_tool_config_class() default_config = config_class() else: config_class = BaseToolConfig default_config = BaseToolConfig() user_overrides = self._config.tools.get(tool_name) if user_overrides is None: return config_class() merged_dict = {**default_config.model_dump(), **user_overrides} return config_class.model_validate(merged_dict) def get(self, tool_name: str) -> BaseTool: """Get a tool instance, creating it lazily on first call. Raises: NoSuchToolError: If the requested tool is not available. """ if tool_name in self._instances: return self._instances[tool_name] with self._lock: if tool_name not in self._available: raise NoSuchToolError( f"Unknown tool: {tool_name}. Available: {list(self._available.keys())}" ) tool_class = self._available[tool_name] self._instances[tool_name] = tool_class.from_config( lambda: self.get_tool_config(tool_name) ) return self._instances[tool_name] def reset_all(self) -> None: self._instances.clear()