From bfae7b45b9421fe54a72497509638a9d49f70a8a Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 29 Oct 2026 17:24:43 -0400 Subject: [PATCH] feat(mcp): add automatic client and server instrumentation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > 🤖 **This text was AI-generated.** `auto_instrument()` previously skipped MCP, requiring separate setup: ```python client = wrap_mcp_client(client) server = wrap_mcp_server(server) ``` Enable both sides automatically, with independent controls and content capture enabled by default. Applications can opt out of payload capture: ```python auto_instrument(mcp_client=True, mcp_server=True, mcp_capture_content=False) ``` Repeated setup updates defaults without adding duplicate spans. Explicit wrappers retain their per-object settings, and MCP dependencies stay optional. Builds on [#1025](https://github.com/braintrustdata/braintrust-sdk-python/pull/1025) for MCP 1.x/2.x and FastMCP 2.x–4.x. Refs #1023 --- README.md | 1 + py/src/braintrust/auto.py | 20 ++- .../auto_test_scripts/test_auto_mcp.py | 114 ++++++++++++++++++ py/src/braintrust/integrations/mcp/_client.py | 27 ++++- .../integrations/mcp/_test_harness.py | 9 +- .../integrations/mcp/integration.py | 20 ++- .../integrations/mcp/test_client.py | 51 ++++++-- .../test_lazy_integration_imports.py | 17 +++ py/src/braintrust/test_span_customizer.py | 4 +- 9 files changed, 234 insertions(+), 29 deletions(-) create mode 100644 py/src/braintrust/integrations/auto_test_scripts/test_auto_mcp.py diff --git a/README.md b/README.md index 685adfe4a..835f21a3e 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,7 @@ BRAINTRUST_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` | diff --git a/py/src/braintrust/auto.py b/py/src/braintrust/auto.py index e7f39ed23..6c8b14b65 100644 --- a/py/src/braintrust/auto.py +++ b/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 diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_auto_mcp.py b/py/src/braintrust/integrations/auto_test_scripts/test_auto_mcp.py new file mode 100644 index 000000000..e1c8e279c --- /dev/null +++ b/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") diff --git a/py/src/braintrust/integrations/mcp/_client.py b/py/src/braintrust/integrations/mcp/_client.py index 9b88c8212..9a4eff370 100644 --- a/py/src/braintrust/integrations/mcp/_client.py +++ b/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) diff --git a/py/src/braintrust/integrations/mcp/_test_harness.py b/py/src/braintrust/integrations/mcp/_test_harness.py index c79a855af..075182bfe 100644 --- a/py/src/braintrust/integrations/mcp/_test_harness.py +++ b/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: diff --git a/py/src/braintrust/integrations/mcp/integration.py b/py/src/braintrust/integrations/mcp/integration.py index a23498d8b..e421360b8 100644 --- a/py/src/braintrust/integrations/mcp/integration.py +++ b/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 diff --git a/py/src/braintrust/integrations/mcp/test_client.py b/py/src/braintrust/integrations/mcp/test_client.py index 18543b4ec..a35fb1bd5 100644 --- a/py/src/braintrust/integrations/mcp/test_client.py +++ b/py/src/braintrust/integrations/mcp/test_client.py @@ -1,14 +1,17 @@ """Real MCP client operations, independently and connected to server traces.""" import asyncio +import inspect import json import re import anyio import pytest -from braintrust import logger +from braintrust import auto_instrument, logger from braintrust.integrations.mcp import _test_harness as harness from braintrust.integrations.mcp import wrap_mcp_client, wrap_mcp_server +from braintrust.integrations.mcp._client import set_global_client_config +from braintrust.integrations.test_utils import verify_autoinstrument_script from braintrust.logger import current_span, start_span from braintrust.test_helpers import find_span_by_name, init_test_logger @@ -18,6 +21,7 @@ def memory_logger(): init_test_logger("test-mcp-client") with logger._internal_with_memory_background_logger() as bgl: yield bgl + set_global_client_config(None) def client_spans(spans): @@ -30,6 +34,19 @@ def assert_client(span, method): assert span["metrics"]["end"] >= span["metrics"]["start"] +def _setup_auto_client(capture_content): + options = { + name: False for name, param in inspect.signature(auto_instrument).parameters.items() if param.default is True + } + assert auto_instrument(**{**options, "mcp_client": True, "mcp_capture_content": capture_content}) == {"mcp": True} + + +@pytest.mark.parametrize("import_order", ["before_import", "after_import", "after_connect"]) +def test_auto_instrument(monkeypatch, import_order): + monkeypatch.setenv("BRAINTRUST_MCP_TEST_IMPORT_ORDER", import_order) + verify_autoinstrument_script("test_auto_mcp.py", timeout=60) + + @pytest.mark.asyncio async def test_client_rpc_error(memory_logger): async with harness.connect_in_memory(harness.build_server()) as client: @@ -110,15 +127,19 @@ async def test_connected_http(memory_logger): @pytest.mark.asyncio -@pytest.mark.parametrize("server_traced", [False, True]) -async def test_connected_stdio(memory_logger, tmp_path, server_traced): +@pytest.mark.parametrize("server_mode", ["none", "wrap", "auto"]) +async def test_connected_stdio(memory_logger, tmp_path, server_mode): spans_file = tmp_path / "server-spans.json" env = { harness.SPANS_FILE_ENV: str(spans_file), - harness.WRAP_ENV: "wrap" if server_traced else "none", + harness.WRAP_ENV: server_mode, + harness.CAPTURE_ENV: str(server_mode != "auto").lower(), } async with harness.connect_stdio(env) as client: - wrap_mcp_client(client.raw) + if server_mode == "auto": + _setup_auto_client(capture_content=False) + else: + wrap_mcp_client(client.raw) with start_span(name="agent") as agent: await client.call_tool("add", {"a": 2, "b": 3}) with anyio.fail_after(10): @@ -128,25 +149,33 @@ async def test_connected_stdio(memory_logger, tmp_path, server_traced): outgoing = find_span_by_name(client_spans(memory_logger.pop()), "tools/call add") assert outgoing["span_parents"] == [agent.span_id] assert outgoing["root_span_id"] == agent.root_span_id - assert outgoing["output"]["content"][0]["text"] == "5" - if server_traced: + if server_mode != "auto": + assert outgoing["output"]["content"][0]["text"] == "5" + if server_mode != "none": incoming = find_span_by_name(remote, "tools/call add") assert incoming["span_parents"] == [outgoing["span_id"]] assert incoming["root_span_id"] == outgoing["root_span_id"] assert find_span_by_name(remote, "nested-work")["span_parents"] == [incoming["span_id"]] + if server_mode == "auto": + for span in (outgoing, incoming): + assert span.get("input") is None and span.get("output") is None else: assert not [s for s in remote if "mcp.method.name" in s.get("metadata", {})] @pytest.mark.asyncio @pytest.mark.parametrize("capture_content", [True, False]) -async def test_client_http_operations(memory_logger, capture_content): +@pytest.mark.parametrize("entrypoint", ["manual", "auto"]) +async def test_client_http_operations(memory_logger, capture_content, entrypoint): + if entrypoint == "auto": + _setup_auto_client(capture_content) request_meta = [] server = harness.build_server(request_meta=request_meta) async with harness.serve_http(harness.http_app(server)) as url: - async with harness.connect_http(url, trace_client=True) as client: - assert wrap_mcp_client(client.raw, capture_content=capture_content) is client.raw - wrap_mcp_client(client.raw, capture_content=capture_content) + async with harness.connect_http(url, trace_client=entrypoint == "manual") as client: + if entrypoint == "manual": + assert wrap_mcp_client(client.raw, capture_content=capture_content) is client.raw + wrap_mcp_client(client.raw, capture_content=capture_content) with start_span(name="agent") as agent: await client.call_tool("add", {"a": 4, "b": 5}, meta={"Baggage": "tenant=demo", "custom": "untouched"}) await client.call_tool("fail") diff --git a/py/src/braintrust/test_lazy_integration_imports.py b/py/src/braintrust/test_lazy_integration_imports.py index 9f8f3afe0..ba326488e 100644 --- a/py/src/braintrust/test_lazy_integration_imports.py +++ b/py/src/braintrust/test_lazy_integration_imports.py @@ -1,3 +1,4 @@ +import importlib.util import json import subprocess import sys @@ -91,6 +92,22 @@ def test_integrations_package_lazily_resolves_public_classes(): } +def test_mcp_auto_without_optional_dependency(): + if importlib.util.find_spec("mcp") is not None: + pytest.skip("Requires the core session without optional MCP dependencies") + loaded = _run_import_script( + """ +import inspect +import json +from braintrust import auto_instrument + +options = {name: False for name, param in inspect.signature(auto_instrument).parameters.items() if param.default is True} +print(json.dumps(auto_instrument(**{**options, "mcp_client": True, "mcp_server": True}))) +""" + ) + assert loaded == {"mcp": False} + + @requires_lazy_imports def test_import_braintrust_defers_lazy_modules(): loaded = _run_import_script( diff --git a/py/src/braintrust/test_span_customizer.py b/py/src/braintrust/test_span_customizer.py index 8a922fc29..bfde3c463 100644 --- a/py/src/braintrust/test_span_customizer.py +++ b/py/src/braintrust/test_span_customizer.py @@ -96,7 +96,7 @@ def test_registration_rejects_misconfigured_customizers(invalid): def test_registration_supported_in_otel_compat_mode(monkeypatch, caplog): - monkeypatch.setattr("braintrust.auto._instrument_integration", lambda _module, _class: False) + monkeypatch.setattr("braintrust.auto._instrument_integration", lambda _module, _class, **_kwargs: False) monkeypatch.setenv("BRAINTRUST_OTEL_COMPAT", "true") set_span_customizers([SpanCustomizer()]) auto_instrument(span_customizers=[_Valid()]) @@ -878,7 +878,7 @@ def mask(value): def test_auto_instrument_registration_and_disable(monkeypatch, with_memory_logger, test_logger): # Registration should work without importing optional provider libraries. - monkeypatch.setattr("braintrust.auto._instrument_integration", lambda _module, _class: False) + monkeypatch.setattr("braintrust.auto._instrument_integration", lambda _module, _class, **_kwargs: False) class Redact(SpanCustomizer): def on_span_export(self, data):