fix(mcp): honor TRACELOOP_TRACE_CONTENT on the client path

packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/fastmcp_instrumentation.py

@@ -9,13 +9,14 @@
 from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 from wrapt import register_post_import_hook, wrap_function_wrapper
 
-from .utils import dont_throw
+from .utils import dont_throw, should_send_prompts
 
 
 class FastMCPInstrumentor:
     """Handles FastMCP-specific instrumentation logic."""
 
     def __init__(self):
+        """Create the instrumentor with no tracer or server name bound yet."""
         self._tracer = None
         self._server_name = None
 
@@ -50,6 +51,7 @@ def _fastmcp_init_wrapper(self):
         @dont_throw
         def traced_method(wrapped, instance, args, kwargs):
             # Call the original __init__ first
+            """Record the server name from FastMCP's constructor arguments."""
             result = wrapped(*args, **kwargs)
 
             if args and len(args) > 0:
@@ -63,6 +65,7 @@ def traced_method(wrapped, instance, args, kwargs):
     def _fastmcp_tool_wrapper(self):
         """Create wrapper for FastMCP tool execution."""
         async def traced_method(wrapped, instance, args, kwargs):
+            """Wrap a FastMCP tool call in server and tool spans."""
             if not self._tracer:
                 return await wrapped(*args, **kwargs)
 
@@ -155,9 +158,7 @@ async def traced_method(wrapped, instance, args, kwargs):
 
     def _should_send_prompts(self):
         """Check if content tracing is enabled (matches traceloop SDK)"""
-        return (
-            os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
-        ).lower() == "true"
+        return should_send_prompts()
 
     def _get_json_encoder(self):
         """Get JSON encoder class (simplified - traceloop SDK uses custom JSONEncoder)"""

packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py

@@ -15,7 +15,11 @@
 from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 
 from opentelemetry.instrumentation.mcp.version import __version__
-from opentelemetry.instrumentation.mcp.utils import dont_throw, Config
+from opentelemetry.instrumentation.mcp.utils import (
+    Config,
+    dont_throw,
+    should_send_prompts,
+)
 from opentelemetry.instrumentation.mcp.fastmcp_instrumentation import (
     FastMCPInstrumentor,
 )
@@ -24,15 +28,19 @@
 
 
 class McpInstrumentor(BaseInstrumentor):
+    """Instrument the MCP client, server and transports with OpenTelemetry spans."""
     def __init__(self, exception_logger=None):
+        """Store the exception logger and build the FastMCP sub-instrumentor."""
         super().__init__()
         Config.exception_logger = exception_logger
         self._fastmcp_instrumentor = FastMCPInstrumentor()
 
     def instrumentation_dependencies(self) -> Collection[str]:
+        """Return the package versions this instrumentation supports."""
         return _instruments
 
     def _instrument(self, **kwargs):
+        """Wrap the MCP client, server sessions and every supported transport."""
         tracer_provider = kwargs.get("tracer_provider")
         tracer = get_tracer(__name__, __version__, tracer_provider)
 
@@ -114,11 +122,13 @@ def _instrument(self, **kwargs):
         )
 
     def _uninstrument(self, **kwargs):
+        """Unwrap the transports this instrumentation replaced."""
         unwrap("mcp.client.stdio", "stdio_client")
         unwrap("mcp.server.stdio", "stdio_server")
         self._fastmcp_instrumentor.uninstrument()
 
     def _transport_wrapper(self, tracer):
+        """Wrap a transport so its read and write streams are instrumented."""
         @asynccontextmanager
         async def traced_method(
             wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any
@@ -129,6 +139,7 @@ async def traced_method(
             ],
             None,
         ]:
+            """Yield the transport's streams wrapped in instrumented proxies."""
             async with wrapped(*args, **kwargs) as result:
                 try:
                     read_stream, write_stream = result
@@ -157,9 +168,11 @@ async def traced_method(
         return traced_method
 
     def _base_session_init_wrapper(self, tracer):
+        """Wrap a server session's incoming message streams to carry trace context."""
         def traced_method(
             wrapped: Callable[..., None], instance: Any, args: Any, kwargs: Any
         ) -> None:
+            """Replace the session's incoming stream pair with context-propagating proxies."""
             wrapped(*args, **kwargs)
             reader = getattr(instance, "_incoming_message_stream_reader", None)
             writer = getattr(instance, "_incoming_message_stream_writer", None)
@@ -178,8 +191,10 @@ def traced_method(
         return traced_method
 
     def patch_mcp_client(self, tracer: Tracer):
+        """Wrap BaseSession.send_request so each MCP request becomes a span."""
         @dont_throw
         async def traced_method(wrapped, instance, args, kwargs):
+            """Start a span for the outgoing request and inject trace context into its meta."""
             meta = None
             method = None
             params = None
@@ -216,6 +231,7 @@ def _fastmcp_client_enter_wrapper(self, tracer):
         @dont_throw
         async def traced_method(wrapped, instance, args, kwargs):
             # Start a root span for the MCP client session and make it current
+            """Wrap a FastMCP client session enter to open a session span."""
             span_context_manager = tracer.start_as_current_span("mcp.client.session")
             span = span_context_manager.__enter__()
             span.set_attribute(SpanAttributes.TRACELOOP_SPAN_KIND, "session")
@@ -243,6 +259,7 @@ def _fastmcp_client_exit_wrapper(self, tracer):
 
         @dont_throw
         async def traced_method(wrapped, instance, args, kwargs):
+            """Close the session span when the FastMCP client exits."""
             try:
                 # Call the original method first
                 result = await wrapped(*args, **kwargs)
@@ -289,8 +306,13 @@ async def _handle_tool_call(self, tracer, method, params, args, kwargs, wrapped)
             )
             span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, entity_name)
 
-            # Add input
-            clean_input = self._extract_clean_input(method, params)
+            # Add input. Tool arguments are request content, so they are
+            # recorded only when content capture is enabled.
+            clean_input = (
+                self._extract_clean_input(method, params)
+                if should_send_prompts()
+                else None
+            )
             if clean_input:
                 try:
                     span.set_attribute(
@@ -308,9 +330,12 @@ async def _handle_tool_call(self, tracer, method, params, args, kwargs, wrapped)
     async def _handle_mcp_method(self, tracer, method, args, kwargs, wrapped):
         """Handle non-tool MCP methods with simple serialization"""
         with tracer.start_as_current_span(f"{method}.mcp") as span:
-            span.set_attribute(
-                SpanAttributes.TRACELOOP_ENTITY_INPUT, f"{serialize(args[0])}"
-            )
+            # The serialized request is content: it carries caller-supplied
+            # params, so it is recorded only when content capture is enabled.
+            if should_send_prompts():
+                span.set_attribute(
+                    SpanAttributes.TRACELOOP_ENTITY_INPUT, f"{serialize(args[0])}"
+                )
             return await self._execute_and_handle_result(
                 span, method, args, kwargs, wrapped, clean_output=False
             )
@@ -321,8 +346,11 @@ async def _execute_and_handle_result(
         """Execute the wrapped function and handle the result"""
         try:
             result = await wrapped(*args, **kwargs)
-            # Add output
-            if clean_output:
+            # Add output. The response body is content, so it is recorded only
+            # when content capture is enabled.
+            if not should_send_prompts():
+                pass
+            elif clean_output:
                 clean_output_data = self._extract_clean_output(method, result)
                 if clean_output_data:
                     try:
@@ -457,6 +485,7 @@ def serialize(request, depth=0, max_depth=4):
     depth += 1
 
     def is_serializable(request):
+        """Return whether a value can be JSON-encoded without a fallback."""
         try:
             json.dumps(request)
             return True
@@ -492,18 +521,23 @@ def is_serializable(request):
 
 class InstrumentedStreamReader(ObjectProxy):  # type: ignore
     # ObjectProxy missing context manager - https://github.com/GrahamDumpleton/wrapt/issues/73
+    """Stream reader proxy that extracts trace context from incoming messages."""
     def __init__(self, wrapped, tracer):
+        """Wrap a stream reader and keep the tracer used for its spans."""
         super().__init__(wrapped)
         self._tracer = tracer
 
     async def __aenter__(self) -> Any:
+        """Enter the wrapped stream reader."""
         return await self.__wrapped__.__aenter__()
 
     async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> Any:
+        """Exit the wrapped stream reader."""
         return await self.__wrapped__.__aexit__(exc_type, exc_value, traceback)
 
     @dont_throw
     async def __aiter__(self) -> AsyncGenerator[Any, None]:
+        """Iterate the wrapped reader, restoring trace context from each message."""
         from mcp.types import JSONRPCMessage, JSONRPCRequest
 
         async for item in self.__wrapped__:
@@ -538,18 +572,23 @@ async def __aiter__(self) -> AsyncGenerator[Any, None]:
 
 class InstrumentedStreamWriter(ObjectProxy):  # type: ignore
     # ObjectProxy missing context manager - https://github.com/GrahamDumpleton/wrapt/issues/73
+    """Stream writer proxy that records outgoing responses on a span."""
     def __init__(self, wrapped, tracer):
+        """Wrap a stream writer and keep the tracer used for its spans."""
         super().__init__(wrapped)
         self._tracer = tracer
 
     async def __aenter__(self) -> Any:
+        """Enter the wrapped stream writer."""
         return await self.__wrapped__.__aenter__()
 
     async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> Any:
+        """Exit the wrapped stream writer."""
         return await self.__wrapped__.__aexit__(exc_type, exc_value, traceback)
 
     @dont_throw
     async def send(self, item: Any) -> Any:
+        """Record the outgoing response on a span and forward it to the wrapped stream."""
         from mcp.types import JSONRPCMessage, JSONRPCRequest
 
         # Handle different item types based on what's available
@@ -565,9 +604,13 @@ async def send(self, item: Any) -> Any:
 
         with self._tracer.start_as_current_span("ResponseStreamWriter") as span:
             if hasattr(request, "result"):
-                span.set_attribute(
-                    SpanAttributes.MCP_RESPONSE_VALUE, f"{serialize(request.result)}"
-                )
+                # The response body is content; the error status below is not,
+                # so only the value itself is gated.
+                if should_send_prompts():
+                    span.set_attribute(
+                        SpanAttributes.MCP_RESPONSE_VALUE,
+                        f"{serialize(request.result)}",
+                    )
                 if "isError" in request.result:
                     if request.result["isError"] is True:
                         span.set_status(
@@ -592,42 +635,53 @@ async def send(self, item: Any) -> Any:
 
 @dataclass(slots=True, frozen=True)
 class ItemWithContext:
+    """A stream item paired with the OpenTelemetry context it was written under."""
     item: Any
     ctx: context.Context
 
 
 class ContextSavingStreamWriter(ObjectProxy):  # type: ignore
     # ObjectProxy missing context manager - https://github.com/GrahamDumpleton/wrapt/issues/73
+    """Stream writer proxy that attaches the current context to each item."""
     def __init__(self, wrapped, tracer):
+        """Wrap a stream writer and keep the tracer used for its spans."""
         super().__init__(wrapped)
         self._tracer = tracer
 
     async def __aenter__(self) -> Any:
+        """Enter the wrapped stream writer."""
         return await self.__wrapped__.__aenter__()
 
     async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> Any:
+        """Exit the wrapped stream writer."""
         return await self.__wrapped__.__aexit__(exc_type, exc_value, traceback)
 
     @dont_throw
     async def send(self, item: Any) -> Any:
         # Removed RequestStreamWriter span creation - we don't need low-level protocol spans
+        """Forward the item together with the context it was sent under."""
         ctx = context.get_current()
         return await self.__wrapped__.send(ItemWithContext(item, ctx))
 
 
 class ContextAttachingStreamReader(ObjectProxy):  # type: ignore
     # ObjectProxy missing context manager - https://github.com/GrahamDumpleton/wrapt/issues/73
+    """Stream reader proxy that restores each item's saved context while it is handled."""
     def __init__(self, wrapped, tracer):
+        """Wrap a stream reader and keep the tracer used for its spans."""
         super().__init__(wrapped)
         self._tracer = tracer
 
     async def __aenter__(self) -> Any:
+        """Enter the wrapped stream reader."""
         return await self.__wrapped__.__aenter__()
 
     async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> Any:
+        """Exit the wrapped stream reader."""
         return await self.__wrapped__.__aexit__(exc_type, exc_value, traceback)
 
     async def __aiter__(self) -> AsyncGenerator[Any, None]:
+        """Yield each item with its saved context attached, detaching it afterwards."""
         async for item in self.__wrapped__:
             item_with_context = cast(ItemWithContext, item)
             restore = context.attach(item_with_context.ctx)

packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/utils.py

@@ -2,13 +2,26 @@
 
 import asyncio
 import logging
+import os
 import traceback
 
 
 class Config:
+    """Module-level configuration for the MCP instrumentation."""
     exception_logger = None
 
 
+def should_send_prompts() -> bool:
+    """Whether request/response content may be recorded on spans.
+
+    Mirrors the traceloop SDK's ``TRACELOOP_TRACE_CONTENT`` switch: content
+    capture is on unless an operator explicitly turns it off. Shared by the
+    FastMCP server wrapper and the MCP client path so a single environment
+    variable governs both, which is what the package README documents.
+    """
+    return (os.getenv("TRACELOOP_TRACE_CONTENT") or "true").lower() == "true"
+
+
 def dont_throw(func):
     """
     A decorator that wraps the passed in function and logs exceptions instead of throwing them.
@@ -17,18 +30,21 @@ def dont_throw(func):
     logger = logging.getLogger(func.__module__)
 
     async def async_wrapper(*args, **kwargs):
+        """Await the wrapped coroutine, logging instead of raising on failure."""
         try:
             return await func(*args, **kwargs)
         except Exception as e:
             _handle_exception(e, func, logger)
 
     def sync_wrapper(*args, **kwargs):
+        """Call the wrapped function, logging instead of raising on failure."""
         try:
             return func(*args, **kwargs)
         except Exception as e:
             _handle_exception(e, func, logger)
 
     def _handle_exception(e, func, logger):
+        """Log a tracing failure and hand it to the configured exception logger."""
         logger.debug(
             "OpenLLMetry failed to trace in %s, error: %s",
             func.__name__,