Build an MCP client with discovery, calls and tracing

"""The docs server the traced client talks to. Run with `python docs_server.py`.

If MCP_SPANS_FILE is set, the server writes its OpenTelemetry spans there. mcp 2.3.0
creates a server span for every request and continues the trace the client sends in _meta.
"""
import os
import time
from typing import TypedDict

from opentelemetry import trace

if os.environ.get("MCP_SPANS_FILE"):
    from span_export import configure
    configure(os.environ["MCP_SPANS_FILE"])

from mcp.server import MCPServer
from mcp.server.caching import CacheHint
from mcp.server.mcpserver.exceptions import ToolError

PAGES = {"quickstart": "Start a worker with taskq worker emails.",
         "retries": "A failed job is retried 3 times with exponential backoff."}
server = MCPServer("docs", version="1.0.0", instructions="Search the taskq docs.",
                   cache_hints={"tools/list": CacheHint(ttl_ms=60_000, scope="public")})
tracer = trace.get_tracer("docs")


class SearchHit(TypedDict):
    slug: str
    text: str


@server.tool()
def search_docs(query: str) -> list[SearchHit]:
    """Return the docs pages that contain the query."""
    with tracer.start_as_current_span("docs.search") as span:
        if not query.strip():
            raise ToolError("The query is empty.")
        hits = [{"slug": slug, "text": text} for slug, text in PAGES.items() if query.lower() in text.lower()]
        span.set_attribute("docs.hits", len(hits))
        return hits


@server.tool()
def rebuild_index(seconds: float) -> str:
    """Rebuild the search index; takes `seconds` seconds."""
    time.sleep(seconds)
    return "index rebuilt"


if __name__ == "__main__":
    server.run()

"""An MCP client that keeps what discovery tells it, checks calls before sending and traces each one.

Usage: python traced_client.py [--write-log execution-log.json] [--write-trace trace-tree.json]
"""
import argparse
import asyncio
import json
import os
import sys
import tempfile
from pathlib import Path

import jsonschema
from opentelemetry import trace

from mcp import Client, MCPError, StdioServerParameters
from mcp.types import REQUEST_TIMEOUT
from span_export import configure, span_tree

HERE = Path(__file__).resolve().parent
tracer = trace.get_tracer("traced-client")


class TracedClient:
    """Wraps an SDK Client: remembers discovery, lists tools only if offered, and classifies every call."""

    def __init__(self, client):
        self.client, self.tools, self.log = client, {}, []

    def discovery(self):
        caps, info = self.client.server_capabilities, self.client.server_info
        return {"protocol_version": self.client.protocol_version, "server": f"{info.name} {info.version}",
                "capabilities": [name for name in ("tools", "resources", "prompts") if getattr(caps, name) is not None]}

    async def load_tools(self):
        if self.client.server_capabilities.tools is None:
            self.log.append({"request": "tools/list", "outcome": "skipped", "detail": "server does not offer tools"})
            return
        self.tools = {tool.name: tool for tool in (await self.client.list_tools()).tools}
        self.log.append({"request": "tools/list", "outcome": "result", "detail": sorted(self.tools)})

    async def call(self, name, arguments, timeout=5.0):
        with tracer.start_as_current_span(f"app call {name}") as span:
            outcome, detail = await self._call(name, arguments, timeout)
            span.set_attribute("app.outcome", outcome)
        self.log.append({"request": f"tools/call {name}", "arguments": arguments, "outcome": outcome, "detail": detail})

    async def _call(self, name, arguments, timeout):
        tool = self.tools.get(name)
        if tool is None:
            return "refused_locally", "tool is not in the server's tools/list"
        problem = jsonschema.exceptions.best_match(jsonschema.Draft202012Validator(tool.input_schema).iter_errors(arguments))
        if problem:
            return "refused_locally", f"arguments do not match inputSchema: {problem.message}"
        try:
            result = await self.client.call_tool(name, arguments, read_timeout_seconds=timeout)
        except MCPError as error:
            return ("timeout" if error.error.code == REQUEST_TIMEOUT else "protocol_error"), error.error.message
        except RuntimeError as error:  # mcp 2.3.0 checks structuredContent against outputSchema and raises this
            return "invalid_output", str(getattr(error.__cause__, "message", error))
        if result.is_error:
            return "tool_error", result.content[0].text
        return "result", result.structured_content


def server(script, **env):
    return StdioServerParameters(command=sys.executable, args=[str(HERE / script)], env={**os.environ, **env})


async def docs_session(server_spans):
    async with Client(server("docs_server.py", MCP_SPANS_FILE=str(server_spans))) as client:
        with tracer.start_as_current_span("answer docs question"):
            app = TracedClient(client)
            discovery = app.discovery()
            await app.load_tools()
            await app.load_tools()  # served from the SDK's cache: the listing has ttlMs 60000
            await app.call("search_docs", {"query": "worker"})
            await app.call("search_docs", {"query": 42})
            await app.call("search_docs", {"query": " "})
            await app.call("rebuild_index", {"seconds": 1.0}, timeout=0.2)
            await app.call("search_docs", {"query": "retried"})
            await app.call("delete_page", {"slug": "retries"})
    return {"discovery": discovery, "log": app.log}


async def broken_session():
    async with Client(server("broken_server.py")) as client:
        app = TracedClient(client)
        discovery = app.discovery()
        await app.load_tools()
        await app.call("search_docs", {"query": "retries"})
    return {"discovery": discovery, "log": app.log}


def run(directory):
    """Run both sessions with span files in `directory`; return the execution log and the joined trace."""
    client_spans, server_spans = Path(directory) / "client-spans.jsonl", Path(directory) / "server-spans.jsonl"
    configure(client_spans)
    docs = asyncio.run(docs_session(server_spans))
    tree = span_tree(client_spans, server_spans)
    return {"docs": docs, "broken": asyncio.run(broken_session())}, tree


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--write-log", type=Path)
    parser.add_argument("--write-trace", type=Path)
    options = parser.parse_args(argv)
    with tempfile.TemporaryDirectory() as directory:
        log, tree = run(directory)
    for data in log.values():
        print(f"== {data['discovery']['server']} offers {', '.join(data['discovery']['capabilities'])}")
        for entry in data["log"]:
            print(f"  {entry['request']:26} {entry['outcome']:16} {json.dumps(entry['detail'])[:64]}")
    print(f"trace: {tree['spans']} spans in {tree['traces']} trace(s)")
    for path, data in ((options.write_log, log), (options.write_trace, tree)):
        if path:
            path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8")
            print(f"wrote {path}")


if __name__ == "__main__":
    main()

"""Write this process's OpenTelemetry spans to a JSON Lines file, and join two files into one trace tree."""
import json
import os
from pathlib import Path

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter, SpanExportResult


class JsonLinesExporter(SpanExporter):
    def __init__(self, path):
        self.path = path

    def export(self, spans):
        with open(self.path, "a", encoding="utf-8") as out:
            for span in spans:
                out.write(json.dumps({
                    "name": span.name, "kind": span.kind.name, "status": span.status.status_code.name,
                    "trace_id": format(span.context.trace_id, "032x"), "span_id": format(span.context.span_id, "016x"),
                    "parent_id": format(span.parent.span_id, "016x") if span.parent else None,
                    "start": span.start_time, "attributes": dict(span.attributes or {})}) + "\n")
        return SpanExportResult.SUCCESS


_exporter = None


def configure(path):
    """Send this process's spans to `path`. A later call only changes the file."""
    global _exporter
    if _exporter is None:
        os.environ["OTEL_SDK_DISABLED"] = "false"  # the repository's test runner disables OpenTelemetry; this exporter only writes a file
        _exporter = JsonLinesExporter(path)
        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(_exporter))
        trace.set_tracer_provider(provider)
    _exporter.path = path


def span_tree(*paths):
    """Nest the spans from every file by parent, without IDs and times, so two runs compare equal."""
    spans = sorted((json.loads(line) for path in paths for line in Path(path).read_text().splitlines()), key=lambda s: s["start"])
    ids = {s["span_id"] for s in spans}

    def node(span):
        attributes = {k: v for k, v in span["attributes"].items() if k != "jsonrpc.request.id"}
        children = [node(s) for s in spans if s["parent_id"] == span["span_id"]]
        return {"name": span["name"], "kind": span["kind"], "status": span["status"], "attributes": attributes, "children": children}

    return {"traces": len({s["trace_id"] for s in spans}), "spans": len(spans),
            "roots": [node(s) for s in spans if s["parent_id"] not in ids]}