fast-agent-mcp 0.1.12__py3-none-any.whl → 0.1.13__py3-none-any.whl

mcp_agent/mcp/mcp_aggregator.py

@@ -1,31 +1,31 @@
 from asyncio import Lock, gather
 from typing import (
-    List,
-    Dict,
-    Optional,
     TYPE_CHECKING,
     Any,
     Callable,
+    Dict,
+    List,
+    Optional,
     TypeVar,
 )
+
 from mcp import GetPromptResult, ReadResourceResult
-from pydantic import AnyUrl, BaseModel, ConfigDict
 from mcp.client.session import ClientSession
 from mcp.server.lowlevel.server import Server
 from mcp.server.stdio import stdio_server
 from mcp.types import (
     CallToolResult,
     ListToolsResult,
+    Prompt,
     TextContent,
     Tool,
-    Prompt,
 )
+from pydantic import AnyUrl, BaseModel, ConfigDict
 
+from mcp_agent.context_dependent import ContextDependent
 from mcp_agent.event_progress import ProgressAction
 from mcp_agent.logging.logger import get_logger
 from mcp_agent.mcp.gen_client import gen_client
-
-from mcp_agent.context_dependent import ContextDependent
 from mcp_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
 from mcp_agent.mcp.mcp_connection_manager import MCPConnectionManager
 
@@ -33,9 +33,7 @@ if TYPE_CHECKING:
     from mcp_agent.context import Context
 
 
-logger = get_logger(
-    __name__
-)  # This will be replaced per-instance when agent_name is available
+logger = get_logger(__name__)  # This will be replaced per-instance when agent_name is available
 
 SEP = "-"
 
@@ -79,9 +77,7 @@ class MCPAggregator(ContextDependent):
         if self.connection_persistence:
             # Try to get existing connection manager from context
             if not hasattr(self.context, "_connection_manager"):
-                self.context._connection_manager = MCPConnectionManager(
-                    self.context.server_registry
-                )
+                self.context._connection_manager = MCPConnectionManager(self.context.server_registry)
                 await self.context._connection_manager.__aenter__()
             self._persistent_connection_manager = self.context._connection_manager
 
@@ -99,7 +95,7 @@ class MCPAggregator(ContextDependent):
         context: Optional["Context"] = None,
         name: str = None,
         **kwargs,
-    ):
+    ) -> None:
         """
         :param server_names: A list of server names to connect to.
         :param connection_persistence: Whether to maintain persistent connections to servers (default: True).
@@ -130,23 +126,17 @@ class MCPAggregator(ContextDependent):
         self._prompt_cache: Dict[str, List[Prompt]] = {}
         self._prompt_cache_lock = Lock()
 
-    async def close(self):
+    async def close(self) -> None:
         """
         Close all persistent connections when the aggregator is deleted.
         """
         if self.connection_persistence and self._persistent_connection_manager:
             try:
                 # Only attempt cleanup if we own the connection manager
-                if (
-                    hasattr(self.context, "_connection_manager")
-                    and self.context._connection_manager
-                    == self._persistent_connection_manager
-                ):
+                if hasattr(self.context, "_connection_manager") and self.context._connection_manager == self._persistent_connection_manager:
                     logger.info("Shutting down all persistent connections...")
                     await self._persistent_connection_manager.disconnect_all()
-                    await self._persistent_connection_manager.__aexit__(
-                        None, None, None
-                    )
+                    await self._persistent_connection_manager.__aexit__(None, None, None)
                     delattr(self.context, "_connection_manager")
                 self.initialized = False
             except Exception as e:
@@ -185,7 +175,7 @@ class MCPAggregator(ContextDependent):
             logger.error(f"Error creating MCPAggregator: {e}")
             await instance.__aexit__(None, None, None)
 
-    async def load_servers(self):
+    async def load_servers(self) -> None:
         """
         Discover tools from each server in parallel and build an index of namespaced tool names.
         Also populate the prompt cache.
@@ -212,9 +202,7 @@ class MCPAggregator(ContextDependent):
                     },
                 )
 
-                await self._persistent_connection_manager.get_server(
-                    server_name, client_session_factory=MCPAgentClientSession
-                )
+                await self._persistent_connection_manager.get_server(server_name, client_session_factory=MCPAgentClientSession)
 
             logger.info(
                 f"MCP Servers initialized for agent '{self.agent_name}'",
@@ -232,9 +220,7 @@ class MCPAggregator(ContextDependent):
                 logger.error(f"Error loading tools from server '{server_name}'", data=e)
                 return []
 
-        async def fetch_prompts(
-            client: ClientSession, server_name: str
-        ) -> List[Prompt]:
+        async def fetch_prompts(client: ClientSession, server_name: str) -> List[Prompt]:
             # Only fetch prompts if the server supports them
             capabilities = await self.get_capabilities(server_name)
             if not capabilities or not capabilities.prompts:
@@ -253,17 +239,11 @@ class MCPAggregator(ContextDependent):
             prompts: List[Prompt] = []
 
             if self.connection_persistence:
-                server_connection = (
-                    await self._persistent_connection_manager.get_server(
-                        server_name, client_session_factory=MCPAgentClientSession
-                    )
-                )
+                server_connection = await self._persistent_connection_manager.get_server(server_name, client_session_factory=MCPAgentClientSession)
                 tools = await fetch_tools(server_connection.session)
                 prompts = await fetch_prompts(server_connection.session, server_name)
             else:
-                async with gen_client(
-                    server_name, server_registry=self.context.server_registry
-                ) as client:
+                async with gen_client(server_name, server_registry=self.context.server_registry) as client:
                     tools = await fetch_tools(client)
                     prompts = await fetch_prompts(client, server_name)
 
@@ -319,9 +299,7 @@ class MCPAggregator(ContextDependent):
             return None
 
         try:
-            server_conn = await self._persistent_connection_manager.get_server(
-                server_name, client_session_factory=MCPAgentClientSession
-            )
+            server_conn = await self._persistent_connection_manager.get_server(server_name, client_session_factory=MCPAgentClientSession)
             # server_capabilities is a property, not a coroutine
             return server_conn.server_capabilities
         except Exception as e:
@@ -383,9 +361,7 @@ class MCPAggregator(ContextDependent):
                 return error_factory(error_msg) if error_factory else None
 
         if self.connection_persistence:
-            server_connection = await self._persistent_connection_manager.get_server(
-                server_name, client_session_factory=MCPAgentClientSession
-            )
+            server_connection = await self._persistent_connection_manager.get_server(server_name, client_session_factory=MCPAgentClientSession)
             return await try_execute(server_connection.session)
         else:
             logger.debug(
@@ -396,9 +372,7 @@ class MCPAggregator(ContextDependent):
                     "agent_name": self.agent_name,
                 },
             )
-            async with gen_client(
-                server_name, server_registry=self.context.server_registry
-            ) as client:
+            async with gen_client(server_name, server_registry=self.context.server_registry) as client:
                 result = await try_execute(client)
                 logger.debug(
                     f"Closing temporary connection to server: {server_name}",
@@ -410,9 +384,7 @@ class MCPAggregator(ContextDependent):
                 )
                 return result
 
-    async def _parse_resource_name(
-        self, name: str, resource_type: str
-    ) -> tuple[str, str]:
+    async def _parse_resource_name(self, name: str, resource_type: str) -> tuple[str, str]:
         """
         Parse a possibly namespaced resource name into server name and local resource name.
 
@@ -447,9 +419,7 @@ class MCPAggregator(ContextDependent):
 
         return server_name, local_name
 
-    async def call_tool(
-        self, name: str, arguments: dict | None = None
-    ) -> CallToolResult:
+    async def call_tool(self, name: str, arguments: dict | None = None) -> CallToolResult:
         """
         Call a namespaced tool, e.g., 'server_name.tool_name'.
         """
@@ -481,14 +451,10 @@ class MCPAggregator(ContextDependent):
             operation_name=local_tool_name,
             method_name="call_tool",
             method_args={"name": local_tool_name, "arguments": arguments},
-            error_factory=lambda msg: CallToolResult(
-                isError=True, content=[TextContent(type="text", text=msg)]
-            ),
+            error_factory=lambda msg: CallToolResult(isError=True, content=[TextContent(type="text", text=msg)]),
         )
 
-    async def get_prompt(
-        self, prompt_name: str = None, arguments: dict[str, str] = None
-    ) -> GetPromptResult:
+    async def get_prompt(self, prompt_name: str | None, arguments: dict[str, str] | None) -> GetPromptResult:
         """
         Get a prompt from a server.
 
@@ -540,13 +506,9 @@ class MCPAggregator(ContextDependent):
                 async with self._prompt_cache_lock:
                     if server_name in self._prompt_cache:
                         # Check if any prompt in the cache has this name
-                        prompt_names = [
-                            prompt.name for prompt in self._prompt_cache[server_name]
-                        ]
+                        prompt_names = [prompt.name for prompt in self._prompt_cache[server_name]]
                         if local_prompt_name not in prompt_names:
-                            logger.debug(
-                                f"Prompt '{local_prompt_name}' not found in cache for server '{server_name}'"
-                            )
+                            logger.debug(f"Prompt '{local_prompt_name}' not found in cache for server '{server_name}'")
                             return GetPromptResult(
                                 description=f"Prompt '{local_prompt_name}' not found on server '{server_name}'",
                                 messages=[],
@@ -568,9 +530,7 @@ class MCPAggregator(ContextDependent):
 
             # Add namespaced name and source server to the result
             if result and result.messages:
-                result.namespaced_name = (
-                    namespaced_name or f"{server_name}{SEP}{local_prompt_name}"
-                )
+                result.namespaced_name = namespaced_name or f"{server_name}{SEP}{local_prompt_name}"
 
                 # Store the arguments in the result for display purposes
                 if arguments:
@@ -590,18 +550,14 @@ class MCPAggregator(ContextDependent):
                     potential_servers.append(s_name)
 
         if potential_servers:
-            logger.debug(
-                f"Found prompt '{local_prompt_name}' in cache for servers: {potential_servers}"
-            )
+            logger.debug(f"Found prompt '{local_prompt_name}' in cache for servers: {potential_servers}")
 
             # Try each server from the cache
             for s_name in potential_servers:
                 # Check if this server supports prompts
                 capabilities = await self.get_capabilities(s_name)
                 if not capabilities or not capabilities.prompts:
-                    logger.debug(
-                        f"Server '{s_name}' does not support prompts, skipping"
-                    )
+                    logger.debug(f"Server '{s_name}' does not support prompts, skipping")
                     continue
 
                 try:
@@ -620,9 +576,7 @@ class MCPAggregator(ContextDependent):
 
                     # If we got a successful result with messages, return it
                     if result and result.messages:
-                        logger.debug(
-                            f"Successfully retrieved prompt '{local_prompt_name}' from server '{s_name}'"
-                        )
+                        logger.debug(f"Successfully retrieved prompt '{local_prompt_name}' from server '{s_name}'")
                         # Add namespaced name using the actual server where found
                         result.namespaced_name = f"{s_name}{SEP}{local_prompt_name}"
 
@@ -635,9 +589,7 @@ class MCPAggregator(ContextDependent):
                 except Exception as e:
                     logger.debug(f"Error retrieving prompt from server '{s_name}': {e}")
         else:
-            logger.debug(
-                f"Prompt '{local_prompt_name}' not found in any server's cache"
-            )
+            logger.debug(f"Prompt '{local_prompt_name}' not found in any server's cache")
 
             # If not in cache, perform a full search as fallback (cache might be outdated)
             # First identify servers that support prompts
@@ -647,9 +599,7 @@ class MCPAggregator(ContextDependent):
                 if capabilities and capabilities.prompts:
                     supported_servers.append(s_name)
                 else:
-                    logger.debug(
-                        f"Server '{s_name}' does not support prompts, skipping from fallback search"
-                    )
+                    logger.debug(f"Server '{s_name}' does not support prompts, skipping from fallback search")
 
             # Try all supported servers in order
             for s_name in supported_servers:
@@ -670,9 +620,7 @@ class MCPAggregator(ContextDependent):
 
                     # If we got a successful result with messages, return it
                     if result and result.messages:
-                        logger.debug(
-                            f"Found prompt '{local_prompt_name}' on server '{s_name}' (not in cache)"
-                        )
+                        logger.debug(f"Found prompt '{local_prompt_name}' on server '{s_name}' (not in cache)")
                         # Add namespaced name using the actual server where found
                         result.namespaced_name = f"{s_name}{SEP}{local_prompt_name}"
 
@@ -691,21 +639,15 @@ class MCPAggregator(ContextDependent):
                             )
 
                             prompts = getattr(prompt_list_result, "prompts", [])
-                            matching_prompts = [
-                                p for p in prompts if p.name == local_prompt_name
-                            ]
+                            matching_prompts = [p for p in prompts if p.name == local_prompt_name]
                             if matching_prompts:
                                 async with self._prompt_cache_lock:
                                     if s_name not in self._prompt_cache:
                                         self._prompt_cache[s_name] = []
                                     # Add if not already in the cache
-                                    prompt_names_in_cache = [
-                                        p.name for p in self._prompt_cache[s_name]
-                                    ]
+                                    prompt_names_in_cache = [p.name for p in self._prompt_cache[s_name]]
                                     if local_prompt_name not in prompt_names_in_cache:
-                                        self._prompt_cache[s_name].append(
-                                            matching_prompts[0]
-                                        )
+                                        self._prompt_cache[s_name].append(matching_prompts[0])
                         except Exception:
                             # Ignore errors when updating cache
                             pass
@@ -754,9 +696,7 @@ class MCPAggregator(ContextDependent):
                 async with self._prompt_cache_lock:
                     if server_name in self._prompt_cache:
                         results[server_name] = self._prompt_cache[server_name]
-                        logger.debug(
-                            f"Returning cached prompts for server '{server_name}'"
-                        )
+                        logger.debug(f"Returning cached prompts for server '{server_name}'")
                         return results
 
                 # Check if server supports prompts
@@ -790,9 +730,7 @@ class MCPAggregator(ContextDependent):
                 if capabilities and capabilities.prompts:
                     supported_servers.append(s_name)
                 else:
-                    logger.debug(
-                        f"Server '{s_name}' does not support prompts, skipping"
-                    )
+                    logger.debug(f"Server '{s_name}' does not support prompts, skipping")
                     # Add empty list to results for this server
                     results[s_name] = []
 
@@ -828,9 +766,7 @@ class MCPAggregator(ContextDependent):
         logger.debug(f"Available prompts across servers: {results}")
         return results
 
-    async def get_resource(
-        self, server_name: str, resource_uri: str
-    ) -> ReadResourceResult:
+    async def get_resource(self, server_name: str, resource_uri: str) -> ReadResourceResult:
         """
         Get a resource directly from an MCP server by URI.
 
@@ -881,7 +817,7 @@ class MCPCompoundServer(Server):
     A compound server (server-of-servers) that aggregates multiple MCP servers and is itself an MCP server
     """
 
-    def __init__(self, server_names: List[str], name: str = "MCPCompoundServer"):
+    def __init__(self, server_names: List[str], name: str = "MCPCompoundServer") -> None:
         super().__init__(name)
         self.aggregator = MCPAggregator(server_names)
 
@@ -896,9 +832,7 @@ class MCPCompoundServer(Server):
         tools_result = await self.aggregator.list_tools()
         return tools_result.tools
 
-    async def _call_tool(
-        self, name: str, arguments: dict | None = None
-    ) -> CallToolResult:
+    async def _call_tool(self, name: str, arguments: dict | None = None) -> CallToolResult:
         """Call a specific tool from the aggregated servers."""
         try:
             result = await self.aggregator.call_tool(name=name, arguments=arguments)
@@ -909,9 +843,7 @@ class MCPCompoundServer(Server):
                 content=[TextContent(type="text", text=f"Error calling tool: {e}")],
             )
 
-    async def _get_prompt(
-        self, name: str = None, arguments: dict[str, str] = None
-    ) -> GetPromptResult:
+    async def _get_prompt(self, name: str = None, arguments: dict[str, str] = None) -> GetPromptResult:
         """
         Get a prompt from the aggregated servers.
 
@@ -920,14 +852,10 @@ class MCPCompoundServer(Server):
             arguments: Optional dictionary of string arguments for prompt templating
         """
         try:
-            result = await self.aggregator.get_prompt(
-                prompt_name=name, arguments=arguments
-            )
+            result = await self.aggregator.get_prompt(prompt_name=name, arguments=arguments)
             return result
         except Exception as e:
-            return GetPromptResult(
-                description=f"Error getting prompt: {e}", messages=[]
-            )
+            return GetPromptResult(description=f"Error getting prompt: {e}", messages=[])
 
     async def _list_prompts(self, server_name: str = None) -> Dict[str, List[str]]:
         """List available prompts from the aggregated servers."""

mcp_agent/mcp/mcp_connection_manager.py

@@ -2,38 +2,37 @@
 Manages the lifecycle of multiple MCP server connections.
 """
 
-from datetime import timedelta
 import asyncio
+from datetime import timedelta
 from typing import (
+    TYPE_CHECKING,
     AsyncGenerator,
     Callable,
     Dict,
     Optional,
-    TYPE_CHECKING,
 )
 
-from anyio import Event, create_task_group, Lock
+from anyio import Event, Lock, create_task_group
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-
 from mcp import ClientSession
+from mcp.client.sse import sse_client
 from mcp.client.stdio import (
     StdioServerParameters,
     get_default_environment,
 )
-from mcp.client.sse import sse_client
 from mcp.types import JSONRPCMessage, ServerCapabilities
 
 from mcp_agent.config import MCPServerSettings
+from mcp_agent.context_dependent import ContextDependent
 from mcp_agent.core.exceptions import ServerInitializationError
 from mcp_agent.event_progress import ProgressAction
 from mcp_agent.logging.logger import get_logger
 from mcp_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
 from mcp_agent.mcp.stdio import stdio_client_with_rich_stderr
-from mcp_agent.context_dependent import ContextDependent
 
 if TYPE_CHECKING:
-    from mcp_agent.mcp_server_registry import InitHookCallable, ServerRegistry
     from mcp_agent.context import Context
+    from mcp_agent.mcp_server_registry import InitHookCallable, ServerRegistry
 
 logger = get_logger(__name__)
 
@@ -64,7 +63,7 @@ class ServerConnection:
             ClientSession,
         ],
         init_hook: Optional["InitHookCallable"] = None,
-    ):
+    ) -> None:
         self.server_name = server_name
         self.server_config = server_config
         self.session: ClientSession | None = None
@@ -134,11 +133,7 @@ class ServerConnection:
         Create a new session instance for this server connection.
         """
 
-        read_timeout = (
-            timedelta(seconds=self.server_config.read_timeout_seconds)
-            if self.server_config.read_timeout_seconds
-            else None
-        )
+        read_timeout = timedelta(seconds=self.server_config.read_timeout_seconds) if self.server_config.read_timeout_seconds else None
 
         session = self._client_session_factory(read_stream, send_stream, read_timeout)
 
@@ -192,9 +187,7 @@ class MCPConnectionManager(ContextDependent):
     Integrates with the application context system for proper resource management.
     """
 
-    def __init__(
-        self, server_registry: "ServerRegistry", context: Optional["Context"] = None
-    ):
+    def __init__(self, server_registry: "ServerRegistry", context: Optional["Context"] = None) -> None:
         super().__init__(context=context)
         self.server_registry = server_registry
         self.running_servers: Dict[str, ServerConnection] = {}
@@ -255,9 +248,7 @@ class MCPConnectionManager(ContextDependent):
         if not config:
             raise ValueError(f"Server '{server_name}' not found in registry.")
 
-        logger.debug(
-            f"{server_name}: Found server configuration=", data=config.model_dump()
-        )
+        logger.debug(f"{server_name}: Found server configuration=", data=config.model_dump())
 
         def transport_context_factory():
             if config.transport == "stdio":
@@ -309,9 +300,7 @@ class MCPConnectionManager(ContextDependent):
 
             # If server exists but isn't healthy, remove it so we can create a new one
             if server_conn:
-                logger.info(
-                    f"{server_name}: Server exists but is unhealthy, recreating..."
-                )
+                logger.info(f"{server_name}: Server exists but is unhealthy, recreating...")
                 self.running_servers.pop(server_name)
                 server_conn.request_shutdown()
 
@@ -328,19 +317,13 @@ class MCPConnectionManager(ContextDependent):
         # Check if the server is healthy after initialization
         if not server_conn.is_healthy():
             error_msg = server_conn._error_message or "Unknown error"
-            raise ServerInitializationError(
-                f"MCP Server: '{server_name}': Failed to initialize with error: '{error_msg}'. Check fastagent.config.yaml"
-            )
+            raise ServerInitializationError(f"MCP Server: '{server_name}': Failed to initialize with error: '{error_msg}'. Check fastagent.config.yaml")
 
         return server_conn
 
-    async def get_server_capabilities(
-        self, server_name: str
-    ) -> ServerCapabilities | None:
+    async def get_server_capabilities(self, server_name: str) -> ServerCapabilities | None:
         """Get the capabilities of a specific server."""
-        server_conn = await self.get_server(
-            server_name, client_session_factory=MCPAgentClientSession
-        )
+        server_conn = await self.get_server(server_name, client_session_factory=MCPAgentClientSession)
         return server_conn.server_capabilities if server_conn else None
 
     async def disconnect_server(self, server_name: str) -> None:
@@ -353,13 +336,9 @@ class MCPConnectionManager(ContextDependent):
             server_conn = self.running_servers.pop(server_name, None)
         if server_conn:
             server_conn.request_shutdown()
-            logger.info(
-                f"{server_name}: Shutdown signal sent (lifecycle task will exit)."
-            )
+            logger.info(f"{server_name}: Shutdown signal sent (lifecycle task will exit).")
         else:
-            logger.info(
-                f"{server_name}: No persistent connection found. Skipping server shutdown"
-            )
+            logger.info(f"{server_name}: No persistent connection found. Skipping server shutdown")
 
     async def disconnect_all(self) -> None:
         """Disconnect all servers that are running under this connection manager."""

mcp_agent/mcp/prompt_message_multipart.py

@@ -1,14 +1,14 @@
 from typing import List, Union
-from pydantic import BaseModel
 
 from mcp.types import (
-    PromptMessage,
-    TextContent,
-    ImageContent,
     EmbeddedResource,
-    Role,
     GetPromptResult,
+    ImageContent,
+    PromptMessage,
+    Role,
+    TextContent,
 )
+from pydantic import BaseModel
 
 
 class PromptMessageMultipart(BaseModel):
@@ -21,9 +21,7 @@ class PromptMessageMultipart(BaseModel):
     content: List[Union[TextContent, ImageContent, EmbeddedResource]]
 
     @classmethod
-    def from_prompt_messages(
-        cls, messages: List[PromptMessage]
-    ) -> List["PromptMessageMultipart"]:
+    def to_multipart(cls, messages: List[PromptMessage]) -> List["PromptMessageMultipart"]:
         """Convert a sequence of PromptMessages into PromptMessageMultipart objects."""
         if not messages:
             return []
@@ -41,7 +39,8 @@ class PromptMessageMultipart(BaseModel):
                 current_group = cls(role=msg.role, content=[msg.content])
             else:
                 # Same role, add to current message
-                current_group.content.append(msg.content)
+                if current_group is not None:
+                    current_group.content.append(msg.content)
 
         # Add the last group
         if current_group is not None:
@@ -49,16 +48,11 @@ class PromptMessageMultipart(BaseModel):
 
         return result
 
-    def to_prompt_messages(self) -> List[PromptMessage]:
+    def from_multipart(self) -> List[PromptMessage]:
         """Convert this PromptMessageMultipart to a sequence of standard PromptMessages."""
-        return [
-            PromptMessage(role=self.role, content=content_part)
-            for content_part in self.content
-        ]
+        return [PromptMessage(role=self.role, content=content_part) for content_part in self.content]
 
     @classmethod
-    def parse_get_prompt_result(
-        cls, result: GetPromptResult
-    ) -> List["PromptMessageMultipart"]:
+    def parse_get_prompt_result(cls, result: GetPromptResult) -> List["PromptMessageMultipart"]:
         """Parse a GetPromptResult into PromptMessageMultipart objects."""
-        return cls.from_prompt_messages(result.messages)
+        return cls.to_multipart(result.messages)