feat(mcp): add automatic client and server instrumentation

README.md

@@ -57,6 +57,7 @@ BRAINTRUST_API_KEY=<YOUR_API_KEY> braintrust eval tutorial_eval.py
 | --- | --- | --- |
 | [OpenAI](py/src/braintrust/integrations/openai/) | Yes | `openai>=1.71.0` |
 | [Anthropic](py/src/braintrust/integrations/anthropic/) | Yes | `anthropic>=0.48.0` |
+| [MCP / FastMCP](py/src/braintrust/integrations/mcp/) | Yes (client and server) | `mcp>=1.28.1,<3`; `fastmcp>=2.14.7,<5` |
 | [Cohere](py/src/braintrust/integrations/cohere/) | Yes | `cohere>=5.0.0` |
 | [Mistral](py/src/braintrust/integrations/mistral/) | Yes | `mistralai>=1.12.4` |
 | [LiteLLM](py/src/braintrust/integrations/litellm/) | Yes | `litellm>=1.74.0` |

py/src/braintrust/auto.py

@@ -61,12 +61,16 @@ def auto_instrument(
     livekit_agents: bool = True,
     pipecat: bool = True,
     typesafe: bool = True,
+    mcp_client: bool = True,
+    mcp_server: bool = True,
+    mcp_capture_content: bool = True,
     span_customizers: Sequence[SpanCustomizer] | None = None,
 ) -> dict[str, bool]:
     """
     Auto-instrument supported AI/ML libraries for Braintrust tracing.
 
-    Safe to call multiple times - already instrumented libraries are skipped.
+    Safe to call multiple times - hooks are installed once. MCP settings are
+    updated on each call; explicit per-client/server wrapping takes precedence.
 
     Note on import order: If you use `from openai import OpenAI` style imports,
     call auto_instrument() first. If you use `import openai` style imports,
@@ -103,6 +107,12 @@ def auto_instrument(
         livekit_agents: Enable LiveKit Agents instrumentation (default: True)
         pipecat: Enable Pipecat AI instrumentation (default: True)
         typesafe: Enable TypeSafe instrumentation (default: True)
+        mcp_client: Enable MCP client instrumentation (default: True)
+        mcp_server: Enable MCP server instrumentation (default: True)
+        mcp_capture_content: Capture MCP arguments, results, resource contents,
+            and server notification parameters (default: True). False keeps
+            names, timing, identifiers, and error classification. Applies to
+            both enabled MCP sides; does not affect other integrations.
         span_customizers: Ordered synchronous export customizers for native SDK spans
             and OTel spans sent through OtelExporter or BraintrustSpanProcessor,
             including manual spans; generic OTLP exporters, dataset rows, and
@@ -221,12 +231,16 @@ def auto_instrument(
         results["pipecat"] = _instrument_integration("pipecat", "PipecatIntegration")
     if typesafe:
         results["typesafe"] = _instrument_integration("typesafe", "TypeSafeIntegration")
+    if mcp_client or mcp_server:
+        results["mcp"] = _instrument_integration(
+            "mcp", "MCPIntegration", client=mcp_client, server=mcp_server, capture_content=mcp_capture_content
+        )
 
     return results
 
 
-def _instrument_integration(module_name: str, class_name: str) -> bool:
+def _instrument_integration(module_name: str, class_name: str, **kwargs: Any) -> bool:
     with _try_patch():
         integration: Any = getattr(import_module(f"braintrust.integrations.{module_name}"), class_name)
-        return integration.setup()
+        return integration.setup(**kwargs)
     return False

py/src/braintrust/integrations/auto_test_scripts/test_auto_mcp.py

@@ -0,0 +1,114 @@
+"""Exercise auto setup, independent controls, capture settings and manual overrides."""
+
+import inspect
+import os
+import sys
+
+from braintrust.auto import auto_instrument
+
+
+# Keep this subprocess focused on MCP; no optional model instrumentation is needed.
+OPTIONS = {
+    name: False for name, param in inspect.signature(auto_instrument).parameters.items() if param.default is True
+}
+ORDER = os.environ["BRAINTRUST_MCP_TEST_IMPORT_ORDER"]
+assert auto_instrument(**OPTIONS) == {}
+assert "mcp" not in sys.modules and "fastmcp" not in sys.modules
+
+
+def setup(client=True, server=True, capture=True):
+    result = auto_instrument(**{**OPTIONS, "mcp_client": client, "mcp_server": server, "mcp_capture_content": capture})
+    assert result["mcp"] is True
+
+
+if ORDER == "before_import":
+    # Exercise the public defaults as well as the explicit controls below.
+    assert auto_instrument(**{name: value for name, value in OPTIONS.items() if not name.startswith("mcp_")}) == {
+        "mcp": True
+    }
+
+import anyio  # noqa: E402
+from braintrust.integrations.mcp import _test_harness as harness  # noqa: E402
+from braintrust.integrations.mcp import wrap_mcp_client, wrap_mcp_server  # noqa: E402
+from braintrust.integrations.test_utils import autoinstrument_test_context  # noqa: E402
+from braintrust.logger import start_span  # noqa: E402
+from braintrust.test_helpers import find_span_by_name  # noqa: E402
+
+
+async def main():
+    with autoinstrument_test_context("test_auto_mcp", use_vcr=False) as memory_logger:
+        # Server runners retain callbacks when they start; automatic server
+        # setup belongs before serving. Client setup also covers active sessions.
+        modes = {
+            "before_import": [(True, True)],
+            "after_import": [(True, False), (False, True)],
+            "after_connect": [(True, False)],
+        }[ORDER]
+        for client_enabled, server_enabled in modes:
+            server = harness.build_server()
+            if ORDER != "after_connect":
+                setup(client_enabled, server_enabled)
+            async with harness.serve_http(harness.http_app(server)) as url:
+                async with harness.connect_http(url) as client:
+                    negotiation = [
+                        s
+                        for s in memory_logger.pop()
+                        if s.get("metadata", {}).get("mcp.method.name") in ("initialize", "server/discover")
+                    ]
+                    if ORDER != "after_connect":
+                        assert (
+                            any(s["metadata"].get("mcp.span.kind") == "client" for s in negotiation) == client_enabled
+                        )
+                        assert (
+                            any(s["metadata"].get("mcp.span.kind") != "client" for s in negotiation) == server_enabled
+                        )
+                    # The combined case checks live capture updates; the
+                    # independent cases only need to verify content opt-out.
+                    for capture in (True, False) if ORDER == "before_import" else (False,):
+                        setup(client_enabled, server_enabled, capture)
+                        setup(client_enabled, server_enabled, capture)
+                        with start_span(name="agent") as agent:
+                            await client.call_tool("add", {"a": 2, "b": 3})
+                        spans = memory_logger.pop()
+                        mcp = [s for s in spans if s["span_attributes"]["name"] == "tools/call add"]
+                        outgoing = [s for s in mcp if s["metadata"].get("mcp.span.kind") == "client"]
+                        incoming = [s for s in mcp if s not in outgoing]
+                        assert len(outgoing) == int(client_enabled), spans
+                        assert len(incoming) == int(server_enabled), spans
+                        for span in mcp:
+                            if capture:
+                                assert span["input"] == {"a": 2, "b": 3}
+                                assert span["output"]["content"][0]["text"] == "5"
+                            else:
+                                assert span.get("input") is None and span.get("output") is None
+                        if outgoing:
+                            assert outgoing[0]["span_parents"] == [agent.span_id]
+                        if outgoing and incoming:
+                            assert incoming[0]["span_parents"] == [outgoing[0]["span_id"]]
+                            assert incoming[0]["root_span_id"] == outgoing[0]["root_span_id"] == agent.root_span_id
+                        if incoming:
+                            assert find_span_by_name(spans, "nested-work")["span_parents"] == [incoming[0]["span_id"]]
+
+        # Global capture changes leave explicit per-object settings in control.
+        if ORDER != "before_import":
+            return
+        for manual_capture in (True, False):
+            setup(capture=not manual_capture)
+            server = wrap_mcp_server(harness.build_server(), capture_content=manual_capture)
+            async with harness.serve_http(harness.http_app(server)) as url:
+                async with harness.connect_http(url) as client:
+                    wrap_mcp_client(client.raw, capture_content=manual_capture)
+                    setup(capture=not manual_capture)
+                    memory_logger.pop()
+                    await client.call_tool("add", {"a": 2, "b": 3})
+                    spans = memory_logger.pop()
+                    tools = [s for s in spans if s["span_attributes"]["name"] == "tools/call add"]
+                    assert len(tools) == 2
+                    if manual_capture:
+                        assert all(s["input"] == {"a": 2, "b": 3} for s in tools)
+                    else:
+                        assert all(s.get("input") is None and s.get("output") is None for s in tools)
+
+
+anyio.run(main)
+print("SUCCESS")

py/src/braintrust/integrations/mcp/_client.py

@@ -38,6 +38,17 @@
 
 
 CLIENT_CONFIG_ATTR = "__braintrust_mcp_client_config__"
+_global_client_config: ClientConfig | None = None
+
+
+def set_global_client_config(config: ClientConfig | None) -> None:
+    global _global_client_config
+    _global_client_config = config
+
+
+def client_config(client: Any) -> ClientConfig | None:
+    config = getattr(client, CLIENT_CONFIG_ATTR, None)
+    return _global_client_config if config is None else config
 
 
 # Pydantic include keys are Python field names, which differ across MCP majors.
@@ -75,7 +86,9 @@ def configure_session(client: Any, config: ClientConfig) -> None:
     set_session_config(session, config)
 
 
-def set_session_config(session: Any, config: ClientConfig) -> None:
+def set_session_config(session: Any, config: ClientConfig | None) -> None:
+    # None inherits automatic defaults. Mark only client-owned dispatchers;
+    # server dispatchers share the same class and must not become traced clients.
     setattr(session, CLIENT_CONFIG_ATTR, config)
     dispatcher = getattr(session, "_dispatcher", None)
     if dispatcher is not None:
@@ -97,9 +110,7 @@ async def client_enter_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: A
 
 def session_init_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
     result = wrapped(*args, **kwargs)
-    config = _enter_config.get()
-    if config is not None:
-        contain(set_session_config, instance, config)
+    contain(set_session_config, instance, _enter_config.get())
     return result
 
 
@@ -154,9 +165,13 @@ def _record_exception(operation: ServerOperation, error: BaseException) -> None:
 
 
 async def send_request_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
-    config = getattr(instance, CLIENT_CONFIG_ATTR, None)
+    config = client_config(instance)
     if not isinstance(config, ClientConfig):
         return await wrapped(*args, **kwargs)
+    # Also cover sessions connected before automatic setup installed the hooks.
+    dispatcher = getattr(instance, "_dispatcher", None)
+    if dispatcher is not None and not hasattr(dispatcher, CLIENT_CONFIG_ATTR):
+        contain(set_session_config, instance, None)
     if upstream_owns_operations():
         # MCP 1.x inherits this method. Resolve the live base method so an OTel
         # hook installed after our subclass hook remains visible as well.
@@ -210,7 +225,7 @@ def _prepare_raw(method: str, params: Any, config: ClientConfig) -> tuple[Server
 async def raw_request_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
     # 2.x negotiation and direct dispatcher calls can bypass send_request. The
     # typed wrapper remains the owner when it is already tracing this dispatcher.
-    config = getattr(instance, CLIENT_CONFIG_ATTR, None)
+    config = client_config(instance) if hasattr(instance, CLIENT_CONFIG_ATTR) else None
     active = _outgoing_operation.get()
     if not isinstance(config, ClientConfig) or (active is not None and active[1] is instance):
         return await wrapped(*args, **kwargs)

py/src/braintrust/integrations/mcp/_test_harness.py

@@ -35,6 +35,7 @@
 PNG_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
 SPANS_FILE_ENV = "BRAINTRUST_MCP_TEST_SPANS_FILE"
 WRAP_ENV = "BRAINTRUST_MCP_TEST_WRAP"
+CAPTURE_ENV = "BRAINTRUST_MCP_TEST_CAPTURE_CONTENT"
 
 # Set by the ``wait_for_cancel`` tool so tests can cancel a running request.
 cancel_tool_started: anyio.Event | None = None
@@ -283,10 +284,16 @@ async def _serve_stdio() -> None:
 
     init_test_logger("test-mcp-stdio")
     with logger._internal_with_memory_background_logger() as memory_logger:
+        if os.environ.get(WRAP_ENV) == "auto":
+            from braintrust import auto_instrument
+
+            capture = os.environ.get(CAPTURE_ENV, "true") == "true"
+            auto_instrument(mcp_client=False, mcp_capture_content=capture)
+            auto_instrument(mcp_client=False, mcp_capture_content=capture)
         server = build_server()
         if os.environ.get(WRAP_ENV) == "setup":
             MCPIntegration.setup()
-        elif os.environ.get(WRAP_ENV) != "none":
+        elif os.environ.get(WRAP_ENV) not in ("none", "auto"):
             wrap_mcp_server(server)
         run = server.run_stdio_async
         if "show_banner" in inspect.signature(run).parameters:

py/src/braintrust/integrations/mcp/integration.py

@@ -2,6 +2,7 @@
 
 from braintrust.integrations.base import BaseIntegration
 
+from ._client import ClientConfig, set_global_client_config
 from .patchers import CLIENT_PATCHERS, SERVER_PATCHERS
 from .tracing import MCPServerConfig, set_global_server_config
 
@@ -11,8 +12,9 @@ class MCPIntegration(BaseIntegration):
 
     Covers low-level servers and the server abstractions built on them: the
     bundled ``FastMCP`` (1.x) / ``MCPServer`` (2.x) and standalone ``fastmcp``.
-    ``setup()`` enables server tracing; ``wrap_mcp_client()`` opts individual
-    clients or sessions into tracing independently.
+    ``setup()`` enables server tracing by default. Pass ``client=True`` to
+    enable clients too, and ``server=False`` as well for clients only. Manual
+    wrappers override the process defaults for individual clients and servers.
     """
 
     name = "mcp"
@@ -22,16 +24,22 @@ class MCPIntegration(BaseIntegration):
     patchers = (*SERVER_PATCHERS, *CLIENT_PATCHERS)
 
     @classmethod
-    def setup(cls, *, target: Any | None = None, capture_content: bool = True) -> bool:
-        """Trace every MCP server in the process.
-
+    def setup(
+        cls, *, target: Any | None = None, client: bool = False, server: bool = True, capture_content: bool = True
+    ) -> bool:
+        """Configure process-wide MCP tracing, independently for each side.
+
+        Repeated calls update defaults without adding hooks or duplicate spans.
+        :param client: Enable outgoing client round trips.
+        :param server: Enable incoming server requests and notifications.
         :param capture_content: Log arguments, results, and notification
             parameters. When False, spans keep names, timing, allowlisted
             metadata, and error classification only.
         """
         if not cls.install_patchers(target=target):
             return False
-        set_global_server_config(MCPServerConfig(capture_content=capture_content))
+        set_global_client_config(ClientConfig(capture_content=capture_content) if client else None)
+        set_global_server_config(MCPServerConfig(capture_content=capture_content) if server else None)
         return True
 
     @classmethod