Add instrumentation for autogen MCP tool adapter. (#1409)

* Add instrumentation for MCP tool adapter.

* [MegaLinter] Apply linters fixes

Author: umaannamalai

newrelic/config.py

@@ -2863,6 +2863,10 @@ 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(
+        "autogen_ext.tools.mcp._base", "newrelic.hooks.mlmodel_autogen", "instrument_autogen_ext_tools_mcp__base"
+    )
+
     _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")

newrelic/hooks/mlmodel_autogen.py

@@ -0,0 +1,38 @@
+# 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 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
+
+
+async def wrap_from_server_params(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("tool_name") or "tool"
+    function_trace_name = f"{func_name}/{tool_name}"
+    with FunctionTrace(name=function_trace_name, group="Llm", source=wrapped):
+        return await wrapped(*args, **kwargs)
+
+
+def instrument_autogen_ext_tools_mcp__base(module):
+    if hasattr(module, "McpToolAdapter"):
+        if hasattr(module.McpToolAdapter, "from_server_params"):
+            wrap_function_wrapper(module, "McpToolAdapter.from_server_params", wrap_from_server_params)

tests/mlmodel_autogen/conftest.py

@@ -0,0 +1,30 @@
+# 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,
+    "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,
+    "ai_monitoring.enabled": True,
+}
+
+collector_agent_registration = collector_agent_registration_fixture(
+    app_name="Python Agent Test (mlmodel_autogen)", default_settings=_default_settings
+)

tests/mlmodel_autogen/test_mcp_tool_adapter.py

@@ -0,0 +1,67 @@
+# 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 unittest.mock import AsyncMock
+
+import pytest
+from autogen_ext.tools.mcp import SseMcpToolAdapter, SseServerParams
+from mcp import ClientSession, Tool
+from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
+
+from newrelic.api.background_task import background_task
+
+# Test setup derived from: https://github.com/microsoft/autogen/blob/main/python/packages/autogen-ext/tests/tools/test_mcp_tools.py
+# autogen MIT license: https://github.com/microsoft/autogen/blob/main/LICENSE and
+# https://github.com/microsoft/autogen/blob/main/LICENSE-CODE
+
+
+@pytest.fixture
+def mock_sse_session():
+    session = AsyncMock(spec=ClientSession)
+    session.initialize = AsyncMock()
+    session.call_tool = AsyncMock()
+    session.list_tools = AsyncMock()
+    return session
+
+
+@pytest.fixture
+def add_exclamation():
+    return Tool(
+        name="add_exclamation",
+        description="A test SSE tool that adds an exclamation mark to a string",
+        inputSchema={"type": "object", "properties": {"input": {"type": "string"}}, "required": ["input"]},
+    )
+
+
+@validate_transaction_metrics(
+    "test_mcp_tool_adapter:test_from_server_params_tracing",
+    scoped_metrics=[("Llm/autogen_ext.tools.mcp._sse:SseMcpToolAdapter.from_server_params/add_exclamation", 1)],
+    rollup_metrics=[("Llm/autogen_ext.tools.mcp._sse:SseMcpToolAdapter.from_server_params/add_exclamation", 1)],
+    background_task=True,
+)
+@background_task()
+def test_from_server_params_tracing(loop, mock_sse_session, monkeypatch, add_exclamation):
+    async def _test():
+        params = SseServerParams(url="http://test-url")
+        mock_context = AsyncMock()
+        mock_context.__aenter__.return_value = mock_sse_session
+        monkeypatch.setattr(
+            "autogen_ext.tools.mcp._base.create_mcp_server_session", lambda *args, **kwargs: mock_context
+        )
+
+        mock_sse_session.list_tools.return_value.tools = [add_exclamation]
+
+        adapter = await SseMcpToolAdapter.from_server_params(params, "add_exclamation")
+
+    loop.run_until_complete(_test())

tox.ini

@@ -155,6 +155,7 @@ envlist =
     python-logger_logging-{py37,py38,py39,py310,py311,py312,py313,pypy310},
     python-logger_loguru-{py37,py38,py39,py310,py311,py312,py313,pypy310}-logurulatest,
     python-logger_structlog-{py37,py38,py39,py310,py311,py312,py313,pypy310}-structloglatest,
+    python-mlmodel_autogen-{py310,py311,py312,py313,pypy310},
     python-mlmodel_gemini-{py39,py310,py311,py312,py313},
     python-mlmodel_langchain-{py39,py310,py311,py312},
     ;; Package not ready for Python 3.13 (uses an older version of numpy)
@@ -406,6 +407,8 @@ deps =
     framework_tornado: pycurl
     framework_tornado-tornadolatest: tornado
     framework_tornado-tornadomaster: https://github.com/tornadoweb/tornado/archive/master.zip
+    mlmodel_autogen: autogen-ext
+    mlmodel_autogen: mcp
     mlmodel_gemini: google-genai
     mlmodel_openai-openai0: openai[datalib]<1.0
     mlmodel_openai-openai107: openai[datalib]<1.8
@@ -562,6 +565,7 @@ changedir =
     messagebroker_kafkapython: tests/messagebroker_kafkapython
     messagebroker_kombu: tests/messagebroker_kombu
     messagebroker_pika: tests/messagebroker_pika
+    mlmodel_autogen: tests/mlmodel_autogen
     mlmodel_gemini: tests/mlmodel_gemini
     mlmodel_langchain: tests/mlmodel_langchain
     mlmodel_openai: tests/mlmodel_openai