flock-core 0.4.511__py3-none-any.whl → 0.4.513__py3-none-any.whl

flock/core/mcp/mcp_client.py

@@ -4,9 +4,10 @@ import asyncio
 import random
 from abc import ABC, abstractmethod
 from asyncio import Lock
+from collections.abc import Callable
 from contextlib import (
-    AbstractAsyncContextManager,
     AsyncExitStack,
+    asynccontextmanager,
 )
 from datetime import timedelta
 from typing import (
@@ -27,7 +28,7 @@ from mcp import (
     McpError,
     ServerCapabilities,
 )
-from mcp.types import CallToolResult, JSONRPCMessage
+from mcp.types import CallToolResult
 from opentelemetry import trace
 from pydantic import (
     BaseModel,
@@ -54,9 +55,11 @@ from flock.core.mcp.types.types import (
 )
 from flock.core.mcp.util.helpers import cache_key_generator
 
-logger = get_logger("core.mcp.client_base")
+logger = get_logger("mcp.client")
 tracer = trace.get_tracer(__name__)
 
+GetSessionIdCallback = Callable[[], str | None]
+
 
 class FlockMCPClientBase(BaseModel, ABC):
     """Wrapper for mcp ClientSession.
@@ -159,11 +162,9 @@ class FlockMCPClientBase(BaseModel, ABC):
                     max_tries = cfg.connection_config.max_retries or 1
                     base_delay = 0.1
                     span.set_attribute("client.name", client.config.name)
+                    span.set_attribute("max_tries", max_tries)
 
                     for attempt in range(1, max_tries + 2):
-                        span.set_attribute(
-                            "max_tries", max_tries
-                        )  # TODO: shift outside of loop
                         span.set_attribute("base_delay", base_delay)
                         span.set_attribute("attempt", attempt)
                         await client._ensure_connected()
@@ -364,12 +365,7 @@ class FlockMCPClientBase(BaseModel, ABC):
         self,
         params: ServerParameters,
         additional_params: dict[str, Any] | None = None,
-    ) -> AbstractAsyncContextManager[
-        tuple[
-            MemoryObjectReceiveStream[JSONRPCMessage | Exception],
-            MemoryObjectSendStream[JSONRPCMessage],
-        ]
-    ]:
+    ) -> Any:
         """Given your custom ServerParameters, return an async-contextmgr whose __aenter yields (read_stream, write_stream)."""
         ...
 
@@ -390,6 +386,7 @@ class FlockMCPClientBase(BaseModel, ABC):
                 return []
 
             async def _get_tools_internal() -> list[FlockMCPToolBase]:
+                # TODO: Crash
                 response: ListToolsResult = await self.session.list_tools()
                 flock_tools = []
 
@@ -520,13 +517,31 @@ class FlockMCPClientBase(BaseModel, ABC):
             if self.session_stack:
                 # manually __aexit__
                 await self.session_stack.aclose()
-                self.session_stack = None
-                self.client_session = None
+                self.client_session = None  # remove the reference
 
     # --- Private Methods ---
+    @asynccontextmanager
+    async def _safe_transport_ctx(self, cm: Any):
+        """Enter the real transport ctxmg, yield its value, but on __aexit__ always swallow all errors."""
+        val = await cm.__aenter__()
+        try:
+            yield val
+        finally:
+            try:
+                await cm.__aexit__(None, None, None)
+            except Exception as e:
+                logger.debug(
+                    f"Suppressed transport-ctx exit error "
+                    f"for server '{self.config.name}': {e!r}"
+                )
+
     async def _create_session(self) -> None:
-        """Create and hol onto a single ClientSession + ExitStack."""
+        """Create and hold onto a single ClientSession + ExitStack."""
         logger.debug(f"Creating Client Session for server '{self.config.name}'")
+        if self.session_stack:
+            await self.session_stack.aclose()
+        if self.client_session:
+            self.client_session = None
         stack = AsyncExitStack()
         await stack.__aenter__()
 
@@ -536,7 +551,30 @@ class FlockMCPClientBase(BaseModel, ABC):
         transport_ctx = await self.create_transport(
             server_params, self.additional_params
         )
-        read, write = await stack.enter_async_context(transport_ctx)
+        safe_transport = self._safe_transport_ctx(transport_ctx)
+        result = await stack.enter_async_context(safe_transport)
+
+        # support old (read, write) or new (read, write, get_sesssion_id_callback)
+        read: MemoryObjectReceiveStream | None = None
+        write: MemoryObjectSendStream | None = None
+        get_session_id_callback: GetSessionIdCallback | None = None
+        if isinstance(result, tuple) and len(result) == 2:
+            # old type
+            read, write = result
+            get_session_id_callback = None
+        elif isinstance(result, tuple) and len(result) == 3:
+            # new type
+            read, write, get_session_id_callback = result
+        else:
+            raise RuntimeError(
+                f"create_transport returned unexpected tuple of {result}"
+            )
+
+        if read is None or write is None:
+            raise RuntimeError(
+                f"create_transport did not create any read or write streams."
+            )
+
         read_timeout = self.config.connection_config.read_timeout_seconds
 
         if (
@@ -553,6 +591,8 @@ class FlockMCPClientBase(BaseModel, ABC):
             else timedelta(seconds=float(read_timeout))
         )
 
+        # TODO: get_session_id_callback is currently ignored.
+
         session = await stack.enter_async_context(
             ClientSession(
                 read_stream=read,
@@ -603,18 +643,7 @@ class FlockMCPClientBase(BaseModel, ABC):
 
         self.connected_server_capabilities = init
 
-        init_report = f"""
-            Server Init Handshake completed Server '{self.config.name}'
-            Lists the following Capabilities:
-
-            - Protocol Version: {init.protocolVersion}
-            - Instructions: {init.instructions or "No specific Instructions"}
-            - MCP Implementation:
-                - Name: {init.serverInfo.name}
-                - Version: {init.serverInfo.version}
-            - Capabilities:
-                {init.capabilities}
-            """
+        init_report = f"Server: '{self.config.name}': Protocol-Version: {init.protocolVersion}, Instructions: {init.instructions or 'No specific instructions'}, MCP_Implementation: Name: {init.serverInfo.name}, Version: {init.serverInfo.version}, Capabilities: {init.capabilities}"
 
         logger.debug(init_report)
 

flock/core/mcp/mcp_client_manager.py

@@ -19,7 +19,7 @@ from flock.core.mcp.mcp_client import (
 )
 from flock.core.mcp.mcp_config import FlockMCPConfigurationBase
 
-logger = get_logger("core.mcp.connection_manager_base")
+logger = get_logger("mcp.client_manager")
 tracer = trace.get_tracer(__name__)
 
 TClient = TypeVar("TClient", bound="FlockMCPClientBase")

flock/core/mcp/mcp_config.py

@@ -1,7 +1,9 @@
 """Base Config for MCP Clients."""
 
-from typing import Literal, TypeVar
+import importlib
+from typing import Any, Literal, TypeVar
 
+import httpx
 from pydantic import BaseModel, ConfigDict, Field, create_model
 
 from flock.core.mcp.types.types import (
@@ -11,6 +13,15 @@ from flock.core.mcp.types.types import (
     FlockSamplingMCPCallback,
     MCPRoot,
     ServerParameters,
+    SseServerParameters,
+    StdioServerParameters,
+    StreamableHttpServerParameters,
+    WebsocketServerParameters,
+)
+from flock.core.serialization.serializable import Serializable
+from flock.core.serialization.serialization_utils import (
+    deserialize_item,
+    serialize_item,
 )
 
 LoggingLevel = Literal[
@@ -32,7 +43,7 @@ D = TypeVar("D", bound="FlockMCPCachingConfigurationBase")
 E = TypeVar("E", bound="FlockMCPFeatureConfigurationBase")
 
 
-class FlockMCPCachingConfigurationBase(BaseModel):
+class FlockMCPCachingConfigurationBase(BaseModel, Serializable):
     """Configuration for Caching in Clients."""
 
     tool_cache_max_size: float = Field(
@@ -79,6 +90,18 @@ class FlockMCPCachingConfigurationBase(BaseModel):
         extra="allow",
     )
 
+    def to_dict(self, path_type: str = "relative"):
+        """Serialize the config object."""
+        return self.model_dump(
+            exclude_none=True,
+            mode="json",
+        )
+
+    @classmethod
+    def from_dict(cls: type[D], data: dict[str, Any]) -> D:
+        """Deserialize from a dict."""
+        return cls(**{k: v for k, v in data.items()})
+
     @classmethod
     def with_fields(cls: type[D], **field_definitions) -> type[D]:
         """Create a new config class with additional fields."""
@@ -87,7 +110,7 @@ class FlockMCPCachingConfigurationBase(BaseModel):
         )
 
 
-class FlockMCPCallbackConfigurationBase(BaseModel):
+class FlockMCPCallbackConfigurationBase(BaseModel, Serializable):
     """Base Configuration Class for Callbacks for Clients."""
 
     sampling_callback: FlockSamplingMCPCallback | None = Field(
@@ -111,6 +134,52 @@ class FlockMCPCallbackConfigurationBase(BaseModel):
 
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
+    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
+
+    @classmethod
+    def from_dict(cls: type[A], data: dict[str, Any]) -> A:
+        """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
+
     @classmethod
     def with_fields(cls: type[A], **field_definitions) -> type[A]:
         """Create a new config class with additional fields."""
@@ -119,7 +188,7 @@ class FlockMCPCallbackConfigurationBase(BaseModel):
         )
 
 
-class FlockMCPConnectionConfigurationBase(BaseModel):
+class FlockMCPConnectionConfigurationBase(BaseModel, Serializable):
     """Base Configuration Class for Connection Parameters for a client."""
 
     max_retries: int = Field(
@@ -131,9 +200,9 @@ class FlockMCPConnectionConfigurationBase(BaseModel):
         ..., description="Connection parameters for the server."
     )
 
-    transport_type: Literal["stdio", "websockets", "sse", "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."
@@ -150,6 +219,77 @@ class FlockMCPConnectionConfigurationBase(BaseModel):
 
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
+    def to_dict(self, path_type: str = "relative") -> dict[str, Any]:
+        """Serialize object to a dict."""
+        exclude = ["connection_parameters"]
+
+        data = self.model_dump(
+            exclude=exclude,
+            exclude_defaults=False,
+            exclude_none=True,
+            mode="json",
+        )
+
+        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:
+        """Deserialize from dict."""
+        connection_params = data.get("connection_parameters")
+        connection_params_obj = None
+        auth_obj: httpx.Auth | None = None
+        if connection_params:
+            kind = connection_params.get("transport_type", None)
+            auth_spec = connection_params.get("auth", None)
+            if auth_spec:
+                # find the concrete implementation and
+                # instantiate 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()})
+
+                if auth_obj:
+                    connection_params["auth"] = auth_obj
+                else:
+                    # just to be sure
+                    connection_params.pop("auth", None)
+                match kind:
+                    case "stdio":
+                        connection_params_obj = StdioServerParameters(
+                            **{k: v for k, v in connection_params.items()}
+                        )
+                    case "websockets":
+                        connection_params_obj = WebsocketServerParameters(
+                            **{k: v for k, v in connection_params.items()}
+                        )
+                    case "streamable_http":
+                        connection_params_obj = StreamableHttpServerParameters(
+                            **{k: v for k, v in connection_params.items()}
+                        )
+                    case "sse":
+                        connection_params_obj = SseServerParameters(
+                            **{k: v for k, v in connection_params.items()}
+                        )
+                    case _:
+                        # handle custom server params
+                        connection_params_obj = ServerParameters(
+                            **{k: v for k, v in 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.")
+
     @classmethod
     def with_fields(cls: type[B], **field_definitions) -> type[B]:
         """Create a new config class with additional fields."""
@@ -158,7 +298,7 @@ class FlockMCPConnectionConfigurationBase(BaseModel):
         )
 
 
-class FlockMCPFeatureConfigurationBase(BaseModel):
+class FlockMCPFeatureConfigurationBase(BaseModel, Serializable):
     """Base Configuration Class for switching MCP Features on and off."""
 
     roots_enabled: bool = Field(
@@ -186,6 +326,18 @@ class FlockMCPFeatureConfigurationBase(BaseModel):
         extra="allow",
     )
 
+    def to_dict(self, path_type: str = "relative"):
+        """Serialize the object."""
+        return self.model_dump(
+            mode="json",
+            exclude_none=True,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]):
+        """Deserialize from a dict."""
+        return cls(**{k: v for k, v in data.items()})
+
     @classmethod
     def with_fields(cls: type[E], **field_definitions) -> type[E]:
         """Create a new config class with additional fields."""
@@ -194,7 +346,7 @@ class FlockMCPFeatureConfigurationBase(BaseModel):
         )
 
 
-class FlockMCPConfigurationBase(BaseModel):
+class FlockMCPConfigurationBase(BaseModel, Serializable):
     """Base Configuration Class for MCP Clients.
 
     Each Client should implement their own config
@@ -229,6 +381,90 @@ class FlockMCPConfigurationBase(BaseModel):
         extra="allow",
     )
 
+    def to_dict(self, path_type: str = "relative") -> dict[str, Any]:
+        """Serialize the object to a dict."""
+        # each built-in type should serialize, deserialize it self.
+        exclude = [
+            "connection_config",
+            "caching_config",
+            "callback_config",
+            "feature_config",
+        ]
+
+        data = self.model_dump(
+            exclude=exclude,
+            exclude_defaults=False,
+            exclude_none=True,
+            mode="json",
+        )
+
+        # add the core properties
+        data["connection_config"] = self.connection_config.to_dict(path_type)
+        data["caching_config"] = self.caching_config.to_dict(path_type)
+        data["callback_config"] = self.callback_config.to_dict(path_type)
+        data["feature_config"] = self.feature_config.to_dict(path_type)
+
+        return data
+
+    @classmethod
+    def from_dict(cls: type[C], data: dict[str, Any]) -> C:
+        """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"]
+        }
+
+        if connection_config:
+            # Forcing a square into a round hole
+            try:
+                config_field = cls.model_fields["connection_config"]
+                config_cls = config_field.annotation
+            except (AttributeError, KeyError):
+                # fallback
+                config_cls = FlockMCPConnectionConfigurationBase
+            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')}")
+
+        if caching_config:
+            try:
+                config_field = cls.model_fields["caching_config"]
+                config_cls = config_field.annotation
+            except (AttributeError, KeyError):
+                # fallback
+                config_cls = FlockMCPCachingConfigurationBase
+            instance_data["caching_config"] = config_cls.from_dict(caching_config)
+        else:
+            instance_data["caching_config"] = FlockMCPCachingConfigurationBase()
+
+        if feature_config:
+            try:
+                config_field = cls.model_fields["feature_config"]
+                config_cls = config_field.annotation
+            except (AttributeError, KeyError):
+                # fallback
+                config_cls = FlockMCPFeatureConfigurationBase
+            instance_data["feature_config"] = config_cls.from_dict(feature_config)
+        else:
+            instance_data["feature_config"] = FlockMCPFeatureConfigurationBase()
+
+        if callback_config:
+            try:
+                config_field = cls.model_fields["callback_config"]
+                config_cls = config_field.annotation
+            except (AttributeError, KeyError):
+                # fallback
+                config_cls = FlockMCPCallbackConfigurationBase
+            instance_data["callback_config"] = config_cls.from_dict(callback_config)
+        else:
+            instance_data["callback_config"] = FlockMCPCallbackConfigurationBase()
+
+        return cls(**{k: v for k, v in instance_data.items()})
+
     @classmethod
     def with_fields(cls: type[C], **field_definitions) -> type[C]:
         """Create a new config class with additional fields."""

flock/core/mcp/types/callbacks.py

@@ -8,6 +8,8 @@ from mcp.types import (
     CreateMessageRequestParams,
     ErrorData,
     ListRootsResult,
+    LoggingMessageNotificationParams,
+    ServerNotification,
     ServerRequest,
 )
 
@@ -19,10 +21,6 @@ from flock.core.mcp.types.handlers import (
     handle_incoming_server_notification,
     handle_logging_message,
 )
-from flock.core.mcp.types.types import (
-    FlockLoggingMessageNotificationParams,
-    ServerNotification,
-)
 
 
 async def default_sampling_callback(
@@ -76,7 +74,7 @@ async def default_list_roots_callback(
 
 
 async def default_logging_callback(
-    params: FlockLoggingMessageNotificationParams,
+    params: LoggingMessageNotificationParams,
     logger: FlockLogger,
     server_name: str,
 ) -> None:

flock/core/mcp/types/factories.py

@@ -1,6 +1,6 @@
 """Factories for default MCP Callbacks."""
 
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from mcp.shared.context import RequestContext
 from mcp.types import (
@@ -8,6 +8,12 @@ from mcp.types import (
 )
 
 from flock.core.logging.logging import FlockLogger, get_logger
+from flock.core.mcp.types.callbacks import (
+    default_list_roots_callback,
+    default_logging_callback,
+    default_message_handler,
+    default_sampling_callback,
+)
 from flock.core.mcp.types.types import (
     FlockListRootsMCPCallback,
     FlockLoggingMCPCallback,
@@ -17,18 +23,10 @@ from flock.core.mcp.types.types import (
     ServerNotification,
 )
 
-if TYPE_CHECKING:
-    from flock.core.mcp.types.callbacks import (
-        default_list_roots_callback,
-        default_logging_callback,
-        default_message_handler,
-        default_sampling_callback,
-    )
-
-default_logging_callback_logger = get_logger("core.mcp.callback.logging")
-default_sampling_callback_logger = get_logger("core.mcp.callback.sampling")
-default_list_roots_callback_logger = get_logger("core.mcp.callback.sampling")
-default_message_handler_logger = get_logger("core.mcp.callback.message")
+default_logging_callback_logger = get_logger("mcp.callback.logging")
+default_sampling_callback_logger = get_logger("mcp.callback.sampling")
+default_list_roots_callback_logger = get_logger("mcp.callback.roots")
+default_message_handler_logger = get_logger("mcp.callback.message")
 
 
 def default_flock_mcp_logging_callback_factory(
@@ -88,7 +86,7 @@ def default_flock_mcp_message_handler_callback_factory(
     ) -> None:
         await default_message_handler(
             req=n,
-            logger_to_use=logger_to_use,
+            logger=logger_to_use,
             associated_client=associated_client,
         )
 

flock/core/mcp/types/handlers.py

@@ -8,26 +8,23 @@ from mcp.shared.context import RequestContext
 from mcp.shared.session import RequestResponder
 from mcp.types import (
     INTERNAL_ERROR,
+    CancelledNotification,
     ClientResult,
     ErrorData,
     ListRootsRequest,
-    ServerNotification as _MCPServerNotification,
-    ServerRequest,
-)
-
-from flock.core.logging.logging import FlockLogger
-from flock.core.mcp.mcp_client import Any
-from flock.core.mcp.types.types import (
-    CancelledNotification,
-    FlockLoggingMessageNotificationParams,
     LoggingMessageNotification,
+    LoggingMessageNotificationParams,
     ProgressNotification,
     ResourceListChangedNotification,
     ResourceUpdatedNotification,
     ServerNotification,
+    ServerRequest,
     ToolListChangedNotification,
 )
 
+from flock.core.logging.logging import FlockLogger
+from flock.core.mcp.mcp_client import Any
+
 
 async def handle_incoming_exception(
     e: Exception,
@@ -129,11 +126,11 @@ async def handle_tool_list_changed_notification(
     await associated_client.invalidate_tool_cache()
 
 
-_SERVER_NOTIFICATION_MAP: dict[type[_MCPServerNotification], Callable] = {
+_SERVER_NOTIFICATION_MAP: dict[type[ServerNotification], Callable] = {
     ResourceListChangedNotification: handle_resource_list_changed_notification,
     ResourceUpdatedNotification: handle_resource_update_notification,
     LoggingMessageNotification: lambda n, log, client: handle_logging_message(
-        params=n,
+        params=n.params,
         logger=log,
         server_name=client.config.name,
     ),
@@ -154,7 +151,7 @@ async def handle_incoming_server_notification(
 
 
 async def handle_logging_message(
-    params: FlockLoggingMessageNotificationParams,
+    params: LoggingMessageNotificationParams,
     logger: FlockLogger,
     server_name: str,
 ) -> None: