fast-agent-mcp 0.0.7__py3-none-any.whl

mcp_agent/mcp/mcp_agent_client_session.py

@@ -0,0 +1,242 @@
+"""
+A derived client session for the MCP Agent framework.
+It adds logging and supports sampling requests.
+"""
+
+from typing import Optional
+
+from mcp import ClientSession
+from mcp.shared.session import (
+    RequestResponder,
+    ReceiveResultT,
+    ReceiveNotificationT,
+    RequestId,
+    SendNotificationT,
+    SendRequestT,
+    SendResultT,
+)
+from mcp.types import (
+    ClientResult,
+    CreateMessageRequest,
+    CreateMessageResult,
+    ErrorData,
+    JSONRPCNotification,
+    JSONRPCRequest,
+    ServerRequest,
+    TextContent,
+    ListRootsRequest,
+    ListRootsResult,
+    Root,
+)
+
+from mcp_agent.config import MCPServerSettings
+from mcp_agent.context_dependent import ContextDependent
+from mcp_agent.logging.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class MCPAgentClientSession(ClientSession, ContextDependent):
+    """
+    MCP Agent framework acts as a client to the servers providing tools/resources/prompts for the agent workloads.
+    This is a simple client session for those server connections, and supports
+        - handling sampling requests
+        - notifications
+        - MCP root configuration
+
+    Developers can extend this class to add more custom functionality as needed
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.server_config: Optional[MCPServerSettings] = None
+
+    async def _received_request(
+        self, responder: RequestResponder[ServerRequest, ClientResult]
+    ) -> None:
+        logger.debug("Received request:", data=responder.request.model_dump())
+        request = responder.request.root
+
+        if isinstance(request, CreateMessageRequest):
+            return await self.handle_sampling_request(request, responder)
+        elif isinstance(request, ListRootsRequest):
+            # Handle list_roots request by returning configured roots
+            if hasattr(self, "server_config") and self.server_config.roots:
+                roots = [
+                    Root(
+                        uri=root.server_uri_alias or root.uri,
+                        name=root.name,
+                    )
+                    for root in self.server_config.roots
+                ]
+
+                await responder.respond(ListRootsResult(roots=roots))
+            else:
+                await responder.respond(ListRootsResult(roots=[]))
+            return
+
+        # Handle other requests as usual
+        await super()._received_request(responder)
+
+    async def send_request(
+        self,
+        request: SendRequestT,
+        result_type: type[ReceiveResultT],
+    ) -> ReceiveResultT:
+        logger.debug("send_request: request=", data=request.model_dump())
+        try:
+            result = await super().send_request(request, result_type)
+            logger.debug("send_request: response=", data=result.model_dump())
+            return result
+        except Exception as e:
+            logger.error(f"send_request failed: {e}")
+            raise
+
+    async def send_notification(self, notification: SendNotificationT) -> None:
+        logger.debug("send_notification:", data=notification.model_dump())
+        try:
+            return await super().send_notification(notification)
+        except Exception as e:
+            logger.error("send_notification failed", data=e)
+            raise
+
+    async def _send_response(
+        self, request_id: RequestId, response: SendResultT | ErrorData
+    ) -> None:
+        logger.debug(
+            f"send_response: request_id={request_id}, response=",
+            data=response.model_dump(),
+        )
+        return await super()._send_response(request_id, response)
+
+    async def _received_notification(self, notification: ReceiveNotificationT) -> None:
+        """
+        Can be overridden by subclasses to handle a notification without needing
+        to listen on the message stream.
+        """
+        logger.info(
+            "_received_notification: notification=",
+            data=notification.model_dump(),
+        )
+        return await super()._received_notification(notification)
+
+    async def send_progress_notification(
+        self, progress_token: str | int, progress: float, total: float | None = None
+    ) -> None:
+        """
+        Sends a progress notification for a request that is currently being
+        processed.
+        """
+        logger.debug(
+            "send_progress_notification: progress_token={progress_token}, progress={progress}, total={total}"
+        )
+        return await super().send_progress_notification(
+            progress_token=progress_token, progress=progress, total=total
+        )
+
+    async def _receive_loop(self) -> None:
+        async with (
+            self._read_stream,
+            self._write_stream,
+            self._incoming_message_stream_writer,
+        ):
+            async for message in self._read_stream:
+                if isinstance(message, Exception):
+                    await self._incoming_message_stream_writer.send(message)
+                elif isinstance(message.root, JSONRPCRequest):
+                    validated_request = self._receive_request_type.model_validate(
+                        message.root.model_dump(
+                            by_alias=True, mode="json", exclude_none=True
+                        )
+                    )
+                    responder = RequestResponder(
+                        request_id=message.root.id,
+                        request_meta=validated_request.root.params.meta
+                        if validated_request.root.params
+                        else None,
+                        request=validated_request,
+                        session=self,
+                    )
+
+                    await self._received_request(responder)
+                    if not responder._responded:
+                        await self._incoming_message_stream_writer.send(responder)
+                elif isinstance(message.root, JSONRPCNotification):
+                    notification = self._receive_notification_type.model_validate(
+                        message.root.model_dump(
+                            by_alias=True, mode="json", exclude_none=True
+                        )
+                    )
+
+                    await self._received_notification(notification)
+                    await self._incoming_message_stream_writer.send(notification)
+                else:  # Response or error
+                    stream = self._response_streams.pop(message.root.id, None)
+                    if stream:
+                        await stream.send(message.root)
+                    else:
+                        await self._incoming_message_stream_writer.send(
+                            RuntimeError(
+                                "Received response with an unknown "
+                                f"request ID: {message}"
+                            )
+                        )
+
+    async def handle_sampling_request(
+        self,
+        request: CreateMessageRequest,
+        responder: RequestResponder[ServerRequest, ClientResult],
+    ):
+        logger.info("Handling sampling request: %s", request)
+        config = self.context.config
+        session = self.context.upstream_session
+        if session is None:
+            # TODO: saqadri - consider whether we should be handling the sampling request here as a client
+            logger.warning(
+                "Error: No upstream client available for sampling requests. Request:",
+                data=request,
+            )
+            try:
+                from anthropic import AsyncAnthropic
+
+                client = AsyncAnthropic(api_key=config.anthropic.api_key)
+
+                params = request.params
+                response = await client.messages.create(
+                    model="claude-3-sonnet-20240229",
+                    max_tokens=params.maxTokens,
+                    messages=[
+                        {
+                            "role": m.role,
+                            "content": m.content.text
+                            if hasattr(m.content, "text")
+                            else m.content.data,
+                        }
+                        for m in params.messages
+                    ],
+                    system=getattr(params, "systemPrompt", None),
+                    temperature=getattr(params, "temperature", 0.7),
+                    stop_sequences=getattr(params, "stopSequences", None),
+                )
+
+                await responder.respond(
+                    CreateMessageResult(
+                        model="claude-3-sonnet-20240229",
+                        role="assistant",
+                        content=TextContent(type="text", text=response.content[0].text),
+                    )
+                )
+            except Exception as e:
+                logger.error(f"Error handling sampling request: {e}")
+                await responder.respond(ErrorData(code=-32603, message=str(e)))
+        else:
+            try:
+                # If a session is available, we'll pass-through the sampling request to the upstream client
+                result = await session.send_request(
+                    request=ServerRequest(request), result_type=CreateMessageResult
+                )
+
+                # Pass the result from the upstream client back to the server. We just act as a pass-through client here.
+                await responder.send_result(result)
+            except Exception as e:
+                await responder.send_error(code=-32603, message=str(e))

mcp_agent/mcp/mcp_agent_server.py

@@ -0,0 +1,56 @@
+import asyncio
+from mcp.server import NotificationOptions
+from mcp.server.fastmcp import FastMCP
+from mcp.server.stdio import stdio_server
+from mcp_agent.executor.temporal import get_temporal_client
+from mcp_agent.telemetry.tracing import setup_tracing
+
+app = FastMCP("mcp-agent-server")
+
+setup_tracing("mcp-agent-server")
+
+
+async def run():
+    async with stdio_server() as (read_stream, write_stream):
+        await app._mcp_server.run(
+            read_stream,
+            write_stream,
+            app._mcp_server.create_initialization_options(
+                notification_options=NotificationOptions(
+                    tools_changed=True, resources_changed=True
+                )
+            ),
+        )
+
+
+@app.tool
+async def run_workflow(query: str):
+    """Run the workflow given its name or id"""
+    pass
+
+
+@app.tool
+async def pause_workflow(workflow_id: str):
+    """Pause a running workflow."""
+    temporal_client = await get_temporal_client()
+    handle = temporal_client.get_workflow_handle(workflow_id)
+    await handle.signal("pause")
+
+
+@app.tool
+async def resume_workflow(workflow_id: str):
+    """Resume a paused workflow."""
+    temporal_client = await get_temporal_client()
+    handle = temporal_client.get_workflow_handle(workflow_id)
+    await handle.signal("resume")
+
+
+async def provide_user_input(workflow_id: str, input_data: str):
+    """Provide user/human input to a waiting workflow step."""
+    temporal_client = await get_temporal_client()
+    handle = temporal_client.get_workflow_handle(workflow_id)
+    await handle.signal("human_input", input_data)
+
+
+if __name__ == "__main__":
+    asyncio.run(run())

mcp_agent/mcp/mcp_aggregator.py

@@ -0,0 +1,394 @@
+from asyncio import Lock, gather
+from typing import List, Dict, Optional, TYPE_CHECKING
+
+from pydantic import 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,
+    Tool,
+)
+
+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
+
+if TYPE_CHECKING:
+    from mcp_agent.context import Context
+
+
+logger = get_logger(
+    __name__
+)  # This will be replaced per-instance when agent_name is available
+
+SEP = "-"
+
+
+class NamespacedTool(BaseModel):
+    """
+    A tool that is namespaced by server name.
+    """
+
+    tool: Tool
+    server_name: str
+    namespaced_tool_name: str
+
+
+class MCPAggregator(ContextDependent):
+    """
+    Aggregates multiple MCP servers. When a developer calls, e.g. call_tool(...),
+    the aggregator searches all servers in its list for a server that provides that tool.
+    """
+
+    initialized: bool = False
+    """Whether the aggregator has been initialized with tools and resources from all servers."""
+
+    connection_persistence: bool = False
+    """Whether to maintain a persistent connection to the server."""
+
+    server_names: List[str]
+    """A list of server names to connect to."""
+
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
+
+    async def __aenter__(self):
+        if self.initialized:
+            return self
+
+        # Keep a connection manager to manage persistent connections for this aggregator
+        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
+                )
+                await self.context._connection_manager.__aenter__()
+            self._persistent_connection_manager = self.context._connection_manager
+
+        await self.load_servers()
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+    def __init__(
+        self,
+        server_names: List[str],
+        connection_persistence: bool = False,
+        context: Optional["Context"] = None,
+        name: str = None,
+        **kwargs,
+    ):
+        """
+        :param server_names: A list of server names to connect to.
+        Note: The server names must be resolvable by the gen_client function, and specified in the server registry.
+        """
+        super().__init__(
+            context=context,
+            **kwargs,
+        )
+
+        self.server_names = server_names
+        self.connection_persistence = connection_persistence
+        self.agent_name = name
+        self._persistent_connection_manager: MCPConnectionManager = None
+
+        # Set up logger with agent name in namespace if available
+        global logger
+        logger_name = f"{__name__}.{name}" if name else __name__
+        logger = get_logger(logger_name)
+
+        # Maps namespaced_tool_name -> namespaced tool info
+        self._namespaced_tool_map: Dict[str, NamespacedTool] = {}
+        # Maps server_name -> list of tools
+        self._server_to_tool_map: Dict[str, List[NamespacedTool]] = {}
+        self._tool_map_lock = Lock()
+
+        # TODO: saqadri - add resources and prompt maps as well
+
+    async def close(self):
+        """
+        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
+                ):
+                    logger.info("Shutting down all persistent connections...")
+                    await self._persistent_connection_manager.disconnect_all()
+                    await self._persistent_connection_manager.__aexit__(
+                        None, None, None
+                    )
+                    delattr(self.context, "_connection_manager")
+                self.initialized = False
+            except Exception as e:
+                logger.error(f"Error during connection manager cleanup: {e}")
+
+    @classmethod
+    async def create(
+        cls,
+        server_names: List[str],
+        connection_persistence: bool = False,
+    ) -> "MCPAggregator":
+        """
+        Factory method to create and initialize an MCPAggregator.
+        Use this instead of constructor since we need async initialization.
+        If connection_persistence is True, the aggregator will maintain a
+        persistent connection to the servers for as long as this aggregator is around.
+        By default we do not maintain a persistent connection.
+        """
+
+        logger.info(f"Creating MCPAggregator with servers: {server_names}")
+
+        instance = cls(
+            server_names=server_names,
+            connection_persistence=connection_persistence,
+        )
+
+        try:
+            await instance.__aenter__()
+
+            logger.debug("Loading servers...")
+            await instance.load_servers()
+
+            logger.debug("MCPAggregator created and initialized.")
+            return instance
+        except Exception as e:
+            logger.error(f"Error creating MCPAggregator: {e}")
+            await instance.__aexit__(None, None, None)
+
+    async def load_servers(self):
+        """
+        Discover tools from each server in parallel and build an index of namespaced tool names.
+        """
+        if self.initialized:
+            logger.debug("MCPAggregator already initialized.")
+            return
+
+        async with self._tool_map_lock:
+            self._namespaced_tool_map.clear()
+            self._server_to_tool_map.clear()
+
+        for server_name in self.server_names:
+            if self.connection_persistence:
+                logger.info(
+                    f"Creating persistent connection to server: {server_name}",
+                    data={
+                        "progress_action": ProgressAction.STARTING,
+                        "server_name": server_name,
+                        "agent_name": self.agent_name,
+                    },
+                )
+                await self._persistent_connection_manager.get_server(
+                    server_name, client_session_factory=MCPAgentClientSession
+                )
+
+            logger.info(
+                f"MCP Servers initialized for agent '{self.agent_name}'",
+                data={
+                    "progress_action": ProgressAction.INITIALIZED,
+                    "agent_name": self.agent_name,
+                },
+            )
+
+        async def fetch_tools(client: ClientSession):
+            try:
+                result: ListToolsResult = await client.list_tools()
+                return result.tools or []
+            except Exception as e:
+                logger.error(f"Error loading tools from server '{server_name}'", data=e)
+                return []
+
+        async def load_server_tools(server_name: str):
+            tools: List[Tool] = []
+            if self.connection_persistence:
+                server_connection = (
+                    await self._persistent_connection_manager.get_server(
+                        server_name, client_session_factory=MCPAgentClientSession
+                    )
+                )
+                tools = await fetch_tools(server_connection.session)
+            else:
+                async with gen_client(
+                    server_name, server_registry=self.context.server_registry
+                ) as client:
+                    tools = await fetch_tools(client)
+
+            return server_name, tools
+
+        # Gather tools from all servers concurrently
+        results = await gather(
+            *(load_server_tools(server_name) for server_name in self.server_names),
+            return_exceptions=True,
+        )
+
+        for result in results:
+            if isinstance(result, BaseException):
+                continue
+            server_name, tools = result
+
+            self._server_to_tool_map[server_name] = []
+            for tool in tools:
+                namespaced_tool_name = f"{server_name}{SEP}{tool.name}"
+                namespaced_tool = NamespacedTool(
+                    tool=tool,
+                    server_name=server_name,
+                    namespaced_tool_name=namespaced_tool_name,
+                )
+
+                self._namespaced_tool_map[namespaced_tool_name] = namespaced_tool
+                self._server_to_tool_map[server_name].append(namespaced_tool)
+            logger.debug(
+                "MCP Aggregator initialized",
+                data={
+                    "progress_action": ProgressAction.INITIALIZED,
+                    "server_name": server_name,
+                    "agent_name": self.agent_name,
+                },
+            )
+        self.initialized = True
+
+    async def list_servers(self) -> List[str]:
+        """Return the list of server names aggregated by this agent."""
+        if not self.initialized:
+            await self.load_servers()
+
+        return self.server_names
+
+    async def list_tools(self) -> ListToolsResult:
+        """
+        :return: Tools from all servers aggregated, and renamed to be dot-namespaced by server name.
+        """
+        if not self.initialized:
+            await self.load_servers()
+
+        return ListToolsResult(
+            tools=[
+                namespaced_tool.tool.model_copy(update={"name": namespaced_tool_name})
+                for namespaced_tool_name, namespaced_tool in self._namespaced_tool_map.items()
+            ]
+        )
+
+    async def call_tool(
+        self, name: str, arguments: dict | None = None
+    ) -> CallToolResult:
+        """
+        Call a namespaced tool, e.g., 'server_name.tool_name'.
+        """
+        if not self.initialized:
+            await self.load_servers()
+
+        server_name: str = None
+        local_tool_name: str = None
+
+        if SEP in name:  # Namespaced tool name
+            server_name, local_tool_name = name.split(SEP, 1)
+        else:
+            # Assume un-namespaced, loop through all servers to find the tool. First match wins.
+            for _, tools in self._server_to_tool_map.items():
+                for namespaced_tool in tools:
+                    if namespaced_tool.tool.name == name:
+                        server_name = namespaced_tool.server_name
+                        local_tool_name = name
+                        break
+
+            if server_name is None or local_tool_name is None:
+                logger.error(f"Error: Tool '{name}' not found")
+                return CallToolResult(isError=True, message=f"Tool '{name}' not found")
+
+        logger.info(
+            "Requesting tool call",
+            data={
+                "progress_action": ProgressAction.CALLING_TOOL,
+                "tool_name": local_tool_name,
+                "server_name": server_name,
+                "agent_name": self.agent_name,
+            },
+        )
+
+        async def try_call_tool(client: ClientSession):
+            try:
+                return await client.call_tool(name=local_tool_name, arguments=arguments)
+            except Exception as e:
+                return CallToolResult(
+                    isError=True,
+                    message=f"Failed to call tool '{local_tool_name}' on server '{server_name}': {e}",
+                )
+
+        if self.connection_persistence:
+            server_connection = await self._persistent_connection_manager.get_server(
+                server_name, client_session_factory=MCPAgentClientSession
+            )
+            return await try_call_tool(server_connection.session)
+        else:
+            logger.debug(
+                f"Creating temporary connection to server: {server_name}",
+                data={
+                    "progress_action": ProgressAction.STARTING,
+                    "server_name": server_name,
+                    "agent_name": self.agent_name,
+                },
+            )
+            async with gen_client(
+                server_name, server_registry=self.context.server_registry
+            ) as client:
+                result = await try_call_tool(client)
+                logger.debug(
+                    f"Closing temporary connection to server: {server_name}",
+                    data={
+                        "progress_action": ProgressAction.SHUTDOWN,
+                        "server_name": server_name,
+                        "agent_name": self.agent_name,
+                    },
+                )
+                return result
+
+
+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"):
+        super().__init__(name)
+        self.aggregator = MCPAggregator(server_names)
+
+        # Register handlers
+        # TODO: saqadri - once we support resources and prompts, add handlers for those as well
+        self.list_tools()(self._list_tools)
+        self.call_tool()(self._call_tool)
+
+    async def _list_tools(self) -> List[Tool]:
+        """List all tools aggregated from connected MCP servers."""
+        tools_result = await self.aggregator.list_tools()
+        return tools_result.tools
+
+    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)
+            return result.content
+        except Exception as e:
+            return CallToolResult(isError=True, message=f"Error calling tool: {e}")
+
+    async def run_stdio_async(self) -> None:
+        """Run the server using stdio transport."""
+        async with stdio_server() as (read_stream, write_stream):
+            await self.run(
+                read_stream=read_stream,
+                write_stream=write_stream,
+                initialization_options=self.create_initialization_options(),
+            )