Update MCP instrumentation to check if AIM is enabled. (#1456)

* Update MCP instrumentation to check if AIM is enabled.

* Linting

---------

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

Author: 3 people

newrelic/hooks/adapter_mcp.py

@@ -21,7 +21,6 @@
 from newrelic.common.signature import bind_args
 from newrelic.core.config import global_settings
 
-
 _logger = logging.getLogger(__name__)
 
 
@@ -30,7 +29,7 @@ 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()
+    settings = transaction.settings or global_settings()
     if not settings.ai_monitoring.enabled:
         return await wrapped(*args, **kwargs)
 
@@ -48,7 +47,7 @@ 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()
+    settings = transaction.settings or global_settings()
     if not settings.ai_monitoring.enabled:
         return await wrapped(*args, **kwargs)
 
@@ -74,7 +73,7 @@ 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()
+    settings = transaction.settings or global_settings()
     if not settings.ai_monitoring.enabled:
         return await wrapped(*args, **kwargs)
 

tests/adapter_mcp/test_mcp.py

@@ -13,11 +13,17 @@
 # limitations under the License.
 
 import pytest
+
 from fastmcp.client import Client
 from fastmcp.client.transports import FastMCPTransport
 from fastmcp.server.server import FastMCP
 from mcp.server.fastmcp.tools import ToolManager
 
+from testing_support.fixtures import function_not_called
+from testing_support.ml_testing_utils import disabled_ai_monitoring_settings
+from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
+
+
 from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
 from newrelic.api.background_task import background_task
 
@@ -61,7 +67,6 @@ 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.content[0])
         assert "Python is awesome!" in content
 
@@ -135,3 +140,38 @@ async def _test():
         assert "Python is cool" in content
 
     loop.run_until_complete(_test())
+
+
+@disabled_ai_monitoring_settings
+@function_not_called("newrelic.api.function_trace", "FunctionTrace.__enter__")
+@background_task()
+def test_tool_tracing_aim_disabled(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.content[0])
+        assert "Python is awesome!" in content
+
+    loop.run_until_complete(_test())
+
+
+@disabled_ai_monitoring_settings
+@function_not_called("newrelic.api.function_trace", "FunctionTrace.__enter__")
+@background_task()
+def test_resource_tracing_aim_disabled(loop, fastmcp_server):
+    run_read_resources(loop, fastmcp_server, "greeting://Python")
+
+
+@disabled_ai_monitoring_settings
+@function_not_called("newrelic.api.function_trace", "FunctionTrace.__enter__")
+@background_task()
+def test_prompt_tracing_aim_disabled(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())