flock-core 0.5.0b28__py3-none-any.whl → 0.5.56b0__py3-none-any.whl

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

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

@@ -15,32 +15,53 @@ from anyio.streams.memory import (
 from mcp import (
     ClientSession,
     CreateMessageResult,
+)
+from mcp import (
     StdioServerParameters as _MCPStdioServerParameters,
 )
 from mcp.shared.context import RequestContext
 from mcp.shared.session import RequestResponder
 from mcp.types import (
     CancelledNotification as _MCPCancelledNotification,
+)
+from mcp.types import (
     ClientResult,
     CreateMessageRequestParams,
     ErrorData,
     JSONRPCMessage,
     ListRootsResult,
+    ServerRequest,
+)
+from mcp.types import (
     LoggingMessageNotification as _MCPLoggingMessageNotification,
+)
+from mcp.types import (
     LoggingMessageNotificationParams as _MCPLoggingMessageNotificationParams,
+)
+from mcp.types import (
     ProgressNotification as _MCPProgressNotification,
+)
+from mcp.types import (
     PromptListChangedNotification as _MCPPromptListChangedNotification,
+)
+from mcp.types import (
     ResourceListChangedNotification as _MCPResourceListChangedNotification,
+)
+from mcp.types import (
     ResourceUpdatedNotification as _MCPResourceUpdateNotification,
+)
+from mcp.types import (
     Root as _MCPRoot,
+)
+from mcp.types import (
     ServerNotification as _MCPServerNotification,
-    ServerRequest,
+)
+from mcp.types import (
     ToolListChangedNotification as _MCPToolListChangedNotification,
 )
 from pydantic import AnyUrl, BaseModel, ConfigDict, Field
 
-from flock.core.mcp.util.helpers import get_default_env
-from flock.core.serialization.serializable import Serializable
+from flock.mcp.util.helpers import get_default_env
 
 
 class ServerNotification(_MCPServerNotification):
@@ -75,9 +96,7 @@ class PromptListChangedNotification(_MCPPromptListChangedNotification):
     """A notification message sent by the server side informing a client about a change in the list of offered Prompts."""
 
 
-class FlockLoggingMessageNotificationParams(
-    _MCPLoggingMessageNotificationParams
-):
+class FlockLoggingMessageNotificationParams(_MCPLoggingMessageNotificationParams):
     """Parameters contained within a Logging Message Notification."""
 
 
@@ -85,7 +104,7 @@ class MCPRoot(_MCPRoot):
     """Wrapper for mcp.types.Root."""
 
 
-class ServerParameters(BaseModel, Serializable):
+class ServerParameters(BaseModel):
     """Base Type for server parameters."""
 
     model_config = ConfigDict(
@@ -93,17 +112,12 @@ class ServerParameters(BaseModel, Serializable):
     )
 
     transport_type: Literal["stdio", "websockets", "sse", "streamable_http"] = Field(
-        ...,
-        description="which type of transport these connection params are used for."
+        ..., description="which type of transport these connection params are used for."
     )
 
     def to_dict(self, path_type: str = "relative"):
         """Serialize."""
-        return self.model_dump(
-            exclude_defaults=False,
-            exclude_none=True,
-            mode="json"
-        )
+        return self.model_dump(exclude_defaults=False, exclude_none=True, mode="json")
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]):
@@ -114,10 +128,7 @@ class ServerParameters(BaseModel, Serializable):
 class StdioServerParameters(_MCPStdioServerParameters, ServerParameters):
     """Base Type for Stdio Server parameters."""
 
-    transport_type: Literal["stdio"] = Field(
-        default="stdio",
-        description="Use stdio params."
-    )
+    transport_type: Literal["stdio"] = Field(default="stdio", description="Use stdio params.")
 
     env: dict[str, str] | None = Field(
         default_factory=get_default_env,
@@ -129,28 +140,23 @@ class WebsocketServerParameters(ServerParameters):
     """Base Type for Websocket Server params."""
 
     transport_type: Literal["websockets"] = Field(
-        default="websockets",
-        description="Use websocket params."
+        default="websockets", description="Use websocket params."
     )
 
     url: str | AnyUrl = Field(..., description="Url the server listens at.")
 
+
 class StreamableHttpServerParameters(ServerParameters):
     """Base Type for StreamableHttp params."""
 
     transport_type: Literal["streamable_http"] = Field(
-        default="streamable_http",
-        description="Use streamable http params."
+        default="streamable_http", description="Use streamable http params."
     )
 
-    url: str | AnyUrl = Field(
-        ...,
-        description="The url the server listens at."
-    )
+    url: str | AnyUrl = Field(..., description="The url the server listens at.")
 
     headers: dict[str, Any] | None = Field(
-        default=None,
-        description="Additional headers to pass to the client."
+        default=None, description="Additional headers to pass to the client."
     )
 
     timeout: float | int = Field(
@@ -159,19 +165,13 @@ class StreamableHttpServerParameters(ServerParameters):
     )
 
     sse_read_timeout: float | int = Field(
-        default=60*5,
-        description="How long the client will wait before disconnecting from the server."
+        default=60 * 5,
+        description="How long the client will wait before disconnecting from the server.",
     )
 
-    terminate_on_close: bool = Field(
-        default=True,
-        description="Terminate connection on close"
-    )
+    terminate_on_close: bool = Field(default=True, description="Terminate connection on close")
 
-    auth: httpx.Auth | None = Field(
-        default=None,
-        description="Httpx Auth Scheme"
-    )
+    auth: httpx.Auth | None = Field(default=None, description="Httpx Auth Scheme")
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]):
@@ -188,7 +188,7 @@ class StreamableHttpServerParameters(ServerParameters):
                 mod = importlib.import_module(impl["module_path"])
                 real_cls = getattr(mod, impl["classname"])
                 if params:
-                    auth_obj = real_cls(**{k: v for k, v in params.items()})
+                    auth_obj = real_cls(**dict(params.items()))
                 else:
                     # assume that the implementation handles it.
                     auth_obj = real_cls()
@@ -196,9 +196,9 @@ class StreamableHttpServerParameters(ServerParameters):
                 raise ValueError("No concrete implementation for auth provided.")
 
         data["auth"] = auth_obj
-        return cls(**{k: v for k, v in data.items()})
+        return cls(**dict(data.items()))
 
-    def to_dict(self, path_type = "relative"):
+    def to_dict(self, path_type="relative"):
         """Serialize the object."""
         exclude = ["auth"]
 
@@ -224,24 +224,24 @@ class StreamableHttpServerParameters(ServerParameters):
             except Exception:
                 params = None
 
-        data["auth"] = {
-            "implementation": {
-                "class_name": type(self.auth).__name__,
-                "module_path": type(self.auth).__module__,
-                "file_path": file_path,
-            },
-            "params": params,
-        }
+            data["auth"] = {
+                "implementation": {
+                    "class_name": type(self.auth).__name__,
+                    "module_path": type(self.auth).__module__,
+                    "file_path": file_path,
+                },
+                "params": params,
+            }
+        else:
+            data["auth"] = None
 
         return data
 
+
 class SseServerParameters(ServerParameters):
     """Base Type for SSE Server params."""
 
-    transport_type: Literal["sse"] = Field(
-        default="sse",
-        description="Use sse server params."
-    )
+    transport_type: Literal["sse"] = Field(default="sse", description="Use sse server params.")
 
     url: str | AnyUrl = Field(..., description="The url the server listens at.")
 
@@ -256,10 +256,7 @@ class SseServerParameters(ServerParameters):
         description="How long the client will wait before disconnecting from the server.",
     )
 
-    auth: httpx.Auth | None = Field(
-        default=None,
-        description="Httpx Auth Scheme."
-    )
+    auth: httpx.Auth | None = Field(default=None, description="Httpx Auth Scheme.")
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]):
@@ -267,7 +264,7 @@ class SseServerParameters(ServerParameters):
         # find and import the concrete implementation for
         # the auth object.
         auth_obj: httpx.Auth | None = None
-        auth_impl = data.pop("auth", None) # get the specs for the auth class
+        auth_impl = data.pop("auth", None)  # get the specs for the auth class
         if auth_impl:
             # find the concrete implementation
             impl = auth_impl.pop("implementation", None)
@@ -276,7 +273,7 @@ class SseServerParameters(ServerParameters):
                 mod = importlib.import_module(impl["module_path"])
                 real_cls = getattr(mod, impl["class_name"])
                 if params:
-                    auth_obj = real_cls(**{k: v for k, v in params.items()})
+                    auth_obj = real_cls(**dict(params.items()))
                 else:
                     # assume that implementation handles it
                     auth_obj = real_cls()
@@ -284,9 +281,9 @@ class SseServerParameters(ServerParameters):
                 raise ValueError("No concrete implementation for auth provided.")
 
         data["auth"] = auth_obj
-        return cls(**{k: v for k, v in data.items()})
+        return cls(**dict(data.items()))
 
-    def to_dict(self, path_type = "relative"):
+    def to_dict(self, path_type="relative"):
         """Serialize the object."""
         exclude = ["auth"]
 
@@ -312,17 +309,20 @@ class SseServerParameters(ServerParameters):
             except Exception:
                 params = None
 
-        data["auth"] = {
-            "implementation": {
-                "class_name": type(self.auth).__name__,
-                "module_path": type(self.auth).__module__,
-                "file_path": file_path,
-            },
-            "params": params,
-        }
+            data["auth"] = {
+                "implementation": {
+                    "class_name": type(self.auth).__name__,
+                    "module_path": type(self.auth).__module__,
+                    "file_path": file_path,
+                },
+                "params": params,
+            }
+        else:
+            data["auth"] = None
 
         return data
 
+
 MCPCLientInitFunction = Callable[
     ...,
     AbstractAsyncContextManager[
@@ -351,10 +351,6 @@ FlockLoggingMCPCallback = Callable[
 ]
 
 FlockMessageHandlerMCPCallback = Callable[
-    [
-        RequestResponder[ServerRequest, ClientResult]
-        | ServerNotification
-        | Exception
-    ],
+    [RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception],
     Awaitable[None],
 ]

flock/{core/mcp → mcp}/util/helpers.py

@@ -18,6 +18,6 @@ def get_default_env() -> dict[str, str]:
 def cache_key_generator(agent_id: str, run_id: str, *args, **kwargs) -> str:
     """Helper function to generate cache keys for Flock MCP caches."""
     args_digest = hashlib.md5(
-        json.dumps(kwargs, sort_keys=True).encode()
+        json.dumps(kwargs, sort_keys=True).encode(), usedforsecurity=False
     ).hexdigest()
     return f"{agent_id}:{run_id}:{args_digest}"