From ba84850a725b20729b751f2d11cf7ba071595a0d Mon Sep 17 00:00:00 2001 From: Michael <7780875+pandego@users.noreply.github.com> Date: Mon, 9 Mar 2026 09:05:23 +0100 Subject: [PATCH 1/2] fix(mcp): instrument default streamable http client for tracing --- .../adk/tools/mcp_tool/mcp_session_manager.py | 22 +++++++++- .../mcp_tool/test_mcp_session_manager.py | 44 +++++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index e0cd1ebc89..99d36d4fcd 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -38,7 +38,7 @@ from mcp.client.session import SamplingFnT from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client -from mcp.client.streamable_http import create_mcp_http_client +from mcp.client.streamable_http import create_mcp_http_client as _create_mcp_http_client from mcp.client.streamable_http import McpHttpClientFactory from mcp.client.streamable_http import streamablehttp_client from pydantic import BaseModel @@ -49,6 +49,26 @@ logger = logging.getLogger('google_adk.' + __name__) +def create_mcp_http_client( + headers=None, + timeout=None, + auth=None, +): + """Creates MCP HTTP client and instruments it when OTel is available.""" + client = _create_mcp_http_client( + headers=headers, + timeout=timeout, + auth=auth, + ) + try: + from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor + + HTTPXClientInstrumentor.instrument_client(client) + except ImportError: + pass + return client + + def _has_cancelled_error_context(exc: BaseException) -> bool: """Returns True if `exc` is/was caused by `asyncio.CancelledError`. diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index 327df114a8..098c4a6fc5 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -18,12 +18,14 @@ from io import StringIO import json import sys +import builtins from unittest.mock import ANY from unittest.mock import AsyncMock from unittest.mock import Mock from unittest.mock import patch from google.adk.platform import thread as platform_thread +from google.adk.tools.mcp_tool.mcp_session_manager import create_mcp_http_client from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_errors from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams @@ -191,6 +193,48 @@ def test_init_with_streamable_http_default_httpx_factory( ].get_default(), ) + @patch("google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client") + def test_default_httpx_factory_instruments_client_when_available( + self, mock_base_factory + ): + """Test default MCP HTTP factory instruments HTTPX client when available.""" + client = Mock() + mock_base_factory.return_value = client + + mock_instrumentor = Mock() + with patch.dict( + sys.modules, + { + "opentelemetry.instrumentation.httpx": Mock( + HTTPXClientInstrumentor=mock_instrumentor + ) + }, + ): + result = create_mcp_http_client() + + assert result is client + mock_instrumentor.instrument_client.assert_called_once_with(client) + + @patch("google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client") + def test_default_httpx_factory_handles_missing_opentelemetry( + self, mock_base_factory + ): + """Test default MCP HTTP factory works without OTel instrumentation.""" + client = Mock() + mock_base_factory.return_value = client + + original_import = builtins.__import__ + + def import_with_missing_otel(name, *args, **kwargs): + if name == "opentelemetry.instrumentation.httpx": + raise ImportError("missing test dependency") + return original_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=import_with_missing_otel): + result = create_mcp_http_client() + + assert result is client + def test_generate_session_key_stdio(self): """Test session key generation for stdio connections.""" manager = MCPSessionManager(self.mock_stdio_connection_params) From 4da844e87da19cb214e8bd31a6b6ffe60cf55fef Mon Sep 17 00:00:00 2001 From: Michael <7780875+pandego@users.noreply.github.com> Date: Mon, 9 Mar 2026 21:35:36 +0100 Subject: [PATCH 2/2] fix(mcp): satisfy CI for trace context client factory --- src/google/adk/tools/mcp_tool/mcp_session_manager.py | 9 +++++---- .../tools/mcp_tool/test_mcp_session_manager.py | 10 +++++++--- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index 99d36d4fcd..12816de33d 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -32,6 +32,7 @@ from typing import TextIO from typing import Union +import httpx from mcp import ClientSession from mcp import SamplingCapability from mcp import StdioServerParameters @@ -50,10 +51,10 @@ def create_mcp_http_client( - headers=None, - timeout=None, - auth=None, -): + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, +) -> httpx.AsyncClient: """Creates MCP HTTP client and instruments it when OTel is available.""" client = _create_mcp_http_client( headers=headers, diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index 098c4a6fc5..6d0ecab331 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -13,12 +13,12 @@ # limitations under the License. import asyncio +import builtins from datetime import timedelta import hashlib from io import StringIO import json import sys -import builtins from unittest.mock import ANY from unittest.mock import AsyncMock from unittest.mock import Mock @@ -193,7 +193,9 @@ def test_init_with_streamable_http_default_httpx_factory( ].get_default(), ) - @patch("google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client") + @patch( + "google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client" + ) def test_default_httpx_factory_instruments_client_when_available( self, mock_base_factory ): @@ -215,7 +217,9 @@ def test_default_httpx_factory_instruments_client_when_available( assert result is client mock_instrumentor.instrument_client.assert_called_once_with(client) - @patch("google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client") + @patch( + "google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client" + ) def test_default_httpx_factory_handles_missing_opentelemetry( self, mock_base_factory ):