diff --git a/python/instrumentation/openinference-instrumentation-mcp/src/openinference/instrumentation/mcp/__init__.py b/python/instrumentation/openinference-instrumentation-mcp/src/openinference/instrumentation/mcp/__init__.py index a22d05d890..2930633da3 100644 --- a/python/instrumentation/openinference-instrumentation-mcp/src/openinference/instrumentation/mcp/__init__.py +++ b/python/instrumentation/openinference-instrumentation-mcp/src/openinference/instrumentation/mcp/__init__.py @@ -61,6 +61,14 @@ def _instrument(self, **kwargs: Any) -> None: ), "mcp.server.stdio", ) + register_post_import_hook( + lambda _: wrap_function_wrapper( + "mcp.server.lowlevel.server", + "Server._handle_request", + self._wrap_handle_request, + ), + "mcp.server.lowlevel.server", + ) # While we prefer to instrument the lowest level primitive, the transports above, it doesn't # mean context will be propagated to handlers automatically. Notably, the MCP SDK passes @@ -79,6 +87,7 @@ def _instrument(self, **kwargs: Any) -> None: def _uninstrument(self, **kwargs: Any) -> None: unwrap("mcp.client.stdio", "stdio_client") unwrap("mcp.server.stdio", "stdio_server") + unwrap("mcp.client.session", "ClientSession.call_tool") @asynccontextmanager async def _wrap_transport_with_callback( @@ -98,6 +107,21 @@ async def _wrap_plain_transport( async with wrapped(*args, **kwargs) as (read_stream, write_stream): yield InstrumentedStreamReader(read_stream), InstrumentedStreamWriter(write_stream) + async def _wrap_handle_request( + self, wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any + ) -> Any: + token = None + try: + # Message has been deserialized, we need to extract the traceparent + _meta = {"traceparent": args[1].params.meta.traceparent} + ctx = propagate.extract(_meta) + token = context.attach(ctx) + finally: + res = await wrapped(*args, **kwargs) + if token: + context.detach(token) + return res + def _base_session_init_wrapper( self, wrapped: Callable[..., None], instance: Any, args: Any, kwargs: Any ) -> None: diff --git a/python/instrumentation/openinference-instrumentation-mcp/src/openinference/instrumentation/mcp/version.py b/python/instrumentation/openinference-instrumentation-mcp/src/openinference/instrumentation/mcp/version.py index 67bc602abf..9c73af26be 100644 --- a/python/instrumentation/openinference-instrumentation-mcp/src/openinference/instrumentation/mcp/version.py +++ b/python/instrumentation/openinference-instrumentation-mcp/src/openinference/instrumentation/mcp/version.py @@ -1 +1 @@ -__version__ = "1.3.0" +__version__ = "1.3.1"