MCP Streamable HTTP Spans (DO NOT MERGE YET) (#1462)

* Add tool manager span.

* Add tool manager spans to MCP instrumentation to support streaming.

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: umaannamalai

newrelic/config.py

@@ -2851,6 +2851,8 @@ def _process_module_builtin_defaults():
         "instrument_autogen_agentchat_agents__assistant_agent",
     )
     _process_module_definition("mcp.client.session", "newrelic.hooks.adapter_mcp", "instrument_mcp_client_session")
+    _process_module_definition("mcp.server.fastmcp.tools.tool_manager", "newrelic.hooks.adapter_mcp", "instrument_mcp_server_fastmcp_tools_tool_manager")
+
 
     _process_module_definition("structlog._base", "newrelic.hooks.logger_structlog", "instrument_structlog__base")
     _process_module_definition("structlog._frames", "newrelic.hooks.logger_structlog", "instrument_structlog__frames")

newrelic/hooks/adapter_mcp.py

@@ -19,6 +19,8 @@
 from newrelic.common.object_names import callable_name
 from newrelic.common.object_wrapper import wrap_function_wrapper
 from newrelic.common.signature import bind_args
+from newrelic.core.config import global_settings
+
 
 _logger = logging.getLogger(__name__)
 
@@ -28,6 +30,10 @@ async def wrap_call_tool(wrapped, instance, args, kwargs):
     if not transaction:
         return await wrapped(*args, **kwargs)
 
+    settings = transaction.settings if transaction.settings is not None else global_settings()
+    if not settings.ai_monitoring.enabled:
+        return await wrapped(*args, **kwargs)
+
     func_name = callable_name(wrapped)
     bound_args = bind_args(wrapped, args, kwargs)
     tool_name = bound_args.get("name") or "tool"
@@ -42,6 +48,10 @@ async def wrap_read_resource(wrapped, instance, args, kwargs):
     if not transaction:
         return await wrapped(*args, **kwargs)
 
+    settings = transaction.settings if transaction.settings is not None else global_settings()
+    if not settings.ai_monitoring.enabled:
+        return await wrapped(*args, **kwargs)
+
     func_name = callable_name(wrapped)
     bound_args = bind_args(wrapped, args, kwargs)
     # Set a default value in case we can't parse out the URI scheme successfully
@@ -64,6 +74,10 @@ async def wrap_get_prompt(wrapped, instance, args, kwargs):
     if not transaction:
         return await wrapped(*args, **kwargs)
 
+    settings = transaction.settings if transaction.settings is not None else global_settings()
+    if not settings.ai_monitoring.enabled:
+        return await wrapped(*args, **kwargs)
+
     func_name = callable_name(wrapped)
     bound_args = bind_args(wrapped, args, kwargs)
     prompt_name = bound_args.get("name") or "prompt"
@@ -81,3 +95,9 @@ def instrument_mcp_client_session(module):
             wrap_function_wrapper(module, "ClientSession.read_resource", wrap_read_resource)
         if hasattr(module.ClientSession, "get_prompt"):
             wrap_function_wrapper(module, "ClientSession.get_prompt", wrap_get_prompt)
+
+
+def instrument_mcp_server_fastmcp_tools_tool_manager(module):
+    if hasattr(module, "ToolManager"):
+        if hasattr(module.ToolManager, "call_tool"):
+            wrap_function_wrapper(module, "ToolManager.call_tool", wrap_call_tool)

tests/adapter_mcp/conftest.py

@@ -22,6 +22,7 @@
     "transaction_tracer.stack_trace_threshold": 0.0,
     "debug.log_data_collector_payloads": True,
     "debug.record_transaction_failure": True,
+    "ai_monitoring.enabled": True,
 }
 
 collector_agent_registration = collector_agent_registration_fixture(

tests/adapter_mcp/test_mcp.py

@@ -16,8 +16,9 @@
 from fastmcp.client import Client
 from fastmcp.client.transports import FastMCPTransport
 from fastmcp.server.server import FastMCP
-from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
+from mcp.server.fastmcp.tools import ToolManager
 
+from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
 from newrelic.api.background_task import background_task
 
 
@@ -49,13 +50,13 @@ def echo_prompt(message: str):
 
 
 @validate_transaction_metrics(
-    "test_mcp:test_tool_tracing",
+    "test_mcp:test_tool_tracing_via_client_session",
     scoped_metrics=[("Llm/tool/MCP/mcp.client.session:ClientSession.call_tool/add_exclamation", 1)],
     rollup_metrics=[("Llm/tool/MCP/mcp.client.session:ClientSession.call_tool/add_exclamation", 1)],
     background_task=True,
 )
 @background_task()
-def test_tool_tracing(loop, fastmcp_server):
+def test_tool_tracing_via_client_session(loop, fastmcp_server):
     async def _test():
         async with Client(transport=FastMCPTransport(fastmcp_server)) as client:
             # Call the MCP tool, so we can validate the trace naming is correct.
@@ -67,6 +68,26 @@ async def _test():
     loop.run_until_complete(_test())
 
 
+@validate_transaction_metrics(
+    "test_mcp:test_tool_tracing_via_tool_manager",
+    scoped_metrics=[("Llm/tool/MCP/mcp.server.fastmcp.tools.tool_manager:ToolManager.call_tool/add_exclamation", 1)],
+    rollup_metrics=[("Llm/tool/MCP/mcp.server.fastmcp.tools.tool_manager:ToolManager.call_tool/add_exclamation", 1)],
+    background_task=True,
+)
+@background_task()
+def test_tool_tracing_via_tool_manager(loop):
+    async def _test():
+        def add_exclamation(phrase):
+            return f"{phrase}!"
+
+        manager = ToolManager()
+        manager.add_tool(add_exclamation)
+        result = await manager.call_tool("add_exclamation", {"phrase": "Python is awesome"})
+        assert result == "Python is awesome!"
+
+    loop.run_until_complete(_test())
+
+
 # Separate out the test function to work with the transaction metrics validator
 def run_read_resources(loop, fastmcp_server, resource_uri):
     async def _test():