MCP Servers

@abc.abstractmethod
async def connect(self):
    """Connect to the server. For example, this might mean spawning a subprocess or
    opening a network connection. The server is expected to remain connected until
    `cleanup()` is called.
    """
    pass

@abc.abstractmethod
async def cleanup(self):
    """Cleanup the server. For example, this might mean closing a subprocess or
    closing a network connection.
    """
    pass

@abc.abstractmethod
async def list_tools(self) -> list[MCPTool]:
    """List the tools available on the server."""
    pass

@abc.abstractmethod
async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> CallToolResult:
    """Invoke a tool on the server."""
    pass

async def connect(self):
    """Connect to the server."""
    try:
        transport = await self.exit_stack.enter_async_context(self.create_streams())
        read, write = transport
        session = await self.exit_stack.enter_async_context(ClientSession(read, write))
        await session.initialize()
        self.session = session
    except Exception as e:
        logger.error(f"Error initializing MCP server: {e}")
        await self.cleanup()
        raise

async def cleanup(self):
    """Cleanup the server."""
    async with self._cleanup_lock:
        try:
            await self.exit_stack.aclose()
            self.session = None
        except Exception as e:
            logger.error(f"Error cleaning up server: {e}")

async def list_tools(self) -> list[MCPTool]:
    """List the tools available on the server."""
    if not self.session:
        raise UserError("Server not initialized. Make sure you call `connect()` first.")

    # Return from cache if caching is enabled, we have tools, and the cache is not dirty
    if self.cache_tools_list and not self._cache_dirty and self._tools_list:
        return self._tools_list

    # Reset the cache dirty to False
    self._cache_dirty = False

    # Fetch the tools from the server
    self._tools_list = (await self.session.list_tools()).tools
    return self._tools_list

async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> CallToolResult:
    """Invoke a tool on the server."""
    if not self.session:
        raise UserError("Server not initialized. Make sure you call `connect()` first.")

    return await self.session.call_tool(tool_name, arguments)

def invalidate_tools_cache(self):
    """Invalidate the tools cache."""
    self._cache_dirty = True

def __init__(
    self,
    params: MCPServerStdioParams,
    cache_tools_list: bool = False,
    name: str | None = None,
):
    """Create a new MCP server based on the stdio transport.

    Args:
        params: The params that configure the server. This includes the command to run to
            start the server, the args to pass to the command, the environment variables to
            set for the server, the working directory to use when spawning the process, and
            the text encoding used when sending/receiving messages to the server.
        cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be
            cached and only fetched from the server once. If `False`, the tools list will be
            fetched from the server on each call to `list_tools()`. The cache can be
            invalidated by calling `invalidate_tools_cache()`. You should set this to `True`
            if you know the server will not change its tools list, because it can drastically
            improve latency (by avoiding a round-trip to the server every time).
        name: A readable name for the server. If not provided, we'll create one from the
            command.
    """
    super().__init__(cache_tools_list)

    self.params = StdioServerParameters(
        command=params["command"],
        args=params.get("args", []),
        env=params.get("env"),
        cwd=params.get("cwd"),
        encoding=params.get("encoding", "utf-8"),
        encoding_error_handler=params.get("encoding_error_handler", "strict"),
    )

    self._name = name or f"stdio: {self.params.command}"

def create_streams(
    self,
) -> AbstractAsyncContextManager[
    tuple[
        MemoryObjectReceiveStream[JSONRPCMessage | Exception],
        MemoryObjectSendStream[JSONRPCMessage],
    ]
]:
    """Create the streams for the server."""
    return stdio_client(self.params)

async def connect(self):
    """Connect to the server."""
    try:
        transport = await self.exit_stack.enter_async_context(self.create_streams())
        read, write = transport
        session = await self.exit_stack.enter_async_context(ClientSession(read, write))
        await session.initialize()
        self.session = session
    except Exception as e:
        logger.error(f"Error initializing MCP server: {e}")
        await self.cleanup()
        raise

async def cleanup(self):
    """Cleanup the server."""
    async with self._cleanup_lock:
        try:
            await self.exit_stack.aclose()
            self.session = None
        except Exception as e:
            logger.error(f"Error cleaning up server: {e}")

async def list_tools(self) -> list[MCPTool]:
    """List the tools available on the server."""
    if not self.session:
        raise UserError("Server not initialized. Make sure you call `connect()` first.")

    # Return from cache if caching is enabled, we have tools, and the cache is not dirty
    if self.cache_tools_list and not self._cache_dirty and self._tools_list:
        return self._tools_list

    # Reset the cache dirty to False
    self._cache_dirty = False

    # Fetch the tools from the server
    self._tools_list = (await self.session.list_tools()).tools
    return self._tools_list

async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> CallToolResult:
    """Invoke a tool on the server."""
    if not self.session:
        raise UserError("Server not initialized. Make sure you call `connect()` first.")

    return await self.session.call_tool(tool_name, arguments)

def invalidate_tools_cache(self):
    """Invalidate the tools cache."""
    self._cache_dirty = True

def __init__(
    self,
    params: MCPServerSseParams,
    cache_tools_list: bool = False,
    name: str | None = None,
):
    """Create a new MCP server based on the HTTP with SSE transport.

    Args:
        params: The params that configure the server. This includes the URL of the server,
            the headers to send to the server, the timeout for the HTTP request, and the
            timeout for the SSE connection.

        cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be
            cached and only fetched from the server once. If `False`, the tools list will be
            fetched from the server on each call to `list_tools()`. The cache can be
            invalidated by calling `invalidate_tools_cache()`. You should set this to `True`
            if you know the server will not change its tools list, because it can drastically
            improve latency (by avoiding a round-trip to the server every time).

        name: A readable name for the server. If not provided, we'll create one from the
            URL.
    """
    super().__init__(cache_tools_list)

    self.params = params
    self._name = name or f"sse: {self.params['url']}"

def create_streams(
    self,
) -> AbstractAsyncContextManager[
    tuple[
        MemoryObjectReceiveStream[JSONRPCMessage | Exception],
        MemoryObjectSendStream[JSONRPCMessage],
    ]
]:
    """Create the streams for the server."""
    return sse_client(
        url=self.params["url"],
        headers=self.params.get("headers", None),
        timeout=self.params.get("timeout", 5),
        sse_read_timeout=self.params.get("sse_read_timeout", 60 * 5),
    )