MCP Clients & Hosts

# stdio_client_example.py
# Connects to the document_server.py from MCP01 and calls its tools.
from __future__ import annotations

import asyncio
import json
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


async def run_stdio_client() -> None:
    # ── CHUNK 1: Define how to launch the server ──
    # StdioServerParameters tells the SDK how to spawn the server subprocess.
    # Use absolute paths in production to avoid "file not found" errors.
    server_params = StdioServerParameters(
        command="python",
        args=["document_server.py"],    # path to the MCP01 server
        env=None,                        # inherit parent environment
    )

    # ── CHUNK 2: Open transport and session ──
    # stdio_client() spawns the subprocess and returns (read_stream, write_stream).
    # ClientSession wraps those streams with the full MCP protocol.
    # Both context managers are nested: session closes before transport closes.
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:

            # ── CHUNK 3: Mandatory initialize handshake ──
            # Sends initialize → waits for server_info + capabilities.
            # Returns an InitializeResult with server name, version, capabilities.
            init_result = await session.initialize()
            print(f"Connected to: {init_result.serverInfo.name} "
                  f"v{init_result.serverInfo.version}")
            print(f"Capabilities: {init_result.capabilities}")

            # ── CHUNK 4: Discover available tools ──
            # Only call this after initialize() succeeds.
            # tools_result.tools is a list of Tool objects with name, description,
            # and inputSchema (the JSON Schema for each tool's parameters).
            tools_result = await session.list_tools()
            print(f"\nAvailable tools ({len(tools_result.tools)}):")
            for tool in tools_result.tools:
                print(f"  - {tool.name}: {tool.description[:60]}...")

            # ── CHUNK 5: Call a tool ──
            # call_tool() sends tools/call and returns a CallToolResult.
            # result.content is a list of content blocks (text, image, etc.)
            # result.isError is True if the tool raised an exception.
            result = await session.call_tool(
                "search_documents",
                {"query": "MCP", "limit": 3}
            )

            if result.isError:
                print(f"\nTool error: {result.content[0].text}")
            else:
                # Content blocks are typed — text blocks have a .text attribute
                for block in result.content:
                    if block.type == "text":
                        docs = json.loads(block.text)
                        print(f"\nSearch results ({len(docs)} found):")
                        for doc in docs:
                            print(f"  [{doc['score']}] {doc['id']}: {doc['preview'][:60]}")

            # Also list resources (if the server supports them)
            # GOTCHA: Only call list_resources() if init_result.capabilities.resources
            # is set. Calling it on a tools-only server returns an error.
            if hasattr(init_result.capabilities, 'resources') and init_result.capabilities.resources:
                resources = await session.list_resources()
                print(f"\nResources: {len(resources.resources)}")


if __name__ == "__main__":
    asyncio.run(run_stdio_client())

# http_client_example.py
# Connects to the HTTP/SSE MCP server from MCP03.
# Start the MCP03 server first: uvicorn server:app --port 8000
from __future__ import annotations

import asyncio
import json
import os
from mcp import ClientSession
from mcp.client.sse import sse_client


SERVER_URL = os.getenv("MCP_SERVER_URL", "http://localhost:8000/sse")
API_TOKEN  = os.getenv("MCP_API_TOKEN", "")  # set in production


async def run_http_client() -> None:
    # ── CHUNK 1: Auth headers (production pattern) ──
    # Pass as dict to sse_client. These headers are sent on all requests.
    # For development with no auth, pass headers={} or omit the argument.
    headers: dict[str, str] = {}
    if API_TOKEN:
        headers["Authorization"] = f"Bearer {API_TOKEN}"

    # ── CHUNK 2: Open SSE transport ──
    # sse_client opens the long-lived SSE stream to SERVER_URL.
    # The server must be running before this call — no retry on initial connect.
    async with sse_client(SERVER_URL, headers=headers) as (read, write):
        async with ClientSession(read, write) as session:

            # ── CHUNK 3: Same initialize/discover/call pattern as stdio ──
            # The protocol is transport-agnostic — everything after this point
            # is identical whether you used stdio_client or sse_client.
            init_result = await session.initialize()
            print(f"Connected to: {init_result.serverInfo.name}")

            tools = await session.list_tools()
            print(f"Tools available: {[t.name for t in tools.tools]}")

            # ── CHUNK 4: Call a tool over HTTP ──
            # Under the hood: HTTP POST to /messages with the tools/call body.
            # The server processes the call and streams the result back via SSE.
            result = await session.call_tool(
                "read_file",
                {"path": "/tmp/hello.txt"}
            )

            if result.isError:
                print(f"Error: {result.content[0].text}")
            else:
                for block in result.content:
                    if block.type == "text":
                        print(f"File contents:\n{block.text}")

            # ── CHUNK 5: Resources work the same way ──
            # If the server supports resources, fetch them identically to stdio.
            if getattr(init_result.capabilities, 'resources', None):
                resources = await session.list_resources()
                print(f"Resources: {[str(r.uri) for r in resources.resources]}")

                if resources.resources:
                    first_uri = resources.resources[0].uri
                    content = await session.read_resource(first_uri)
                    print(f"First resource content: {str(content.contents[0].text)[:200]}")


if __name__ == "__main__":
    asyncio.run(run_http_client())

# mcp_host.py — multi-server host with Claude integration
from __future__ import annotations

import asyncio
import json
from contextlib import AsyncExitStack
from typing import Any

import anthropic
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.sse import sse_client


class MCPHost:
    """Connects to multiple MCP servers and routes LLM tool calls."""

    def __init__(self) -> None:
        self.clients: dict[str, ClientSession] = {}
        self.tool_map: dict[str, str] = {}      # tool_name → server_name
        self.all_tools: list[dict] = []          # flat list for the LLM API
        self._stack = AsyncExitStack()           # manages all context managers

    # ── CHUNK 1: Connect to a stdio server ──
    async def add_stdio_server(self, name: str, params: StdioServerParameters) -> None:
        """Add a local MCP server (spawned as a subprocess)."""
        read, write = await self._stack.enter_async_context(stdio_client(params))
        session = await self._stack.enter_async_context(ClientSession(read, write))
        await session.initialize()
        await self._register_server(name, session)

    # ── CHUNK 2: Connect to an HTTP/SSE server ──
    async def add_http_server(self, name: str, url: str, headers: dict[str, str] | None = None) -> None:
        """Add a remote MCP server (connected via HTTP/SSE)."""
        read, write = await self._stack.enter_async_context(sse_client(url, headers=headers or {}))
        session = await self._stack.enter_async_context(ClientSession(read, write))
        await session.initialize()
        await self._register_server(name, session)

    async def _register_server(self, name: str, session: ClientSession) -> None:
        """Store the session and build the tool_map for this server."""
        self.clients[name] = session
        tools_result = await session.list_tools()
        for tool in tools_result.tools:
            self.tool_map[tool.name] = name
            # Convert MCP tool schema to Anthropic API tool format
            self.all_tools.append({
                "name": tool.name,
                "description": tool.description or "",
                "input_schema": tool.inputSchema or {"type": "object", "properties": {}},
            })
        print(f"  [{name}] registered {len(tools_result.tools)} tools: "
              f"{[t.name for t in tools_result.tools]}")

    # ── CHUNK 3: Route a tool call to the right server ──
    async def call_tool(self, tool_name: str, args: dict[str, Any]) -> Any:
        """Route a tool call to the server that owns it."""
        server_name = self.tool_map.get(tool_name)
        if server_name is None:
            raise ValueError(f"Tool '{tool_name}' not found in any connected server. "
                             f"Available: {list(self.tool_map.keys())}")
        session = self.clients[server_name]
        result = await session.call_tool(tool_name, args)
        if result.isError:
            raise RuntimeError(f"Tool '{tool_name}' on '{server_name}' returned error: "
                               f"{result.content[0].text if result.content else 'unknown'}")
        return result.content

    # ── CHUNK 4: Run an agentic conversation with Claude ──
    async def run_agent(self, user_message: str) -> str:
        """Run a full agentic loop: prompt → tool calls → final answer."""
        client = anthropic.Anthropic()
        messages = [{"role": "user", "content": user_message}]

        while True:
            response = client.messages.create(
                model="claude-opus-4-5",
                max_tokens=4096,
                tools=self.all_tools,
                messages=messages,
            )

            # If Claude isn't calling a tool, we're done
            if response.stop_reason == "end_turn":
                text_blocks = [b.text for b in response.content if hasattr(b, 'text')]
                return " ".join(text_blocks)

            # Append Claude's response (may contain text + tool_use blocks)
            messages.append({"role": "assistant", "content": response.content})

            # Process all tool calls in parallel for speed
            tool_use_blocks = [b for b in response.content if b.type == "tool_use"]
            if not tool_use_blocks:
                text_blocks = [b.text for b in response.content if hasattr(b, 'text')]
                return " ".join(text_blocks)

            tool_results = []
            for block in tool_use_blocks:
                try:
                    content_blocks = await self.call_tool(block.name, block.input)
                    result_text = " ".join(
                        b.text for b in content_blocks if hasattr(b, 'text')
                    )
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result_text,
                    })
                except Exception as e:
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": f"Error: {e}",
                        "is_error": True,
                    })

            messages.append({"role": "user", "content": tool_results})

    async def close(self) -> None:
        """Shut down all server connections and subprocesses."""
        await self._stack.aclose()


# ── CHUNK 5: Usage — connect to MCP01 + MCP02 simultaneously ──
async def main() -> None:
    host = MCPHost()
    try:
        print("Connecting to MCP servers...")

        # Add both servers concurrently for faster startup
        await asyncio.gather(
            host.add_stdio_server(
                "doc-server",
                StdioServerParameters(command="python", args=["document_server.py"])
            ),
            host.add_stdio_server(
                "file-server",
                StdioServerParameters(command="python", args=["file_server.py"])
            ),
        )

        print(f"\nTotal tools available: {len(host.all_tools)}")
        print(f"Tool routing map: {host.tool_map}")

        # Ask Claude a question that requires tools from both servers
        print("\nRunning agent...")
        answer = await host.run_agent(
            "Search my documents for anything about MCP, then save a note "
            "summarizing what you found."
        )
        print(f"\nAgent answer:\n{answer}")

    finally:
        await host.close()


if __name__ == "__main__":
    asyncio.run(main())

# robust_mcp_client.py — production-grade client wrapper
from __future__ import annotations

import asyncio
import logging
from contextlib import AsyncExitStack
from typing import Any

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.sse import sse_client

logger = logging.getLogger(__name__)


class MCPClientError(Exception):
    """Base class for client-side MCP errors."""


class ToolNotFoundError(MCPClientError):
    """Raised when the requested tool doesn't exist on the server."""


class InvalidParamsError(MCPClientError):
    """Raised when the tool rejected the provided arguments."""


class ServerUnavailableError(MCPClientError):
    """Raised when the server process died or the HTTP connection dropped."""


class RobustMCPClient:
    """
    MCP client wrapper that handles the five common failure scenarios.

    Usage:
        client = RobustMCPClient(StdioServerParameters(command="python", args=["server.py"]))
        async with client:
            result = await client.call_tool("search_documents", {"query": "hello"})
    """

    MAX_RETRIES = 3
    BACKOFF_BASE = 1.0   # seconds
    BACKOFF_MAX  = 30.0  # seconds

    def __init__(
        self,
        params: StdioServerParameters | None = None,
        http_url: str | None = None,
        http_headers: dict[str, str] | None = None,
    ) -> None:
        if params is None and http_url is None:
            raise ValueError("Provide either StdioServerParameters or http_url")
        self._params = params
        self._http_url = http_url
        self._http_headers = http_headers or {}
        self._session: ClientSession | None = None
        self._stack = AsyncExitStack()
        self._connected = False
        self._known_tools: set[str] = set()

    async def __aenter__(self) -> "RobustMCPClient":
        await self._connect()
        return self

    async def __aexit__(self, *_: Any) -> None:
        await self._stack.aclose()

    # ── Error Scenario 1 & 2: Connect with retry ──
    async def _connect(self) -> None:
        """Connect to the server, retrying with exponential backoff on failure."""
        for attempt in range(self.MAX_RETRIES):
            try:
                # Start fresh each attempt
                if self._stack._exit_callbacks:
                    await self._stack.aclose()
                    self._stack = AsyncExitStack()

                if self._params is not None:
                    # stdio: spawn subprocess
                    read, write = await self._stack.enter_async_context(
                        stdio_client(self._params)
                    )
                else:
                    # HTTP: open SSE stream — Scenario 2: timeout handled here
                    read, write = await self._stack.enter_async_context(
                        sse_client(self._http_url, headers=self._http_headers)
                    )

                self._session = await self._stack.enter_async_context(
                    ClientSession(read, write)
                )
                init_result = await self._session.initialize()
                tools = await self._session.list_tools()
                self._known_tools = {t.name for t in tools.tools}
                self._connected = True
                logger.info(f"Connected to {init_result.serverInfo.name} "
                            f"(attempt {attempt + 1})")
                return

            except Exception as e:
                self._connected = False
                if attempt < self.MAX_RETRIES - 1:
                    delay = min(self.BACKOFF_BASE * (2 ** attempt), self.BACKOFF_MAX)
                    logger.warning(f"Connect failed (attempt {attempt + 1}): {e}. "
                                   f"Retrying in {delay:.1f}s...")
                    await asyncio.sleep(delay)
                else:
                    raise ServerUnavailableError(
                        f"Failed to connect after {self.MAX_RETRIES} attempts: {e}"
                    ) from e

    # ── Error Scenario 1: Detect and restart crashed stdio server ──
    async def _ensure_connected(self) -> None:
        """Reconnect if the server process died (stdio only)."""
        if not self._connected:
            logger.warning("Server disconnected — attempting reconnect...")
            await self._connect()

    # ── Core call_tool with all error scenarios handled ──
    async def call_tool(
        self,
        tool_name: str,
        args: dict[str, Any],
        *,
        retry_on_disconnect: bool = True,
    ) -> list:
        """
        Call a tool with full error handling.

        Args:
            tool_name: Name of the tool to call.
            args: Tool arguments. These are logged on InvalidParams for debugging.
            retry_on_disconnect: If True, reconnect and retry once on disconnect.

        Returns:
            List of MCP content blocks from the tool result.

        Raises:
            ToolNotFoundError: Tool doesn't exist on this server.
            InvalidParamsError: Server rejected the arguments.
            ServerUnavailableError: Cannot reach the server after retries.
            MCPClientError: Any other MCP-level error.
        """
        await self._ensure_connected()

        # Scenario 3: tool_not_found — fail fast, no retry
        if tool_name not in self._known_tools:
            # Try refreshing the tool list first (server may have restarted)
            try:
                tools = await self._session.list_tools()
                self._known_tools = {t.name for t in tools.tools}
            except Exception:
                pass  # If list_tools fails, fall through to the error below

            if tool_name not in self._known_tools:
                raise ToolNotFoundError(
                    f"Tool '{tool_name}' not found. "
                    f"Available: {sorted(self._known_tools)}"
                )

        try:
            result = await self._session.call_tool(tool_name, args)
        except Exception as e:
            err_str = str(e).lower()

            # Scenario 1/2: Connection dropped → reconnect and retry once
            if retry_on_disconnect and any(
                kw in err_str for kw in ("connection", "closed", "eof", "timeout", "broken")
            ):
                logger.warning(f"Connection error on tool call: {e}. Reconnecting...")
                self._connected = False
                await self._connect()
                result = await self._session.call_tool(tool_name, args)
            else:
                raise MCPClientError(f"Tool call failed: {e}") from e

        # Scenario 4: invalid_params — the tool itself rejected the args
        if result.isError:
            error_text = result.content[0].text if result.content else "unknown error"

            if "invalid" in error_text.lower() or "param" in error_text.lower():
                logger.error(
                    f"InvalidParams for tool '{tool_name}'. "
                    f"Args sent: {args}. Error: {error_text}"
                )
                raise InvalidParamsError(
                    f"Tool '{tool_name}' rejected arguments: {error_text}. "
                    f"Args provided: {args}"
                )

            raise MCPClientError(f"Tool '{tool_name}' error: {error_text}")

        return result.content

    # ── Scenario 5: Refresh HTTP auth token ──
    async def refresh_auth(self, new_token: str) -> None:
        """Re-establish the HTTP connection with a new auth token."""
        if self._http_url is None:
            raise ValueError("refresh_auth is only applicable to HTTP clients")

        self._http_headers["Authorization"] = f"Bearer {new_token}"
        self._connected = False
        logger.info("Auth token refreshed — reconnecting...")
        await self._connect()


# ── Usage example ──
async def main() -> None:
    params = StdioServerParameters(command="python", args=["document_server.py"])

    async with RobustMCPClient(params=params) as client:
        try:
            # Normal call
            content = await client.call_tool("search_documents", {"query": "MCP"})
            print("Search result:", content[0].text[:100] if content else "empty")

        except ToolNotFoundError as e:
            print(f"Tool missing: {e}")

        except InvalidParamsError as e:
            print(f"Bad args (check LLM output): {e}")

        except ServerUnavailableError as e:
            print(f"Server down (alert on-call): {e}")

        except MCPClientError as e:
            print(f"MCP error: {e}")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    asyncio.run(main())