autobyteus 1.1.2__py3-none-any.whl → 1.1.3__py3-none-any.whl

autobyteus/tools/mcp/server/stdio_managed_mcp_server.py

@@ -0,0 +1,33 @@
+# file: autobyteus/autobyteus/tools/mcp/server/stdio_managed_mcp_server.py
+import logging
+from typing import cast
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from .base_managed_mcp_server import BaseManagedMcpServer
+from ..types import StdioMcpServerConfig
+
+logger = logging.getLogger(__name__)
+
+class StdioManagedMcpServer(BaseManagedMcpServer):
+    """Manages the lifecycle of a stdio-based MCP server."""
+
+    def __init__(self, config: StdioMcpServerConfig):
+        super().__init__(config)
+
+    async def _create_client_session(self) -> ClientSession:
+        """Starts a subprocess and establishes a client session over its stdio."""
+        config = cast(StdioMcpServerConfig, self._config)
+        stdio_params = StdioServerParameters(
+            command=config.command,
+            args=config.args,
+            env=config.env,
+            cwd=config.cwd
+        )
+        
+        logger.debug(f"Establishing stdio connection for server '{self.server_id}' with command: {config.command}")
+        read_stream, write_stream = await self._exit_stack.enter_async_context(stdio_client(stdio_params))
+        session = await self._exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
+        logger.debug(f"ClientSession established for stdio server '{self.server_id}'.")
+        return session

autobyteus/tools/mcp/server_instance_manager.py

@@ -0,0 +1,93 @@
+# file: autobyteus/autobyteus/tools/mcp/server_instance_manager.py
+import logging
+from typing import Dict, List, AsyncIterator
+from contextlib import asynccontextmanager
+
+from autobyteus.utils.singleton import SingletonMeta
+from .config_service import McpConfigService
+from .server import BaseManagedMcpServer, StdioManagedMcpServer, HttpManagedMcpServer
+from .types import McpTransportType, McpServerInstanceKey, BaseMcpConfig
+
+logger = logging.getLogger(__name__)
+
+class McpServerInstanceManager(metaclass=SingletonMeta):
+    """
+    Manages the lifecycle of BaseManagedMcpServer instances, providing
+    isolated server connections on a per-agent, per-server_id basis.
+    """
+    def __init__(self):
+        self._config_service = McpConfigService()
+        self._active_servers: Dict[McpServerInstanceKey, BaseManagedMcpServer] = {}
+        logger.info("McpServerInstanceManager initialized.")
+    
+    def _create_server_instance(self, server_config: BaseMcpConfig) -> BaseManagedMcpServer:
+        """Factory method to create a server instance from a config."""
+        if server_config.transport_type == McpTransportType.STDIO:
+            return StdioManagedMcpServer(server_config)
+        elif server_config.transport_type == McpTransportType.STREAMABLE_HTTP:
+            return HttpManagedMcpServer(server_config)
+        else:
+            raise NotImplementedError(f"No ManagedMcpServer implementation for transport type '{server_config.transport_type}'.")
+
+    def get_server_instance(self, agent_id: str, server_id: str) -> BaseManagedMcpServer:
+        """
+        Retrieves or creates a dedicated, long-lived managed server instance
+        for a given agent and server ID.
+        """
+        instance_key = McpServerInstanceKey(agent_id=agent_id, server_id=server_id)
+        
+        if instance_key in self._active_servers:
+            return self._active_servers[instance_key]
+
+        logger.info(f"Creating new persistent server instance for {instance_key}.")
+        config = self._config_service.get_config(server_id)
+        if not config:
+            raise ValueError(f"No configuration found for server_id '{server_id}'.")
+
+        server_instance = self._create_server_instance(config)
+        self._active_servers[instance_key] = server_instance
+        return server_instance
+
+    @asynccontextmanager
+    async def managed_discovery_session(self, server_config: BaseMcpConfig) -> AsyncIterator[BaseManagedMcpServer]:
+        """
+        Provides a temporary server instance for a one-shot operation like discovery.
+        Guarantees the instance is closed upon exiting the context.
+        This method uses the provided config object directly and does not look it up
+        in the config service, making it suitable for stateless previews.
+        """
+        if not server_config:
+            raise ValueError("A valid BaseMcpConfig object must be provided to managed_discovery_session.")
+        
+        logger.debug(f"Creating temporary discovery instance for server '{server_config.server_id}'.")
+        temp_server_instance = self._create_server_instance(server_config)
+        try:
+            yield temp_server_instance
+        finally:
+            logger.debug(f"Closing temporary discovery instance for server '{server_config.server_id}'.")
+            await temp_server_instance.close()
+
+    async def cleanup_mcp_server_instances_for_agent(self, agent_id: str):
+        """
+        Closes all active MCP server instances and removes them from the cache for a specific agent.
+        """
+        logger.info(f"Cleaning up all MCP server instances for agent '{agent_id}'.")
+        keys_to_remove: List[McpServerInstanceKey] = []
+        for instance_key, server_instance in self._active_servers.items():
+            if instance_key.agent_id == agent_id:
+                try:
+                    await server_instance.close()
+                except Exception as e:
+                    logger.error(f"Error closing MCP server '{instance_key.server_id}' for agent '{instance_key.agent_id}': {e}", exc_info=True)
+                keys_to_remove.append(instance_key)
+        
+        for key in keys_to_remove:
+            del self._active_servers[key]
+        logger.info(f"Finished cleaning up MCP server instances for agent '{agent_id}'.")
+
+    async def cleanup_all_mcp_server_instances(self):
+        """Closes all active MCP server instances for all agents and clears the cache."""
+        logger.info("Cleaning up all active MCP server instances.")
+        agent_ids = {key.agent_id for key in self._active_servers.keys()}
+        for agent_id in agent_ids:
+            await self.cleanup_mcp_server_instances_for_agent(agent_id)

autobyteus/tools/mcp/tool.py

@@ -4,98 +4,80 @@ from typing import Any, Optional, TYPE_CHECKING
 
 from autobyteus.tools.base_tool import BaseTool
 from autobyteus.tools.parameter_schema import ParameterSchema
+from .server.proxy import McpServerProxy
 
 if TYPE_CHECKING:
     from autobyteus.agent.context import AgentContext
-    from .call_handlers.base_handler import McpCallHandler
-    from .types import BaseMcpConfig
 
 logger = logging.getLogger(__name__)
 
 class GenericMcpTool(BaseTool):
     """
     A generic tool wrapper for executing tools on a remote MCP server.
-    This tool delegates the entire remote call, including connection and
-    protocol specifics, to an MCP call handler.
+    This tool is a lightweight blueprint. At execution time, it uses a proxy
+    to interact with the server, completely hiding the connection management logic.
     """
 
     def __init__(self,
-                 mcp_server_config: 'BaseMcpConfig',
-                 mcp_remote_tool_name: str,
-                 mcp_call_handler: 'McpCallHandler',
+                 server_id: str,
+                 remote_tool_name: str,
                  name: str, 
                  description: str,
                  argument_schema: ParameterSchema):
         """
-        Initializes the GenericMcpTool instance.
+        Initializes the GenericMcpTool instance with identifiers.
         """
         super().__init__() 
         
-        self._mcp_server_config = mcp_server_config
-        self._mcp_remote_tool_name = mcp_remote_tool_name
-        self._mcp_call_handler = mcp_call_handler
+        self._server_id = server_id
+        self._remote_tool_name = remote_tool_name
         
+        # Instance-specific properties for schema generation
         self._instance_name = name
         self._instance_description = description
         self._instance_argument_schema = argument_schema
         
-        # Override the base class's schema-related methods with instance-specific
-        # versions for correct validation and usage generation.
         self.get_name = self.get_instance_name
         self.get_description = self.get_instance_description
         self.get_argument_schema = self.get_instance_argument_schema
         
-        logger.info(f"GenericMcpTool instance created for remote tool '{mcp_remote_tool_name}' on server '{self._mcp_server_config.server_id}'. "
+        logger.info(f"GenericMcpTool instance created for remote tool '{remote_tool_name}' on server '{self._server_id}'. "
                     f"Registered in AutoByteUs as '{self._instance_name}'.")
 
     # --- Getters for instance-specific data ---
+    def get_instance_name(self) -> str: return self._instance_name
+    def get_instance_description(self) -> str: return self._instance_description
+    def get_instance_argument_schema(self) -> ParameterSchema: return self._instance_argument_schema
 
-    def get_instance_name(self) -> str:
-        return self._instance_name
-
-    def get_instance_description(self) -> str:
-        return self._instance_description
-
-    def get_instance_argument_schema(self) -> ParameterSchema:
-        return self._instance_argument_schema
-
-    # --- Base class methods that are NOT overridden at instance level ---
-
+    # --- Base class methods (class-level, not instance-level) ---
     @classmethod
-    def get_name(cls) -> str:
-        return "GenericMcpTool" 
-
+    def get_name(cls) -> str: return "GenericMcpTool"
     @classmethod
-    def get_description(cls) -> str:
-        return "A generic wrapper for executing tools on a remote MCP server. Specifics are instance-based."
-
+    def get_description(cls) -> str: return "A generic wrapper for executing remote MCP tools."
     @classmethod
-    def get_argument_schema(cls) -> Optional[ParameterSchema]:
-        return None 
+    def get_argument_schema(cls) -> Optional[ParameterSchema]: return None 
 
     async def _execute(self, context: 'AgentContext', **kwargs: Any) -> Any:
         """
-        Executes the remote MCP tool by delegating to the injected call handler.
+        Creates a proxy for the remote server and executes the tool call.
+        The proxy handles all interaction with the McpServerInstanceManager.
         """
+        agent_id = context.agent_id
         tool_name_for_log = self.get_instance_name()
-        logger.info(f"GenericMcpTool '{tool_name_for_log}': Delegating call for remote tool '{self._mcp_remote_tool_name}' "
-                    f"on server '{self._mcp_server_config.server_id}' to handler.")
         
-        if not self._mcp_call_handler:
-             logger.error(f"GenericMcpTool '{tool_name_for_log}': McpCallHandler is not set. Cannot execute.")
-             raise RuntimeError("McpCallHandler not available in GenericMcpTool instance.")
-
+        logger.info(f"GenericMcpTool '{tool_name_for_log}': Creating proxy for agent '{agent_id}' and server '{self._server_id}'.")
+        
         try:
-            # The handler is responsible for the entire end-to-end call.
-            return await self._mcp_call_handler.handle_call(
-                config=self._mcp_server_config,
-                remote_tool_name=self._mcp_remote_tool_name,
+            # The proxy is created on-demand for each execution.
+            proxy = McpServerProxy(agent_id=agent_id, server_id=self._server_id)
+            
+            return await proxy.call_tool(
+                tool_name=self._remote_tool_name,
                 arguments=kwargs
             )
         except Exception as e:
             logger.error(
-                f"The MCP call handler for tool '{tool_name_for_log}' raised an exception: {e}",
+                f"Execution failed for tool '{tool_name_for_log}' on server '{self._server_id}' for agent '{agent_id}': {e}",
                 exc_info=True
             )
-            # Re-raise to ensure the agent knows the tool call failed.
             raise

autobyteus/tools/mcp/tool_registrar.py

@@ -0,0 +1,177 @@
+# file: autobyteus/autobyteus/tools/mcp/tool_registrar.py
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+# Consolidated imports from the autobyteus.autobyteus.mcp package public API
+from .config_service import McpConfigService
+from .factory import McpToolFactory
+from .schema_mapper import McpSchemaMapper
+from .server_instance_manager import McpServerInstanceManager
+from .types import BaseMcpConfig
+from .server import BaseManagedMcpServer
+
+from autobyteus.tools.registry import ToolRegistry, ToolDefinition
+from autobyteus.tools.tool_category import ToolCategory
+from autobyteus.utils.singleton import SingletonMeta
+from mcp import types as mcp_types
+
+
+logger = logging.getLogger(__name__)
+
+class McpToolRegistrar(metaclass=SingletonMeta):
+    """
+    Orchestrates the discovery of remote MCP tools and their registration
+    with the AutoByteUs ToolRegistry.
+    """
+    def __init__(self):
+        """
+        Initializes the McpToolRegistrar singleton.
+        """
+        self._config_service: McpConfigService = McpConfigService()
+        self._tool_registry: ToolRegistry = ToolRegistry()
+        self._instance_manager: McpServerInstanceManager = McpServerInstanceManager()
+        self._registered_tools_by_server: Dict[str, List[ToolDefinition]] = {}
+        logger.info("McpToolRegistrar initialized.")
+
+    async def _fetch_tools_from_server(self, server_config: BaseMcpConfig) -> List[mcp_types.Tool]:
+        """
+        Uses the instance manager to get a temporary, managed session for discovery.
+        """
+        async with self._instance_manager.managed_discovery_session(server_config) as discovery_server:
+            # The context manager guarantees the server is connected and will be closed.
+            remote_tools = await discovery_server.list_remote_tools()
+            return remote_tools
+
+    def _create_tool_definition_from_remote(
+        self,
+        remote_tool: mcp_types.Tool,
+        server_config: BaseMcpConfig,
+        schema_mapper: McpSchemaMapper
+    ) -> ToolDefinition:
+        """
+        Maps a single remote tool from an MCP server to an AutoByteUs ToolDefinition.
+        """
+        actual_arg_schema = schema_mapper.map_to_autobyteus_schema(remote_tool.inputSchema)
+        actual_desc = remote_tool.description
+        
+        registered_name = remote_tool.name
+        if server_config.tool_name_prefix:
+            registered_name = f"{server_config.tool_name_prefix.rstrip('_')}_{remote_tool.name}"
+
+        tool_factory = McpToolFactory(
+            server_id=server_config.server_id,
+            remote_tool_name=remote_tool.name,
+            registered_tool_name=registered_name,
+            tool_description=actual_desc,
+            tool_argument_schema=actual_arg_schema
+        )
+        
+        return ToolDefinition(
+            name=registered_name,
+            description=actual_desc,
+            argument_schema=actual_arg_schema,
+            category=ToolCategory.MCP,
+            metadata={"mcp_server_id": server_config.server_id}, # Store origin in generic metadata
+            custom_factory=tool_factory.create_tool,
+            config_schema=None,
+            tool_class=None
+        )
+
+    async def discover_and_register_tools(self, mcp_config: Optional[Union[BaseMcpConfig, Dict[str, Any]]] = None) -> List[ToolDefinition]:
+        """
+        Discovers tools from MCP servers and registers them.
+        """
+        configs_to_process: List[BaseMcpConfig]
+        
+        if mcp_config:
+            validated_config: BaseMcpConfig
+            if isinstance(mcp_config, dict):
+                try:
+                    validated_config = self._config_service.load_config(mcp_config)
+                except ValueError as e:
+                    logger.error(f"Failed to parse provided MCP config dictionary: {e}")
+                    raise
+            elif isinstance(mcp_config, BaseMcpConfig):
+                validated_config = self._config_service.add_config(mcp_config)
+            else:
+                raise TypeError(f"mcp_config must be a BaseMcpConfig object or a dictionary, not {type(mcp_config)}.")
+            
+            logger.info(f"Starting targeted MCP tool discovery for server: {validated_config.server_id}")
+            self.unregister_tools_from_server(validated_config.server_id)
+            configs_to_process = [validated_config]
+        else:
+            logger.info("Starting full MCP tool discovery. Unregistering all existing MCP tools first.")
+            all_server_ids = list(self._registered_tools_by_server.keys())
+            for server_id in all_server_ids:
+                self.unregister_tools_from_server(server_id)
+            self._registered_tools_by_server.clear()
+            configs_to_process = self._config_service.get_all_configs()
+
+        if not configs_to_process:
+            logger.info("No MCP server configurations to process. Skipping discovery.")
+            return []
+
+        schema_mapper = McpSchemaMapper()
+        registered_tool_definitions: List[ToolDefinition] = []
+        for server_config in configs_to_process:
+            if not server_config.enabled:
+                logger.info(f"MCP server '{server_config.server_id}' is disabled. Skipping.")
+                continue
+
+            logger.info(f"Discovering tools from MCP server: '{server_config.server_id}'")
+            
+            try:
+                remote_tools = await self._fetch_tools_from_server(server_config)
+                logger.info(f"Discovered {len(remote_tools)} tools from server '{server_config.server_id}'.")
+
+                for remote_tool in remote_tools:
+                    try:
+                        tool_def = self._create_tool_definition_from_remote(remote_tool, server_config, schema_mapper)
+                        self._tool_registry.register_tool(tool_def)
+                        self._registered_tools_by_server.setdefault(server_config.server_id, []).append(tool_def)
+                        registered_tool_definitions.append(tool_def)
+                    except Exception as e_tool:
+                        logger.error(f"Failed to process or register remote tool '{remote_tool.name}': {e_tool}", exc_info=True)
+            
+            except Exception as e_server:
+                logger.error(f"Failed to discover tools from MCP server '{server_config.server_id}': {e_server}", exc_info=True)
+        
+        logger.info(f"MCP tool discovery and registration process completed. Total tools registered: {len(registered_tool_definitions)}.")
+        return registered_tool_definitions
+
+    async def list_remote_tools(self, mcp_config: Union[BaseMcpConfig, Dict[str, Any]]) -> List[ToolDefinition]:
+        validated_config: BaseMcpConfig
+        if isinstance(mcp_config, dict):
+            validated_config = McpConfigService.parse_mcp_config_dict(mcp_config)
+        elif isinstance(mcp_config, BaseMcpConfig):
+            validated_config = mcp_config
+        else:
+            raise TypeError(f"mcp_config must be a BaseMcpConfig object or a dictionary, not {type(mcp_config)}.")
+        
+        logger.info(f"Previewing tools from MCP server: '{validated_config.server_id}'")
+        schema_mapper = McpSchemaMapper()
+        tool_definitions: List[ToolDefinition] = []
+
+        try:
+            remote_tools = await self._fetch_tools_from_server(validated_config)
+            logger.info(f"Discovered {len(remote_tools)} tools from server '{validated_config.server_id}' for preview.")
+            for remote_tool in remote_tools:
+                tool_def = self._create_tool_definition_from_remote(remote_tool, validated_config, schema_mapper)
+                tool_definitions.append(tool_def)
+        except Exception as e_server:
+            logger.error(f"Failed to discover tools for preview from MCP server '{validated_config.server_id}': {e_server}", exc_info=True)
+            raise
+            
+        logger.info(f"MCP tool preview completed. Found {len(tool_definitions)} tools.")
+        return tool_definitions
+    
+    def unregister_tools_from_server(self, server_id: str) -> bool:
+        if not self.is_server_registered(server_id):
+            return False
+        tools_to_unregister = self._registered_tools_by_server.pop(server_id, [])
+        for tool_def in tools_to_unregister:
+            self._tool_registry.unregister_tool(tool_def.name)
+        return True
+        
+    def is_server_registered(self, server_id: str) -> bool:
+        return server_id in self._registered_tools_by_server

autobyteus/tools/mcp/types.py

@@ -1,4 +1,4 @@
-# file: autobyteus/autobyteus/mcp/types.py
+# file: autobyteus/autobyteus/tools/mcp/types.py
 import logging
 from typing import List, Dict, Any, Optional, Type
 from dataclasses import dataclass, field, InitVar
@@ -9,9 +9,17 @@ logger = logging.getLogger(__name__)
 class McpTransportType(str, Enum):
     """Enumeration of supported MCP transport types."""
     STDIO = "stdio"
-    SSE = "sse" 
     STREAMABLE_HTTP = "streamable_http"
 
+@dataclass(frozen=True)
+class McpServerInstanceKey:
+    """
+    A dedicated, hashable key for identifying a unique server instance.
+    An instance is unique for a given agent and a specific server configuration.
+    """
+    agent_id: str
+    server_id: str
+
 @dataclass
 class BaseMcpConfig:
     """
@@ -59,25 +67,6 @@ class StdioMcpServerConfig(BaseMcpConfig):
         if self.cwd is not None and not isinstance(self.cwd, str):
             raise ValueError(f"StdioMcpServerConfig '{self.server_id}' 'cwd' must be a string if provided.") 
 
-@dataclass
-class SseMcpServerConfig(BaseMcpConfig):
-    """Configuration parameters for an MCP server using SSE transport."""
-    url: Optional[str] = None # Changed: Added default None
-    token: Optional[str] = None
-    headers: Dict[str, str] = field(default_factory=dict)
-
-    def __post_init__(self):
-        super().__post_init__()
-        self.transport_type = McpTransportType.SSE
-
-        if self.url is None or not isinstance(self.url, str) or not self.url.strip():
-            raise ValueError(f"SseMcpServerConfig '{self.server_id}' 'url' must be a non-empty string.")
-        
-        if self.token is not None and not isinstance(self.token, str):
-            raise ValueError(f"SseMcpServerConfig '{self.server_id}' 'token' must be a string if provided.") 
-        if not isinstance(self.headers, dict) or not all(isinstance(k, str) and isinstance(v, str) for k, v in self.headers.items()):
-            raise ValueError(f"SseMcpServerConfig '{self.server_id}' 'headers' must be a Dict[str, str].") 
-
 @dataclass
 class StreamableHttpMcpServerConfig(BaseMcpConfig):
     """Configuration parameters for an MCP server using Streamable HTTP transport."""

autobyteus/tools/registry/tool_definition.py

@@ -31,7 +31,8 @@ class ToolDefinition:
                  category: ToolCategory,
                  config_schema: Optional['ParameterSchema'] = None,
                  tool_class: Optional[Type['BaseTool']] = None,
-                 custom_factory: Optional[Callable[['ToolConfig'], 'BaseTool']] = None):
+                 custom_factory: Optional[Callable[['ToolConfig'], 'BaseTool']] = None,
+                 metadata: Optional[Dict[str, Any]] = None):
         """
         Initializes the ToolDefinition.
         """
@@ -56,6 +57,10 @@ class ToolDefinition:
              raise TypeError(f"ToolDefinition '{name}' received an invalid 'config_schema'. Expected ParameterSchema or None.")
         if not isinstance(category, ToolCategory):
             raise TypeError(f"ToolDefinition '{name}' requires a ToolCategory for 'category'. Got {type(category)}")
+        
+        # Validation for MCP-specific metadata
+        if category == ToolCategory.MCP and not (metadata and metadata.get("mcp_server_id")):
+            raise ValueError(f"ToolDefinition '{name}' with category MCP must provide a 'mcp_server_id' in its metadata.")
 
         self._name = name
         self._description = description
@@ -64,6 +69,7 @@ class ToolDefinition:
         self._tool_class = tool_class
         self._custom_factory = custom_factory
         self._category = category
+        self._metadata = metadata or {}
         
         logger.debug(f"ToolDefinition created for tool '{self.name}'.")
 
@@ -82,6 +88,8 @@ class ToolDefinition:
     def config_schema(self) -> Optional['ParameterSchema']: return self._config_schema
     @property
     def category(self) -> ToolCategory: return self._category
+    @property
+    def metadata(self) -> Dict[str, Any]: return self._metadata
     
     # --- Schema Generation API ---
     def get_usage_xml(self, provider: Optional[LLMProvider] = None) -> str:
@@ -136,4 +144,5 @@ class ToolDefinition:
 
     def __repr__(self) -> str:
         creator_repr = f"class='{self._tool_class.__name__}'" if self._tool_class else "factory=True"
-        return (f"ToolDefinition(name='{self.name}', category='{self.category.value}', {creator_repr})")
+        metadata_repr = f", metadata={self.metadata}" if self.metadata else ""
+        return (f"ToolDefinition(name='{self.name}', category='{self.category.value}'{metadata_repr}, {creator_repr})")

autobyteus/tools/registry/tool_registry.py

@@ -5,6 +5,7 @@ from typing import Dict, List, Optional, Type, TYPE_CHECKING
 from autobyteus.tools.registry.tool_definition import ToolDefinition
 from autobyteus.utils.singleton import SingletonMeta
 from autobyteus.tools.tool_config import ToolConfig
+from autobyteus.tools.tool_category import ToolCategory
 
 if TYPE_CHECKING:
     from autobyteus.tools.base_tool import BaseTool
@@ -78,19 +79,8 @@ class ToolRegistry(metaclass=SingletonMeta):
 
     def create_tool(self, name: str, config: Optional[ToolConfig] = None) -> 'BaseTool':
         """
-        Creates a tool instance using its definition, either from a factory or a class.
-
-        Args:
-            name: The name of the tool to create.
-            config: Optional ToolConfig with constructor parameters for class-based tools
-                    or to be passed to a custom factory.
-
-        Returns:
-            The tool instance if the definition exists.
-
-        Raises:
-            ValueError: If the tool definition is not found or is invalid.
-            TypeError: If tool instantiation fails.
+        Creates a tool instance using its definition and injects the definition
+        back into the instance.
         """
         definition = self.get_tool_definition(name)
         if not definition:
@@ -98,24 +88,20 @@ class ToolRegistry(metaclass=SingletonMeta):
             raise ValueError(f"No tool definition found for name '{name}'")
         
         try:
-            # Prefer the custom factory if it exists
+            tool_instance: 'BaseTool'
             if definition.custom_factory:
                 logger.info(f"Creating tool instance for '{name}' using its custom factory.")
-                # Pass the config to the factory. The factory can choose to use it or not.
                 tool_instance = definition.custom_factory(config)
-            
-            # Fall back to instantiating the tool_class
             elif definition.tool_class:
-                # For class-based tools, the convention is to pass the ToolConfig object
-                # itself to the constructor under the 'config' keyword argument.
                 logger.info(f"Creating tool instance for '{name}' using class '{definition.tool_class.__name__}' and passing ToolConfig.")
                 tool_instance = definition.tool_class(config=config)
-            
             else:
-                # This case should be prevented by ToolDefinition's validation
                 raise ValueError(f"ToolDefinition for '{name}' is invalid: missing both tool_class and custom_factory.")
 
-            logger.debug(f"Successfully created tool instance for '{name}'")
+            # Inject the definition into the newly created instance.
+            tool_instance.definition = definition
+            logger.debug(f"Injected ToolDefinition into instance of '{name}'.")
+
             return tool_instance
 
         except Exception as e:
@@ -126,23 +112,36 @@ class ToolRegistry(metaclass=SingletonMeta):
     def list_tools(self) -> List[ToolDefinition]:
         """
         Returns a list of all registered tool definitions.
-
-        Returns:
-            A list of ToolDefinition objects.
         """
         return list(self._definitions.values())
 
     def list_tool_names(self) -> List[str]:
         """
         Returns a list of the names of all registered tools.
-
-        Returns:
-            A list of tool name strings.
         """
         return list(self._definitions.keys())
 
     def get_all_definitions(self) -> Dict[str, ToolDefinition]:
         """Returns the internal dictionary of definitions."""
         return dict(ToolRegistry._definitions)
+        
+    def get_tools_by_mcp_server(self, server_id: str) -> List[ToolDefinition]:
+        """
+        Returns a list of all registered tool definitions that originated
+        from a specific MCP server.
+
+        Args:
+            server_id: The unique ID of the MCP server to query for.
+
+        Returns:
+            A list of matching ToolDefinition objects.
+        """
+        if not server_id:
+            return []
+        
+        return [
+            td for td in self._definitions.values()
+            if td.category == ToolCategory.MCP and td.metadata.get("mcp_server_id") == server_id
+        ]
 
 default_tool_registry = ToolRegistry()

{autobyteus-1.1.2.dist-info → autobyteus-1.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: autobyteus
-Version: 1.1.2
+Version: 1.1.3
 Summary: Multi-Agent framework
 Home-page: https://github.com/AutoByteus/autobyteus
 Author: Ryan Zheng
@@ -33,6 +33,7 @@ Requires-Dist: numpy==2.2.5
 Requires-Dist: aiohttp
 Requires-Dist: autobyteus-llm-client==1.1.1
 Requires-Dist: brui-core==1.0.8
+Requires-Dist: mcp[cli]==1.9.1
 Provides-Extra: dev
 Requires-Dist: coverage; extra == "dev"
 Requires-Dist: flake8; extra == "dev"