flock-core 0.5.0b28__py3-none-any.whl → 0.5.0b51__py3-none-any.whl

flock/{core/mcp/mcp_config.py → mcp/config.py}

@@ -1,12 +1,12 @@
 """Base Config for MCP Clients."""
 
 import importlib
-from typing import Any, Literal, TypeVar
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
-import httpx
 from pydantic import BaseModel, ConfigDict, Field, create_model
+from typing_extensions import Self
 
-from flock.core.mcp.types.types import (
+from flock.mcp.types import (
     FlockListRootsMCPCallback,
     FlockLoggingMCPCallback,
     FlockMessageHandlerMCPCallback,
@@ -18,11 +18,11 @@ from flock.core.mcp.types.types import (
     StreamableHttpServerParameters,
     WebsocketServerParameters,
 )
-from flock.core.serialization.serializable import Serializable
-from flock.core.serialization.serialization_utils import (
-    deserialize_item,
-    serialize_item,
-)
+
+
+if TYPE_CHECKING:
+    import httpx
+
 
 LoggingLevel = Literal[
     "debug",
@@ -43,7 +43,7 @@ D = TypeVar("D", bound="FlockMCPCachingConfiguration")
 E = TypeVar("E", bound="FlockMCPFeatureConfiguration")
 
 
-class FlockMCPCachingConfiguration(BaseModel, Serializable):
+class FlockMCPCachingConfiguration(BaseModel):
     """Configuration for Caching in Clients."""
 
     tool_cache_max_size: float = Field(
@@ -98,19 +98,17 @@ class FlockMCPCachingConfiguration(BaseModel, Serializable):
         )
 
     @classmethod
-    def from_dict(cls: type[D], data: dict[str, Any]) -> D:
+    def from_dict(cls, data: dict[str, Any]) -> Self:
         """Deserialize from a dict."""
-        return cls(**{k: v for k, v in data.items()})
+        return cls(**dict(data.items()))
 
     @classmethod
-    def with_fields(cls: type[D], **field_definitions) -> type[D]:
+    def with_fields(cls, **field_definitions) -> type[Self]:
         """Create a new config class with additional fields."""
-        return create_model(
-            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
-        )
+        return create_model(f"Dynamic{cls.__name__}", __base__=cls, **field_definitions)
 
 
-class FlockMCPCallbackConfiguration(BaseModel, Serializable):
+class FlockMCPCallbackConfiguration(BaseModel):
     """Base Configuration Class for Callbacks for Clients."""
 
     sampling_callback: FlockSamplingMCPCallback | None = Field(
@@ -136,59 +134,21 @@ class FlockMCPCallbackConfiguration(BaseModel, Serializable):
 
     def to_dict(self, path_type: str = "relative"):
         """Serialize the object."""
-        # we need to register callables.
-        data: dict[str, Any] = {}
-        if self.sampling_callback:
-            sampling_callback_data = serialize_item(self.sampling_callback)
-            data["sampling_callback"] = sampling_callback_data
-
-        if self.list_roots_callback:
-            list_roots_callback_data = serialize_item(self.list_roots_callback)
-            data["list_roots_callback"] = list_roots_callback_data
-
-        if self.logging_callback:
-            logging_callback_data = serialize_item(self.logging_callback)
-            data["logging_callback"] = logging_callback_data
-
-        if self.message_handler:
-            message_handler_data = serialize_item(self.message_handler)
-            data["message_handler"] = message_handler_data
-
-        return data
+        # Callbacks are set programmatically and cannot be serialized
+        return {}
 
     @classmethod
-    def from_dict(cls: type[A], data: dict[str, Any]) -> A:
+    def from_dict(cls, data: dict[str, Any]) -> Self:
         """Deserialize from a dict."""
-        instance = cls()
-        if data:
-            if "sampling_callback" in data:
-                instance.sampling_callback = deserialize_item(
-                    data["sampling_callback"]
-                )
-            if "list_roots_callback" in data:
-                instance.list_roots_callback = deserialize_item(
-                    data["list_roots_callback"]
-                )
-            if "logging_callback" in data:
-                instance.logging_callback = deserialize_item(
-                    data["logging_callback"]
-                )
-            if "message_handler" in data:
-                instance.message_handler = deserialize_item(
-                    data["message_handler"]
-                )
-
-        return instance
+        return cls()
 
     @classmethod
-    def with_fields(cls: type[A], **field_definitions) -> type[A]:
+    def with_fields(cls, **field_definitions) -> type[Self]:
         """Create a new config class with additional fields."""
-        return create_model(
-            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
-        )
+        return create_model(f"Dynamic{cls.__name__}", __base__=cls, **field_definitions)
 
 
-class FlockMCPConnectionConfiguration(BaseModel, Serializable):
+class FlockMCPConnectionConfiguration(BaseModel):
     """Base Configuration Class for Connection Parameters for a client."""
 
     max_retries: int = Field(
@@ -200,17 +160,15 @@ class FlockMCPConnectionConfiguration(BaseModel, Serializable):
         ..., description="Connection parameters for the server."
     )
 
-    transport_type: Literal[
-        "stdio", "websockets", "sse", "streamable_http", "custom"
-    ] = Field(..., description="Type of transport to use.")
+    transport_type: Literal["stdio", "websockets", "sse", "streamable_http", "custom"] = Field(
+        ..., description="Type of transport to use."
+    )
 
     mount_points: list[MCPRoot] | None = Field(
         default=None, description="Initial Mountpoints to operate under."
     )
 
-    read_timeout_seconds: float | int = Field(
-        default=60 * 5, description="Read Timeout."
-    )
+    read_timeout_seconds: float | int = Field(default=60 * 5, description="Read Timeout.")
 
     server_logging_level: LoggingLevel = Field(
         default="error",
@@ -230,14 +188,12 @@ class FlockMCPConnectionConfiguration(BaseModel, Serializable):
             mode="json",
         )
 
-        data["connection_parameters"] = self.connection_parameters.to_dict(
-            path_type=path_type
-        )
+        data["connection_parameters"] = self.connection_parameters.to_dict(path_type=path_type)
 
         return data
 
     @classmethod
-    def from_dict(cls: type[B], data: dict[str, Any]) -> B:
+    def from_dict(cls, data: dict[str, Any]) -> Self:
         """Deserialize from dict."""
         connection_params = data.get("connection_parameters")
         connection_params_obj = None
@@ -248,13 +204,13 @@ class FlockMCPConnectionConfiguration(BaseModel, Serializable):
             if auth_spec:
                 # find the concrete implementation and
                 # instantiate it.
-                    # find the concrete implementation for auth and instatiate it.
+                # find the concrete implementation for auth and instatiate it.
                 impl = auth_spec.get("implementation", None)
                 params = auth_spec.get("params", None)
                 if impl and params:
-                        mod = importlib.import_module(impl["module_path"])
-                        real_cls = getattr(mod, impl["class_name"])
-                        auth_obj = real_cls(**{k: v for k, v in params.items()})
+                    mod = importlib.import_module(impl["module_path"])
+                    real_cls = getattr(mod, impl["class_name"])
+                    auth_obj = real_cls(**dict(params.items()))
 
                 if auth_obj:
                     connection_params["auth"] = auth_obj
@@ -264,41 +220,36 @@ class FlockMCPConnectionConfiguration(BaseModel, Serializable):
                 match kind:
                     case "stdio":
                         connection_params_obj = StdioServerParameters(
-                            **{k: v for k, v in connection_params.items()}
+                            **dict(connection_params.items())
                         )
                     case "websockets":
                         connection_params_obj = WebsocketServerParameters(
-                            **{k: v for k, v in connection_params.items()}
+                            **dict(connection_params.items())
                         )
                     case "streamable_http":
                         connection_params_obj = StreamableHttpServerParameters(
-                            **{k: v for k, v in connection_params.items()}
+                            **dict(connection_params.items())
                         )
                     case "sse":
                         connection_params_obj = SseServerParameters(
-                            **{k: v for k, v in connection_params.items()}
+                            **dict(connection_params.items())
                         )
                     case _:
                         # handle custom server params
-                        connection_params_obj = ServerParameters(
-                            **{k: v for k, v in connection_params.items()}
-                        )
+                        connection_params_obj = ServerParameters(**dict(connection_params.items()))
 
         if connection_params_obj:
             data["connection_parameters"] = connection_params_obj
-            return cls(**{k: v for k, v in data.items()})
-        else:
-            raise ValueError("No connection parameters provided.")
+            return cls(**dict(data.items()))
+        raise ValueError("No connection parameters provided.")
 
     @classmethod
-    def with_fields(cls: type[B], **field_definitions) -> type[B]:
+    def with_fields(cls, **field_definitions) -> type[Self]:
         """Create a new config class with additional fields."""
-        return create_model(
-            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
-        )
+        return create_model(f"Dynamic{cls.__name__}", __base__=cls, **field_definitions)
 
 
-class FlockMCPFeatureConfiguration(BaseModel, Serializable):
+class FlockMCPFeatureConfiguration(BaseModel):
     """Base Configuration Class for switching MCP Features on and off."""
 
     roots_enabled: bool = Field(
@@ -319,10 +270,10 @@ class FlockMCPFeatureConfiguration(BaseModel, Serializable):
     tool_whitelist: list[str] | None = Field(
         default=None,
         description="Whitelist of tool names that are enabled for this MCP server. "
-                   "If provided, only tools with names in this list will be available "
-                   "from this server. Used in conjunction with allow_all_tools setting. "
-                   "Note: Agent-level tool filtering is generally preferred over "
-                   "server-level filtering for better granular control."
+        "If provided, only tools with names in this list will be available "
+        "from this server. Used in conjunction with allow_all_tools setting. "
+        "Note: Agent-level tool filtering is generally preferred over "
+        "server-level filtering for better granular control.",
     )
 
     prompts_enabled: bool = Field(
@@ -345,34 +296,30 @@ class FlockMCPFeatureConfiguration(BaseModel, Serializable):
     @classmethod
     def from_dict(cls, data: dict[str, Any]):
         """Deserialize from a dict."""
-        return cls(**{k: v for k, v in data.items()})
+        return cls(**dict(data.items()))
 
     @classmethod
-    def with_fields(cls: type[E], **field_definitions) -> type[E]:
+    def with_fields(cls, **field_definitions) -> type[Self]:
         """Create a new config class with additional fields."""
-        return create_model(
-            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
-        )
+        return create_model(f"Dynamic{cls.__name__}", __base__=cls, **field_definitions)
 
 
-class FlockMCPConfiguration(BaseModel, Serializable):
+class FlockMCPConfiguration(BaseModel):
     """Base Configuration Class for MCP Clients.
 
     Each Client should implement their own config
     model by inheriting from this class.
     """
 
-    name: str = Field(
-        ..., description="Name of the server the client connects to."
-    )
+    name: str = Field(..., description="Name of the server the client connects to.")
 
     allow_all_tools: bool = Field(
         default=True,
         description="Whether to allow usage of all tools from this MCP server. "
-                   "When True (default), all tools are available unless restricted "
-                   "by tool_whitelist. When False, tool access is controlled entirely "
-                   "by tool_whitelist (if provided). Setting to False with no whitelist "
-                   "will block all tools from this server."
+        "When True (default), all tools are available unless restricted "
+        "by tool_whitelist. When False, tool access is controlled entirely "
+        "by tool_whitelist (if provided). Setting to False with no whitelist "
+        "will block all tools from this server.",
     )
 
     connection_config: FlockMCPConnectionConfiguration = Field(
@@ -425,16 +372,14 @@ class FlockMCPConfiguration(BaseModel, Serializable):
         return data
 
     @classmethod
-    def from_dict(cls: type[C], data: dict[str, Any]) -> C:
+    def from_dict(cls, data: dict[str, Any]) -> Self:
         """Deserialize the class."""
         connection_config = data.pop("connection_config", None)
         caching_config = data.pop("caching_config", None)
         feature_config = data.pop("feature_config", None)
         callback_config = data.pop("callback_config", None)
 
-        instance_data: dict[str, Any] = {
-            "name": data["name"]
-        }
+        instance_data: dict[str, Any] = {"name": data["name"]}
 
         if connection_config:
             # Forcing a square into a round hole
@@ -446,7 +391,9 @@ class FlockMCPConfiguration(BaseModel, Serializable):
                 config_cls = FlockMCPConnectionConfiguration
             instance_data["connection_config"] = config_cls.from_dict(connection_config)
         else:
-            raise ValueError(f"connection_config MUST be specified for '{data.get('name', 'unknown_server')}")
+            raise ValueError(
+                f"connection_config MUST be specified for '{data.get('name', 'unknown_server')}"
+            )
 
         if caching_config:
             try:
@@ -481,11 +428,9 @@ class FlockMCPConfiguration(BaseModel, Serializable):
         else:
             instance_data["callback_config"] = FlockMCPCallbackConfiguration()
 
-        return cls(**{k: v for k, v in instance_data.items()})
+        return cls(**dict(instance_data.items()))
 
     @classmethod
-    def with_fields(cls: type[C], **field_definitions) -> type[C]:
+    def with_fields(cls, **field_definitions) -> type[Self]:
         """Create a new config class with additional fields."""
-        return create_model(
-            f"Dynamic{cls.__name__}", __base__=cls, **field_definitions
-        )
+        return create_model(f"Dynamic{cls.__name__}", __base__=cls, **field_definitions)

flock/mcp/manager.py

@@ -0,0 +1,255 @@
+"""MCP Client Manager for Connection Pooling and Lifecycle Management.
+
+This module provides the FlockMCPClientManager class which manages MCP client
+connections with per-(agent_id, run_id) isolation and lazy connection establishment.
+
+Architecture Decisions:
+- AD004: Per-(agent_id, run_id) Connection Isolation
+  Each unique (agent_id, run_id) pair gets its own set of isolated MCP connections
+- AD005: Lazy Connection Establishment
+  Connections are only established when first requested
+- AD007: Graceful Degradation on MCP Failures
+  Failures to connect to individual servers don't prevent agent execution
+"""
+
+import asyncio
+from typing import Any
+
+from flock.logging.logging import get_logger
+from flock.mcp.client import FlockMCPClient
+from flock.mcp.config import FlockMCPConfiguration
+
+
+logger = get_logger(__name__)
+
+
+class FlockMCPClientManager:
+    """Manages MCP client connections with per-(agent_id, run_id) isolation.
+
+    This manager implements lazy connection pooling where clients are only
+    initialized when first requested. Each unique (agent_id, run_id) pair
+    gets its own set of isolated MCP connections.
+
+    Architecture Decision: AD004 - Per-(agent_id, run_id) Connection Isolation
+
+    Example:
+        ```python
+        # Initialize manager with configurations
+        configs = {
+            "filesystem": FlockMCPConfiguration(
+                name="filesystem",
+                connection_config=connection_config,
+                feature_config=feature_config
+            )
+        }
+        manager = FlockMCPClientManager(configs)
+
+        # Get client for specific agent and run
+        client = await manager.get_client("filesystem", "agent_1", "run_123")
+
+        # Get all tools for agent
+        tools = await manager.get_tools_for_agent(
+            "agent_1", "run_123", {"filesystem"}
+        )
+
+        # Cleanup after run completes
+        await manager.cleanup_run("agent_1", "run_123")
+        ```
+    """
+
+    def __init__(self, configs: dict[str, FlockMCPConfiguration]):
+        """Initialize the manager with MCP server configurations.
+
+        Args:
+            configs: Dictionary mapping server names to their configurations
+        """
+        self._configs = configs
+        # Pool structure: (agent_id, run_id) → server_name → FlockMCPClient
+        self._pool: dict[tuple[str, str], dict[str, FlockMCPClient]] = {}
+        self._lock = asyncio.Lock()
+
+    async def get_client(self, server_name: str, agent_id: str, run_id: str) -> FlockMCPClient:
+        """Get or create an MCP client for the given context.
+
+        Architecture Decision: AD005 - Lazy Connection Establishment
+        Connections are only established when first requested.
+
+        Args:
+            server_name: Name of the MCP server
+            agent_id: Agent requesting the client
+            run_id: Current run identifier
+
+        Returns:
+            FlockMCPClient instance ready for use
+
+        Raises:
+            ValueError: If server_name not registered
+        """
+        if server_name not in self._configs:
+            raise ValueError(
+                f"MCP server '{server_name}' not registered. "
+                f"Available servers: {list(self._configs.keys())}"
+            )
+
+        key = (agent_id, run_id)
+
+        async with self._lock:
+            # Check if we have a client pool for this (agent, run)
+            if key not in self._pool:
+                self._pool[key] = {}
+
+            # Check if we have a client for this server
+            if server_name not in self._pool[key]:
+                logger.info(
+                    f"Creating new MCP client for server '{server_name}' "
+                    f"(agent={agent_id}, run={run_id})"
+                )
+                config = self._configs[server_name]
+
+                # Instantiate the correct concrete client class based on transport type
+                # Lazy import to avoid requiring all dependencies
+                transport_type = config.connection_config.transport_type
+                if transport_type == "stdio":
+                    from flock.mcp.servers.stdio.flock_stdio_server import (
+                        FlockStdioClient,
+                    )
+
+                    client = FlockStdioClient(config=config)
+                elif transport_type == "sse":
+                    from flock.mcp.servers.sse.flock_sse_server import (
+                        FlockSSEClient,
+                    )
+
+                    client = FlockSSEClient(config=config)
+                elif transport_type == "websocket":
+                    from flock.mcp.servers.websockets.flock_websocket_server import (
+                        FlockWSClient,
+                    )
+
+                    client = FlockWSClient(config=config)
+                elif transport_type == "streamable_http":
+                    from flock.mcp.servers.streamable_http.flock_streamable_http_server import (
+                        FlockStreamableHttpClient,
+                    )
+
+                    client = FlockStreamableHttpClient(config=config)
+                else:
+                    raise ValueError(
+                        f"Unsupported transport type: {transport_type}. "
+                        f"Supported types: stdio, sse, websocket, streamable_http"
+                    )
+
+                await client._connect()
+                self._pool[key][server_name] = client
+
+            return self._pool[key][server_name]
+
+    async def get_tools_for_agent(
+        self, agent_id: str, run_id: str, server_names: set[str]
+    ) -> dict[str, Any]:
+        """Get all tools from specified servers for an agent.
+
+        Architecture Decision: AD003 - Tool Namespacing
+        All tools are returned with format: {server}__{tool}
+
+        Args:
+            agent_id: Agent requesting tools
+            run_id: Current run identifier
+            server_names: Set of MCP server names to fetch tools from
+
+        Returns:
+            Dictionary mapping namespaced tool names to tool definitions
+
+        Note:
+            Architecture Decision: AD007 - Graceful Degradation
+            If a server fails to load, we log the error and continue with
+            other servers rather than failing the entire operation.
+        """
+        tools = {}
+
+        for server_name in server_names:
+            try:
+                client = await self.get_client(server_name, agent_id, run_id)
+                server_tools = await client.get_tools(agent_id, run_id)
+
+                # Apply namespacing: AD003
+                for tool in server_tools:
+                    namespaced_name = f"{server_name}__{tool.name}"
+                    tools[namespaced_name] = {
+                        "server_name": server_name,
+                        "original_name": tool.name,
+                        "tool": tool,
+                        "client": client,
+                    }
+
+                logger.debug(
+                    f"Loaded {len(server_tools)} tools from server '{server_name}' "
+                    f"for agent {agent_id}"
+                )
+
+            except Exception as e:
+                # Architecture Decision: AD007 - Graceful Degradation
+                logger.exception(
+                    f"Failed to load tools from MCP server '{server_name}': {e}. "
+                    f"Agent {agent_id} will continue without these tools."
+                )
+                # Continue loading other servers
+                continue
+
+        return tools
+
+    async def cleanup_run(self, agent_id: str, run_id: str) -> None:
+        """Clean up all MCP connections for a completed run.
+
+        Args:
+            agent_id: Agent identifier
+            run_id: Run identifier to clean up
+        """
+        key = (agent_id, run_id)
+
+        async with self._lock:
+            if key in self._pool:
+                logger.info(f"Cleaning up MCP connections for run {run_id}")
+                clients = self._pool[key]
+
+                # Disconnect all clients for this run
+                for server_name, client in clients.items():
+                    try:
+                        await client.disconnect()
+                        logger.debug(f"Disconnected from MCP server '{server_name}'")
+                    except Exception as e:
+                        logger.warning(f"Error disconnecting from '{server_name}': {e}")
+
+                # Remove from pool
+                del self._pool[key]
+
+    async def cleanup_all(self) -> None:
+        """Clean up all MCP connections (orchestrator shutdown)."""
+        async with self._lock:
+            logger.info("Shutting down all MCP connections")
+
+            for _key, clients in self._pool.items():
+                for server_name, client in clients.items():
+                    try:
+                        await client.disconnect()
+                    except Exception as e:
+                        logger.warning(f"Error disconnecting from '{server_name}': {e}")
+
+            self._pool.clear()
+            logger.info("All MCP connections closed")
+
+    def list_servers(self) -> dict[str, dict[str, Any]]:
+        """List all registered MCP servers with their configurations.
+
+        Returns:
+            Dictionary mapping server names to configuration metadata
+        """
+        return {
+            name: {
+                "transport_type": config.connection_config.transport_type,
+                "tools_enabled": config.feature_config.tools_enabled,
+                "prompts_enabled": config.feature_config.prompts_enabled,
+                "sampling_enabled": config.feature_config.sampling_enabled,
+            }
+            for name, config in self._configs.items()
+        }

flock/mcp/servers/sse/__init__.py

@@ -1 +1 @@
-"""Default SSE Server Implementation for Flock."""
+"""Default SSE Server Implementation for flock."""