fast-agent-mcp 0.1.12__py3-none-any.whl → 0.2.0__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 = "-"
 
@@ -99,7 +97,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,7 +128,7 @@ 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.
         """
@@ -139,14 +137,11 @@ class MCPAggregator(ContextDependent):
                 # Only attempt cleanup if we own the connection manager
                 if (
                     hasattr(self.context, "_connection_manager")
-                    and self.context._connection_manager
-                    == self._persistent_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 +180,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.
@@ -232,9 +227,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,10 +246,8 @@ 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)
@@ -378,7 +369,9 @@ class MCPAggregator(ContextDependent):
                 method = getattr(client, method_name)
                 return await method(**method_args)
             except Exception as e:
-                error_msg = f"Failed to {method_name} '{operation_name}' on server '{server_name}': {e}"
+                error_msg = (
+                    f"Failed to {method_name} '{operation_name}' on server '{server_name}': {e}"
+                )
                 logger.error(error_msg)
                 return error_factory(error_msg) if error_factory else None
 
@@ -410,9 +403,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 +438,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'.
         """
@@ -487,7 +476,7 @@ class MCPAggregator(ContextDependent):
         )
 
     async def get_prompt(
-        self, prompt_name: str = None, arguments: dict[str, str] = None
+        self, prompt_name: str | None, arguments: dict[str, str] | None
     ) -> GetPromptResult:
         """
         Get a prompt from a server.
@@ -540,9 +529,7 @@ 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}'"
@@ -568,9 +555,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:
@@ -599,9 +584,7 @@ class MCPAggregator(ContextDependent):
                 # 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:
@@ -635,9 +618,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
@@ -691,9 +672,7 @@ 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:
@@ -703,9 +682,7 @@ class MCPAggregator(ContextDependent):
                                         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
@@ -723,114 +700,102 @@ class MCPAggregator(ContextDependent):
             messages=[],
         )
 
-    async def list_prompts(self, server_name: str = None):
+    async def list_prompts(self, server_name: str = None) -> Dict[str, List[Prompt]]:
         """
         List available prompts from one or all servers.
 
         :param server_name: Optional server name to list prompts from. If not provided,
                            lists prompts from all servers.
-        :return: Dictionary mapping server names to lists of available prompts
+        :return: Dictionary mapping server names to lists of Prompt objects
         """
         if not self.initialized:
             await self.load_servers()
 
-        results = {}
+        results: Dict[str, List[Prompt]] = {}
+
+        # If specific server requested
+        if server_name:
+            if server_name not in self.server_names:
+                logger.error(f"Server '{server_name}' not found")
+                return results
 
-        # If we already have the data in cache and not requesting a specific server,
-        # we can use the cache directly
-        if not server_name:
+            # Check cache first
             async with self._prompt_cache_lock:
-                if all(s_name in self._prompt_cache for s_name in self.server_names):
-                    # Return the cached prompt objects
-                    for s_name, prompt_list in self._prompt_cache.items():
-                        results[s_name] = prompt_list
-                    logger.debug("Returning cached prompts for all servers")
+                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}'")
                     return results
 
-        # If server_name is provided, only list prompts from that server
-        if server_name:
-            if server_name in self.server_names:
-                # Check if we can use the cache
-                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}'"
-                        )
-                        return results
+            # Check if server supports prompts
+            capabilities = await self.get_capabilities(server_name)
+            if not capabilities or not capabilities.prompts:
+                logger.debug(f"Server '{server_name}' does not support prompts")
+                results[server_name] = []
+                return results
 
-                # Check if server supports prompts
-                capabilities = await self.get_capabilities(server_name)
-                if not capabilities or not capabilities.prompts:
-                    logger.debug(f"Server '{server_name}' does not support prompts")
-                    results[server_name] = []
-                    return results
+            # Fetch from server
+            result = await self._execute_on_server(
+                server_name=server_name,
+                operation_type="prompts-list",
+                operation_name="",
+                method_name="list_prompts",
+                error_factory=lambda _: None,
+            )
+
+            # Get prompts from result
+            prompts = getattr(result, "prompts", [])
+
+            # Update cache
+            async with self._prompt_cache_lock:
+                self._prompt_cache[server_name] = prompts
+
+            results[server_name] = prompts
+            return results
+
+        # No specific server - check if we can use the cache for all servers
+        async with self._prompt_cache_lock:
+            if all(s_name in self._prompt_cache for s_name in self.server_names):
+                for s_name, prompt_list in self._prompt_cache.items():
+                    results[s_name] = prompt_list
+                logger.debug("Returning cached prompts for all servers")
+                return results
+
+        # Identify servers that support prompts
+        supported_servers = []
+        for s_name in self.server_names:
+            capabilities = await self.get_capabilities(s_name)
+            if capabilities and capabilities.prompts:
+                supported_servers.append(s_name)
+            else:
+                logger.debug(f"Server '{s_name}' does not support prompts, skipping")
+                results[s_name] = []
 
-                # If not in cache and server supports prompts, fetch from server
+        # Fetch prompts from supported servers
+        for s_name in supported_servers:
+            try:
                 result = await self._execute_on_server(
-                    server_name=server_name,
+                    server_name=s_name,
                     operation_type="prompts-list",
                     operation_name="",
                     method_name="list_prompts",
-                    error_factory=lambda _: [],
+                    error_factory=lambda _: None,
                 )
 
-                # Update cache with the result
-                async with self._prompt_cache_lock:
-                    self._prompt_cache[server_name] = getattr(result, "prompts", [])
+                prompts = getattr(result, "prompts", [])
 
-                results[server_name] = result
-            else:
-                logger.error(f"Server '{server_name}' not found")
-        else:
-            # We need to filter the servers that support prompts
-            supported_servers = []
-            for s_name in self.server_names:
-                capabilities = await self.get_capabilities(s_name)
-                if capabilities and capabilities.prompts:
-                    supported_servers.append(s_name)
-                else:
-                    logger.debug(
-                        f"Server '{s_name}' does not support prompts, skipping"
-                    )
-                    # Add empty list to results for this server
-                    results[s_name] = []
-
-            # Process servers sequentially to ensure proper resource cleanup
-            # This helps prevent resource leaks especially on Windows
-            if supported_servers:
-                server_results = []
-                for s_name in supported_servers:
-                    try:
-                        result = await self._execute_on_server(
-                            server_name=s_name,
-                            operation_type="prompts-list",
-                            operation_name="",
-                            method_name="list_prompts",
-                            error_factory=lambda _: [],
-                        )
-                        server_results.append(result)
-                    except Exception as e:
-                        logger.debug(f"Error fetching prompts from {s_name}: {e}")
-                        server_results.append(e)
-
-                for i, result in enumerate(server_results):
-                    if isinstance(result, BaseException):
-                        continue
-
-                    s_name = supported_servers[i]
-                    results[s_name] = result
+                # Update cache and results
+                async with self._prompt_cache_lock:
+                    self._prompt_cache[s_name] = prompts
 
-                    # Update cache with the result
-                    async with self._prompt_cache_lock:
-                        self._prompt_cache[s_name] = getattr(result, "prompts", [])
+                results[s_name] = prompts
+            except Exception as e:
+                logger.debug(f"Error fetching prompts from {s_name}: {e}")
+                results[s_name] = []
 
         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 +846,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 +861,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)
@@ -920,16 +883,12 @@ 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]]:
+    async def _list_prompts(self, server_name: str = None) -> Dict[str, List[Prompt]]:
         """List available prompts from the aggregated servers."""
         try:
             return await self.aggregator.list_prompts(server_name=server_name)

mcp_agent/mcp/mcp_connection_manager.py

@@ -2,38 +2,38 @@
 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,
+    stdio_client,
 )
-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.logger_textio import get_stderr_handler
 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 +64,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
@@ -194,7 +194,7 @@ class MCPConnectionManager(ContextDependent):
 
     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 +255,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":
@@ -266,8 +264,11 @@ class MCPConnectionManager(ContextDependent):
                     args=config.args,
                     env={**get_default_environment(), **(config.env or {})},
                 )
-                # Create stdio client config with redirected stderr
-                return stdio_client_with_rich_stderr(server_params)
+                # Create custom error handler to ensure all output is captured
+                error_handler = get_stderr_handler(server_name)
+                # Explicitly ensure we're using our custom logger for stderr
+                logger.debug(f"{server_name}: Creating stdio client with custom error handler")
+                return stdio_client(server_params, errlog=error_handler)
             elif config.transport == "sse":
                 return sse_client(config.url)
             else:
@@ -309,9 +310,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()
 
@@ -334,9 +333,7 @@ class MCPConnectionManager(ContextDependent):
 
         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
@@ -353,13 +350,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."""