From e955392742152c294e063665ed8c99e141cd1f4a Mon Sep 17 00:00:00 2001 From: Sebastian Sosa <1sebastian1sosa1@gmail.com> Date: Tue, 20 May 2025 02:36:35 -0400 Subject: [PATCH 1/6] fix: propagate context to MCP server --- .../instrumentation/mcp/__init__.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) 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 a22d05d89..c2af2fa3a 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 @@ -98,6 +106,18 @@ async def _wrap_plain_transport( async with wrapped(*args, **kwargs) as (read_stream, write_stream): yield InstrumentedStreamReader(read_stream), InstrumentedStreamWriter(write_stream) + def _wrap_handle_request( + self, wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any + ) -> Any: + try: + # Message has been deserialized, we need to extract the traceparent + _meta = {"traceparent": args[1].params.meta.traceparent} + ctx = propagate.extract(_meta) + context.attach(ctx) + return wrapped(*args, **kwargs) + finally: + pass + def _base_session_init_wrapper( self, wrapped: Callable[..., None], instance: Any, args: Any, kwargs: Any ) -> None: From 53b2c22b56f22eb42139a3810674d58efc730756 Mon Sep 17 00:00:00 2001 From: Sebastian Sosa <1sebastian1sosa1@gmail.com> Date: Tue, 20 May 2025 16:38:56 -0400 Subject: [PATCH 2/6] mcp instrumentation version --- .../src/openinference/instrumentation/mcp/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 a955fdae1..9c73af26b 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.2.1" +__version__ = "1.3.1" From ad60941835e00fe4198d194f675c556e35e2a1bf Mon Sep 17 00:00:00 2001 From: Sebastian Sosa <1sebastian1sosa1@gmail.com> Date: Wed, 21 May 2025 23:08:40 -0400 Subject: [PATCH 3/6] fix: always run wrapped function --- .../src/openinference/instrumentation/mcp/__init__.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) 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 c2af2fa3a..d0741270b 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 @@ -114,9 +114,8 @@ def _wrap_handle_request( _meta = {"traceparent": args[1].params.meta.traceparent} ctx = propagate.extract(_meta) context.attach(ctx) - return wrapped(*args, **kwargs) finally: - pass + return wrapped(*args, **kwargs) def _base_session_init_wrapper( self, wrapped: Callable[..., None], instance: Any, args: Any, kwargs: Any From dd7f0fdd36e55b2b16ac052c82b0b6ec28712094 Mon Sep 17 00:00:00 2001 From: Sebastian Sosa <1sebastian1sosa1@gmail.com> Date: Thu, 22 May 2025 13:28:02 -0400 Subject: [PATCH 4/6] fix: must detach content to avoid trace bleed --- .../src/openinference/instrumentation/mcp/__init__.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) 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 d0741270b..75f3935ea 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 @@ -106,16 +106,20 @@ async def _wrap_plain_transport( async with wrapped(*args, **kwargs) as (read_stream, write_stream): yield InstrumentedStreamReader(read_stream), InstrumentedStreamWriter(write_stream) - def _wrap_handle_request( + 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) - context.attach(ctx) + token = context.attach(ctx) finally: - return wrapped(*args, **kwargs) + 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 From 3a2ffc927896cf537722190629f6c19d5c413573 Mon Sep 17 00:00:00 2001 From: Sebastian Sosa <1sebastian1sosa1@gmail.com> Date: Thu, 22 May 2025 22:42:40 -0400 Subject: [PATCH 5/6] uninstrument --- .../src/openinference/instrumentation/mcp/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 75f3935ea..9de6c083f 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 @@ -87,7 +87,8 @@ 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( self, wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any From 147f75d4f24f45a093d798676acaf201ff63cf13 Mon Sep 17 00:00:00 2001 From: Sebastian Sosa <1sebastian1sosa1@gmail.com> Date: Fri, 23 May 2025 14:03:18 -0400 Subject: [PATCH 6/6] fix: formatting --- .../src/openinference/instrumentation/mcp/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 9de6c083f..2930633da 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 @@ -88,7 +88,7 @@ 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( self, wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any