feat(mcp): trace client operations and propagate context

py/noxfile.py

@@ -337,26 +337,26 @@ def run(session, version, spec=spec):
         },
         {
             "name": "test_mcp",
-            "doc": "MCP server instrumentation on the bundled FastMCP (1.x) / MCPServer (2.x).",
+            "doc": "MCP client/server instrumentation on bundled FastMCP (1.x) / MCPServer (2.x).",
             "versions": MCP_VERSIONS,
             "dependencies": [("mcp", "$version")],
-            "tests": f"{INTEGRATION_DIR}/mcp/test_mcp.py",
+            "tests": [f"{INTEGRATION_DIR}/mcp/test_mcp.py", f"{INTEGRATION_DIR}/mcp/test_client.py"],
             "env": {"BRAINTRUST_MCP_TEST_FLAVOR": "mcp"},
         },
         {
             "name": "test_fastmcp",
-            "doc": "MCP server instrumentation on standalone FastMCP.",
+            "doc": "MCP client/server instrumentation on standalone FastMCP.",
             "versions": FASTMCP_VERSIONS,
             "dependencies": [("fastmcp", "$version")],
-            "tests": f"{INTEGRATION_DIR}/mcp/test_mcp.py",
+            "tests": [f"{INTEGRATION_DIR}/mcp/test_mcp.py", f"{INTEGRATION_DIR}/mcp/test_client.py"],
             "env": {"BRAINTRUST_MCP_TEST_FLAVOR": "fastmcp"},
         },
         {
             "name": "test_fastmcp_slim",
-            "doc": "MCP server instrumentation on fastmcp-slim without the fastmcp meta-package.",
+            "doc": "MCP client/server instrumentation on fastmcp-slim without the fastmcp meta-package.",
             "versions": FASTMCP_SLIM_VERSIONS,
             "dependencies": [("fastmcp-slim", "$version")],
-            "tests": f"{INTEGRATION_DIR}/mcp/test_mcp.py",
+            "tests": [f"{INTEGRATION_DIR}/mcp/test_mcp.py", f"{INTEGRATION_DIR}/mcp/test_client.py"],
             "env": {"BRAINTRUST_MCP_TEST_FLAVOR": "fastmcp"},
         },
         {
@@ -417,6 +417,22 @@ def test_openai_agents(session, version):
 
 LITELLM_VERSIONS = _get_matrix_versions("litellm")
 
+
+@nox.session()
+def test_mcp_interop(session):
+    """MCP 1.x coexistence with the upstream OTel instrumentor and LangChain tools."""
+    _install_test_deps(session, "test-mcp-interop")
+    # The upstream OTel instrumentor targets the MCP 1.x BaseSession API.
+    _install_matrix_dep(session, "mcp", "1.28.1")
+    _install_matrix_dep(session, "langchain-core", LATEST)
+    _run_tests(
+        session,
+        f"{INTEGRATION_DIR}/mcp/test_interop.py",
+        version="1.28.1",
+        env={"BRAINTRUST_MCP_TEST_FLAVOR": "mcp"},
+    )
+
+
 # LiteLLM >= 1.102.0 fetches model_prices_and_context_window.json from GitHub on
 # import. Recording that lands a ~3MB blob in the cassette, so force the bundled
 # copy instead and keep the fetch off the wire.

py/pyproject.toml

@@ -205,6 +205,11 @@ test-langchain = [
     "langgraph==1.1.6",
 ]
 
+test-mcp-interop = [
+    "opentelemetry-instrumentation-mcp==0.62.4",
+    "langchain-mcp-adapters==0.3.2",
+]
+
 test-deepagents = [
     {include-group = "test"},
     "langchain-openai==1.1.14",
@@ -436,6 +441,10 @@ conflicts = [
         {group = "test-litellm"},
         {group = "test-mcp-v1"},
     ],
+    [
+        {group = "test-litellm"},
+        {group = "test-mcp-interop"},
+    ],
 ]
 
 # ---------------------------------------------------------------------------

py/src/braintrust/integrations/mcp/__init__.py

@@ -1,12 +1,14 @@
-"""Braintrust instrumentation for Model Context Protocol (MCP) servers."""
+"""Braintrust instrumentation for Model Context Protocol clients and servers."""
 
 import logging
+import sys
 from typing import Any, TypeVar
 
 from braintrust.integrations.versioning import version_satisfies
 
+from ._client import CLIENT_CONFIG_ATTR, ClientConfig, configure_session, set_session_config
 from .integration import MCPIntegration
-from .patchers import ServerRunPatcher
+from .patchers import ClientEnterPatcher, ServerRunPatcher
 from .tracing import SERVER_CONFIG_ATTR, MCPServerConfig
 
 
@@ -59,4 +61,48 @@ def wrap_mcp_server(server: ServerT, capture_content: bool = True) -> ServerT:
     return server
 
 
-__all__ = ["MCPIntegration", "wrap_mcp_server"]
+def wrap_mcp_client(client_or_session: ServerT, capture_content: bool = True) -> ServerT:
+    """Trace outgoing MCP round trips and propagate W3C context in ``params._meta``.
+
+    Accepts an official ``mcp.ClientSession`` (1.x/2.x), ``mcp.Client`` (2.x),
+    or standalone ``fastmcp.Client`` (2.x–4.x). Returns the same object; wrap
+    before connecting to include initialization, or wrap an active session.
+    Repeated calls update content settings without adding operation spans.
+
+    Content capture defaults to True. False omits arguments and results;
+    method, timing, identifiers, and error classification remain available.
+    Each client and server can be instrumented independently.
+
+    One-way notifications and detached task APIs have no client operation span.
+    If ``opentelemetry-instrumentation-mcp`` is enabled, its instrumentor owns
+    client spans, propagation, and capture settings; no duplicate span is added.
+    """
+    try:
+        import mcp
+        from mcp.client.session import ClientSession
+    except ImportError:
+        return client_or_session
+    supported = [ClientSession]
+    if hasattr(mcp, "Client"):
+        supported.append(mcp.Client)
+    fastmcp = sys.modules.get("fastmcp")
+    if fastmcp is not None and hasattr(fastmcp, "Client"):
+        supported.append(fastmcp.Client)
+    if not isinstance(client_or_session, tuple(supported)):
+        logger.warning("Unsupported MCP client %r. Not wrapping.", type(client_or_session).__name__)
+        return client_or_session
+    if not MCPIntegration.install_patchers():
+        logger.warning("Unsupported mcp version. Not wrapping.")
+        return client_or_session
+    config = ClientConfig(capture_content=capture_content)
+    if isinstance(client_or_session, ClientSession):
+        set_session_config(client_or_session, config)
+    else:
+        setattr(client_or_session, CLIENT_CONFIG_ATTR, config)
+        configure_session(client_or_session, config)
+        # Python resolves special methods on the class, not the instance.
+        ClientEnterPatcher.wrap_target(type(client_or_session))
+    return client_or_session
+
+
+__all__ = ["MCPIntegration", "wrap_mcp_client", "wrap_mcp_server"]

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

@@ -0,0 +1,235 @@
+"""Client hooks shared by MCP 1.x/2.x and standalone FastMCP.
+
+ClientSession.send_request owns typed operation spans. MCP 2.x raw dispatcher
+hooks also cover negotiation/custom requests that bypass that method, and defer
+to it when already active. High-level client hooks copy settings to each new
+session, including reconnects and initialization.
+
+Client span contract:
+* One span per eligible request round trip, through response parsing. Session
+  lifetimes, one-way notifications, and detached task APIs have no client span.
+* Names/types and content match the server contract: tools/call <name> is a
+  tool; prompts/get <name>, resources/read, and other methods are functions.
+  Discovery payloads are omitted. Content capture defaults to enabled.
+* metadata.mcp.span.kind = client distinguishes the round trip from execution;
+  protocol correlation/identity fields use the server's allowlist. Provenance
+  is mcp-auto. Client spans do not own model/token usage.
+* Parent is the active caller. Its outgoing span injects W3C traceparent,
+  tracestate and Braintrust baggage in params._meta, preserving other entries.
+* JSON-RPC errors (including caller errors) fail the client span. Returned tool
+  errors preserve the result and never change application exception behavior.
+  Runtime cancellation ends the span with mcp.cancelled and no error.
+* A framework tool span may parent this distinct network round trip. An enabled
+  opentelemetry-instrumentation-mcp owns client operations and propagation,
+  including its own capture settings; Braintrust adds no duplicate client span.
+"""
+
+import contextvars
+import sys
+from collections.abc import Mapping
+from typing import Any, get_args
+
+import anyio
+from braintrust import logger as bt_logger
+from braintrust.logger import inject_trace_context
+
+from .tracing import MCPServerConfig as ClientConfig
+from .tracing import ServerOperation, call_arg, contain, start_client_operation
+
+
+CLIENT_CONFIG_ATTR = "__braintrust_mcp_client_config__"
+
+
+# Pydantic include keys are Python field names, which differ across MCP majors.
+# Content-free spans need only outcome, protocol identity, and result metadata.
+_RESULT_METADATA_FIELDS = frozenset(
+    {"meta", "isError", "is_error", "protocolVersion", "protocol_version", "serverInfo", "server_info"}
+)
+
+
+_enter_config: contextvars.ContextVar[ClientConfig | None] = contextvars.ContextVar(
+    "braintrust_mcp_client_enter_config", default=None
+)
+_outgoing_operation: contextvars.ContextVar[tuple[ServerOperation, Any] | None] = contextvars.ContextVar(
+    "braintrust_mcp_client_operation", default=None
+)
+
+
+def upstream_owns_operations() -> bool:
+    """Defer to an explicitly enabled OTel MCP instrumentor, in either setup order.
+
+    Its client hooks own operation spans and propagation. Do not import optional
+    OTel packages or change their patch chains/configuration on the user's behalf.
+    """
+    module = sys.modules.get("opentelemetry.instrumentation.mcp")
+    instrumentor = getattr(module, "McpInstrumentor", None)
+    return bool(instrumentor is not None and instrumentor().is_instrumented_by_opentelemetry)
+
+
+def configure_session(client: Any, config: ClientConfig) -> None:
+    try:
+        session = client.session
+    except (AttributeError, RuntimeError):
+        # A client that has not connected has no session yet.
+        return
+    set_session_config(session, config)
+
+
+def set_session_config(session: Any, config: ClientConfig) -> None:
+    setattr(session, CLIENT_CONFIG_ATTR, config)
+    dispatcher = getattr(session, "_dispatcher", None)
+    if dispatcher is not None:
+        setattr(dispatcher, CLIENT_CONFIG_ATTR, config)
+
+
+async def client_enter_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
+    config = getattr(instance, CLIENT_CONFIG_ATTR, None)
+    if not isinstance(config, ClientConfig):
+        return await wrapped(*args, **kwargs)
+    token = _enter_config.set(config)
+    try:
+        result = await wrapped(*args, **kwargs)
+        contain(configure_session, instance, config)
+        return result
+    finally:
+        _enter_config.reset(token)
+
+
+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)
+    return result
+
+
+def _prepare(session: Any, request: Any, config: ClientConfig) -> tuple[ServerOperation, Any] | None:
+    data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
+    params = data.get("params") or {}
+    operation = start_client_operation(
+        method=data["method"],
+        params=params,
+        request_id=getattr(session, "_request_id", None),  # MCP 1.x allocates synchronously on entry.
+        config=config,
+        protocol_version=getattr(session, "_negotiated_version", None),
+    )
+    if operation is None:
+        return None
+    try:
+        # Work on telemetry/transport copies, never the caller's protocol object.
+        meta = dict(params.get("_meta") or {})
+        inject_trace_context(meta, span=operation.span)
+        from mcp import types
+
+        root = getattr(request, "root", request)
+        protocol_params = root.params
+        if protocol_params is None:
+            annotation = type(root).model_fields["params"].annotation
+            params_type = next(
+                (candidate for candidate in get_args(annotation) if hasattr(candidate, "model_fields")), annotation
+            )
+            protocol_params = params_type()
+        meta_type = getattr(types, "RequestParamsMeta", None) or types.RequestParams.Meta
+        # 1.x Meta is a Pydantic model; 2.x RequestParamsMeta is a TypedDict.
+        protocol_params = protocol_params.model_copy(update={"meta": meta_type(**meta)})
+        copied_root = root.model_copy(update={"params": protocol_params})
+        copied = request.model_copy(update={"root": copied_root}) if root is not request else copied_root
+    except Exception:
+        # An unfamiliar extension shape must still execute unchanged.
+        copied = request
+    return operation, copied
+
+
+def _record_exception(operation: ServerOperation, error: BaseException) -> None:
+    if isinstance(error, anyio.get_cancelled_exc_class()):
+        operation.mark_cancelled()
+        return
+    from mcp.shared import exceptions
+
+    error_class = getattr(exceptions, "MCPError", None) or getattr(exceptions, "McpError")
+    if isinstance(error, error_class):
+        operation.record_error_response(error.error.code, error.error.message)
+    else:
+        operation.record_exception(error)
+
+
+async def send_request_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
+    config = getattr(instance, CLIENT_CONFIG_ATTR, None)
+    if not isinstance(config, ClientConfig):
+        return await wrapped(*args, **kwargs)
+    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.
+        from mcp.shared.session import BaseSession  # pylint: disable=import-error,no-name-in-module  # MCP 1.x
+
+        return await BaseSession.send_request(instance, *args, **kwargs)
+    prepared = contain(_prepare, instance, call_arg(args, kwargs, 0, "request"), config)
+    if prepared is None:
+        return await wrapped(*args, **kwargs)
+    operation, request = prepared
+    if args:
+        args = (request, *args[1:])
+    else:
+        kwargs = {**kwargs, "request": request}
+    return await _run_request(operation, getattr(instance, "_dispatcher", None), wrapped, args, kwargs)
+
+
+async def _run_request(operation: ServerOperation, dispatcher: Any, wrapped: Any, args: Any, kwargs: Any) -> Any:
+    token = _outgoing_operation.set((operation, dispatcher))
+    manager = bt_logger._state.context_manager
+    span_token = manager.set_current_span(operation.span)
+    try:
+        result = await wrapped(*args, **kwargs)
+        include = None if operation.captures_output else _RESULT_METADATA_FIELDS
+        payload = (
+            result
+            if isinstance(result, Mapping)
+            else contain(lambda: result.model_dump(by_alias=True, mode="json", exclude_none=True, include=include))
+        )
+        contain(operation.record_result, payload)
+        return result
+    except BaseException as error:
+        contain(_record_exception, operation, error)
+        raise
+    finally:
+        manager.unset_current_span(span_token)
+        _outgoing_operation.reset(token)
+        contain(operation.finish)
+
+
+def _prepare_raw(method: str, params: Any, config: ClientConfig) -> tuple[ServerOperation, Any] | None:
+    params = params or {}
+    meta = dict(params.get("_meta") or {})
+    operation = start_client_operation(method=method, params=params, request_id=None, config=config)
+    if operation is None:
+        return None
+    inject_trace_context(meta, span=operation.span)
+    return operation, {**params, "_meta": meta}
+
+
+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)
+    active = _outgoing_operation.get()
+    if not isinstance(config, ClientConfig) or (active is not None and active[1] is instance):
+        return await wrapped(*args, **kwargs)
+    prepared = contain(_prepare_raw, call_arg(args, kwargs, 0, "method"), call_arg(args, kwargs, 1, "params"), config)
+    if prepared is None:
+        return await wrapped(*args, **kwargs)
+    operation, params = prepared
+    if len(args) > 1:
+        args = (args[0], params, *args[2:])
+    else:
+        kwargs = {**kwargs, "params": params}
+    return await _run_request(operation, instance, wrapped, args, kwargs)
+
+
+def allocate_id_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
+    request_id = wrapped(*args, **kwargs)
+    active = _outgoing_operation.get()
+    if active is not None and active[1] is instance:
+        operation = active[0]
+        operation.request_id = request_id
+        operation.span.log(metadata={"jsonrpc.request.id": str(request_id)})
+    return request_id