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

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

@@ -17,7 +17,7 @@ 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()
-    ).hexdigest()
+    # merge args and kwargs into a single string
+    merged = json.dumps({"args": args, "kwargs": kwargs}, sort_keys=True)
+    args_digest = hashlib.md5(merged.encode(), usedforsecurity=False).hexdigest()
     return f"{agent_id}:{run_id}:{args_digest}"