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

flock/{core/mcp/mcp_client.py → mcp/client.py}

@@ -11,16 +11,13 @@ from contextlib import (
 )
 from datetime import timedelta
 from typing import (
+    TYPE_CHECKING,
     Any,
 )
 
 import httpx
 from anyio import ClosedResourceError
-from anyio.streams.memory import (
-    MemoryObjectReceiveStream,
-    MemoryObjectSendStream,
-)
-from cachetools import TTLCache, cached
+from cachetools import TTLCache
 from mcp import (
     ClientSession,
     InitializeResult,
@@ -36,16 +33,10 @@ from pydantic import (
     Field,
 )
 
-from flock.core.logging.logging import get_logger
-from flock.core.mcp.flock_mcp_tool_base import FlockMCPToolBase
-from flock.core.mcp.mcp_config import FlockMCPConfigurationBase
-from flock.core.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.core.mcp.types.types import (
+from flock.logging.logging import get_logger
+from flock.mcp.config import FlockMCPConfiguration
+from flock.mcp.tool import FlockMCPTool
+from flock.mcp.types import (
     FlockListRootsMCPCallback,
     FlockLoggingMCPCallback,
     FlockMessageHandlerMCPCallback,
@@ -53,7 +44,21 @@ from flock.core.mcp.types.types import (
     MCPRoot,
     ServerParameters,
 )
-from flock.core.mcp.util.helpers import cache_key_generator
+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.util.helpers import cache_key_generator
+
+
+if TYPE_CHECKING:
+    from anyio.streams.memory import (
+        MemoryObjectReceiveStream,
+        MemoryObjectSendStream,
+    )
+
 
 logger = get_logger("mcp.client")
 tracer = trace.get_tracer(__name__)
@@ -61,7 +66,7 @@ tracer = trace.get_tracer(__name__)
 GetSessionIdCallback = Callable[[], str | None]
 
 
-class FlockMCPClientBase(BaseModel, ABC):
+class FlockMCPClient(BaseModel, ABC):
     """Wrapper for mcp ClientSession.
 
     Class will attempt to re-establish connection if possible.
@@ -71,9 +76,7 @@ class FlockMCPClientBase(BaseModel, ABC):
     """
 
     # --- Properties ---
-    config: FlockMCPConfigurationBase = Field(
-        ..., description="The config for this client instance."
-    )
+    config: FlockMCPConfiguration = Field(..., description="The config for this client instance.")
 
     tool_cache: TTLCache | None = Field(
         default=None,
@@ -150,13 +153,14 @@ class FlockMCPClientBase(BaseModel, ABC):
     class _SessionProxy:
         def __init__(self, client: Any):
             self._client = client
+            # Check if roots are specified in the config
+            if not self.current_roots and self.config.connection_config.mount_points:
+                self.current_roots = self.config.connection_config.mount_points
 
         def __getattr__(self, name: str):
             # return an async function that auto-reconnects, then calls through.
             async def _method(*args, **kwargs):
-                with tracer.start_as_current_span(
-                    "session_proxy.__getattr__"
-                ) as span:
+                with tracer.start_as_current_span("session_proxy.__getattr__") as span:
                     client = self._client
                     cfg = client.config
                     max_tries = cfg.connection_config.max_retries or 1
@@ -170,18 +174,14 @@ class FlockMCPClientBase(BaseModel, ABC):
                         await client._ensure_connected()
                         try:
                             # delegate the real session
-                            return await getattr(client.client_session, name)(
-                                *args, **kwargs
-                            )
+                            return await getattr(client.client_session, name)(*args, **kwargs)
                         except McpError as e:
                             # only retry on a transport timeout
                             if e.error.code == httpx.codes.REQUEST_TIMEOUT:
                                 kind = "timeout"
                             else:
                                 # application-level MCP error -> give up immediately
-                                logger.error(
-                                    f"MCP error in session.{name}: {e.error}"
-                                )
+                                logger.exception(f"MCP error in session.{name}: {e.error}")
                                 return None
                         except (BrokenPipeError, ClosedResourceError) as e:
                             kind = type(e).__name__
@@ -199,9 +199,7 @@ class FlockMCPClientBase(BaseModel, ABC):
                             try:
                                 await client.disconnect()
                             except Exception as e:
-                                logger.warning(
-                                    f"Error tearing down stale session: {e}"
-                                )
+                                logger.warning(f"Error tearing down stale session: {e}")
                                 span.record_exception(e)
                             return None
 
@@ -213,19 +211,20 @@ class FlockMCPClientBase(BaseModel, ABC):
                             await client.disconnect()
                             await client._connect()
                         except Exception as e:
-                            logger.error(f"Reconnect failed: {e}")
+                            logger.exception(f"Reconnect failed: {e}")
                             span.record_exception(e)
 
                         # Exponential backoff + 10% jitter
                         delay = base_delay ** (2 ** (attempt - 1))
                         delay += random.uniform(0, delay * 0.1)
                         await asyncio.sleep(delay)
+                return None
 
             return _method
 
     def __init__(
         self,
-        config: FlockMCPConfigurationBase,
+        config: FlockMCPConfiguration,
         lock: Lock | None = None,
         tool_cache: TTLCache | None = None,
         tool_result_cache: TTLCache | None = None,
@@ -262,10 +261,7 @@ class FlockMCPClientBase(BaseModel, ABC):
         )
 
         # Check if roots are specified in the config:
-        if (
-            not self.current_roots
-            and self.config.connection_config.mount_points
-        ):
+        if not self.current_roots and self.config.connection_config.mount_points:
             # That means that the roots are set in the config
             self.current_roots = self.config.connection_config.mount_points
 
@@ -296,55 +292,39 @@ class FlockMCPClientBase(BaseModel, ABC):
         # set up callbacks
         if not self.logging_callback:
             if not self.config.callback_config.logging_callback:
-                self.logging_callback = (
-                    default_flock_mcp_logging_callback_factory(
-                        associated_client=self,
-                        logger=logger,
-                    )
+                self.logging_callback = default_flock_mcp_logging_callback_factory(
+                    associated_client=self,
+                    logger=logger,
                 )
             else:
-                self.logging_callback = (
-                    self.config.callback_config.logging_callback
-                )
+                self.logging_callback = self.config.callback_config.logging_callback
 
         if not self.message_handler:
             if not self.config.callback_config.message_handler:
-                self.message_handler = (
-                    default_flock_mcp_message_handler_callback_factory(
-                        associated_client=self,
-                        logger=logger,
-                    )
+                self.message_handler = default_flock_mcp_message_handler_callback_factory(
+                    associated_client=self,
+                    logger=logger,
                 )
             else:
-                self.message_handler = (
-                    self.config.callback_config.message_handler
-                )
+                self.message_handler = self.config.callback_config.message_handler
 
         if not self.list_roots_callback:
             if not self.config.callback_config.list_roots_callback:
-                self.list_roots_callback = (
-                    default_flock_mcp_list_roots_callback_factory(
-                        associated_client=self,
-                        logger=logger,
-                    )
+                self.list_roots_callback = default_flock_mcp_list_roots_callback_factory(
+                    associated_client=self,
+                    logger=logger,
                 )
             else:
-                self.list_roots_callback = (
-                    self.config.callback_config.list_roots_callback
-                )
+                self.list_roots_callback = self.config.callback_config.list_roots_callback
 
         if not self.sampling_callback:
             if not self.config.callback_config.sampling_callback:
-                self.sampling_callback = (
-                    default_flock_mcp_sampling_callback_factory(
-                        associated_client=self,
-                        logger=logger,
-                    )
+                self.sampling_callback = default_flock_mcp_sampling_callback_factory(
+                    associated_client=self,
+                    logger=logger,
                 )
             else:
-                self.sampling_callback = (
-                    self.config.callback_config.sampling_callback
-                )
+                self.sampling_callback = self.config.callback_config.sampling_callback
 
     @property
     def session(self) -> _SessionProxy:
@@ -374,61 +354,69 @@ class FlockMCPClientBase(BaseModel, ABC):
         self,
         agent_id: str,
         run_id: str,
-    ) -> list[FlockMCPToolBase]:
+    ) -> list[FlockMCPTool]:
         """Gets a list of available tools from the server."""
 
-        @cached(cache=self.tool_cache, key=cache_key_generator)
-        async def _get_tools_cached(
-            agent_id: str,
-            run_id: str,
-        ) -> list[FlockMCPToolBase]:
-            if not self.config.feature_config.tools_enabled:
-                return []
-
-            async def _get_tools_internal() -> list[FlockMCPToolBase]:
-                # TODO: Crash
-                response: ListToolsResult = await self.session.list_tools()
-                flock_tools = []
-
-                for tool in response.tools:
-                    converted_tool = FlockMCPToolBase.from_mcp_tool(
-                        tool,
-                        agent_id=agent_id,
-                        run_id=run_id,
-                    )
-                    if converted_tool:
-                        flock_tools.append(converted_tool)
-                return flock_tools
-
-            return await _get_tools_internal()
+        cache_key = cache_key_generator(agent_id, run_id)
+
+        # Check cache first
+        if cache_key in self.tool_cache:
+            return self.tool_cache[cache_key]
+
+        if not self.config.feature_config.tools_enabled:
+            result = []
+            self.tool_cache[cache_key] = result
+            return result
+
+        async def _get_tools_internal() -> list[FlockMCPTool]:
+            response: ListToolsResult = await self.session.list_tools()
+            flock_tools = []
+            tool_whitelist = self.config.feature_config.tool_whitelist
+
+            for tool in response.tools:
+                # Skip tools that are not whitelisted
+                # IF a whitelist is present
+                if (
+                    tool_whitelist is not None
+                    and isinstance(tool_whitelist, list)
+                    and len(tool_whitelist) > 0
+                    and tool.name not in tool_whitelist
+                ):
+                    continue
+                converted_tool = FlockMCPTool.from_mcp_tool(
+                    tool,
+                    agent_id=agent_id,
+                    run_id=run_id,
+                )
+                if converted_tool:
+                    flock_tools.append(converted_tool)
+            return flock_tools
 
-        return await _get_tools_cached(agent_id=agent_id, run_id=run_id)
+        result = await _get_tools_internal()
+        self.tool_cache[cache_key] = result
+        return result
 
     async def call_tool(
         self, agent_id: str, run_id: str, name: str, arguments: dict[str, Any]
     ) -> CallToolResult:
         """Call a tool via the MCP Protocol on the client's server."""
 
-        @cached(cache=self.tool_result_cache, key=cache_key_generator)
-        async def _call_tool_cached(
-            agent_id: str, run_id: str, name: str, arguments: dict[str, Any]
-        ) -> CallToolResult:
-            async def _call_tool_internal(
-                name: str, arguments: dict[str, Any]
-            ) -> CallToolResult:
-                logger.debug(
-                    f"Calling tool '{name}' with arguments {arguments}"
-                )
-                return await self.session.call_tool(
-                    name=name,
-                    arguments=arguments,
-                )
+        cache_key = cache_key_generator(agent_id, run_id, name, arguments)
 
-            return await _call_tool_internal(name=name, arguments=arguments)
+        # Check cache first
+        if cache_key in self.tool_result_cache:
+            return self.tool_result_cache[cache_key]
 
-        return await _call_tool_cached(
-            agent_id=agent_id, run_id=run_id, name=name, arguments=arguments
-        )
+        async def _call_tool_internal(name: str, arguments: dict[str, Any]) -> CallToolResult:
+            logger.debug(f"Calling tool '{name}' with arguments {arguments}")
+            return await self.session.call_tool(
+                name=name,
+                arguments=arguments,
+            )
+
+        result = await _call_tool_internal(name=name, arguments=arguments)
+        self.tool_result_cache[cache_key] = result
+        return result
 
     async def get_server_name(self) -> str:
         """Return the server_name.
@@ -446,6 +434,14 @@ class FlockMCPClientBase(BaseModel, ABC):
         async with self.lock:
             return self.current_roots
 
+    def _get_roots_no_lock(self) -> list[MCPRoot] | None:
+        """Get the currently set roots without acquiring a lock.
+
+        WARNING: Only use this internally when you're sure there's no race condition.
+        This is primarily for use during initialization when the lock is already held.
+        """
+        return self.current_roots
+
     async def set_roots(self, new_roots: list[MCPRoot]) -> None:
         """Set the current roots of the client.
 
@@ -465,33 +461,23 @@ class FlockMCPClientBase(BaseModel, ABC):
         async with self.lock:
             if self.tool_cache:
                 self.tool_cache.clear()
-                logger.debug(
-                    f"Invalidated tool_cache for server '{self.config.name}'"
-                )
+                logger.debug(f"Invalidated tool_cache for server '{self.config.name}'")
 
     async def invalidate_resource_list_cache(self) -> None:
         """Invalidate the entries in the resource list cache."""
-        logger.debug(
-            f"Invalidating resource_list_cache for server '{self.config.name}'"
-        )
+        logger.debug(f"Invalidating resource_list_cache for server '{self.config.name}'")
         async with self.lock:
             if self.resource_list_cache:
                 self.resource_list_cache.clear()
-                logger.debug(
-                    f"Invalidated resource_list_cache for server '{self.config.name}'"
-                )
+                logger.debug(f"Invalidated resource_list_cache for server '{self.config.name}'")
 
     async def invalidate_resource_contents_cache(self) -> None:
         """Invalidate the entries in the resource contents cache."""
-        logger.debug(
-            f"Invalidating resource_contents_cache for server '{self.config.name}'."
-        )
+        logger.debug(f"Invalidating resource_contents_cache for server '{self.config.name}'.")
         async with self.lock:
             if self.resource_contents_cache:
                 self.resource_contents_cache.clear()
-                logger.debug(
-                    f"Invalidated resource_contents_cache for server '{self.config.name}'"
-                )
+                logger.debug(f"Invalidated resource_contents_cache for server '{self.config.name}'")
 
     async def invalidate_resource_contents_cache_entry(self, key: str) -> None:
         """Invalidate a single entry in the resource contents cache."""
@@ -531,8 +517,7 @@ class FlockMCPClientBase(BaseModel, ABC):
                 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}"
+                    f"Suppressed transport-ctx exit error for server '{self.config.name}': {e!r}"
                 )
 
     async def _create_session(self) -> None:
@@ -548,42 +533,29 @@ class FlockMCPClientBase(BaseModel, ABC):
         server_params = self.config.connection_config.connection_parameters
 
         # Single Hook
-        transport_ctx = await self.create_transport(
-            server_params, self.additional_params
-        )
+        transport_ctx = await self.create_transport(server_params, self.additional_params)
         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
+            read, write, _get_session_id_callback = result
         else:
-            raise RuntimeError(
-                f"create_transport returned unexpected tuple of {result}"
-            )
+            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."
-            )
+            raise RuntimeError("create_transport did not create any read or write streams.")
 
         read_timeout = self.config.connection_config.read_timeout_seconds
 
-        if (
-            self.additional_params
-            and "read_timeout_seconds" in self.additional_params
-        ):
-            read_timeout = self.additional_params.get(
-                "read_timeout_seconds", read_timeout
-            )
+        if self.additional_params and "read_timeout_seconds" in self.additional_params:
+            read_timeout = self.additional_params.get("read_timeout_seconds", read_timeout)
 
         timeout_seconds = (
             read_timeout
@@ -622,11 +594,10 @@ class FlockMCPClientBase(BaseModel, ABC):
                 )
                 return self.client_session
 
-            else:
-                logger.debug(
-                    f"Client Session for Server '{self.config.name}' does not exist yet. Connecting..."
-                )
-                await self._create_session()
+            logger.debug(
+                f"Client Session for Server '{self.config.name}' does not exist yet. Connecting..."
+            )
+            await self._create_session()
 
             if not self.connected_server_capabilities:
                 # This means we never asked the server to initialize the connection.
@@ -636,9 +607,7 @@ class FlockMCPClientBase(BaseModel, ABC):
     async def _perform_initial_handshake(self) -> None:
         """Tell the server who we are, what capabilities we have, and what roots we're interested in."""
         # 1) do the LSP-style initialize handshake
-        logger.debug(
-            f"Performing intialize handshake with server '{self.config.name}'"
-        )
+        logger.debug(f"Performing intialize handshake with server '{self.config.name}'")
         init: InitializeResult = await self.client_session.initialize()
 
         self.connected_server_capabilities = init
@@ -650,6 +619,10 @@ class FlockMCPClientBase(BaseModel, ABC):
         # 2) if we already know our current roots, notify the server
         #    so that it will follow up with a ListRootsRequest
         if self.current_roots and self.config.feature_config.roots_enabled:
+            logger.debug(
+                f"Notifying server '{self.config.name}' of {len(self.current_roots)} root(s): "
+                f"{[r.uri for r in self.current_roots]}"
+            )
             await self.client_session.send_roots_list_changed()
 
         # 3) Tell the server, what logging level we would like to use