Merge pull request #1403 from newrelic/feature-mcp-instrumentation

Add MCP Instrumentation

Author: TimPansino

newrelic/config.py

@@ -2861,6 +2861,9 @@ def _process_module_builtin_defaults():
 
     _process_module_definition("loguru", "newrelic.hooks.logger_loguru", "instrument_loguru")
     _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("structlog._base", "newrelic.hooks.logger_structlog", "instrument_structlog__base")
     _process_module_definition("structlog._frames", "newrelic.hooks.logger_structlog", "instrument_structlog__frames")
     _process_module_definition("paste.httpserver", "newrelic.hooks.adapter_paste", "instrument_paste_httpserver")

newrelic/hooks/adapter_mcp.py

@@ -0,0 +1,83 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from newrelic.api.function_trace import FunctionTrace
+from newrelic.api.transaction import current_transaction
+from newrelic.common.object_names import callable_name
+from newrelic.common.object_wrapper import wrap_function_wrapper
+from newrelic.common.signature import bind_args
+
+_logger = logging.getLogger(__name__)
+
+
+async def wrap_call_tool(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/MCP", source=wrapped):
+        return await wrapped(*args, **kwargs)
+
+
+async def wrap_read_resource(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)
+    # Set a default value in case we can't parse out the URI scheme successfully
+    resource_scheme = "resource"
+
+    try:
+        resource_uri = bound_args.get("uri")
+        resource_scheme = getattr(resource_uri, "scheme", "resource")
+    except Exception:
+        _logger.warning("Unable to parse resource URI scheme for MCP read_resource call")
+
+    function_trace_name = f"{func_name}/{resource_scheme}"
+
+    with FunctionTrace(name=function_trace_name, group="Llm/resource/MCP", source=wrapped):
+        return await wrapped(*args, **kwargs)
+
+
+async def wrap_get_prompt(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)
+    prompt_name = bound_args.get("name") or "prompt"
+    function_trace_name = f"{func_name}/{prompt_name}"
+
+    with FunctionTrace(name=function_trace_name, group="Llm/prompt/MCP", source=wrapped):
+        return await wrapped(*args, **kwargs)
+
+
+def instrument_mcp_client_session(module):
+    if hasattr(module, "ClientSession"):
+        if hasattr(module.ClientSession, "call_tool"):
+            wrap_function_wrapper(module, "ClientSession.call_tool", wrap_call_tool)
+        if hasattr(module.ClientSession, "read_resource"):
+            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)

tests/adapter_mcp/conftest.py

@@ -0,0 +1,29 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from testing_support.fixture.event_loop import event_loop as loop
+from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture
+
+_default_settings = {
+    "package_reporting.enabled": False,  # Turn off package reporting for testing as it causes slowdowns.
+    "transaction_tracer.explain_threshold": 0.0,
+    "transaction_tracer.transaction_threshold": 0.0,
+    "transaction_tracer.stack_trace_threshold": 0.0,
+    "debug.log_data_collector_payloads": True,
+    "debug.record_transaction_failure": True,
+}
+
+collector_agent_registration = collector_agent_registration_fixture(
+    app_name="Python Agent Test (adapter_mcp)", default_settings=_default_settings
+)

tests/adapter_mcp/test_mcp.py

@@ -0,0 +1,116 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+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 newrelic.api.background_task import background_task
+
+
+@pytest.fixture
+def fastmcp_server():
+    server = FastMCP("Test Tools")
+
+    @server.tool()
+    def add_exclamation(phrase):
+        return f"{phrase}!"
+
+    @server.resource("greeting://{name}")
+    def get_greeting(name):
+        return f"Hello, {name}!"
+
+    @server.resource("files:///home/user/documents/{filename}")
+    def get_file(filename):
+        return f"Hello, Python! You requested the file: {filename}."
+
+    @server.resource("postgres://python/customers/schema")
+    def get_db_info():
+        return "Hello, Python!"
+
+    @server.prompt()
+    def echo_prompt(message: str):
+        return f"Echo this message: {message}"
+
+    return server
+
+
+@validate_transaction_metrics(
+    "test_mcp:test_tool_tracing",
+    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):
+    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.
+            result = await client.call_tool("add_exclamation", {"phrase": "Python is awesome"})
+
+            content = str(result[0])
+        assert "Python is awesome!" in content
+
+    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():
+        async with Client(transport=FastMCPTransport(fastmcp_server)) as client:
+            result = await client.read_resource(resource_uri)
+            content = str(result[0])
+        assert "Hello, Python!" in content
+
+    loop.run_until_complete(_test())
+
+
+@pytest.mark.parametrize(
+    "resource_uri",
+    ("greeting://Python", "files:///home/user/documents/python.pdf", "postgres://python/customers/schema"),
+)
+def test_resource_tracing(loop, fastmcp_server, resource_uri):
+    resource_scheme = resource_uri.split(":/")[0]
+
+    @validate_transaction_metrics(
+        "test_mcp:test_resource_tracing.<locals>._test",
+        scoped_metrics=[(f"Llm/resource/MCP/mcp.client.session:ClientSession.read_resource/{resource_scheme}", 1)],
+        rollup_metrics=[(f"Llm/resource/MCP/mcp.client.session:ClientSession.read_resource/{resource_scheme}", 1)],
+        background_task=True,
+    )
+    @background_task()
+    def _test():
+        run_read_resources(loop, fastmcp_server, resource_uri)
+
+    _test()
+
+
+@validate_transaction_metrics(
+    "test_mcp:test_prompt_tracing",
+    scoped_metrics=[("Llm/prompt/MCP/mcp.client.session:ClientSession.get_prompt/echo_prompt", 1)],
+    rollup_metrics=[("Llm/prompt/MCP/mcp.client.session:ClientSession.get_prompt/echo_prompt", 1)],
+    background_task=True,
+)
+@background_task()
+def test_prompt_tracing(loop, fastmcp_server):
+    async def _test():
+        async with Client(transport=FastMCPTransport(fastmcp_server)) as client:
+            result = await client.get_prompt("echo_prompt", {"message": "Python is cool"})
+
+            content = str(result)
+        assert "Python is cool" in content
+
+    loop.run_until_complete(_test())

tox.ini

@@ -83,6 +83,7 @@ envlist =
     python-adapter_gunicorn-{py37,py38,py39,py310,py311,py312,py313}-aiohttp03-gunicornlatest,
     python-adapter_hypercorn-{py38,py39,py310,py311,py312,py313}-hypercornlatest,
     python-adapter_hypercorn-py38-hypercorn{0010,0011,0012,0013},
+    python-adapter_mcp-{py310,py311,py312,py313,pypy310},
     python-adapter_uvicorn-{py37,py38,py39,py310,py311,py312,py313}-uvicornlatest,
     python-adapter_uvicorn-py38-uvicorn014,
     python-adapter_waitress-{py37,py38,py39,py310,py311,py312,py313}-waitresslatest,
@@ -210,6 +211,7 @@ deps =
     adapter_hypercorn-hypercorn0011: hypercorn[h3]<0.12
     adapter_hypercorn-hypercorn0010: hypercorn[h3]<0.11
     adapter_hypercorn: niquests
+    adapter_mcp: fastmcp
     adapter_uvicorn-uvicorn014: uvicorn<0.15
     adapter_uvicorn-uvicornlatest: uvicorn
     adapter_uvicorn: typing-extensions
@@ -477,6 +479,7 @@ changedir =
     adapter_gevent: tests/adapter_gevent
     adapter_gunicorn: tests/adapter_gunicorn
     adapter_hypercorn: tests/adapter_hypercorn
+    adapter_mcp: tests/adapter_mcp
     adapter_uvicorn: tests/adapter_uvicorn
     adapter_waitress: tests/adapter_waitress
     agent_features: tests/agent_features