flock-core 0.4.543__py3-none-any.whl → 0.5.0__py3-none-any.whl

flock/mcp/servers/websockets/flock_websocket_server.py

@@ -13,22 +13,21 @@ from mcp.shared.message import SessionMessage
 from opentelemetry import trace
 from pydantic import Field
 
-from flock.core.logging.logging import get_logger
-from flock.core.mcp.flock_mcp_server import FlockMCPServerBase
-from flock.core.mcp.mcp_client import FlockMCPClientBase
-from flock.core.mcp.mcp_client_manager import FlockMCPClientManagerBase
-from flock.core.mcp.mcp_config import (
-    FlockMCPConfigurationBase,
-    FlockMCPConnectionConfigurationBase,
+from flock.logging.logging import get_logger
+from flock.mcp.client import FlockMCPClient
+from flock.mcp.config import (
+    FlockMCPConfiguration,
+    FlockMCPConnectionConfiguration,
 )
-from flock.core.mcp.types.types import WebsocketServerParameters
+from flock.mcp.types import WebsocketServerParameters
+
 
 logger = get_logger("mcp.ws.server")
 tracer = trace.get_tracer(__name__)
 
 
 # Optional to provide type hints.
-class FlockWSConnectionConfig(FlockMCPConnectionConfigurationBase):
+class FlockWSConnectionConfig(FlockMCPConnectionConfiguration):
     """Concrete ConnectionConfig for a WS Client."""
 
     # Only thing we need to override here is the concrete transport_type
@@ -44,7 +43,7 @@ class FlockWSConnectionConfig(FlockMCPConnectionConfigurationBase):
 
 
 # Optional to provide type hints.
-class FlockWSConfig(FlockMCPConfigurationBase):
+class FlockWSConfig(FlockMCPConfiguration):
     """Configuration for Websocket clients."""
 
     # The only thing we need to override here is the concrete
@@ -56,10 +55,10 @@ class FlockWSConfig(FlockMCPConfigurationBase):
     )
 
 
-class FlockWSClient(FlockMCPClientBase):
+class FlockWSClient(FlockMCPClient):
     """Client for Websocket servers."""
 
-    config: FlockWSConfig = Field(..., description="Client Configuration")
+    config: FlockMCPConfiguration = Field(..., description="Client Configuration")
 
     # This one we HAVE to specify. This tells Flock
     # how to create the underlying connection.
@@ -84,36 +83,4 @@ class FlockWSClient(FlockMCPClientBase):
             # If present, then apply the changes in "url" to the create_transport logic.
             param_copy.url = additional_params.get("url", params.url)
 
-        return websocket_client(
-            url=param_copy.url
-        )  # return the async context manager
-
-
-# not really needed, but kept for type hints and as an example.
-class FlockWSClientManager(FlockMCPClientManagerBase):
-    """Manager for handling websocket clients."""
-
-    client_config: FlockWSConfig = Field(
-        ..., description="Configuration for clients."
-    )
-
-    async def make_client(self, additional_params=None):
-        """Create a new client instance."""
-        new_client = FlockWSClient(
-            config=self.client_config,
-            additional_params=additional_params,
-        )
-        return new_client
-
-
-class FlockWSServer(FlockMCPServerBase):
-    """Class which represents an MCP Server using the websocket transport type."""
-
-    config: FlockWSConfig = Field(..., description="Config for the server.")
-
-    # Specify the concrete type for the server.
-    async def initialize(self) -> FlockWSClientManager:
-        """Called when initializing the server."""
-        client_manager = FlockWSClientManager(client_config=self.config)
-
-        return client_manager
+        return websocket_client(url=param_copy.url)  # return the async context manager

flock/{core/mcp/flock_mcp_tool_base.py → mcp/tool.py}

@@ -3,17 +3,20 @@
 from typing import Any, TypeVar
 
 from dspy import Tool as DSPyTool
+from dspy.adapters.types.tool import convert_input_schema_to_tool_args
 from mcp import Tool
 from mcp.types import CallToolResult, TextContent, ToolAnnotations
 from opentelemetry import trace
 from pydantic import BaseModel, Field
+from typing_extensions import Self
+
+from flock.logging.logging import get_logger
 
-from flock.core.logging.logging import get_logger
 
 logger = get_logger("mcp.tool")
 tracer = trace.get_tracer(__name__)
 
-T = TypeVar("T", bound="FlockMCPToolBase")
+T = TypeVar("T", bound="FlockMCPTool")
 
 TYPE_MAPPING = {
     "string": str,
@@ -25,22 +28,16 @@ TYPE_MAPPING = {
 }
 
 
-class FlockMCPToolBase(BaseModel):
+class FlockMCPTool(BaseModel):
     """Base Class for MCP Tools for Flock."""
 
     name: str = Field(..., description="Name of the tool")
 
-    agent_id: str = Field(
-        ..., description="Associated agent_id. Used for internal tracking."
-    )
+    agent_id: str = Field(..., description="Associated agent_id. Used for internal tracking.")
 
-    run_id: str = Field(
-        ..., description="Associated run_id. Used for internal tracking."
-    )
+    run_id: str = Field(..., description="Associated run_id. Used for internal tracking.")
 
-    description: str | None = Field(
-        ..., description="A human-readable description of the tool"
-    )
+    description: str | None = Field(..., description="A human-readable description of the tool")
 
     input_schema: dict[str, Any] = Field(
         ...,
@@ -52,9 +49,7 @@ class FlockMCPToolBase(BaseModel):
     )
 
     @classmethod
-    def from_mcp_tool(
-        cls: type[T], tool: Tool, agent_id: str, run_id: str
-    ) -> T:
+    def from_mcp_tool(cls, tool: Tool, agent_id: str, run_id: str) -> Self:
         """Convert MCP Tool to Flock Tool."""
         return cls(
             name=tool.name,
@@ -66,7 +61,7 @@ class FlockMCPToolBase(BaseModel):
         )
 
     @classmethod
-    def to_mcp_tool(cls: type[T], instance: T) -> Tool | None:
+    def to_mcp_tool(cls, instance: Self) -> Tool | None:
         """Convert a flock mcp tool into a mcp tool."""
         return Tool(
             name=instance.name,
@@ -75,63 +70,18 @@ class FlockMCPToolBase(BaseModel):
             annotations=instance.annotations,
         )
 
-    def resolve_json_schema_reference(self, schema: dict) -> dict:
-        """Recursively resolve json model schema, expanding all references."""
-        if "$defs" not in schema and "definitions" not in schema:
-            return schema
-
-        def resolve_refs(obj: Any) -> Any:
-            if not isinstance(obj, dict[list, list]):
-                return obj
-            if isinstance(obj, dict) and "$ref" in obj:
-                # ref_path = obj["$ref"].split("/")[-1]
-                return {resolve_refs(v) for k, v in obj.items()}
-
-            return [resolve_refs(item) for item in obj]
-
-        resolved_schema = resolve_refs(schema)
-
-        resolved_schema.pop("$defs", None)
-        return resolved_schema
-
+    # Use DSPy's converter for JSON Schema → Tool args to stay aligned with DSPy.
     def _convert_input_schema_to_tool_args(
         self, input_schema: dict[str, Any]
-    ) -> tuple[dict[str, Any], dict[str, type], dict[str, str]]:
-        """Convert an input schema to tool arguments compatible with Dspy Tool.
-
-        Args:
-            schema: an input schema describing the tool's input parameters
-
-        Returns:
-            A tuple of (args, arg_types, arg_desc) for Dspy Tool definition
-        """
-        args, arg_types, arg_desc = {}, {}, {}
-        properties = input_schema.get("properties")
-        if properties is None:
-            return args, arg_types, arg_desc
-
-        required = input_schema.get("required", [])
-
-        defs = input_schema.get("$defs", {})
-
-        for name, prop in properties.items():
-            if len(defs) > 0:
-                prop = self.resolve_json_schema_reference(
-                    {"$defs": defs, **prop}
-                )
-
-            args[name] = prop
-
-            arg_types[name] = TYPE_MAPPING.get(prop.get("type"), Any)
-            arg_desc[name] = prop.get("description", "No description provided")
-            if name in required:
-                arg_desc[name] += " (Required)"
-
-        return args, arg_types, arg_desc
-
-    def _convert_mcp_tool_result(
-        self, call_tool_result: CallToolResult
-    ) -> str | list[Any]:
+    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, str]]:
+        try:
+            return convert_input_schema_to_tool_args(input_schema)
+        except Exception as e:  # pragma: no cover - defensive
+            logger.exception("Failed to convert MCP tool schema to DSPy tool args: %s", e)
+            # Fallback to empty definitions to avoid breaking execution
+            return {}, {}, {}
+
+    def _convert_mcp_tool_result(self, call_tool_result: CallToolResult) -> str | list[Any]:
         text_contents: list[TextContent] = []
         non_text_contents = []
 
@@ -158,9 +108,7 @@ class FlockMCPToolBase(BaseModel):
 
     def as_dspy_tool(self, server: Any) -> DSPyTool:
         """Wrap this tool as a DSPyTool for downstream."""
-        args, arg_type, args_desc = self._convert_input_schema_to_tool_args(
-            self.input_schema
-        )
+        args, arg_type, args_desc = self._convert_input_schema_to_tool_args(self.input_schema)
 
         async def func(*args, **kwargs):
             with tracer.start_as_current_span(f"tool.{self.name}.call") as span:
@@ -172,9 +120,7 @@ class FlockMCPToolBase(BaseModel):
                     logger.debug(
                         f"Tool: {self.name}: got client for server '{server_name}' for agent {self.agent_id} on run {self.run_id}"
                     )
-                    logger.debug(
-                        f"Tool: {self.name}: calling server '{server_name}'"
-                    )
+                    logger.debug(f"Tool: {self.name}: calling server '{server_name}'")
                     result = await server.call_tool(
                         agent_id=self.agent_id,
                         run_id=self.run_id,
@@ -186,7 +132,7 @@ class FlockMCPToolBase(BaseModel):
                     )
                     return self._convert_mcp_tool_result(result)
                 except Exception as e:
-                    logger.error(
+                    logger.exception(
                         f"Tool: Exception ocurred when calling tool '{self.name}': {e}"
                     )
                     span.record_exception(e)

flock/mcp/types/__init__.py

@@ -0,0 +1,42 @@
+"""MCP Type Definitions for Flock-Flow."""
+
+from flock.mcp.types.factories import (
+    default_flock_mcp_list_roots_callback_factory,
+    default_flock_mcp_logging_callback_factory,
+    default_flock_mcp_message_handler_callback_factory,
+    default_flock_mcp_sampling_callback_factory,
+)
+from flock.mcp.types.types import (
+    FlockListRootsMCPCallback,
+    FlockLoggingMCPCallback,
+    FlockLoggingMessageNotificationParams,
+    FlockMessageHandlerMCPCallback,
+    FlockSamplingMCPCallback,
+    MCPRoot,
+    ServerNotification,
+    ServerParameters,
+    SseServerParameters,
+    StdioServerParameters,
+    StreamableHttpServerParameters,
+    WebsocketServerParameters,
+)
+
+
+__all__ = [
+    "FlockListRootsMCPCallback",
+    "FlockLoggingMCPCallback",
+    "FlockLoggingMessageNotificationParams",
+    "FlockMessageHandlerMCPCallback",
+    "FlockSamplingMCPCallback",
+    "MCPRoot",
+    "ServerNotification",
+    "ServerParameters",
+    "SseServerParameters",
+    "StdioServerParameters",
+    "StreamableHttpServerParameters",
+    "WebsocketServerParameters",
+    "default_flock_mcp_list_roots_callback_factory",
+    "default_flock_mcp_logging_callback_factory",
+    "default_flock_mcp_message_handler_callback_factory",
+    "default_flock_mcp_sampling_callback_factory",
+]

flock/{core/mcp → mcp}/types/callbacks.py

@@ -1,5 +1,7 @@
 """MCP Callbacks."""
 
+from typing import Any
+
 from mcp.shared.context import RequestContext
 from mcp.shared.session import RequestResponder
 from mcp.types import (
@@ -13,9 +15,8 @@ from mcp.types import (
     ServerRequest,
 )
 
-from flock.core.logging.logging import FlockLogger
-from flock.core.mcp.mcp_client import Any
-from flock.core.mcp.types.handlers import (
+from flock.logging.logging import FlockLogger
+from flock.mcp.types.handlers import (
     handle_incoming_exception,
     handle_incoming_request,
     handle_incoming_server_notification,
@@ -27,14 +28,12 @@ async def default_sampling_callback(
     ctx: RequestContext, params: CreateMessageRequestParams, logger: FlockLogger
 ) -> ErrorData:
     """Default Callback for Sampling."""
-    logger.info(f"Rejecting Sampling Request.")
+    logger.info("Rejecting Sampling Request.")
     return ErrorData(code=INVALID_REQUEST, message="Sampling not supported.")
 
 
 async def default_message_handler(
-    req: RequestResponder[ServerRequest, ClientResult]
-    | ServerNotification
-    | Exception,
+    req: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
     logger: FlockLogger,
     associated_client: Any,
 ) -> None:
@@ -65,12 +64,12 @@ async def default_list_roots_callback(
 ) -> ListRootsResult | ErrorData:
     """Default List Roots Callback."""
     if associated_client.config.feature_config.roots_enabled:
-        current_roots = await associated_client.get_current_roots()
+        # Use lock-free version to avoid deadlock during initialization
+        # when the lock is already held by _connect()
+        current_roots = associated_client._get_roots_no_lock()
+        logger.debug(f"Server requested list/roots. Sending: {current_roots}")
         return ListRootsResult(roots=current_roots)
-    else:
-        return ErrorData(
-            code=INVALID_REQUEST, message="List roots not supported."
-        )
+    return ErrorData(code=INVALID_REQUEST, message="List roots not supported.")
 
 
 async def default_logging_callback(
@@ -79,6 +78,4 @@ async def default_logging_callback(
     server_name: str,
 ) -> None:
     """Default Logging Handling Callback."""
-    await handle_logging_message(
-        params=params, logger=logger, server_name=server_name
-    )
+    await handle_logging_message(params=params, logger=logger, server_name=server_name)

flock/{core/mcp → mcp}/types/factories.py

@@ -7,14 +7,14 @@ from mcp.types import (
     CreateMessageRequestParams,
 )
 
-from flock.core.logging.logging import FlockLogger, get_logger
-from flock.core.mcp.types.callbacks import (
+from flock.logging.logging import FlockLogger, get_logger
+from flock.mcp.types.callbacks import (
     default_list_roots_callback,
     default_logging_callback,
     default_message_handler,
     default_sampling_callback,
 )
-from flock.core.mcp.types.types import (
+from flock.mcp.types.types import (
     FlockListRootsMCPCallback,
     FlockLoggingMCPCallback,
     FlockLoggingMessageNotificationParams,
@@ -23,6 +23,7 @@ from flock.core.mcp.types.types import (
     ServerNotification,
 )
 
+
 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")
@@ -45,6 +46,8 @@ def default_flock_mcp_logging_callback_factory(
             server_name=associated_client.config.name,
         )
 
+    return _method
+
 
 def default_flock_mcp_sampling_callback_factory(
     associated_client: Any,
@@ -60,9 +63,7 @@ def default_flock_mcp_sampling_callback_factory(
         logger_to_use.info(
             f"SAMPLING_REQUEST: server '{associated_client.config.name}' sent a sampling request: {params}"
         )
-        await default_sampling_callback(
-            ctx=ctx, params=params, logger=logger_to_use
-        )
+        await default_sampling_callback(ctx=ctx, params=params, logger=logger_to_use)
 
     return _method
 

flock/{core/mcp → mcp}/types/handlers.py

@@ -1,6 +1,6 @@
 """Handler functions."""
 
-from collections.abc import Callable
+from typing import Any
 
 from mcp import CreateMessageRequest
 from mcp.client.session import ClientResponse
@@ -22,8 +22,7 @@ from mcp.types import (
     ToolListChangedNotification,
 )
 
-from flock.core.logging.logging import FlockLogger
-from flock.core.mcp.mcp_client import Any
+from flock.logging.logging import FlockLogger
 
 
 async def handle_incoming_exception(
@@ -50,7 +49,7 @@ async def handle_progress_notification(
     progress = params.progress
     total = params.total or "Unknown"
     progress_token = params.progressToken
-    metadata = params.meta or {}
+    metadata: dict[str, Any] = params.meta or {}  # type: ignore[assignment]
 
     message = f"PROGRESS_NOTIFICATION: Server '{server_name}' reports Progress: {progress}/{total}. (Token: {progress_token}) (Meta Data: {metadata})"
 
@@ -66,7 +65,7 @@ async def handle_cancellation_notification(
     params = n.params
     request_id_to_cancel = params.requestId
     reason = params.reason or "no reason given"
-    metadata = params.meta or {}
+    metadata: dict[str, Any] = params.meta or {}  # type: ignore[assignment]
 
     message = f"CANCELLATION_REQUEST: Server '{server_name}' requests to cancel request with id: {request_id_to_cancel}. Reason: {reason}. (Metadata: {metadata})"
 
@@ -83,7 +82,7 @@ async def handle_resource_update_notification(
     # its resource_contents_cache
 
     params = n.params
-    metadata = params.meta or {}
+    metadata: dict[str, Any] = params.meta or {}  # type: ignore[assignment]
     uri = params.uri
 
     message = f"RESOURCE_UPDATE: Server '{associated_client.config.name}' reports change on resoure at: {uri}. (Meta Data: {metadata})"
@@ -102,8 +101,8 @@ async def handle_resource_list_changed_notification(
     # This also means that the associated client needs to invalidate
     # its resource_contents_cache
 
-    params = n.params or {}
-    metadata = params.meta or {}
+    params: Any = n.params or {}
+    metadata: dict[str, Any] = params.meta or {}
 
     message = f"TOOLS_LIST_CHANGED: Server '{associated_client.config.name}' reports a change in their tools list: {metadata}. Resetting Tools Cache for associated clients."
 
@@ -117,8 +116,8 @@ async def handle_tool_list_changed_notification(
     associated_client: Any,
 ) -> None:
     """Handle an incoming ToolListChangedNotification."""
-    params = n.params or {}
-    metadata = params.meta or {}
+    params: Any = n.params or {}
+    metadata: dict[str, Any] = params.meta or {}
 
     message = f"TOOLS_LIST_CHANGED: Server '{associated_client.config.name}' reports a change in their tools list: {metadata}. Resetting Tools Cache for associated clients."
 
@@ -126,7 +125,7 @@ async def handle_tool_list_changed_notification(
     await associated_client.invalidate_tool_cache()
 
 
-_SERVER_NOTIFICATION_MAP: dict[type[ServerNotification], Callable] = {
+_SERVER_NOTIFICATION_MAP: Any = {
     ResourceListChangedNotification: handle_resource_list_changed_notification,
     ResourceUpdatedNotification: handle_resource_update_notification,
     LoggingMessageNotification: lambda n, log, client: handle_logging_message(
@@ -159,7 +158,7 @@ async def handle_logging_message(
     level = params.level
     method = logger.debug
     logger_name = params.logger if params.logger else "unknown_remote_logger"
-    metadata = params.meta or {}
+    metadata: dict[str, Any] = params.meta or {}  # type: ignore[assignment]
 
     str_level = "DEBUG"
     prefix = f"Message from Remote MCP Logger '{logger_name}' for server '{server_name}': "
@@ -206,7 +205,6 @@ async def handle_incoming_request(
             case CreateMessageRequest(params=req.request.root.params):
                 with req:
                     # invoke user's sampling callback
-                    # type: ignore
                     response = await associated_client.sampling_callback(
                         ctx, req.request.root.params
                     )
@@ -214,7 +212,6 @@ async def handle_incoming_request(
                     await req.respond(client_resp)
             case ListRootsRequest():
                 with req:
-                    # type: ignore
                     response = await associated_client.list_roots_callback(ctx)
                     client_resp = ClientResponse.validate_python(response)
                     await req.respond(client_resp)
@@ -223,15 +220,13 @@ async def handle_incoming_request(
                 return
     except Exception as e:
         # 1) Log the error and stacktrace
-        logger_to_use.error(
+        logger_to_use.exception(
             f"Error in fallback handle_incoming_request (id={req.request_id}): {e}"
         )
         # 2) If the request wasn't already completed, send a JSON-RPC error back
         if not getattr(req, "_completed", False):
             with req:
-                err = ErrorData(
-                    code=INTERNAL_ERROR, message=f"Client-side error: {e}"
-                )
+                err = ErrorData(code=INTERNAL_ERROR, message=f"Client-side error: {e}")
                 client_err = ClientResponse.validate_python(err)
                 await req.respond(client_err)
     return