diff --git a/packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/fastmcp_instrumentation.py b/packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/fastmcp_instrumentation.py index 64fdb86024..0adaff1abd 100644 --- a/packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/fastmcp_instrumentation.py +++ b/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)""" diff --git a/packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py b/packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py index 688cdb2c7f..3f4366d436 100644 --- a/packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py +++ b/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) diff --git a/packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/utils.py b/packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/utils.py index d4a80e58dc..8968296286 100644 --- a/packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/utils.py +++ b/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__, diff --git a/packages/opentelemetry-instrumentation-mcp/tests/test_content_capture_gate.py b/packages/opentelemetry-instrumentation-mcp/tests/test_content_capture_gate.py new file mode 100644 index 0000000000..d8cbde5a9a --- /dev/null +++ b/packages/opentelemetry-instrumentation-mcp/tests/test_content_capture_gate.py @@ -0,0 +1,124 @@ +"""TRACELOOP_TRACE_CONTENT must gate the MCP client path, not only FastMCP. + +The package documents TRACELOOP_TRACE_CONTENT as the switch that disables content +logging. Before this test, only the FastMCP server-side wrapper consulted it: the +client path (tools/call arguments, non-tool request bodies, response bodies) recorded +content regardless, so an operator who turned the switch off still got request and +response payloads on their spans. + +Each test drives the real client wrapper with a marker value and asserts the marker +is absent from every span attribute when content capture is off, and present when it +is on, so the test fails if either the gate or the capture itself regresses. +""" + +import json + +from fastmcp import Client, FastMCP + +MARKER = "content-capture-marker-9f3a" + + +def _all_attribute_text(span_exporter) -> str: + """Every attribute value across every exported span, as one string.""" + chunks = [] + for span in span_exporter.get_finished_spans(): + for value in (span.attributes or {}).values(): + if isinstance(value, (list, tuple)): + chunks.extend(str(item) for item in value) + else: + chunks.append(str(value)) + return "\n".join(chunks) + + +def _server() -> FastMCP: + """Build a server with one tool that echoes a caller-supplied token.""" + server = FastMCP("content-gate-server") + + @server.tool() + async def echo_secret(token: str) -> str: + """Echo back a caller-supplied token.""" + return f"received {token}" + + return server + + +async def test_tool_arguments_suppressed_when_content_capture_off( + span_exporter, monkeypatch +) -> None: + """With the switch off, tool arguments must not appear on any span.""" + monkeypatch.setenv("TRACELOOP_TRACE_CONTENT", "false") + + async with Client(_server()) as client: + await client.call_tool("echo_secret", {"token": MARKER}) + + assert span_exporter.get_finished_spans(), "expected the tool call to be traced" + assert MARKER not in _all_attribute_text(span_exporter) + + +async def test_tool_arguments_captured_when_content_capture_on( + span_exporter, monkeypatch +) -> None: + """With the switch on, tool arguments are still recorded as before.""" + monkeypatch.setenv("TRACELOOP_TRACE_CONTENT", "true") + + async with Client(_server()) as client: + await client.call_tool("echo_secret", {"token": MARKER}) + + # The gate must not silently disable capture altogether: with the switch on, + # the argument is still recorded. + assert MARKER in _all_attribute_text(span_exporter) + + +async def test_non_tool_response_body_suppressed_when_content_capture_off( + span_exporter, monkeypatch +) -> None: + """list_tools goes through _handle_mcp_method, which serialized the whole response. + + The marker lives in the registered tool description, so it travels back in the + list_tools result and exercises the response-serialization path. + """ + monkeypatch.setenv("TRACELOOP_TRACE_CONTENT", "false") + + server = FastMCP("content-gate-server") + + @server.tool(description=f"A tool whose description carries {MARKER}.") + async def documented(arg: str) -> str: + """A tool that exists only to carry the marker in its description.""" + return arg + + async with Client(server) as client: + tools = await client.list_tools() + + assert any(MARKER in (t.description or "") for t in tools), ( + "the marker must reach the client, otherwise this test proves nothing" + ) + assert span_exporter.get_finished_spans(), "expected the request to be traced" + assert MARKER not in _all_attribute_text(span_exporter) + + +async def test_span_structure_survives_content_capture_off( + span_exporter, monkeypatch +) -> None: + """Turning content off must not remove spans or their non-content attributes.""" + monkeypatch.setenv("TRACELOOP_TRACE_CONTENT", "false") + + async with Client(_server()) as client: + await client.call_tool("echo_secret", {"token": MARKER}) + + spans = span_exporter.get_finished_spans() + tool_spans = [s for s in spans if s.name.endswith(".tool")] + assert tool_spans, f"expected a tool span, got {[s.name for s in spans]}" + + entity_names = [ + (s.attributes or {}).get("traceloop.entity.name") for s in tool_spans + ] + assert "echo_secret" in entity_names + + # Structural attributes stay; only content is withheld. + for span in tool_spans: + attributes = span.attributes or {} + assert "traceloop.span.kind" in attributes + for key in ("traceloop.entity.input", "traceloop.entity.output"): + if key in attributes: + json.loads(attributes[key]) # if present it must still be valid JSON + assert MARKER not in attributes[key]