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

flock/core/flock_agent.py

@@ -384,7 +384,7 @@ class FlockAgent(BaseModel, Serializable, DSPyIntegrationMixin, ABC):
                             )
                         else:
                             logger.warning(
-                                f"No Server with name '{server}' registered! Skipping."
+                                f"No Server with name '{server.config.name}' registered! Skipping."
                             )
                         mcp_tools = mcp_tools + server_tools
 

flock/core/flock_factory.py

@@ -5,7 +5,8 @@ from collections.abc import Callable
 from pathlib import Path
 from typing import Any, Literal
 
-from pydantic import AnyUrl, BaseModel, Field, FileUrl
+import httpx
+from pydantic import AnyUrl, BaseModel, ConfigDict, Field, FileUrl
 
 from flock.core.config.scheduled_agent_config import ScheduledAgentConfig
 from flock.core.flock_agent import FlockAgent, SignatureType
@@ -24,6 +25,7 @@ from flock.core.mcp.types.types import (
     MCPRoot,
     SseServerParameters,
     StdioServerParameters,
+    StreamableHttpServerParameters,
     WebsocketServerParameters,
 )
 from flock.evaluators.declarative.declarative_evaluator import (
@@ -40,6 +42,11 @@ from flock.mcp.servers.stdio.flock_stdio_server import (
     FlockStdioConfig,
     FlockStdioConnectionConfig,
 )
+from flock.mcp.servers.streamable_http.flock_streamable_http_server import (
+    FlockStreamableHttpConfig,
+    FlockStreamableHttpConnectionConfig,
+    FlockStreamableHttpServer,
+)
 from flock.mcp.servers.websockets.flock_websocket_server import (
     FlockWSConfig,
     FlockWSConnectionConfig,
@@ -101,6 +108,44 @@ class FlockFactory:
             description="The text encoding error handler. See https://docs.python.org/3/library/codecs.html#codec-base-classes for explanations of possible values",
         )
 
+    class StreamableHttpParams(BaseModel):
+        """Factory-Params for Streamable Http Servers."""
+
+        url: str | AnyUrl = Field(
+            ...,
+            description="Url the server listens at."
+        )
+
+        headers: dict[str, Any] | None = Field(
+            default=None,
+            description="Additional Headers to pass to the client."
+        )
+
+        auth: httpx.Auth | None = Field(
+            default=None,
+            description="Httpx Auth Schema."
+        )
+
+        timeout_seconds: float | int = Field(
+            default=5,
+            description="Http Timeout in Seconds"
+        )
+
+        sse_read_timeout_seconds: float | int = Field(
+            default=60*5,
+            description="How many seconds to wait for server-sent events until closing the connection."
+        )
+
+        terminate_on_close: bool = Field(
+            default=True,
+            description="Whether or not to terminate the underlying connection on close."
+        )
+
+        model_config = ConfigDict(
+            arbitrary_types_allowed=True,
+            extra="allow",
+        )
+
     class SSEParams(BaseModel):
         """Factory-Params for SSE-Servers."""
 
@@ -123,6 +168,16 @@ class FlockFactory:
             description="How many seconds to wait for server-sent events until closing the connection. (connections will be automatically re-established.)",
         )
 
+        auth: httpx.Auth | None = Field(
+            default=None,
+            description="Httpx Auth Scheme."
+        )
+
+        model_config = ConfigDict(
+            arbitrary_types_allowed=True,
+            extra="allow",
+        )
+
     class WebsocketParams(BaseModel):
         """Factory-Params for Websocket Servers."""
 
@@ -134,7 +189,7 @@ class FlockFactory:
     @staticmethod
     def create_mcp_server(
         name: str,
-        connection_params: SSEParams | StdioParams | WebsocketParams,
+        connection_params: StreamableHttpParams | SSEParams | StdioParams | WebsocketParams,
         max_retries: int = 3,
         mount_points: list[str | MCPRoot] | None = None,
         timeout_seconds: int | float = 10,
@@ -176,6 +231,9 @@ class FlockFactory:
         if isinstance(connection_params, FlockFactory.WebsocketParams):
             server_kind = "websockets"
             concrete_server_cls = FlockWSServer
+        if isinstance(connection_params, FlockFactory.StreamableHttpParams):
+            server_kind = "streamable_http"
+            concrete_server_cls = FlockStreamableHttpServer
 
         # convert mount points.
         mounts: list[MCPRoot] = []
@@ -244,12 +302,37 @@ class FlockFactory:
                 caching_config=caching_config,
                 callback_config=callback_config,
             )
+        elif server_kind == "streamable_http":
+            # build streamable http config
+            connection_config = FlockStreamableHttpConnectionConfig(
+                max_retries=max_retries,
+                connection_parameters=StreamableHttpServerParameters(
+                    url=connection_params.url,
+                    headers=connection_params.headers,
+                    auth=connection_params.auth,
+                    timeout=connection_params.timeout_seconds,
+                    sse_read_timeout=connection_params.sse_read_timeout_seconds,
+                    terminate_on_close=connection_params.terminate_on_close,
+                ),
+                mount_points=mounts,
+                server_logging_level=server_logging_level,
+            )
+
+            server_config = FlockStreamableHttpConfig(
+                name=name,
+                connection_config=connection_config,
+                feature_config=feature_config,
+                caching_config=caching_config,
+                callback_config=callback_config,
+            )
+
         elif server_kind == "sse":
             # build sse config
             connection_config = FlockSSEConnectionConfig(
                 max_retries=max_retries,
                 connection_parameters=SseServerParameters(
                     url=connection_params.url,
+                    auth=connection_params.auth,
                     headers=connection_params.headers,
                     timeout=connection_params.timeout_seconds,
                     sse_read_timeout=connection_params.sse_read_timeout_seconds,

flock/core/logging/logging.py

@@ -68,6 +68,14 @@ COLOR_MAP = {
     "workflow": "cyan",  # Color only
     "activities": "cyan",
     "context": "green",
+    "mcp.server": "blue",
+    "mcp.tool": "cyan",
+    "mcp.client_manager": "light-blue",
+    "mcp.client": "light-cyan",
+    "mcp.callback.logging": "white",
+    "mcp.callback.sampling": "pink",
+    "mcp.callback.root": "light-yellow",
+    "mcp.callback.message": "light-blue",
     # Components & Mechanisms
     "registry": "yellow",  # Color only
     "serialization": "yellow",

flock/core/mcp/flock_mcp_server.py

@@ -26,7 +26,7 @@ from flock.core.serialization.serialization_utils import (
     serialize_item,
 )
 
-logger = get_logger("core.mcp.server_base")
+logger = get_logger("mcp.server")
 tracer = trace.get_tracer(__name__)
 T = TypeVar("T", bound="FlockMCPServerBase")
 
@@ -206,7 +206,6 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
             async with self.condition:
                 try:
                     await self.pre_mcp_call()
-                    # TODO: inject additional params here.
                     additional_params: dict[str, Any] = {}
                     additional_params = await self.before_connect(
                         additional_params=additional_params
@@ -314,7 +313,7 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
     async def post_terminate(self) -> None:
         """Run post-terminate hooks on modules."""
         logger.debug(
-            f"Running post_terminat hooks for modules in server: '{self.config.name}'"
+            f"Running post_terminate hooks for modules in server: '{self.config.name}'"
         )
         with tracer.start_as_current_span("server.post_terminate") as span:
             span.set_attribute("server.name", self.config.name)
@@ -437,7 +436,7 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
 
         FlockRegistry = get_registry()
 
-        exclude = ["modules"]
+        exclude = ["modules", "config"]
 
         logger.debug(f"Serializing server '{self.config.name}' to dict.")
         # Use Pydantic's dump, exclued manually handled fields.
@@ -447,6 +446,11 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
             exclude_none=True,  # Exclude None values for cleaner output
         )
 
+        # --- Let the config handle its own serialization ---
+        config_data = self.config.to_dict(path_type=path_type)
+        data["config"] = config_data
+
+
         builtin_by_transport = {}
 
         try:
@@ -454,12 +458,16 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
             from flock.mcp.servers.stdio.flock_stdio_server import (
                 FlockMCPStdioServer,
             )
+            from flock.mcp.servers.streamable_http.flock_streamable_http_server import (
+                FlockStreamableHttpServer,
+            )
             from flock.mcp.servers.websockets.flock_websocket_server import (
                 FlockWSServer,
             )
 
             builtin_by_transport = {
                 "stdio": FlockMCPStdioServer,
+                "streamable_http": FlockStreamableHttpServer,
                 "sse": FlockSSEServer,
                 "websockets": FlockWSServer,
             }
@@ -570,6 +578,9 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
             from flock.mcp.servers.stdio.flock_stdio_server import (
                 FlockMCPStdioServer,
             )
+            from flock.mcp.servers.streamable_http.flock_streamable_http_server import (
+                FlockStreamableHttpServer,
+            )
             from flock.mcp.servers.websockets.flock_websocket_server import (
                 FlockWSServer,
             )
@@ -577,6 +588,7 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
             builtin_by_transport = {
                 "stdio": FlockMCPStdioServer,
                 "sse": FlockSSEServer,
+                "streamable_http": FlockStreamableHttpServer,
                 "websockets": FlockWSServer,
             }
         except ImportError:
@@ -592,6 +604,20 @@ class FlockMCPServerBase(BaseModel, Serializable, ABC):
             transport = data["config"]["connection_config"]["transport_type"]
             real_cls = builtin_by_transport.get(transport, cls)
 
+        # deserialize the config:
+        config_data = data.pop("config", None)
+        if config_data:
+            # Forcing a square into a round hole
+            # pretty ugly, but gets the job done.
+            try:
+                config_field = real_cls.model_fields["config"]
+                config_cls = config_field.annotation
+            except (AttributeError, KeyError):
+                # fallback if Pydantic v1 or missing
+                config_cls = FlockMCPConfigurationBase
+            config_object = config_cls.from_dict(config_data)
+            data["config"] = config_object
+
         # now construct
         server = real_cls(**{k: v for k, v in data.items() if k != "modules"})
 

flock/core/mcp/flock_mcp_tool_base.py

@@ -10,7 +10,7 @@ from pydantic import BaseModel, Field
 
 from flock.core.logging.logging import get_logger
 
-logger = get_logger("core.mcp.tool_base")
+logger = get_logger("mcp.tool")
 tracer = trace.get_tracer(__name__)
 
 T = TypeVar("T", bound="FlockMCPToolBase")

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")