Skip to content

Commit c34c12f

Browse files
committed
fix(mcp): instrument default streamable http client for tracing
1 parent 9d15517 commit c34c12f

File tree

2 files changed

+65
-1
lines changed

2 files changed

+65
-1
lines changed

src/google/adk/tools/mcp_tool/mcp_session_manager.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@
3636
from mcp import StdioServerParameters
3737
from mcp.client.sse import sse_client
3838
from mcp.client.stdio import stdio_client
39-
from mcp.client.streamable_http import create_mcp_http_client
39+
from mcp.client.streamable_http import create_mcp_http_client as _create_mcp_http_client
4040
from mcp.client.streamable_http import McpHttpClientFactory
4141
from mcp.client.streamable_http import streamablehttp_client
4242
from pydantic import BaseModel
@@ -47,6 +47,26 @@
4747
logger = logging.getLogger('google_adk.' + __name__)
4848

4949

50+
def create_mcp_http_client(
51+
headers=None,
52+
timeout=None,
53+
auth=None,
54+
):
55+
"""Creates MCP HTTP client and instruments it when OTel is available."""
56+
client = _create_mcp_http_client(
57+
headers=headers,
58+
timeout=timeout,
59+
auth=auth,
60+
)
61+
try:
62+
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
63+
64+
HTTPXClientInstrumentor.instrument_client(client)
65+
except ImportError:
66+
pass
67+
return client
68+
69+
5070
def _has_cancelled_error_context(exc: BaseException) -> bool:
5171
"""Returns True if `exc` is/was caused by `asyncio.CancelledError`.
5272

tests/unittests/tools/mcp_tool/test_mcp_session_manager.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,14 @@
1818
from io import StringIO
1919
import json
2020
import sys
21+
import builtins
2122
from unittest.mock import ANY
2223
from unittest.mock import AsyncMock
2324
from unittest.mock import Mock
2425
from unittest.mock import patch
2526

2627
from google.adk.platform import thread as platform_thread
28+
from google.adk.tools.mcp_tool.mcp_session_manager import create_mcp_http_client
2729
from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
2830
from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_errors
2931
from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
@@ -191,6 +193,48 @@ def test_init_with_streamable_http_default_httpx_factory(
191193
].get_default(),
192194
)
193195

196+
@patch("google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client")
197+
def test_default_httpx_factory_instruments_client_when_available(
198+
self, mock_base_factory
199+
):
200+
"""Test default MCP HTTP factory instruments HTTPX client when available."""
201+
client = Mock()
202+
mock_base_factory.return_value = client
203+
204+
mock_instrumentor = Mock()
205+
with patch.dict(
206+
sys.modules,
207+
{
208+
"opentelemetry.instrumentation.httpx": Mock(
209+
HTTPXClientInstrumentor=mock_instrumentor
210+
)
211+
},
212+
):
213+
result = create_mcp_http_client()
214+
215+
assert result is client
216+
mock_instrumentor.instrument_client.assert_called_once_with(client)
217+
218+
@patch("google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client")
219+
def test_default_httpx_factory_handles_missing_opentelemetry(
220+
self, mock_base_factory
221+
):
222+
"""Test default MCP HTTP factory works without OTel instrumentation."""
223+
client = Mock()
224+
mock_base_factory.return_value = client
225+
226+
original_import = builtins.__import__
227+
228+
def import_with_missing_otel(name, *args, **kwargs):
229+
if name == "opentelemetry.instrumentation.httpx":
230+
raise ImportError("missing test dependency")
231+
return original_import(name, *args, **kwargs)
232+
233+
with patch("builtins.__import__", side_effect=import_with_missing_otel):
234+
result = create_mcp_http_client()
235+
236+
assert result is client
237+
194238
def test_generate_session_key_stdio(self):
195239
"""Test session key generation for stdio connections."""
196240
manager = MCPSessionManager(self.mock_stdio_connection_params)

0 commit comments

Comments
 (0)