Add tool manager spans to MCP instrumentation to support streaming.

Author: umaannamalai

newrelic/config.py

@@ -2869,12 +2869,8 @@ def _process_module_builtin_defaults():
     _process_module_definition("loguru._logger", "newrelic.hooks.logger_loguru", "instrument_loguru_logger")
 
     _process_module_definition("mcp.client.session", "newrelic.hooks.adapter_mcp", "instrument_mcp_client_session")
-    _process_module_definition("mcp.server.fastmcp.server", "newrelic.hooks.adapter_mcp", "instrument_mcp_server_fastmcp_server")
-    _process_module_definition("mcp.server.fastmcp.utilities.func_metadata", "newrelic.hooks.adapter_mcp", "instrument_mcp_server_fastmcp_utilities_func_metadata")
-
     _process_module_definition("mcp.server.fastmcp.tools.tool_manager", "newrelic.hooks.adapter_mcp", "instrument_mcp_server_fastmcp_tools_tool_manager")
-    _process_module_definition("mcp.server.fastmcp.tools.base", "newrelic.hooks.adapter_mcp", "instrument_mcp_server_fastmcp_tools_base")
-    _process_module_definition("mcp.server.lowlevel.server", "newrelic.hooks.adapter_mcp", "instrument_mcp_server_lowlevel_server")
+
 
     _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

@@ -43,20 +43,6 @@ async def wrap_call_tool(wrapped, instance, args, kwargs):
         return await wrapped(*args, **kwargs)
 
 
-async def wrap_call_tool1(wrapped, instance, args, kwargs):
-    transaction = current_transaction()
-    if not transaction:
-        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"
-    function_trace_name = f"{func_name}/{tool_name}"
-
-    with FunctionTrace(name=function_trace_name, group="Llm/tool/MCPPPPP", source=wrapped):
-        return await wrapped(*args, **kwargs)
-
-
 async def wrap_read_resource(wrapped, instance, args, kwargs):
     transaction = current_transaction()
     if not transaction:
@@ -101,27 +87,6 @@ async def wrap_get_prompt(wrapped, instance, args, kwargs):
         return await wrapped(*args, **kwargs)
 
 
-async def wrap_call_fn_with_arg_validation(wrapped, instance, args, kwargs):
-    with FunctionTrace(name="call_fn_with_arg_validation", group="Llm/argssss/MCP", source=wrapped):
-        return await wrapped(*args, **kwargs)
-
-
-async def wrap_run(wrapped, instance, args, kwargs):
-    with FunctionTrace(name="call_fn_with_arg_validation", group="Llm/MCP/base/tool/run", source=wrapped):
-        return await wrapped(*args, **kwargs)
-
-
-async def wrap_tool_manager(wrapped, instance, args, kwargs):
-    with FunctionTrace(name="call_tool", group="Llm/MCP/tool_manager", source=wrapped):
-        return await wrapped(*args, **kwargs)
-
-
-async def wrap__handle_request(wrapped, instance, args, kwargs):
-    func_name = callable_name(wrapped)
-    with FunctionTrace(name="func_name", group="Llm/MCP/handle_request", source=wrapped):
-        return await wrapped(*args, **kwargs)
-
-
 def instrument_mcp_client_session(module):
     if hasattr(module, "ClientSession"):
         if hasattr(module.ClientSession, "call_tool"):
@@ -132,31 +97,7 @@ def instrument_mcp_client_session(module):
             wrap_function_wrapper(module, "ClientSession.get_prompt", wrap_get_prompt)
 
 
-def instrument_mcp_server_fastmcp_server(module):
-    if hasattr(module, "FastMCP"):
-        if hasattr(module.FastMCP, "call_tool"):
-            wrap_function_wrapper(module, "FastMCP.call_tool", wrap_call_tool1)
-
-
-def instrument_mcp_server_lowlevel_server(module):
-    wrap_function_wrapper(module, "Server._handle_request", wrap__handle_request)
-
-
-def instrument_mcp_server_fastmcp_utilities_func_metadata(module):
-    pass
-    # if hasattr(module, "FuncMetadata"):
-    #     if hasattr(module.FuncMetadata, "call_fn_with_arg_validation"):
-    #         wrap_function_wrapper(module, "FuncMetadata.call_fn_with_arg_validation", wrap_call_fn_with_arg_validation)
-
-
-def instrument_mcp_server_fastmcp_tools_base(module):
-    pass
-    # if hasattr(module, "Tool"):
-    #     if hasattr(module.Tool, "run"):
-    #         wrap_function_wrapper(module, "Tool.run", wrap_run)
-
-
 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)
+            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():