Skip to content

Commit 431e3c2

Browse files
pandegoharanrk
authored andcommitted
fix(mcp): propagate trace context in default streamable HTTP client
Wraps the default MCP `httpx_client_factory` so that outbound streamable HTTP MCP requests reuse the active OpenTelemetry context when the OTel HTTPX instrumentation is installed. When the instrumentation is not installed the behavior is unchanged (graceful no-op), and the existing customization path via `httpx_client_factory` is preserved. Fixes #4768 Co-authored-by: Haran Rajkumar <haranrk@google.com> PiperOrigin-RevId: 936247683
1 parent 21512ea commit 431e3c2

2 files changed

Lines changed: 66 additions & 1 deletion

File tree

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

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,19 +61,42 @@ class AsyncAuthorizedSession: # pylint: disable=g-bad-classes
6161
from mcp.client.session import SamplingFnT
6262
from mcp.client.sse import sse_client
6363
from mcp.client.stdio import stdio_client
64-
from mcp.client.streamable_http import create_mcp_http_client
64+
from mcp.client.streamable_http import create_mcp_http_client as _create_mcp_http_client
6565
from mcp.client.streamable_http import McpHttpClientFactory
6666
from mcp.client.streamable_http import streamablehttp_client
6767
from pydantic import BaseModel
6868
from pydantic import ConfigDict
6969

70+
try:
71+
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
72+
73+
_HAS_HTTPX_INSTRUMENTOR = True
74+
except (ImportError, AttributeError):
75+
_HAS_HTTPX_INSTRUMENTOR = False
76+
7077
from ...features import FeatureName
7178
from ...features import is_feature_enabled
7279
from .session_context import SessionContext
7380

7481
logger = logging.getLogger('google_adk.' + __name__)
7582

7683

84+
def create_mcp_http_client(
85+
headers: dict[str, str] | None = None,
86+
timeout: httpx.Timeout | None = None,
87+
auth: httpx.Auth | None = None,
88+
) -> httpx.AsyncClient:
89+
"""Creates MCP HTTP client and instruments it when OTel is available."""
90+
client = _create_mcp_http_client(
91+
headers=headers,
92+
timeout=timeout,
93+
auth=auth,
94+
)
95+
if _HAS_HTTPX_INSTRUMENTOR:
96+
HTTPXClientInstrumentor.instrument_client(client)
97+
return client
98+
99+
77100
def _has_cancelled_error_context(exc: BaseException) -> bool:
78101
"""Returns True if `exc` is/was caused by `asyncio.CancelledError`.
79102

‎tests/unittests/tools/mcp_tool/test_mcp_session_manager.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525

2626
from google.adk.platform import thread as platform_thread
2727
from google.adk.tools.mcp_tool.mcp_session_manager import _SharedAsyncTransport
28+
from google.adk.tools.mcp_tool.mcp_session_manager import create_mcp_http_client
2829
from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
2930
from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_errors
3031
from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
@@ -243,6 +244,47 @@ def test_init_with_streamable_http_default_httpx_factory(
243244
].get_default(),
244245
)
245246

247+
@patch(
248+
"google.adk.tools.mcp_tool.mcp_session_manager.HTTPXClientInstrumentor",
249+
create=True,
250+
)
251+
@patch(
252+
"google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client"
253+
)
254+
@patch(
255+
"google.adk.tools.mcp_tool.mcp_session_manager._HAS_HTTPX_INSTRUMENTOR",
256+
True,
257+
)
258+
def test_default_httpx_factory_instruments_client_when_available(
259+
self, mock_base_factory, mock_instrumentor
260+
):
261+
"""Test default MCP HTTP factory instruments HTTPX client when available."""
262+
client = Mock()
263+
mock_base_factory.return_value = client
264+
265+
result = create_mcp_http_client()
266+
267+
assert result is client
268+
mock_instrumentor.instrument_client.assert_called_once_with(client)
269+
270+
@patch(
271+
"google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client"
272+
)
273+
@patch(
274+
"google.adk.tools.mcp_tool.mcp_session_manager._HAS_HTTPX_INSTRUMENTOR",
275+
False,
276+
)
277+
def test_default_httpx_factory_handles_missing_opentelemetry(
278+
self, mock_base_factory
279+
):
280+
"""Test default MCP HTTP factory works without OTel instrumentation."""
281+
client = Mock()
282+
mock_base_factory.return_value = client
283+
284+
result = create_mcp_http_client()
285+
286+
assert result is client
287+
246288
def test_generate_session_key_stdio(self):
247289
"""Test session key generation for stdio connections."""
248290
manager = MCPSessionManager(self.mock_stdio_connection_params)

0 commit comments

Comments
 (0)