feat(mcp): add instrumentation support for Streamable-HTTP transport and bump mcp dependency to ≥1.8.1 (#1640)

Author: rajephon

python/instrumentation/openinference-instrumentation-mcp/pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
 
 [project.optional-dependencies]
 instruments = [
-  "mcp >= 1.8.0",
+  "mcp >= 1.8.1",
 ]
 
 [project.entry-points.opentelemetry_instrumentor]

python/instrumentation/openinference-instrumentation-mcp/src/openinference/instrumentation/mcp/__init__.py

@@ -21,25 +21,43 @@ def instrumentation_dependencies(self) -> Collection[str]:
     def _instrument(self, **kwargs: Any) -> None:
         register_post_import_hook(
             lambda _: wrap_function_wrapper(
-                "mcp.client.sse", "sse_client", self._transport_wrapper
+                "mcp.client.streamable_http",
+                "streamablehttp_client",
+                self._wrap_transport_with_callback,
+            ),
+            "mcp.client.streamable_http",
+        )
+
+        register_post_import_hook(
+            lambda _: wrap_function_wrapper(
+                "mcp.server.streamable_http",
+                "StreamableHTTPServerTransport.connect",
+                self._wrap_plain_transport,
+            ),
+            "mcp.server.streamable_http",
+        )
+
+        register_post_import_hook(
+            lambda _: wrap_function_wrapper(
+                "mcp.client.sse", "sse_client", self._wrap_plain_transport
             ),
             "mcp.client.sse",
         )
         register_post_import_hook(
             lambda _: wrap_function_wrapper(
-                "mcp.server.sse", "SseServerTransport.connect_sse", self._transport_wrapper
+                "mcp.server.sse", "SseServerTransport.connect_sse", self._wrap_plain_transport
             ),
             "mcp.server.sse",
         )
         register_post_import_hook(
             lambda _: wrap_function_wrapper(
-                "mcp.client.stdio", "stdio_client", self._transport_wrapper
+                "mcp.client.stdio", "stdio_client", self._wrap_plain_transport
             ),
             "mcp.client.stdio",
         )
         register_post_import_hook(
             lambda _: wrap_function_wrapper(
-                "mcp.server.stdio", "stdio_server", self._transport_wrapper
+                "mcp.server.stdio", "stdio_server", self._wrap_plain_transport
             ),
             "mcp.server.stdio",
         )
@@ -63,7 +81,18 @@ def _uninstrument(self, **kwargs: Any) -> None:
         unwrap("mcp.server.stdio", "stdio_server")
 
     @asynccontextmanager
-    async def _transport_wrapper(
+    async def _wrap_transport_with_callback(
+        self, wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any
+    ) -> AsyncGenerator[Tuple["InstrumentedStreamReader", "InstrumentedStreamWriter", Any], None]:
+        async with wrapped(*args, **kwargs) as (read_stream, write_stream, get_session_id_callback):
+            yield (
+                InstrumentedStreamReader(read_stream),
+                InstrumentedStreamWriter(write_stream),
+                get_session_id_callback,
+            )
+
+    @asynccontextmanager
+    async def _wrap_plain_transport(
         self, wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any
     ) -> AsyncGenerator[Tuple["InstrumentedStreamReader", "InstrumentedStreamWriter"], None]:
         async with wrapped(*args, **kwargs) as (read_stream, write_stream):

python/instrumentation/openinference-instrumentation-mcp/test-requirements.txt

@@ -1,4 +1,4 @@
-mcp==1.8.0
+mcp==1.8.1
 
 httpx
 opentelemetry-exporter-otlp-proto-http

python/instrumentation/openinference-instrumentation-mcp/tests/mcpserver.py

@@ -7,7 +7,7 @@
 
 from openinference.instrumentation.mcp import MCPInstrumentor
 
-transport = cast(Literal["sse", "stdio"], os.environ.get("MCP_TRANSPORT"))
+transport = cast(Literal["sse", "stdio", "streamable-http"], os.environ.get("MCP_TRANSPORT"))
 otlp_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
 span_exporter = OTLPSpanExporter(f"{otlp_endpoint}/v1/traces")
 tracer_provider = trace_sdk.TracerProvider()

python/instrumentation/openinference-instrumentation-mcp/tests/test_instrumenter.py

@@ -26,6 +26,7 @@ async def mcp_client(
     # instrumentation through fixtures instead.
     from mcp.client.sse import sse_client
     from mcp.client.stdio import StdioServerParameters, stdio_client
+    from mcp.client.streamable_http import streamablehttp_client
 
     async def message_handler(
         message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
@@ -86,9 +87,41 @@ async def message_handler(
             finally:
                 proc.kill()
                 await proc.wait()
+        case "streamable-http":
+            proc = await asyncio.create_subprocess_exec(
+                sys.executable,
+                server_script,
+                env={
+                    "MCP_TRANSPORT": "streamable-http",
+                    "OTEL_EXPORTER_OTLP_ENDPOINT": otlp_endpoint,
+                    "PYTHONPATH": pythonpath,
+                },
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            try:
+                stderr = proc.stderr
+                assert stderr is not None
+                for i in range(100):
+                    line = str(await stderr.readline())
+                    if "Uvicorn running on http://0.0.0.0:" in line:
+                        _, rest = line.split("http://0.0.0.0:", 1)
+                        port, _ = rest.split(" ", 1)
+                        async with streamablehttp_client(f"http://localhost:{port}/mcp") as (
+                            reader,
+                            writer,
+                            _,
+                        ), ClientSession(reader, writer, message_handler=message_handler) as client:
+                            client._receive_request_type = TestServerRequest
+                            await client.initialize()
+                            yield client
+                        break
+            finally:
+                proc.kill()
+                await proc.wait()
 
 
-@pytest.mark.parametrize("transport", ["sse", "stdio"])
+@pytest.mark.parametrize("transport", ["sse", "stdio", "streamable-http"])
 async def test_hello(
     transport: str, tracer: Tracer, telemetry: Telemetry, otlp_collector: OTLPServer
 ) -> None: