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

autobyteus/tools/mcp/factory.py

@@ -2,12 +2,10 @@
 import logging
 from typing import Optional, TYPE_CHECKING
 
-from autobyteus.tools.mcp.tool import GenericMcpTool
+from .tool import GenericMcpTool
 from autobyteus.tools.factory.tool_factory import ToolFactory
 
 if TYPE_CHECKING:
-    from autobyteus.tools.mcp.call_handlers.base_handler import McpCallHandler
-    from autobyteus.tools.mcp.types import BaseMcpConfig
     from autobyteus.tools.parameter_schema import ParameterSchema
     from autobyteus.tools.tool_config import ToolConfig
     from autobyteus.tools.base_tool import BaseTool
@@ -18,52 +16,40 @@ class McpToolFactory(ToolFactory):
     """
     A dedicated factory for creating configured instances of GenericMcpTool.
     
-    This factory captures the necessary context at the time of tool discovery
-    (e.g., server config, remote tool name, call handler) and uses it to 
-    instantiate a GenericMcpTool when requested by the ToolRegistry.
+    This factory captures the key identifiers of a remote tool (server_id,
+    remote_tool_name) and its schema information at the time of discovery.
     """
     def __init__(self,
-                 mcp_server_config: 'BaseMcpConfig',
-                 mcp_remote_tool_name: str,
-                 mcp_call_handler: 'McpCallHandler',
+                 server_id: str,
+                 remote_tool_name: str,
                  registered_tool_name: str,
                  tool_description: str,
                  tool_argument_schema: 'ParameterSchema'):
         """
-        Initializes the factory with the context of a specific remote tool.
+        Initializes the factory with the identifiers and schema of a specific remote tool.
         """
-        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
         self._registered_tool_name = registered_tool_name
         self._tool_description = tool_description
         self._tool_argument_schema = tool_argument_schema
         
         logger.debug(
-            f"McpToolFactory created for remote tool '{self._mcp_remote_tool_name}' "
-            f"on server '{self._mcp_server_config.server_id}' (to be registered as '{self._registered_tool_name}')."
+            f"McpToolFactory created for remote tool '{self._remote_tool_name}' "
+            f"on server '{self._server_id}' (to be registered as '{self._registered_tool_name}')."
         )
 
     def create_tool(self, config: Optional['ToolConfig'] = None) -> 'BaseTool':
         """
         Creates and returns a new instance of GenericMcpTool using the
         configuration captured by this factory.
-
-        Args:
-            config: An optional ToolConfig. This is part of the standard factory
-                    interface but is not used by this specific factory, as all
-                    configuration is provided during initialization.
-        
-        Returns:
-            A configured instance of GenericMcpTool.
         """
         if config:
             logger.debug(f"McpToolFactory for '{self._registered_tool_name}' received a ToolConfig, which will be ignored.")
             
         return GenericMcpTool(
-            mcp_server_config=self._mcp_server_config,
-            mcp_remote_tool_name=self._mcp_remote_tool_name,
-            mcp_call_handler=self._mcp_call_handler,
+            server_id=self._server_id,
+            remote_tool_name=self._remote_tool_name,
             name=self._registered_tool_name,
             description=self._tool_description,
             argument_schema=self._tool_argument_schema

autobyteus/tools/mcp/registrar.py

@@ -2,22 +2,12 @@
 import logging
 from typing import Any, Dict, List, Optional, Union
 
-# Import the new handler architecture components
-from .call_handlers import (
-    McpCallHandler,
-    StdioMcpCallHandler,
-    StreamableHttpMcpCallHandler,
-    SseMcpCallHandler
-)
-
 # Consolidated imports from the autobyteus.autobyteus.mcp package public API
-from autobyteus.tools.mcp import (
-    McpConfigService,
-    McpSchemaMapper,
-    McpToolFactory,
-    McpTransportType,
-    BaseMcpConfig
-)
+from .config_service import McpConfigService
+from .factory import McpToolFactory
+from .schema_mapper import McpSchemaMapper
+from .types import BaseMcpConfig, McpTransportType
+from .server import StdioManagedMcpServer, HttpManagedMcpServer, BaseManagedMcpServer
 
 from autobyteus.tools.registry import ToolRegistry, ToolDefinition
 from autobyteus.tools.tool_category import ToolCategory
@@ -30,40 +20,48 @@ logger = logging.getLogger(__name__)
 class McpToolRegistrar(metaclass=SingletonMeta):
     """
     Orchestrates the discovery of remote MCP tools and their registration
-    with the AutoByteUs ToolRegistry using a handler-based architecture.
-    This class is a singleton.
+    with the AutoByteUs ToolRegistry using a stateful, server-centric architecture.
     """
     def __init__(self):
         """
         Initializes the McpToolRegistrar singleton.
-        It retrieves singleton instances of its service dependencies and initializes
-        its internal state for tracking registered tools.
         """
         self._config_service: McpConfigService = McpConfigService()
         self._tool_registry: ToolRegistry = ToolRegistry()
-        
-        # The handler registry maps a transport type to a reusable handler instance.
-        self._handler_registry: Dict[McpTransportType, McpCallHandler] = {
-            McpTransportType.STDIO: StdioMcpCallHandler(),
-            McpTransportType.STREAMABLE_HTTP: StreamableHttpMcpCallHandler(),
-            McpTransportType.SSE: SseMcpCallHandler(),
-        }
-
-        # Internal state to track which ToolDefinitions were registered from which server.
         self._registered_tools_by_server: Dict[str, List[ToolDefinition]] = {}
-        
-        logger.info(f"McpToolRegistrar initialized with {len(self._handler_registry)} call handlers.")
+        logger.info("McpToolRegistrar initialized.")
+
+    def _create_server_instance_for_discovery(self, server_config: BaseMcpConfig) -> BaseManagedMcpServer:
+        """Creates a server instance based on transport type."""
+        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"Discovery not implemented for transport type: {server_config.transport_type}")
+
+    async def _fetch_tools_from_server(self, server_config: BaseMcpConfig) -> List[mcp_types.Tool]:
+        """
+        Creates a temporary server instance to perform a single, one-shot
+        tool discovery, ensuring resources are properly closed.
+        """
+        discovery_server = self._create_server_instance_for_discovery(server_config)
+        try:
+            # The list_remote_tools method implicitly handles the connection.
+            remote_tools = await discovery_server.list_remote_tools()
+            return remote_tools
+        finally:
+            # The finally block guarantees the temporary server connection is closed.
+            await discovery_server.close()
 
     def _create_tool_definition_from_remote(
         self,
         remote_tool: mcp_types.Tool,
         server_config: BaseMcpConfig,
-        handler: McpCallHandler,
         schema_mapper: McpSchemaMapper
     ) -> ToolDefinition:
         """
         Maps a single remote tool from an MCP server to an AutoByteUs ToolDefinition.
-        This is a helper method to centralize the mapping logic.
         """
         if hasattr(remote_tool, 'model_dump_json'):
             logger.debug(f"Processing remote tool from server '{server_config.server_id}':\n{remote_tool.model_dump_json(indent=2)}")
@@ -75,16 +73,16 @@ class McpToolRegistrar(metaclass=SingletonMeta):
         if server_config.tool_name_prefix:
             registered_name = f"{server_config.tool_name_prefix.rstrip('_')}_{remote_tool.name}"
 
+        # The factory now only needs key identifiers, not live objects.
         tool_factory = McpToolFactory(
-            mcp_server_config=server_config,
-            mcp_remote_tool_name=remote_tool.name,
-            mcp_call_handler=handler,
+            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
         )
         
-        tool_def = ToolDefinition(
+        return ToolDefinition(
             name=registered_name,
             description=actual_desc,
             argument_schema=actual_arg_schema,
@@ -93,54 +91,29 @@ class McpToolRegistrar(metaclass=SingletonMeta):
             config_schema=None,
             tool_class=None
         )
-        return tool_def
 
     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.
-
-        If `mcp_config` is provided (as a validated object or a raw dictionary),
-        it discovers tools only from that specific server, unregistering any
-        of its old tools first.
-        
-        If `mcp_config` is None, it unregisters all existing MCP tools and then
-        discovers tools from all enabled servers found in the McpConfigService.
-        
-        Returns:
-            A list of ToolDefinition objects for all tools that were successfully
-            discovered and registered during this call.
+        This process uses a helper to manage short-lived server instances for discovery.
         """
         configs_to_process: List[BaseMcpConfig]
         
         if mcp_config:
-            validated_config: BaseMcpConfig
-            # If the user provided a raw dictionary, parse and add it via the service.
             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
-            # If a validated object was passed, add it to the service to ensure it's known.
+                configs_to_process = [self._config_service.load_config(mcp_config)]
             elif isinstance(mcp_config, BaseMcpConfig):
-                validated_config = self._config_service.add_config(mcp_config)
+                configs_to_process = [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}")
-            # Unregister any existing tools from this server before re-discovering.
-            self.unregister_tools_from_server(validated_config.server_id)
-            configs_to_process = [validated_config]
+            logger.info(f"Starting targeted MCP tool discovery for server: {configs_to_process[0].server_id}")
+            self.unregister_tools_from_server(configs_to_process[0].server_id)
         else:
-            logger.info("Starting full MCP tool discovery and registration process. Unregistering all existing MCP tools first.")
-            # Get a copy of all server IDs that have tools registered.
+            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)
-            
-            # The registrar's internal state should now be empty, but a clear() is a good safeguard.
             self._registered_tools_by_server.clear()
-            
             configs_to_process = self._config_service.get_all_configs()
 
         if not configs_to_process:
@@ -155,37 +128,21 @@ class McpToolRegistrar(metaclass=SingletonMeta):
                 continue
 
             logger.info(f"Discovering tools from MCP server: '{server_config.server_id}' ({server_config.transport_type.value})")
+            
             try:
-                handler = self._handler_registry.get(server_config.transport_type)
-                if not handler:
-                    logger.error(f"No MCP call handler found for transport type '{server_config.transport_type.value}' on server '{server_config.server_id}'.")
-                    continue
-
-                remote_tools_result = await handler.handle_call(
-                    config=server_config,
-                    remote_tool_name="list_tools",
-                    arguments={}
-                )
-                
-                if not isinstance(remote_tools_result, mcp_types.ListToolsResult):
-                    logger.error(f"Expected ListToolsResult from handler for 'list_tools', but got {type(remote_tools_result)}. Skipping server '{server_config.server_id}'.")
-                    continue
-
-                actual_remote_tools: list[mcp_types.Tool] = remote_tools_result.tools
-                logger.info(f"Discovered {len(actual_remote_tools)} tools from server '{server_config.server_id}'.")
+                # The helper abstracts away the connect/close lifecycle for discovery.
+                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 actual_remote_tools: 
+                for remote_tool in remote_tools:
                     try:
-                        tool_def = self._create_tool_definition_from_remote(remote_tool, server_config, handler, schema_mapper)
-                        
+                        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)
-                        
                         logger.info(f"Successfully registered MCP tool '{remote_tool.name}' from server '{server_config.server_id}' as '{tool_def.name}'.")
-
                     except Exception as e_tool:
-                        logger.error(f"Failed to process or register remote tool '{remote_tool.name}' from server '{server_config.server_id}': {e_tool}", exc_info=True)
+                        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)
@@ -195,23 +152,11 @@ class McpToolRegistrar(metaclass=SingletonMeta):
 
     async def list_remote_tools(self, mcp_config: Union[BaseMcpConfig, Dict[str, Any]]) -> List[ToolDefinition]:
         """
-        Previews tools from a remote MCP server without registering them or storing the configuration.
+        Previews tools from a remote MCP server without registering them.
         This is a stateless "dry-run" or "preview" operation.
-
-        Args:
-            mcp_config: A single MCP server configuration, as a validated object or a raw dictionary.
-        
-        Returns:
-            A list of ToolDefinition objects for all tools discovered on the server.
-        
-        Raises:
-            TypeError: If mcp_config is not a supported type.
-            ValueError: If the provided configuration is invalid.
-            RuntimeError: If the connection or tool call to the remote server fails.
         """
         validated_config: BaseMcpConfig
         if isinstance(mcp_config, dict):
-            # Use the static method to parse the config without storing it in the service.
             validated_config = McpConfigService.parse_mcp_config_dict(mcp_config)
         elif isinstance(mcp_config, BaseMcpConfig):
             validated_config = mcp_config
@@ -224,100 +169,34 @@ class McpToolRegistrar(metaclass=SingletonMeta):
         tool_definitions: List[ToolDefinition] = []
 
         try:
-            handler = self._handler_registry.get(validated_config.transport_type)
-            if not handler:
-                raise ValueError(f"No MCP call handler found for transport type '{validated_config.transport_type.value}'.")
-
-            remote_tools_result = await handler.handle_call(
-                config=validated_config,
-                remote_tool_name="list_tools",
-                arguments={}
-            )
-            
-            if not isinstance(remote_tools_result, mcp_types.ListToolsResult):
-                error_msg = f"Expected ListToolsResult from handler, but got {type(remote_tools_result)}. Cannot preview tools from server '{validated_config.server_id}'."
-                logger.error(error_msg)
-                raise RuntimeError(error_msg)
+            # Use the same helper to fetch tools, abstracting away the connection lifecycle.
+            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.")
 
-            actual_remote_tools: list[mcp_types.Tool] = remote_tools_result.tools
-            logger.info(f"Discovered {len(actual_remote_tools)} tools from server '{validated_config.server_id}' for preview.")
-
-            for remote_tool in actual_remote_tools:
+            for remote_tool in remote_tools:
                 try:
-                    tool_def = self._create_tool_definition_from_remote(remote_tool, validated_config, handler, schema_mapper)
+                    tool_def = self._create_tool_definition_from_remote(remote_tool, validated_config, schema_mapper)
                     tool_definitions.append(tool_def)
                 except Exception as e_tool:
                     logger.error(f"Failed to map remote tool '{remote_tool.name}' from server '{validated_config.server_id}' during preview: {e_tool}", exc_info=True)
-                    # For a preview, we continue to show other valid tools.
             
         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)
-            # Re-raise server-level exceptions to the caller as per AC5.
             raise
             
         logger.info(f"MCP tool preview completed. Found {len(tool_definitions)} tools.")
         return tool_definitions
 
-    def get_registered_tools_for_server(self, server_id: str) -> List[ToolDefinition]:
-        """
-        Returns a list of ToolDefinition objects that were successfully registered 
-        from a specific MCP server during the most recent discovery process.
-
-        Args:
-            server_id: The unique ID of the MCP server to query.
-
-        Returns:
-            A list of ToolDefinition objects. Returns an empty list if the server ID
-            is not found or registered no tools.
-        """
-        return self._registered_tools_by_server.get(server_id, [])
-
-    def get_all_registered_mcp_tools(self) -> List[ToolDefinition]:
-        """
-        Returns a flat list of all ToolDefinitions registered from any MCP server.
-        """
-        all_tools = []
-        for server_tools in self._registered_tools_by_server.values():
-            all_tools.extend(server_tools)
-        return all_tools
-
-    def is_server_registered(self, server_id: str) -> bool:
-        """
-        Checks if any tools from a specific MCP server are currently registered.
-
-        Args:
-            server_id: The unique ID of the MCP server.
-
-        Returns:
-            True if the server has tools registered, False otherwise.
-        """
-        is_registered = server_id in self._registered_tools_by_server
-        logger.debug(f"Checking if server '{server_id}' is registered: {is_registered}")
-        return is_registered
-
     def unregister_tools_from_server(self, server_id: str) -> bool:
-        """
-        Unregisters all tools associated with a given MCP server ID from the
-        main ToolRegistry and removes the server from internal tracking.
-
-        Args:
-            server_id: The unique ID of the MCP server whose tools should be unregistered.
-
-        Returns:
-            True if the server was found and its tools were processed for unregistration,
-            False if the server was not found in the registrar's state.
-        """
         if not self.is_server_registered(server_id):
             logger.info(f"No tools found for server ID '{server_id}'. Nothing to unregister.")
             return False
-
-        tools_to_unregister = self._registered_tools_by_server.get(server_id, [])
+        tools_to_unregister = self._registered_tools_by_server.pop(server_id, [])
         logger.info(f"Unregistering {len(tools_to_unregister)} tools from server ID: '{server_id}'...")
-
         for tool_def in tools_to_unregister:
             self._tool_registry.unregister_tool(tool_def.name)
-
-        # Remove the server from the tracking dictionary
-        del self._registered_tools_by_server[server_id]
         logger.info(f"Successfully unregistered all tools and removed server '{server_id}' from registrar tracking.")
         return True
+        
+    def is_server_registered(self, server_id: str) -> bool:
+        return server_id in self._registered_tools_by_server

autobyteus/tools/mcp/server/__init__.py

@@ -0,0 +1,16 @@
+# file: autobyteus/autobyteus/tools/mcp/server/__init__.py
+"""
+This package contains the core abstractions for managing connections to remote MCP servers.
+"""
+from .base_managed_mcp_server import BaseManagedMcpServer, ServerState
+from .stdio_managed_mcp_server import StdioManagedMcpServer
+from .http_managed_mcp_server import HttpManagedMcpServer
+from .proxy import McpServerProxy
+
+__all__ = [
+    "BaseManagedMcpServer",
+    "ServerState",
+    "StdioManagedMcpServer",
+    "HttpManagedMcpServer",
+    "McpServerProxy",
+]

autobyteus/tools/mcp/server/base_managed_mcp_server.py

@@ -0,0 +1,139 @@
+# file: autobyteus/autobyteus/tools/mcp/server/base_managed_mcp_server.py
+import logging
+import asyncio
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Dict, Any, List, Optional
+from contextlib import AsyncExitStack
+
+from mcp import ClientSession, types as mcp_types
+
+from ..types import BaseMcpConfig
+
+logger = logging.getLogger(__name__)
+
+class ServerState(str, Enum):
+    """Enumerates the possible connection states of a BaseManagedMcpServer."""
+    DISCONNECTED = "disconnected"
+    CONNECTING = "connecting"
+    CONNECTED = "connected"
+    FAILED = "failed"
+    CLOSED = "closed"
+
+class BaseManagedMcpServer(ABC):
+    """
+    Abstract base class representing a connection to a remote MCP server.
+    It manages the entire lifecycle and state of a single server connection.
+    """
+    
+    # --- Attributes ---
+    _config: BaseMcpConfig
+    _state: ServerState
+    _connection_lock: asyncio.Lock
+    _client_session: Optional[ClientSession]
+    _exit_stack: AsyncExitStack
+
+    # --- Initialization ---
+    def __init__(self, config: BaseMcpConfig):
+        self._config = config
+        self._state = ServerState.DISCONNECTED
+        self._connection_lock = asyncio.Lock()
+        self._client_session = None
+        self._exit_stack = AsyncExitStack()
+
+    # --- Public Properties ---
+    @property
+    def server_id(self) -> str:
+        return self._config.server_id
+
+    @property
+    def config(self) -> BaseMcpConfig:
+        return self._config
+        
+    @property
+    def state(self) -> ServerState:
+        return self._state
+
+    # --- Abstract Methods ---
+    @abstractmethod
+    async def _create_client_session(self) -> ClientSession:
+        """
+        Transport-specific logic to establish a connection and return a ClientSession.
+        This method should leverage self._exit_stack to manage resources.
+        """
+        pass
+
+    # --- Public API Methods ---
+    async def connect(self) -> None:
+        """Public method to establish a connection to the server. Idempotent."""
+        if self._state == ServerState.CONNECTED:
+            return
+
+        async with self._connection_lock:
+            # Re-check state after acquiring lock
+            if self._state == ServerState.CONNECTED:
+                return
+            if self._state == ServerState.CONNECTING:
+                logger.debug(f"Connection already in progress for '{self.server_id}'. Awaiting completion.")
+                # A more advanced implementation might use an asyncio.Event to wait here.
+                # For now, serializing via the lock is sufficient.
+                return
+
+            logger.info(f"Connecting to MCP server '{self.server_id}'...")
+            self._state = ServerState.CONNECTING
+            try:
+                # The exit stack must be fresh for each connection attempt.
+                self._exit_stack = AsyncExitStack()
+                self._client_session = await self._create_client_session()
+                self._state = ServerState.CONNECTED
+                logger.info(f"Successfully connected to MCP server '{self.server_id}'.")
+            except Exception as e:
+                self._state = ServerState.FAILED
+                logger.error(f"Failed to connect to MCP server '{self.server_id}': {e}", exc_info=True)
+                # Clean up any partially established resources on failure
+                if self._exit_stack:
+                    await self._exit_stack.aclose()
+                self._client_session = None
+                raise
+
+    async def close(self) -> None:
+        """Public method to gracefully close the connection to the server."""
+        async with self._connection_lock:
+            if self._state in [ServerState.DISCONNECTED, ServerState.CLOSED]:
+                logger.debug(f"Server '{self.server_id}' is already closed or disconnected. No action taken.")
+                return
+
+            logger.info(f"Closing connection to MCP server '{self.server_id}'...")
+            self._state = ServerState.CLOSED
+            
+            try:
+                if self._exit_stack:
+                    await self._exit_stack.aclose()
+            except Exception as e:
+                logger.error(f"Error during resource cleanup for server '{self.server_id}': {e}", exc_info=True)
+            
+            self._client_session = None
+            logger.info(f"Connection to MCP server '{self.server_id}' closed.")
+
+    async def list_remote_tools(self) -> List[mcp_types.Tool]:
+        """Connects if needed and fetches the list of raw tool objects."""
+        if self._state != ServerState.CONNECTED:
+            await self.connect()
+
+        if not self._client_session:
+            raise RuntimeError(f"Cannot list tools: client session not available for server '{self.server_id}'.")
+
+        logger.debug(f"Listing remote tools on server '{self.server_id}'.")
+        result = await self._client_session.list_tools()
+        return result.tools
+
+    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """Connects if needed and executes a tool call on the remote server."""
+        if self._state != ServerState.CONNECTED:
+            await self.connect()
+
+        if not self._client_session:
+            raise RuntimeError(f"Cannot call tool: client session not available for server '{self.server_id}'.")
+
+        logger.debug(f"Calling remote tool '{tool_name}' on server '{self.server_id}'.")
+        return await self._client_session.call_tool(tool_name, arguments)

autobyteus/tools/mcp/server/http_managed_mcp_server.py

@@ -0,0 +1,29 @@
+# file: autobyteus/autobyteus/tools/mcp/server/http_managed_mcp_server.py
+import logging
+from typing import cast
+
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+
+from .base_managed_mcp_server import BaseManagedMcpServer
+from ..types import StreamableHttpMcpServerConfig
+
+logger = logging.getLogger(__name__)
+
+class HttpManagedMcpServer(BaseManagedMcpServer):
+    """Manages the lifecycle of a streamable_http-based MCP server."""
+
+    def __init__(self, config: StreamableHttpMcpServerConfig):
+        super().__init__(config)
+
+    async def _create_client_session(self) -> ClientSession:
+        """Connects to a remote HTTP endpoint and establishes a client session."""
+        config = cast(StreamableHttpMcpServerConfig, self._config)
+        
+        logger.debug(f"Establishing HTTP connection for server '{self.server_id}' to URL: {config.url}")
+        read_stream, write_stream = await self._exit_stack.enter_async_context(
+            streamablehttp_client(config.url, headers=config.headers)
+        )
+        session = await self._exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
+        logger.debug(f"ClientSession established for HTTP server '{self.server_id}'.")
+        return session

autobyteus/tools/mcp/server/proxy.py

@@ -0,0 +1,36 @@
+# file: autobyteus/autobyteus/tools/mcp/server/proxy.py
+import logging
+from typing import Any, Dict
+
+from ..server_instance_manager import McpServerInstanceManager
+
+logger = logging.getLogger(__name__)
+
+class McpServerProxy:
+    """
+    A proxy object that provides the interface of a BaseManagedMcpServer
+    but delegates the actual work to a real server instance retrieved from the
+    McpServerInstanceManager. This decouples the tool from the manager.
+    """
+    def __init__(self, agent_id: str, server_id: str):
+        if not agent_id or not server_id:
+            raise ValueError("McpServerProxy requires both agent_id and server_id.")
+        self._agent_id = agent_id
+        self._server_id = server_id
+        self._instance_manager = McpServerInstanceManager()
+        logger.debug(f"McpServerProxy created for agent '{agent_id}' and server '{server_id}'.")
+
+    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """
+        Gets the real server instance from the manager and delegates the tool call.
+        """
+        logger.debug(f"Proxy: Getting server instance for agent '{self._agent_id}', server '{self._server_id}'.")
+        # The manager handles the logic of creating or returning a cached instance.
+        real_server_instance = self._instance_manager.get_server_instance(
+            agent_id=self._agent_id,
+            server_id=self._server_id
+        )
+        
+        logger.debug(f"Proxy: Delegating 'call_tool({tool_name})' to real server instance.")
+        # The real instance handles its own connection state.
+        return await real_server_instance.call_tool(tool_name, arguments)